using SparseArrays
using IterativeSolvers

struct sparsetensor_vec
    # dims is the array of dimensions of the tensor
    dims :: Array{Int64}
    # spv stores the values of the tensor in a sparse vector
    spv :: SparseVector{Float64}
end

struct sparsetensor_mat
    # dims1 is the array of dimensions of the tensor
    # used for the matrix dimension one
    dims1 :: Array{Int64}
    # dims2 is the array of dimensions of the tensor
    # used for the matrix dimension two
    dims2 :: Array{Int64}
    # spm stores the values of the tensor in a sparse matrix
    spm :: SparseMatrixCSC{Float64}
end

function spvec_to_spmat(x :: sparsetensor_vec, lindex)
    # Constructs the sparse matrix x[1:lindex, lindex+1:end]
    res_dims1 = x.dims[1:lindex]
    res_dims2 = x.dims[lindex+1:end]
    m = prod(res_dims1)
    n = prod(res_dims2)
    n_nz = nnz(x.spv)
    z = map(collect, map(a -> divrem(a,m),x.spv.nzind))
    i = zeros(Int, n_nz)
    j = zeros(Int, n_nz)
    for k = 1:n_nz
        i[k] = Int(z[k][2])
        j[k] = Int(z[k][1]) + 1
		if i[k] == 0
			i[k] = m
			j[k] += (-1)
		end
    end
    res_ma = sparse(i, j, x.spv.nzval, m, n)
    return sparsetensor_mat(res_dims1, res_dims2, res_ma)
end

function spmat_to_spvec(x :: sparsetensor_mat)
	# Transforms the sparse matrix into a sparse vector
	res_dim = vcat(x.dims1,x.dims2)
    return sparsetensor_vec(res_dim, reshape(x.spm,prod(res_dim)))
end

function sppermutedims(x :: sparsetensor_vec, perm)
    # permutes the dimensions of the sparse tensor

    # Prepare the transformation to cartesian coordinates
    n_prod = ones(length(perm))
    for k = 1:(length(perm))
        n_prod[1:(end-k)] = n_prod[1:(end-k)] .* x.dims[k]
    end
    # Prepare the transformation to vector coordinates
    n_prod_perm = ones(length(perm))
    for k = 1:(length(perm))
        n_prod_perm[1:(end-k)] = n_prod_perm[1:(end-k)] .* x.dims[perm[k]]
    end
    # Initialize the cartesian coordinates and their permutation
    x_katind = ones(length(x.spv.nzind), length(perm))
    x_katind_perm = ones(length(x.spv.nzind), length(perm))
    # Initialize the vector indices
    res_ind = ones(length(x.spv.nzind))

    for x_ind = 1:length(x.spv.nzind)
        #Calculate the cartesian coordinates
        x_ind_rem = x.spv.nzind[x_ind]
        for k = 1:(length(perm))
            x_katind[x_ind, length(perm) - k + 1], x_ind_rem = divrem(x_ind_rem, n_prod[k])
			if x_ind_rem == 0
				x_ind_rem = n_prod[k]
			else
				x_katind[x_ind, length(perm) - k + 1] += 1
			end
        end
        # Apply the permutation on the cartesian coordinates
        for k = 1:length(perm)
            x_katind_perm[x_ind, k] = x_katind[x_ind, perm[k]]
        end
        # Retransform the cartesian coordinates into vector indices
        for k = 1:length(perm)
            res_ind[x_ind] = res_ind[x_ind] + n_prod_perm[k]*(x_katind_perm[x_ind,length(perm) - k + 1]-1)
        end
    end
	return sparsetensor_vec(x.dims[perm], sparsevec(res_ind, x.spv.nzval))
end

function test_sppermutedims()
	n = 5
	x = randn(n,n,n,n)
	x_spv = sparsetensor_vec([n,n,n,n],sparse(x[:]))
	@test isapprox(Array(sppermutedims(x_spv,[2,3,1,4]).spv),permutedims(x,[2,3,1,4])[:])
end

# Redefine permutedims for a sparsetensor_mat
function Base.permutedims(A::sparsetensor_mat, perm, i_new)
	return permutedims(A.spm, perm, A.dims1, A.dims2, i_new)
end

# Redefine permutedims for the structs SparseVector
function Base.permutedims(A::SparseVector, perm, dims)
	return sppermutedims(sparsetensor_vec(dims, A), perm)
end

# Redefine permutedims for the structs SparseMatrixCSC
function Base.permutedims(A::SparseMatrixCSC, perm, dims1, dims2, i_new)
	return spvec_to_spmat(sppermutedims(spmat_to_spvec(sparsetensor_mat(dims1, dims2, A)), perm), i_new).spm
end

# Define ttv_decomp for sparse input
function ttv_spdecomp(spv :: sparsetensor_vec, index;tol=1e-10)
	# Decomposes a tensor into its tensor train with core matrices at i=index
	# For index < d it doesn't work yet
	dims = spv.dims #dims = [n_1,...,n_d]
	n_max = maximum(dims)
	d = length(dims)
	ttv_vec = Array{Array{Float64}}(undef,d)
	# ttv_ot[i]= -1 if i < index
	# ttv_ot[i] = 0 if i = index
	# ttv_ot[i] = 1 if i > index
	ttv_ot = -ones(d)
	ttv_ot[index] = 0
	if index < d
		ttv_ot[index+1:d] = ones(d-index)
	end
	rks = ones(Int, d+2) # ttv_rks will be rks[2:d+1]
	tensor_curr = convert(SparseMatrixCSC, reshape(spv.spv, dims[1], :))
	# Calculate ttv_vec[i] for i < index
	for i = 1 : (index - 1)
		# Reshape the currently left tensor
		tensor_curr = convert(SparseMatrixCSC, reshape(tensor_curr, Int(rks[i] * dims[i]), :))
		# Perform a qr decomposition
		qro = qr(tensor_curr)
		# Define the i-th rank
		diag_R = abs.(diag(qro.R))
		rks[i+1] = sum(diag_R.>tol)
		Q = qro.Q[invperm(qro.prow),(1:length(diag_R))[diag_R.>tol]]
		R = qro.R[diag_R.>tol,invperm(qro.pcol)]
		# Initialize ttv_vec[i]
		ttv_vec[i] = zeros(dims[i],rks[i],rks[i+1])
		# Fill in the ttv_vec[i]
		for x = 1 : dims[i]
			ttv_vec[i][x, :, :] = Q[(rks[i]*(x-1) + 1):(rks[i]*x), :]
		end
		# Update the currently left tensor
		tensor_curr = R
	end
	# Calculate ttv_vec[i] for i = index
	# Reshape the currently left tensor
	tensor_curr = reshape(tensor_curr, Int(dims[index]*rks[index]),:)
	# Initialize ttv_vec[i]
	ttv_vec[index] = zeros(dims[index], rks[index], rks[index+1])
	# Fill in the ttv_vec[i]
	for x = 1 : dims[index]
		ttv_vec[index][x, :, :] =
		tensor_curr[Int(rks[index]*(x-1) + 1):Int(rks[index]*x), 1:rks[index+1]]
	end
	# Define the return value as a ttvector
	return ttvector(ttv_vec, dims, rks[2:d+1], ttv_ot)
end

function test_ttv_spdecomp()
	n = 5
	x = randn(n,n,n,n)
	x_spv = sparsetensor_vec([n,n,n,n],sparse(x[:]))
	x_sp_tt = ttv_spdecomp(x_spv,4)
	size(ttv_to_tensor(x_sp_tt))
	@test isapprox(ttv_to_tensor(x_sp_tt),x)
end

function tto_spdecomp(spm :: sparsetensor_mat, index)
	# Decomposes a sparse tensor operator into its tensor train
	# with core matrices at i=index
	# For index < d it doesn't work yet
	# The tensor has to be given as element of R^[n_1*...*n_d,n_1*...*n_d]
	d = length(spm.dims1)
	tto_dims = spm.dims1
	n_max = maximum(tto_dims)
	dims_sq = tto_dims.^2
	# The matrix is reshaped, reordered and reshaped
	# into the matrix [(x_1,y_1),((x_2,y_2),...,(x_d,y_d))]
	# and decomposed into its tensor train with core matrices at i = index
	index_sorted = vec(reshape(Transpose(reshape(1:(2*d),:,2)),1,:))
	dims_sorted = zeros(2*d)
    dims_sorted[1:d] = tto_dims
	dims_sorted[(d+1):(2*d)] = tto_dims
	dims_sorted = dims_sorted[index_sorted]
	ttv = ttv_spdecomp(spmat_to_spvec(sparsetensor_mat([dims_sq[1]], dims_sq[2:end], permutedims(spm.spm, index_sorted, spm.dims1, spm.dims2, d))), index)
	# Define the array of ranks [r_0=1,r_1,...,r_d]
	rks = ones(Int, d+1)
	rks[2:(d+1)] = ttv.ttv_rks
	r_max = maximum(rks)
	# Initialize tto_vec
	tto_vec = Array{Array{Float64}}(undef,d)
	# Fill in tto_vec
	for i = 1:d
		# Initialize tto_vec[i]
		tto_vec[i] = zeros(tto_dims[i], tto_dims[i], rks[i], rks[i+1])
		# Fill in tto_vec[i]
		tto_vec[i][:, :, :, :] =
		reshape(ttv.ttv_vec[i][1:dims_sq[i], 1:rks[i], 1:rks[i+1]],
		tto_dims[i], tto_dims[i], rks[i], :)
	end
	# Define the return value as a ttoperator
	return ttoperator(tto_vec, tto_dims, rks[2:(d+1)], ttv.ttv_ot)
end

function Lap_sp(n::Integer, d::Integer)
	# Returns the tensor of the discrete Laplacian in a box [0,1]^d
	# with n equidistant discretization points in each direction
	# as an element of R^[n^d,n^d]
	I = zeros((n^d) * (2 * d + 1))
	J = zeros((n^d) * (2 * d + 1))
	V = zeros((n^d) * (2 * d + 1))

	for x = 0:n^d-1
		x_dig = digits(x, base = n, pad = d)
		I[(2*d + 1) * x + 1] = x + 1
		J[(2*d + 1) * x + 1] = x + 1
		V[(2*d + 1) * x + 1] = 2*d
		for i = 0:d-1
			if x_dig[i + 1] == 0
				I[(2*d + 1) * x + 2*i + 2] = x + 1
				J[(2*d + 1) * x + 2*i + 2] = x + n^i + 1
				V[(2*d + 1) * x + 2*i + 2] = -1
			elseif x_dig[i + 1] == n - 1
				I[(2*d + 1) * x + 2*i + 2] = x + 1
				J[(2*d + 1) * x + 2*i + 2] = x - n^i + 1
				V[(2*d + 1) * x + 2*i + 2] = -1
			else
				I[(2*d + 1) * x + 2*i + 2] = x + 1
				J[(2*d + 1) * x + 2*i + 2] = x + n^i + 1
				V[(2*d + 1) * x + 2*i + 2] = -1
				I[(2*d + 1) * x + 2*i + 3] = x + 1
				J[(2*d + 1) * x + 2*i + 3] = x - n^i + 1
				V[(2*d + 1) * x + 2*i + 3] = -1
			end
		end
	end
	I = I[I .!= 0]
	J = J[J .!= 0]
	V = V[V .!= 0]
	return sparse(I , J, V, n^d, n^d)
end

"""
# Test
n = 25
d = 2
L = Lap(n, d)
x_o = ones(size(L,1))
x_o_tt = ttv_decomp(reshape(x_o, n * ones(Int,d)...), d)
b = L * x_o
x_s = map(round, 10 * randn(size(L,1)))
x_s_tt = ttv_decomp(reshape(x_s, n * ones(Int, d)...),d)
b_tt = ttv_decomp(reshape(b,n*ones(Int,d)...), d)
L_spm = sparsetensor_mat(n*ones(Int, d), n*ones(Int, d), L)
L_tt = tto_spdecomp(L_spm, d)
x_als_tt = als(L_tt, b_tt, x_s_tt, x_o_tt.ttv_rks)
"""


module Initial

export initial_sod,initial_static,initial_uniform,initial_one_wave
"""
    this function's inputs are nx, xlim, and gamma

        nx is the number of grid points

        xlim is the limits of the domain

        gamma is the heat capacity ratio

    this function's outputs are ux and uu

        ux has 1 row
        the 1st row is x (x cooridnates)

        uu has 3 rows
        the 1st row is u (velocity)
        the 2nd row is m (mass flow rate)
        the 3rd row is ̂E (total energy per volume)

    this function initialize the Sod shock tube problem (Steger and Warming, 1981)
"""

function initial_sod(nx::Int64, xlim::Any, gamma::Float64)

    Lx = xlim[2]-xlim[1];
    Δx = Lx/(nx-1)

    # Initialization of ux and urhop
    ux = xlim[1]:Δx:xlim[2]
    urhop = Array{Float64,2}(UndefInitializer(), 3,length(ux))

    for i in 1:1:nx
    if (ux[i] < (xlim[2]-xlim[1])/2) urhop[:,i] = [0,1,1]
    else urhop[:,i] = [0,0.125,0.1] end
    end

    #From urhop make uu (u rho*u e)
    uu = Array{Float64,2}(UndefInitializer(), 3,length(ux))
    uu[1,:] = urhop[2,:];
    uu[2,:] = urhop[2,:].*urhop[1,:];
    uu[3,:] = urhop[3,:]/(gamma-1) + 0.5urhop[2,:].*urhop[1,:].^2;

    return ux,uu

end

"""
    this function's inputs are nx, xlim, and gamma

        nx is the number of grid points

        xlim is the limits of the domain

        gamma is the heat capacity ratio

    this function's outputs are ux and uu

        ux has 1 row
        the 1st row is x (x cooridnates)

        uu has 3 rows
        the 1st row is u (velocity)
        the 2nd row is m (mass flow rate)
        the 3rd row is ̂E (total energy per volume)

    this function initialize the Sod shock tube problem (Steger and Warming, 1981)
"""
function initial_static(nx::Int64, xlim::Any, gamma::Float64)

    Lx = xlim[2]-xlim[1];
    Δx = Lx/(nx-1)

    # Initialization of ux and urhop
    ux = xlim[1]:Δx:xlim[2]
    urhop = Array{Float64,2}(UndefInitializer(), 3,length(ux))

    for i in 1:1:nx
    if (ux[i] < (xlim[2]-xlim[1])/2) urhop[:,i] = [0,0.125,1]
    else urhop[:,i] = [0,0.125,1] end
    end

    #From urhop make uu (u rho*u e)
    uu = Array{Float64,2}(UndefInitializer(), 3,length(ux))
    uu[1,:] = urhop[2,:];
    uu[2,:] = urhop[2,:].*urhop[1,:];
    uu[3,:] = urhop[3,:]/(gamma-1) + 0.5urhop[2,:].*urhop[1,:].^2;

    return ux,uu

end

"""
    this function's inputs are nx, xlim, and gamma

        nx is the number of grid points

        xlim is the limits of the domain

        gamma is the heat capacity ratio

    this function's outputs are ux and uu

        ux has 1 row
        the 1st row is x (x cooridnates)

        uu has 3 rows
        the 1st row is u (velocity)
        the 2nd row is m (mass flow rate)
        the 3rd row is ̂E (total energy per volume)

    this function initialize the Sod shock tube problem (Steger and Warming, 1981)
"""
function initial_uniform(nx::Int64, xlim::Any, gamma::Float64)

    Lx = xlim[2]-xlim[1];
    Δx = Lx/(nx-1)

    # Initialization of ux and urhop
    ux = xlim[1]:Δx:xlim[2]
    urhop = Array{Float64,2}(UndefInitializer(), 3,length(ux))

    for i in 1:1:nx
    if (ux[i] < (xlim[2]-xlim[1])/2) urhop[:,i] = [1,0.125,1]
    else urhop[:,i] = [1,0.125,1] end
    end

    #From urhop make uu (u rho*u e)
    uu = Array{Float64,2}(UndefInitializer(), 3,length(ux))
    uu[1,:] = urhop[2,:];
    uu[2,:] = urhop[2,:].*urhop[1,:];
    uu[3,:] = urhop[3,:]/(gamma-1) + 0.5urhop[2,:].*urhop[1,:].^2;

    return ux,uu

end

"""
    this function's inputs are nx, xlim, and gamma

        nx is the number of grid points

        xlim is the limits of the domain

        gamma is the heat capacity ratio

    this function's outputs are ux and uu

        ux has 1 row
        the 1st row is x (x cooridnates)

        uu has 3 rows
        the 1st row is u (velocity)
        the 2nd row is m (mass flow rate)
        the 3rd row is ̂E (total energy per volume)

    this function initialize a wave
"""

    """
    still working on that
    """
function initial_one_wave(nx::Int64, xlim::Any, gamma::Float64)

    Lx = xlim[2]-xlim[1];
    Δx = Lx/(nx-1)

    # Initialization of ux and urhop
    ux = xlim[1]:Δx:xlim[2]
    urhop = Array{Float64,2}(UndefInitializer(), 3,length(ux))

    urhop[1,:].=1
    urhop[2,:].=0.001*exp.(-(collect(ux) .- 0.5.*(ux[end]-ux[1])).^2/0.1^2).+0.125
    urhop[3,:].=0.01*exp.(-(collect(ux) .- 0.5.*(ux[end]-ux[1])).^2/0.1^2).+1

    #From urhop make uu (u rho*u e)
    uu = Array{Float64,2}(UndefInitializer(), 3,length(ux))
    uu[1,:] = urhop[2,:];
    uu[2,:] = urhop[2,:].*urhop[1,:];
    uu[3,:] = urhop[3,:]/(gamma-1) + 0.5urhop[2,:].*urhop[1,:].^2;

    return ux,uu

end

"""
still working on that initial one wave at a relative location
"""
function initial_one_wave(nx::Int64, xlim::Any, gamma::Float64, xposi::Float64)

    Lx = xlim[2]-xlim[1];
    Δx = Lx/(nx-1)

    # Initialization of ux and urhop
    ux = xlim[1]:Δx:xlim[2]
    urhop = Array{Float64,2}(UndefInitializer(), 3,length(ux))

    urhop[1,:].=1
    urhop[2,:].=0.001*exp.(-(collect(ux) .- xposi.*(ux[end]-ux[1])).^2/0.1^2).+0.125
    urhop[3,:].=0.01*exp.(-(collect(ux) .- xposi.*(ux[end]-ux[1])).^2/0.1^2).+1

    #From urhop make uu (u rho*u e)
    uu = Array{Float64,2}(UndefInitializer(), 3,length(ux))
    uu[1,:] = urhop[2,:];
    uu[2,:] = urhop[2,:].*urhop[1,:];
    uu[3,:] = urhop[3,:]/(gamma-1) + 0.5urhop[2,:].*urhop[1,:].^2;

    return ux,uu

end

end


module MyOptim

using LinearAlgebra

function conjgrad(fn, grad, x0; maxstep=2.0, niters=50, conv_thr = 1e-2)

    println("CG START")
    x = deepcopy(x0)
    storage = zeros(size(x))

#    println("x0 ", x)

    println("CG RUN GRAD1")
    #iter 1
    f, g = grad(storage, x)

#    println("f ", f, " x   ",  x)

    dx = deepcopy(-g)

    dx1 = deepcopy(dx)

    s = deepcopy(dx)

    step_size = maxstep/10.0

    println("CG RUN LS1")

    x, step_size = linesearch(x, dx, fn, f, step_size)
    step_size = min(step_size, maxstep)

    for i = 1:niters

        f, g = grad(storage, x)

        if sqrt(sum(g.^2)) < conv_thr && f < 0.1
            println("yes conv")
            return x, f, g
        end

        println("MY CONJGRAD $i ", f, " sg ", sqrt(sum(g.^2)), " step_size $step_size --------------------")

        dx1[:] = dx[:]
        dx[:] = -g[:]

#        println("dx ", dx)
#        println("dx1 ", dx1)

        beta = ( dx' * (dx - dx1)) / (dx1' * dx1)  #PR
#        println("beta $beta")
        beta = max(0, beta)

        s = dx + beta * s

        x, step_size, good = linesearch(x, s, fn, f, step_size)
        if good
            step_size =min(step_size, maxstep)
        else
            beta = 0.0
        end
        println("MY CG $i ", f, " sg ", sqrt(sum(g.^2)), " step_size $step_size good $good ")

    end

    println("conv not achieved")
    return x, f, g

end

function linesearch(x, dx, fn, f0, step_size)
#    println("CG FN0 $f0")
#    println("CG dx $dx")
#    println("ls $x $dx $f0 $step_size")
#    println("CG FN1 ss $step_size")
    f0r, flag0r = fn( x + dx * step_size * 0.0)
#    println("CG real $f0r $flag0r")

    f1, flag1 = fn( x + dx * step_size * 0.5)
#    println("CG FN1 $f1 $flag1")
    if f1 > f0 || flag1 == false
#        println("MY LS uphill $f1 > $f0 flag $flag1, reduce step")
        return x, step_size/3.1, false
    end
#    println("CG FN2")
    f2, flag2 = fn( x + dx * step_size)
#    println("CG FN2 $f2 $flag2")

    if flag2 == false
#        println("MY LS flag2 $flag2")
        return x + dx * step_size * 0.5, step_size/1.75, false
    end

    c = f0
    a = 2 * (f2-2*f1+f0)
    b = f2 - a - c

#    println("a b c $a $b $c ", -b/(2*a))

    if a < 0
#        println("MY LS a negative $a, take max step increase stepsize")
        return x + dx * step_size * 1.0, step_size * 2.1, true
    else
        step = min(-b/(2*a), 1.0)
#        println("MY LS NORMAL $step")
        return x + dx * step * step_size, step * 2.0, true
    end

end

end #end module

#=
function myf(x)

    return -0.5*x'*A*x + 0.1*sum(x.^4)

end

using ForwardDiff

function grad(storage, x)
    sleep(0.1)

    return myf(x), ForwardDiff.gradient(myf, x)

end

function grad_only(storage, x)
    storage[:] = ForwardDiff.gradient(myf, x)
    sleep(0.1)
    return storage
end
=#


# # Modal analysis of Argyris frame: effect of prestress

# Source code: [`argyris_frame_modal_tut.jl`](argyris_frame_modal_tut.jl)

# ## Description

# Vibration analysis of a L-shaped frame under a loading.
# The fundamental vibration frequency depends on the prestress force.

# ## Goals

# - Construct an L-shaped frame by merging individual members.
# - Compute the geometric stiffness.
# - Evaluate the effect of prestress on the fundamental frequency of vibration.

# The finite element code relies on the basic functionality implemented in this
# package.
using FinEtools
# The linear deformation code will be needed to evaluate the loading.
using FinEtoolsDeforLinear
# The functionality concerning beam models comes from these modules.
using FinEtoolsFlexBeams.CrossSectionModule: CrossSectionRectangle
using FinEtoolsFlexBeams.RotUtilModule:  update_rotation_field!
using FinEtoolsFlexBeams.MeshFrameMemberModule: frame_member, merge_members
using FinEtoolsFlexBeams.RotUtilModule: initial_Rfield
using FinEtoolsFlexBeams.FEMMCorotBeamModule: FEMMCorotBeam
using FinEtoolsFlexBeams.FEMMCorotBeamModule
CB = FEMMCorotBeamModule

# Parameters:
E = 71240.0 * phun("MPa")
nu = 0.31; # Poisson ratio
rho = 5000 * phun("kg/m^3");
# cross-sectional dimensions and length of each leg in millimeters
b = 0.6 * phun("mm"); h = 30.0 * phun("mm"); L = 240.0 * phun("mm");
# Magnitude of the total applied force, Newton
magn = 1e-5 * phun("N");

# Cross-sectional properties
cs = CrossSectionRectangle(s -> b, s -> h, s -> [0.0, 1.0, 0.0])

##
# ## Generate the discrete model

# Select the number of elements per leg.
n = 8;
members = Tuple{FENodeSet, AbstractFESet}[]
push!(members, frame_member([0 0 L; L 0 L], n, cs))
push!(members, frame_member([L 0 L; L 0 0], n, cs))
fens, fes = merge_members(members; tolerance = L / 10000)

# Construct the requisite fields, geometry and displacement
# Initialize configuration variables
geom0 = NodalField(fens.xyz)
u0 = NodalField(zeros(size(fens.xyz,1), 3))
Rfield0 = initial_Rfield(fens)
dchi = NodalField(zeros(size(fens.xyz,1), 6))

# Apply EBC's
l1 = selectnode(fens; box = [0 0 0 0 L L], tolerance = L / 10000)
for i in [1, 2, 3, 4, 5, 6]
    setebc!(dchi, l1, true, i)
end
applyebc!(dchi)
numberdofs!(dchi);

# Assemble the global discrete system

# Material properties
material = MatDeforElastIso(DeforModelRed3D, rho, E, nu, 0.0)

femm = FEMMCorotBeam(IntegDomain(fes, GaussRule(1, 2)), material)
K = CB.stiffness(femm, geom0, u0, Rfield0, dchi);
M = CB.mass(femm, geom0, u0, Rfield0, dchi);

# Construct force intensity,  loaded boundary, and assemble the load.
tipn = selectnode(fens; box=[L L 0 0  0 0], tolerance=L/n/1000)[1]
loadbdry = FESetP1(reshape([tipn], 1, 1))
lfemm = FEMMBase(IntegDomain(loadbdry, PointRule()))
fi = ForceIntensity(FFlt[-magn, 0, 0, 0, 0, 0]);
F = CB.distribloads(lfemm, geom0, dchi, fi, 3);

# Solve for the displacement under the static load.
scattersysvec!(dchi, K\F);

# Update deflections and rotations so that the initial stress can be computed.
# First the displacements:
u1 = deepcopy(u0)
u1.values .= dchi.values[:, 1:3]
# Then the rotations:
Rfield1 = deepcopy(Rfield0)
update_rotation_field!(Rfield1, dchi)

# The static deflection is now used to compute the internal forces
# which in turn lead to the geometric stiffness.
Kg = CB.geostiffness(femm, geom0, u1, Rfield1, dchi);

##
# ## Solution of the eigenvalue free-vibration problem

using Arpack

# We will solve for this many natural frequencies. Then, since they are ordered
# by magnitude, we will pick the fundamental by taking the first from the
# list.
neigvs = 4

# First we will  sweep through the loading factors that are positive, meaning
# the force points in the direction in which it was defined (towards the
# clamped end of the frame).

lfp = linearspace(0.0, 68000.0, 400)
fsp = let
    fsp = Float64[]
    for load_factor in lfp
        evals, evecs, nconv = eigs(K + load_factor .* Kg, M; nev=neigvs, which=:SM);
        f = evals[1] > 0 ? sqrt(evals[1]) / (2 * pi) : 0;
        push!(fsp, f)
    end
    fsp
end

# Next, we will sweep through arrange of negative load factors: this simply
# turns the force around so that it points away from the clamped end. This can
# also buckle the frame, but the magnitude is higher.
lfm = linearspace(-109000.0, 0.0, 400)
fsm = let
    fsm = Float64[]
    for load_factor in lfm
        evals, evecs, nconv = eigs(K + load_factor .* Kg, M; nev=neigvs, which=:SM);
        f = evals[1] > 0 ? sqrt(evals[1]) / (2 * pi) : 0;
        push!(fsm, f)
    end
    fsm
end

##
# ## Plot of the fundamental frequency is it depends on the loading factor

using Gnuplot

# We concatenate the ranges for the load factors and the calculated fundamental
# frequencies and present them in a single plot.

Gnuplot.gpexec("reset session")
@gp  "set terminal windows 0 "  :-

@gp  :- cat(collect(lfp), collect(lfm); dims=1) cat(fsp, fsm; dims=1) " lw 2 lc rgb 'red' with p title 'Fundamental frequency' "  :-

@gp  :- "set xlabel 'Loading factor P'" :-
@gp  :- "set ylabel 'Frequency(P) [Hz]'" :-
@gp  :- "set title 'Frame fundamental frequency'"

# Clearly, the curve giving the dependence of the fundamental frequency on the
# loading factor consists of two branches. These two branches correspond to two
# different buckling modes: one for the positive orientation of the force and
# one for the negative orientation.

##
# ## Visualize some fundamental mode shapes

# Here we visualize the fundamental vibration modes for different values of the
# loading factor.

using PlotlyJS
using FinEtoolsFlexBeams.VisUtilModule: plot_space_box, plot_solid, render, react!, default_layout_3d, save_to_json
scale = 0.005

vis(loading_factor, evec) = let
    tbox = plot_space_box(reshape(inflatebox!(boundingbox(fens.xyz), 0.3 * L), 2, 3))
    tenv0 = plot_solid(fens, fes; x=geom0.values, u=0.0 .* dchi.values[:, 1:3], R=Rfield0.values, facecolor="rgb(125, 155, 125)", opacity=0.3);
    plots = cat(tbox, tenv0; dims=1)
    layout = default_layout_3d(;width=600, height=600)
    layout[:scene][:aspectmode] = "data"
    pl = render(plots; layout=layout, title = "Loading factor $(loading_factor)")
    for xscale in scale .* sin.(collect(0:1:89) .* (2 * pi / 21))
        scattersysvec!(dchi, xscale .* evec)
        u1 = deepcopy(u0)
        u1.values .= dchi.values[:, 1:3]
        Rfield1 = deepcopy(Rfield0)
        update_rotation_field!(Rfield1, dchi)
        tenv1 = plot_solid(fens, fes; x=geom0.values, u=dchi.values[:, 1:3], R=Rfield1.values, facecolor="rgb(50, 55, 125)");
        plots = cat(tbox, tenv0, tenv1; dims=1)
        react!(pl, plots, pl.plot.layout)
        sleep(0.115)
    end
end

# This is the vibration mode in the lead up to the buckling mode for the
# positive orientation of the force.

loading_factor = 60000
evals, evecs, nconv = eigs(K + loading_factor .* Kg, M; nev=neigvs, which=:SM);
vis(loading_factor, evecs[:, 1])

# This is the same vibration mode for the negative orientation of the force, but
# note that the associated fundamental frequency increased due to the effect of
# the force upon the stiffening of the clamped leg of the frame that is now in
# tension, and therefore stiffer.

loading_factor = -50000
evals, evecs, nconv = eigs(K + loading_factor .* Kg, M; nev=neigvs, which=:SM);
vis(loading_factor, evecs[:, 1])

# Increasing the load factor in the negative orientation further, the
# fundamental frequency will switch: it will be a different mode shape, the one
# that is close to the buckling mode shape for this orientation of the force.

loading_factor = -100000
evals, evecs, nconv = eigs(K + loading_factor .* Kg, M; nev=neigvs, which=:SM);
vis(loading_factor, evecs[:, 1])

true


"""
    LobattoLegendreBasis([RealT=Float64,] polydeg::Integer)

Create a nodal Lobatto-Legendre basis for polynomials of degree `polydeg`.
"""
struct LobattoLegendreBasis{RealT<:Real, NNODES,
                            InverseVandermondeLegendre<:AbstractMatrix{RealT},
                            BoundaryMatrix<:AbstractMatrix{RealT},
                            DerivativeMatrix<:AbstractMatrix{RealT}} <: AbstractBasisSBP{RealT}
  nodes          ::SVector{NNODES, RealT}
  weights        ::SVector{NNODES, RealT}
  inverse_weights::SVector{NNODES, RealT}

  inverse_vandermonde_legendre::InverseVandermondeLegendre
  boundary_interpolation      ::BoundaryMatrix # lhat

  derivative_matrix         ::DerivativeMatrix # dsplit
  derivative_split          ::DerivativeMatrix # dsplit
  derivative_split_transpose::DerivativeMatrix # dsplit_transposed
  derivative_dhat           ::DerivativeMatrix # dhat, neg. adjoint wrt the SBP dot product
end

function LobattoLegendreBasis(RealT, polydeg::Integer)
  nnodes_ = polydeg + 1
  nodes, weights = gauss_lobatto_nodes_weights(nnodes_)
  inverse_weights = inv.(weights)

  _, inverse_vandermonde_legendre = vandermonde_legendre(nodes)

  boundary_interpolation = zeros(nnodes_, 2)
  boundary_interpolation[:, 1] = calc_lhat(-1.0, nodes, weights)
  boundary_interpolation[:, 2] = calc_lhat( 1.0, nodes, weights)

  derivative_matrix          = polynomial_derivative_matrix(nodes)
  derivative_split           = calc_dsplit(nodes, weights)
  derivative_split_transpose = Matrix(derivative_split')
  derivative_dhat            = calc_dhat(nodes, weights)

  # type conversions to make use of StaticArrays etc.
  nodes           = SVector{nnodes_}(convert.(RealT, nodes))
  weights         = SVector{nnodes_}(convert.(RealT, weights))
  inverse_weights = SVector{nnodes_}(convert.(RealT, inverse_weights))

  inverse_vandermonde_legendre = convert.(RealT, inverse_vandermonde_legendre)
  boundary_interpolation       = SMatrix{nnodes_, 2}(convert.(RealT, boundary_interpolation))

  derivative_matrix          = SMatrix{nnodes_, nnodes_}(convert.(RealT, derivative_matrix))
  derivative_split           = SMatrix{nnodes_, nnodes_}(convert.(RealT, derivative_split))
  derivative_split_transpose = SMatrix{nnodes_, nnodes_}(convert.(RealT, derivative_split_transpose))
  derivative_dhat            = SMatrix{nnodes_, nnodes_}(convert.(RealT, derivative_dhat))

  return LobattoLegendreBasis{RealT, nnodes_, typeof(inverse_vandermonde_legendre), typeof(boundary_interpolation), typeof(derivative_matrix)}(
    nodes, weights, inverse_weights,
    inverse_vandermonde_legendre, boundary_interpolation,
    derivative_matrix, derivative_split, derivative_split_transpose, derivative_dhat
  )
end

LobattoLegendreBasis(polydeg::Integer) = LobattoLegendreBasis(Float64, polydeg)

function Base.show(io::IO, basis::LobattoLegendreBasis)
  @nospecialize basis # reduce precompilation time

  print(io, "LobattoLegendreBasis{", real(basis), "}(polydeg=", polydeg(basis), ")")
end
function Base.show(io::IO, ::MIME"text/plain", basis::LobattoLegendreBasis)
  @nospecialize basis # reduce precompilation time

  print(io, "LobattoLegendreBasis{", real(basis), "} with polynomials of degree ", polydeg(basis))
end

@inline Base.real(basis::LobattoLegendreBasis{RealT}) where {RealT} = RealT

@inline nnodes(basis::LobattoLegendreBasis{RealT, NNODES}) where {RealT, NNODES} = NNODES

@inline polydeg(basis::LobattoLegendreBasis) = nnodes(basis) - 1

struct LobattoLegendreMortarL2{RealT<:Real, NNODES, MortarMatrix<:AbstractMatrix{RealT}} <: AbstractMortarL2{RealT}
  forward_upper::MortarMatrix
  forward_lower::MortarMatrix
  reverse_upper::MortarMatrix
  reverse_lower::MortarMatrix
end

function MortarL2(basis::LobattoLegendreBasis)
  RealT = real(basis)
  NNODES = nnodes(basis)

  forward_upper = calc_forward_upper(NNODES)
  forward_lower = calc_forward_lower(NNODES)
  reverse_upper = calc_reverse_upper(NNODES, Val(:gauss))
  reverse_lower = calc_reverse_lower(NNODES, Val(:gauss))

  # type conversions to make use of StaticArrays etc.
  forward_upper = SMatrix{NNODES, NNODES}(convert.(RealT, forward_upper))
  forward_lower = SMatrix{NNODES, NNODES}(convert.(RealT, forward_lower))
  reverse_upper = SMatrix{NNODES, NNODES}(convert.(RealT, reverse_upper))
  reverse_lower = SMatrix{NNODES, NNODES}(convert.(RealT, reverse_lower))

  LobattoLegendreMortarL2{RealT, NNODES, typeof(forward_upper)}(
    forward_upper, forward_lower,
    reverse_upper, reverse_lower)
end

function Base.show(io::IO, mortar::LobattoLegendreMortarL2)
  @nospecialize mortar # reduce precompilation time

  print(io, "LobattoLegendreMortarL2{", real(mortar), "}(polydeg=", polydeg(mortar), ")")
end
function Base.show(io::IO, ::MIME"text/plain", mortar::LobattoLegendreMortarL2)
  @nospecialize mortar # reduce precompilation time

  print(io, "LobattoLegendreMortarL2{", real(mortar), "} with polynomials of degree ", polydeg(mortar))
end

@inline Base.real(mortar::LobattoLegendreMortarL2{RealT}) where {RealT} = RealT

@inline nnodes(mortar::LobattoLegendreMortarL2{RealT, NNODES}) where {RealT, NNODES} = NNODES

@inline polydeg(mortar::LobattoLegendreMortarL2) = nnodes(mortar) - 1

# TODO: We can create EC mortars along the lines of the following implementation.
# abstract type AbstractMortarEC{RealT} <: AbstractMortar{RealT} end

# struct LobattoLegendreMortarEC{RealT<:Real, NNODES, MortarMatrix<:AbstractMatrix{RealT}, SurfaceFlux} <: AbstractMortarEC{RealT}
#   forward_upper::MortarMatrix
#   forward_lower::MortarMatrix
#   reverse_upper::MortarMatrix
#   reverse_lower::MortarMatrix
#   surface_flux::SurfaceFlux
# end

# function MortarEC(basis::LobattoLegendreBasis{RealT}, surface_flux)
#   forward_upper   = calc_forward_upper(n_nodes)
#   forward_lower   = calc_forward_lower(n_nodes)
#   l2reverse_upper = calc_reverse_upper(n_nodes, Val(:gauss_lobatto))
#   l2reverse_lower = calc_reverse_lower(n_nodes, Val(:gauss_lobatto))

#   # type conversions to make use of StaticArrays etc.
#   nnodes_ = nnodes(basis)
#   forward_upper   = SMatrix{nnodes_, nnodes_}(forward_upper)
#   forward_lower   = SMatrix{nnodes_, nnodes_}(forward_lower)
#   l2reverse_upper = SMatrix{nnodes_, nnodes_}(l2reverse_upper)
#   l2reverse_lower = SMatrix{nnodes_, nnodes_}(l2reverse_lower)

#   LobattoLegendreMortarEC{RealT, nnodes_, typeof(forward_upper), typeof(surface_flux)}(
#     forward_upper, forward_lower,
#     l2reverse_upper, l2reverse_lower,
#     surface_flux
#   )
# end

# @inline nnodes(mortar::LobattoLegendreMortarEC{RealT, NNODES}) = NNODES

struct LobattoLegendreAnalyzer{RealT<:Real, NNODES, Vandermonde<:AbstractMatrix{RealT}} <: SolutionAnalyzer{RealT}
  nodes  ::SVector{NNODES, RealT}
  weights::SVector{NNODES, RealT}
  vandermonde::Vandermonde
end

function SolutionAnalyzer(basis::LobattoLegendreBasis{RealT}; analysis_polydeg=2*polydeg(basis)) where {RealT}
  nnodes_ = analysis_polydeg + 1
  nodes, weights = gauss_lobatto_nodes_weights(nnodes_)

  vandermonde = polynomial_interpolation_matrix(basis.nodes, nodes)

  # type conversions to make use of StaticArrays etc.
  nodes   = SVector{nnodes_}(convert.(RealT, nodes))
  weights = SVector{nnodes_}(convert.(RealT, weights))

  vandermonde = convert.(RealT, vandermonde)

  return LobattoLegendreAnalyzer{RealT, nnodes_, typeof(vandermonde)}(
    nodes, weights, vandermonde)
end

function Base.show(io::IO, analyzer::LobattoLegendreAnalyzer)
  @nospecialize analyzer # reduce precompilation time

  print(io, "LobattoLegendreAnalyzer{", real(analyzer), "}(polydeg=", polydeg(analyzer), ")")
end
function Base.show(io::IO, ::MIME"text/plain", analyzer::LobattoLegendreAnalyzer)
  @nospecialize analyzer # reduce precompilation time

  print(io, "LobattoLegendreAnalyzer{", real(analyzer), "} with polynomials of degree ", polydeg(analyzer))
end

@inline Base.real(analyzer::LobattoLegendreAnalyzer{RealT}) where {RealT} = RealT

@inline nnodes(analyzer::LobattoLegendreAnalyzer{RealT, NNODES}) where {RealT, NNODES} = NNODES
@inline eachnode(analyzer::LobattoLegendreAnalyzer) = Base.OneTo(nnodes(analyzer))

@inline polydeg(analyzer::LobattoLegendreAnalyzer) = nnodes(analyzer) - 1

struct LobattoLegendreAdaptorL2{RealT<:Real, NNODES, MortarMatrix<:AbstractMatrix{RealT}} <: AdaptorL2{RealT}
  forward_upper::MortarMatrix
  forward_lower::MortarMatrix
  reverse_upper::MortarMatrix
  reverse_lower::MortarMatrix
end

function AdaptorL2(basis::LobattoLegendreBasis{RealT}) where {RealT}
  nnodes_ = nnodes(basis)
  forward_upper   = calc_forward_upper(nnodes_)
  forward_lower   = calc_forward_lower(nnodes_)
  l2reverse_upper = calc_reverse_upper(nnodes_, Val(:gauss))
  l2reverse_lower = calc_reverse_lower(nnodes_, Val(:gauss))

  # type conversions to make use of StaticArrays etc.
  forward_upper   = SMatrix{nnodes_, nnodes_}(convert.(RealT, forward_upper))
  forward_lower   = SMatrix{nnodes_, nnodes_}(convert.(RealT, forward_lower))
  l2reverse_upper = SMatrix{nnodes_, nnodes_}(convert.(RealT, l2reverse_upper))
  l2reverse_lower = SMatrix{nnodes_, nnodes_}(convert.(RealT, l2reverse_lower))

  LobattoLegendreAdaptorL2{RealT, nnodes_, typeof(forward_upper)}(
    forward_upper, forward_lower,
    l2reverse_upper, l2reverse_lower)
end

function Base.show(io::IO, adaptor::LobattoLegendreAdaptorL2)
  @nospecialize adaptor # reduce precompilation time

  print(io, "LobattoLegendreAdaptorL2{", real(adaptor), "}(polydeg=", polydeg(adaptor), ")")
end
function Base.show(io::IO, ::MIME"text/plain", adaptor::LobattoLegendreAdaptorL2)
  @nospecialize adaptor # reduce precompilation time

  print(io, "LobattoLegendreAdaptorL2{", real(adaptor), "} with polynomials of degree ", polydeg(adaptor))
end

@inline Base.real(adaptor::LobattoLegendreAdaptorL2{RealT}) where {RealT} = RealT

@inline nnodes(adaptor::LobattoLegendreAdaptorL2{RealT, NNODES}) where {RealT, NNODES} = NNODES

@inline polydeg(adaptor::LobattoLegendreAdaptorL2) = nnodes(adaptor) - 1

###############################################################################
# Polynomial derivative and interpolation functions

# TODO: Taal refactor, allow other RealT below and adapt constructors above accordingly

# Calculate the Dhat matrix
function calc_dhat(nodes, weights)
  n_nodes = length(nodes)
  dhat = Matrix(polynomial_derivative_matrix(nodes)')

  for n in 1:n_nodes, j in 1:n_nodes
    dhat[j, n] *= -weights[n] / weights[j]
  end

  return dhat
end

# Calculate the Dsplit matrix for split-form differentiation: dplit = 2D - M⁻¹B
function calc_dsplit(nodes, weights)
  # Start with 2 x the normal D matrix
  dsplit = 2 .* polynomial_derivative_matrix(nodes)

  # Modify to account for
  dsplit[  1,   1] += 1 / weights[1]
  dsplit[end, end] -= 1 / weights[end]

  return dsplit
end

# Calculate the polynomial derivative matrix D
function polynomial_derivative_matrix(nodes)
  n_nodes = length(nodes)
  d = zeros(n_nodes, n_nodes)
  wbary = barycentric_weights(nodes)

  for i in 1:n_nodes, j in 1:n_nodes
    if j != i
      d[i, j] = wbary[j] / wbary[i] * 1 / (nodes[i] - nodes[j])
      d[i, i] -= d[i, j]
    end
  end

  return d
end

# Calculate and interpolation matrix (Vandermonde matrix) between two given sets of nodes
function polynomial_interpolation_matrix(nodes_in, nodes_out)
  n_nodes_in = length(nodes_in)
  n_nodes_out = length(nodes_out)
  wbary_in = barycentric_weights(nodes_in)
  vdm = zeros(n_nodes_out, n_nodes_in)

  for k in 1:n_nodes_out
    match = false
    for j in 1:n_nodes_in
      if isapprox(nodes_out[k], nodes_in[j], rtol=eps())
        match = true
        vdm[k, j] = 1
      end
    end

    if match == false
      s = 0.0
      for j in 1:n_nodes_in
        t = wbary_in[j] / (nodes_out[k] - nodes_in[j])
        vdm[k, j] = t
        s += t
      end
      for j in 1:n_nodes_in
        vdm[k, j] = vdm[k, j] / s
      end
    end
  end

  return vdm
end

# Calculate the barycentric weights for a given node distribution.
function barycentric_weights(nodes)
  n_nodes = length(nodes)
  weights = ones(n_nodes)

  for j = 2:n_nodes, k = 1:(j-1)
    weights[k] *= nodes[k] - nodes[j]
    weights[j] *= nodes[j] - nodes[k]
  end

  for j in 1:n_nodes
    weights[j] = 1 / weights[j]
  end

  return weights
end

# Calculate Lhat.
function calc_lhat(x, nodes, weights)
  n_nodes = length(nodes)
  wbary = barycentric_weights(nodes)

  lhat = lagrange_interpolating_polynomials(x, nodes, wbary)

  for i in 1:n_nodes
    lhat[i] /= weights[i]
  end

  return lhat
end

# Calculate Lagrange polynomials for a given node distribution.
function lagrange_interpolating_polynomials(x, nodes, wbary)
  n_nodes = length(nodes)
  polynomials = zeros(n_nodes)

  for i in 1:n_nodes
    if isapprox(x, nodes[i], rtol=eps(x))
      polynomials[i] = 1
      return polynomials
    end
  end

  for i in 1:n_nodes
    polynomials[i] = wbary[i] / (x - nodes[i])
  end
  total = sum(polynomials)

  for i in 1:n_nodes
    polynomials[i] /= total
  end

  return polynomials
end

# From FLUXO (but really from blue book by Kopriva)
function gauss_lobatto_nodes_weights(n_nodes::Integer)
  # From Kopriva's book
  n_iterations = 10
  tolerance = 1e-15

  # Initialize output
  nodes = zeros(n_nodes)
  weights = zeros(n_nodes)

  # Get polynomial degree for convenience
  N = n_nodes - 1

  # Calculate values at boundary
  nodes[1] = -1.0
  nodes[end] = 1.0
  weights[1] = 2 / (N * (N + 1))
  weights[end] = weights[1]

  # Calculate interior values
  if N > 1
    cont1 = pi/N
    cont2 = 3/(8 * N * pi)

    # Use symmetry -> only left side is computed
    for i in 1:(div(N + 1, 2) - 1)
      # Calculate node
      # Initial guess for Newton method
      nodes[i+1] = -cos(cont1*(i+0.25) - cont2/(i+0.25))

      # Newton iteration to find root of Legendre polynomial (= integration node)
      for k in 0:n_iterations
        q, qder, _ = calc_q_and_l(N, nodes[i+1])
        dx = -q/qder
        nodes[i+1] += dx
        if abs(dx) < tolerance * abs(nodes[i+1])
          break
        end
      end

      # Calculate weight
      _, _, L = calc_q_and_l(N, nodes[i+1])
      weights[i+1] = weights[1] / L^2

      # Set nodes and weights according to symmetry properties
      nodes[N+1-i] = -nodes[i+1]
      weights[N+1-i] = weights[i+1]
    end
  end

  # If odd number of nodes, set center node to origin (= 0.0) and calculate weight
  if n_nodes % 2 == 1
    _, _, L = calc_q_and_l(N, 0)
    nodes[div(N, 2) + 1] = 0.0
    weights[div(N, 2) + 1] = weights[1] / L^2
  end

  return nodes, weights
end

# From FLUXO (but really from blue book by Kopriva)
function calc_q_and_l(N::Integer, x::Float64)
  L_Nm2 = 1.0
  L_Nm1 = x
  Lder_Nm2 = 0.0
  Lder_Nm1 = 1.0

  local L
  for i in 2:N
    L = ((2 * i - 1) * x * L_Nm1 - (i - 1) * L_Nm2) / i
    Lder = Lder_Nm2 + (2 * i - 1) * L_Nm1
    L_Nm2 = L_Nm1
    L_Nm1 = L
    Lder_Nm2 = Lder_Nm1
    Lder_Nm1 = Lder
  end

  q = (2 * N + 1)/(N + 1) * (x * L - L_Nm2)
  qder = (2 * N + 1) * L

  return q, qder, L
end
calc_q_and_l(N::Integer, x::Real) = calc_q_and_l(N, convert(Float64, x))

# From FLUXO (but really from blue book by Kopriva)
function gauss_nodes_weights(n_nodes::Integer)
  # From Kopriva's book
  n_iterations = 10
  tolerance = 1e-15

  # Initialize output
  nodes = ones(n_nodes) * 1000
  weights = zeros(n_nodes)

  # Get polynomial degree for convenience
  N = n_nodes - 1
  if N == 0
    nodes .= 0.0
    weights .= 2.0
    return nodes, weights
  elseif N == 1
    nodes[1] = -sqrt(1/3)
    nodes[end] = -nodes[1]
    weights .= 1.0
    return nodes, weights
  else # N > 1
    # Use symmetry property of the roots of the Legendre polynomials
    for i in 0:(div(N + 1, 2) - 1)
      # Starting guess for Newton method
      nodes[i+1] = -cos(pi / (2 * N + 2) * (2 * i + 1))

      # Newton iteration to find root of Legendre polynomial (= integration node)
      for k in 0:n_iterations
        poly, deriv = legendre_polynomial_and_derivative(N + 1, nodes[i+1])
        dx = -poly / deriv
        nodes[i+1] += dx
        if abs(dx) < tolerance * abs(nodes[i+1])
          break
        end
      end

      # Calculate weight
      poly, deriv = legendre_polynomial_and_derivative(N + 1, nodes[i+1])
      weights[i+1] = (2 * N + 3) / ((1 - nodes[i+1]^2) * deriv^2)

      # Set nodes and weights according to symmetry properties
      nodes[N+1-i] = -nodes[i+1]
      weights[N+1-i] = weights[i+1]
    end

    # If odd number of nodes, set center node to origin (= 0.0) and calculate weight
    if n_nodes % 2 == 1
      poly, deriv = legendre_polynomial_and_derivative(N + 1, 0.0)
      nodes[div(N, 2) + 1] = 0.0
      weights[div(N, 2) + 1] = (2 * N + 3) / deriv^2
    end

    return nodes, weights
  end
end

# From FLUXO (but really from blue book by Kopriva)
function legendre_polynomial_and_derivative(N::Int, x::Real)
  if N == 0
    poly = 1.0
    deriv = 0.0
  elseif N == 1
    poly = convert(Float64, x)
    deriv = 1.0
  else
    poly_Nm2 = 1.0
    poly_Nm1 = convert(Float64, x)
    deriv_Nm2 = 0.0
    deriv_Nm1 = 1.0

    poly = 0.0
    deriv = 0.0
    for i in 2:N
      poly = ((2*i-1) * x * poly_Nm1 - (i-1) * poly_Nm2) / i
      deriv=deriv_Nm2 + (2*i-1)*poly_Nm1
      poly_Nm2=poly_Nm1
      poly_Nm1=poly
      deriv_Nm2=deriv_Nm1
      deriv_Nm1=deriv
    end
  end

  # Normalize
  poly = poly * sqrt(N+0.5)
  deriv = deriv * sqrt(N+0.5)

  return poly, deriv
end

# Calculate Legendre vandermonde matrix and its inverse
function vandermonde_legendre(nodes, N)
  n_nodes = length(nodes)
  n_modes = N + 1
  vandermonde = zeros(n_nodes, n_modes)

  for i in 1:n_nodes
    for m in 1:n_modes
      vandermonde[i, m], _ = legendre_polynomial_and_derivative(m-1, nodes[i])
    end
  end
  # for very high polynomial degree, this is not well conditioned
  inverse_vandermonde = inv(vandermonde)
  return vandermonde, inverse_vandermonde
end
vandermonde_legendre(nodes) = vandermonde_legendre(nodes, length(nodes) - 1)


using Test
import DifferentiableBackwardEuler
import DPFEHM
import PyPlot
import Zygote

mins = [0, 0]; maxs = [3, 2]#size of the domain, in meters
ns = [30, 20]#number of nodes on the grid
coords, neighbors, areasoverlengths, volumes = DPFEHM.regulargrid2d(mins, maxs, ns, 1.0)#build the grid
coords = vcat(coords, zeros(size(coords, 2))')#turn it into a 3d coords

c = 1.0
tspan = [0.0, 1.0 / abs(c)]
u0 = map(i->coords[1, i] < 1 / 2  && coords[1, i] > 0 && coords[2, i] > 0.25 && coords[2, i] < 0.75 ? 1.0 : 0.0, 1:size(coords, 2)) + map(i->exp(-((coords[1, i] - 1.0) ^ 2 + (coords[2, i] - 1.0) ^ 2) * 64), 1:size(coords, 2))
vxs = map(neighbor->c, neighbors)
vys = map(neighbor->c * 0.5, neighbors)
vzs = zeros(length(neighbors))
Ds = 1e-2 * ones(length(neighbors))
Qs = zeros(size(coords, 2))
dirichletnodes = [collect(1:ns[2]); collect(size(coords, 2) - ns[2] + 1:size(coords, 2))]
dirichletus = zeros(size(coords, 2))
f(u, p, t) = DPFEHM.transport_residuals(u, p, vys, vzs, Ds, neighbors, areasoverlengths, dirichletnodes, dirichletus, Qs, volumes, coords)
f_u(u, p, t) = DPFEHM.transport_u(u, p, vys, vzs, Ds, neighbors, areasoverlengths, dirichletnodes, dirichletus, Qs, volumes, coords)
f_p(u, p, t) = DPFEHM.transport_vxs(u, p, vys, vzs, Ds, neighbors, areasoverlengths, dirichletnodes, dirichletus, Qs, volumes, coords)
f_t(u, p, t) = zeros(length(u))
function solveforu(vxs)
	DifferentiableBackwardEuler.steps(u0, f, f_u, f_p, f_t, vxs, tspan[1], tspan[2]; abstol=1e-4, reltol=1e-4)
end
@time u_implicit = solveforu(vxs)

fig, axs = PyPlot.subplots(2, 1; figsize=(8, 8))
img = axs[1].imshow(reshape(u0, reverse(ns)...), extent=[mins[1], maxs[1], mins[2], maxs[2]], origin="lower")
fig.colorbar(img, ax=axs[1])
img = axs[2].imshow(reshape(u_implicit[:, end], reverse(ns)...), extent=[mins[1], maxs[1], mins[2], maxs[2]], origin="lower")
fig.colorbar(img, ax=axs[2])
fig.tight_layout()
display(fig)
PyPlot.println()
PyPlot.close(fig)

_, gradient_node = findmin(map(i->sum((coords[:, i] .- [1.5, 1.0, 0.0]) .^ 2), 1:size(coords, 2)))
g = x->solveforu(x)[gradient_node, end]
@time gradg = Zygote.gradient(g, vxs)[1]
function checkgradientquickly(f, x0, gradf, n; delta::Float64=1e-8, kwargs...)
	indicestocheck = sort(collect(1:length(x0)), by=i->abs(gradf[i]), rev=true)[1:n]
	f0 = f(x0)
	for i in indicestocheck
		x = copy(x0)
		x[i] += delta
		fval = f(x)
		grad_f_i = (fval - f0) / delta
		@test isapprox(gradf[i], grad_f_i; kwargs...)
	end
end
checkgradientquickly(g, vxs, gradg, 3; atol=1e-4, rtol=1e-3)


using PrototypesCriticisms
using Clustering
using Distances
using KernelFunctions
using LinearAlgebra
using Statistics
using Test

@testset verbose=true "PrototypesCriticisms.jl" begin
    @testset "sqmmd (mmd²)" begin
        X = rand(5, 10)
        k = RBFKernel()

        @test sqmmd(X, X, k) ≈ 0

        Y = X
        XX = kernelmatrix(k, X, obsdim=2)
        XY = kernelmatrix(k, X, Y, obsdim=2)
        YY = kernelmatrix(k, Y, obsdim=2)
        @test sqmmd(XX, XY, YY) ≈ 0

        @test mmd²(X, X, k) == sqmmd(X, X, k)
        @test mmd²(XX, XY, YY) == sqmmd(XX, XY, YY)
    end

    @testset "witness" begin
        X = zeros(5, 10)
        Y = ones(5, 10)
        x = X[:, 1]
        y = Y[:, 1]
        k = RBFKernel()
        @test witness(x, X, X, k) ≈ 0
        @test witness(x, X, Y, k) > 0
        @test witness(y, X, Y, k) < 0
    end

    @testset "prototypes" begin
        n = 50
        X = rand(5, n)
        k = RBFKernel()

        @test length(prototypes(X, k, 0)) == 0
        @test Set(prototypes(X, k, n)) == Set(1:n)

        K = kernelmatrix(k, X, obsdim=2)
        @test length(prototypes(K, 0)) == 0
        @test Set(prototypes(K, n)) == Set(1:n)

        k = 2
        c = kmedoids(pairwise(Euclidean(), X, dims=2), k)
        protoids = prototypes(c)
        @test typeof(protoids) == Vector{Vector{Int}}
        @test length(protoids) == k
        for i = 1:k
            @test length(protoids[i]) == 1
        end
        protoids = prototypes(c, 2)
        @test typeof(protoids) == Vector{Vector{Int}}
        @test length(protoids) == k
        for i = 1:k
            @test length(protoids[i]) == 2
        end

        k = 2
        c = kmeans(X, k)
        protoids = prototypes(c)
        @test typeof(protoids) == Vector{Vector{Int}}
        @test length(protoids) == k
        for i = 1:k
            @test length(protoids[i]) == 1
        end
        protoids = prototypes(c, 2)
        @test typeof(protoids) == Vector{Vector{Int}}
        @test length(protoids) == k
        for i = 1:k
            @test length(protoids[i]) == 2
        end

        k = 2
        c = fuzzy_cmeans(X, k, 2)
        protoids = prototypes(c)
        @test typeof(protoids) == Vector{Vector{Int}}
        @test length(protoids) == k
        for i = 1:k
            @test length(protoids[i]) == 1
        end
        protoids = prototypes(c, 2)
        @test typeof(protoids) == Vector{Vector{Int}}
        @test length(protoids) == k
        for i = 1:k
            @test length(protoids[i]) == 2
        end

        S = -pairwise(Euclidean(), X, dims=2)
        S = S - diagm(0 => diag(S)) + median(S) * I
        c = affinityprop(S)
        protoids = prototypes(c, X)
        @test typeof(protoids) == Vector{Vector{Int}}
        @test length(protoids) == nclusters(c)
        for i = 1:nclusters(c)
            @test length(protoids[i]) == 1
        end
        protoids = prototypes(c, X, 2)
        @test typeof(protoids) == Vector{Vector{Int}}
        @test length(protoids) == nclusters(c)
        for i = 1:nclusters(c)
            @test length(protoids[i]) == 2
        end
    end

    @testset "criticisms" begin
        n = 50
        X = rand(5, n)
        k = RBFKernel()

        @test length(criticisms(X, k, [1, 2], 0)) == 0
        @test Set(criticisms(X, k, [1, 2], n - 2)) == Set(3:n)

        K = kernelmatrix(k, X, obsdim=2)
        @test length(criticisms(K, [1, 2], 0)) == 0
        @test Set(criticisms(K, [1, 2], n - 2)) == Set(3:n)

        k = 2
        c = kmedoids(pairwise(Euclidean(), X, dims=2), k)
        critids = criticisms(c)
        @test typeof(critids) == Vector{Vector{Int}}
        @test length(critids) == k
        for i = 1:k
            @test length(critids[i]) == 1
        end
        critids = criticisms(c, 2)
        @test typeof(critids) == Vector{Vector{Int}}
        @test length(critids) == k
        for i = 1:k
            @test length(critids[i]) == 2
        end

        k = 2
        c = kmeans(X, k)
        critids = criticisms(c)
        @test typeof(critids) == Vector{Vector{Int}}
        @test length(critids) == k
        for i = 1:k
            @test length(critids[i]) == 1
        end
        critids = criticisms(c, 2)
        @test typeof(critids) == Vector{Vector{Int}}
        @test length(critids) == k
        for i = 1:k
            @test length(critids[i]) == 2
        end

        k = 2
        c = fuzzy_cmeans(X, k, 2)
        critids = criticisms(c)
        @test typeof(critids) == Vector{Vector{Int}}
        @test length(critids) == k
        for i = 1:k
            @test length(critids[i]) == 1
        end
        critids = criticisms(c, 2)
        @test typeof(critids) == Vector{Vector{Int}}
        @test length(critids) == k
        for i = 1:k
            @test length(critids[i]) == 2
        end

        S = -pairwise(Euclidean(), X, dims=2)
        S = S - diagm(0 => diag(S)) + median(S) * I
        c = affinityprop(S)
        critids = criticisms(c, X)
        @test typeof(critids) == Vector{Vector{Int}}
        @test length(critids) == nclusters(c)
        for i = 1:nclusters(c)
            @test length(critids[i]) == 1
        end
        critids = prototypes(c, X, 2)
        @test typeof(critids) == Vector{Vector{Int}}
        @test length(critids) == nclusters(c)
        for i = 1:nclusters(c)
            @test length(critids[i]) == 2
        end
    end
end


using LinearAlgebra
using JuMP
using Printf

export accp, Halfspace

"""
    Halfspace(slope, constant)
Define a halfspace `{x: dot(slope, x) ≤ constant}`. Note that slope is internally
normalized.
"""
struct Halfspace{Ts<:AbstractVector,Tc<:Number}
    slope::Ts
    constant::Tc
    function Halfspace(slope, constant)
        norm_slope = slope/norm(slope)
        norm_constant = constant/norm(slope)
        new{typeof(norm_slope),typeof(norm_constant)}(norm_slope, norm_constant)
    end
end

mutable struct OuterApproximation
    model::Model
    x::Array{VariableRef,1}
    cutrefs::Array{ConstraintRef,1}
    function OuterApproximation(obj::AbstractVector, A::AbstractMatrix,
    b::AbstractVector, r::Number)
        n = length(obj)
        model = Model(SOCPsolver.Optimizer)
        set_silent(model)
        @variable(model, x[1:n])
        @constraint(model, vcat(r, x) in SecondOrderCone())
        cutrefs = @constraint(model, A * x .<= b)
        @objective(model, Min, dot(obj, x))
        new(model, x, cutrefs)
    end
end

"Compute lower bound by minimizing over the outer approximation"
function lowerbound(oa::OuterApproximation)
    optimize!(oa.model)
    if termination_status(oa.model) == MOI.OPTIMAL
        return objective_value(oa.model)
    else
        error("Could not compute lower bound."*
        " Termination status was $(termination_status(oa.model))")
    end
end
function addcut!(oa::OuterApproximation, hs::Halfspace)
    ref = @constraint(oa.model, hs.slope' * oa.x <= hs.constant)
    push!(oa.cutrefs, ref)
end

"""
    accp(obj::AbstractVector, oracle::Function, r::Number, x0::AbstractVector)
Run an analytic center cutting plane method to minimize `dot(obj, x)` such that
`norm(x) ≤ r` and `oracle(x) == true`. `oracle(x)` should either return
`true` or a `Halfspace`. The starting point `x0` does not have to be feasible.
Returns a near-optimal solution `x` at default settings.
# Arguments
- `opttol=1e-6`: the maximum relative optimality gap at termination
- `maxcons=4`: multiple of the number of variables indicating the number of
  constraints that are kept. The default `maxcons=4` keeps `4n` constraints, where
  `n` is the number of variables.
- `verbose=true`: display output. If the analytic center computation fails, this
  includes the spectrum of the matrix in the linear system that should be solved.
- `trackcalls=false`: the method returns `(x, calls)` when `trackcalls==true`,
  where `calls` is the number of calls to `oracle`.
"""
function accp(obj::AbstractVector, oracle::Function, r::Number,
x0::AbstractVector=zeros(length(obj)); opttol=1e-6, maxcons=4, verbose=true, trackcalls=false)
    n = length(obj)
    nobj = obj / norm(obj)
    x = norm(x0) > 1 ? x0 / norm(x0) : x0
    result = oracle(x)
    if result isa Halfspace
        # x is infeasible
        A = [result.slope';]
        b = [result.constant]
        local best_sol
        best_val = Inf
    else
        # x is feasible. If x ≈ -obj / norm(obj), the feasible set will be
        # (almost) empty after adding a cut.
        if x ≈ -r*nobj
            return x
        end
        A = [nobj';]
        b = [dot(nobj, x)]
        best_sol = copy(x)
        best_val = dot(obj, x)
    end
    trackcalls && (oraclecalls = 1)
    oa = OuterApproximation(obj, A, b, r)
    while true
        # Compute analytic center
        (status, finalnorm) = analytic_center!(A, b, r, x, verbose=verbose)
        if !status
            # Check if the point is at least in the interior of the set we
            # wanted the analytic center of
            if any(A*x .>= b) || norm(x) >= r
                verbose && println()
                error("Could not approximate analytic center: Lagrange gradient norm $finalnorm")
            end
        else
            A, b = prune(A, b, r, x, oa, ceil(Int, maxcons*n))
        end
        # Check if the analytic center is feasible
        result = oracle(x)
        trackcalls && (oraclecalls += 1)
        if result isa Halfspace
            # If the current point is not feasible, add a cut
            A = [A; result.slope']
            b = [b; result.constant]
            addcut!(oa, result)
        else
            # If the current point is feasible, tighten the objective constraint
            best_sol = copy(x)
            best_val = dot(obj, x)
            A = [A; nobj']
            b = [b; dot(nobj, x)]
            addcut!(oa, Halfspace(obj, best_val))
        end
        # Check if we are close enough to the optimum
        lb = lowerbound(oa)
        relgap = (best_val - lb) / (1 + min(abs(best_val), abs(lb)))
        verbose && @printf("\rAbs. gap: %.3e\tRel. gap: %.3e", best_val - lb, relgap)
        if relgap <= opttol
            verbose && println()
            return trackcalls ? (best_sol, oraclecalls) : best_sol
        end
    end
end

"""
    analytic_center!(A::AbstractMatrix, b::AbstractVector, r::Number, x::AbstractVector)
Approximate the analytic center of `{x: Ax ≤ b, norm(x) ≤ r}`, starting from `x`
and storing the result in `x`.
# Arguments
- `verbose=true`: displays the spectrum of the linear system if the method fails.
- `grad_atol=1e-8`: maximum norm of the Lagrangian gradient for successful termination.
- `maxiter=50`: maximum number of Newton steps.
- `α=0.9`: maximum step to boundary - should satisfy `0 < α < 1`.
"""
function analytic_center!(A::AbstractMatrix, b::AbstractVector, r::Number,
x::AbstractVector; verbose=true, grad_atol=1e-8, maxiter=50, α = 0.9)
    # We use an infeasible start Newton method, cf. page 532 in 'Convex
    # Optimization' by Boyd and Vandenberghe.
    (m,n) = size(A)
    if m == 0
        # If there are no constraints yet, the origin is the analytic center
        return (zeros(n), true, 0.)
    end
    d = r - norm(x) > sqrt(eps(Float64)) ? r^2 - x'*x : 1.
    s = [b[i] - A[i,:]'*x > sqrt(eps(Float64)) ? b[i] - A[i,:]'*x : 1. for i=1:m]
    κ = 1.
    λ = zeros(m)
    iter = 1
    LHS = Matrix{Float64}(undef, n, n)
    rhs = Vector{Float64}(undef, n)
    Δx = Vector{Float64}(undef, n)
    Δs = Vector{Float64}(undef, m)
    Δλ = Vector{Float64}(undef, m)
    while iter <= maxiter
        rhs .= ((r^2 - x'*x - d)/d^2 - 1. /d) * 2. * x + A'*(-1. ./ s + (s.^-2) .* (-s + b - A*x))
        LHS .= 4. / d^2 * x*x' + A' * ((s.^-2) .* A)
        for i in 1:n
            LHS[i,i] += 2. * κ
        end
        Δx .= LHS \ rhs
        Δd  = -d + r^2 - x'*x - 2. *x'*Δx
        Δs .= -s + b - A*x - A*Δx
        Δκ  = -κ + 1. / d - Δd / d^2
        Δλ .= -λ + 1. ./ s - Δs ./ (s.^2)
        # Let res(t) be the residual norm if we take step size t
        res = t -> norm_Lagrange_gradient(A,b,r, x+t*Δx, d+t*Δd, s+t*Δs, κ+t*Δκ, λ+t*Δλ)
        res0 = res(0)
        maxt = 1.
        t = min(maxt,
            any(Δs .< 0) ? α*minimum(-s[i]/Δs[i] for i in 1:m if Δs[i] < 0) : maxt,
            Δd < 0 ? -α*d/Δd : maxt,
            Δκ < 0 ? -α*κ/Δκ : maxt
        )

        if res0 <= min(grad_atol, res(t)) && all(λ .>= 0)
            # We have approximated the analytic center as well as we can; going
            # on is pointless
            return (true, res0)
        end
        x .+= t*Δx
        d  += t*Δd
        s .+= t*Δs
        κ  += t*Δκ
        λ .+= t*Δλ
        if iter == maxiter && res(t) <= grad_atol && all(λ .>= 0)
            # We have reached the final iteration, and we are beneath the
            # tolerance. Perhaps we could approximate the analytic center to
            # higher accuracy by continuing, but this is also fine.
            return (true, res(t))
        end
        iter += 1
    end
    verbose && @printf(" Spectrum: [%.3e, %.3e]", eigmin(LHS), eigmax(LHS))
    return (false, norm_Lagrange_gradient(A,b,r, x, d, s, κ, λ))
end
function norm_Lagrange_gradient(A::AbstractMatrix,b::AbstractVector,r::Number,
x::AbstractVector, d::Number, s::AbstractVector, κ::Number, λ::AbstractVector)
    # Returns the norm of the gradient of the Lagrangian of this problem at the
    # point (x, d, s, κ, λ). In the book by Boyd and Vandenberghe, this
    # quantity has the interpretation of the (primal-dual) residual norm.
    v1 = 2*κ * x + A' * λ
    v2 = -1 ./ s + λ
    v3 = s - b + A*x
    normsq = v1'*v1 + (-1. / d + κ)^2 + v2'*v2 + (d - r^2 + x'*x)^2 + v3'*v3
    return sqrt(normsq)
end

"""
    prune(A::AbstractMatrix, b::AbstractVector, r::Number, ac::AbstractVector, oa::OuterApproximation, maxconabs::Int=Inf)
Prune the constraints from `{x: Ax ≤ b, norm(x) ≤ r}` that are redundant, using
the set's analytic center `ac`. Also prune the least relevant constraints
such that at most `maxconabs` constraints remain. Returns `(Ap,bp)`, where `Ap`
and `bp` are the pruned versions of `A` and `b`. Also prunes the constraints
from `oa`.
"""
function prune(A::AbstractMatrix, b::AbstractVector, r::Number, ac::AbstractVector,
oa::OuterApproximation, maxconabs::Int=Inf)
    # The method follows Section 3 in Boyd, Vandenberghe, and Skaf. Returns the
    # pruned A and b.
    (m,n) = size(A)
    if m <= n
        return (A,b)
    end
    slack = b - A*ac
    H = slack.^-2 .* A
    H = sum(A[i,:] * H[i,:]' for i in 1:m)
    H += 2. / (r^2 - ac'*ac) * I + 4. / (r^2 - ac'*ac) * ac*ac'
    H_inv_AT = H \ Matrix(A')
    redcons = Int[]
    η = Float64[]
    for i = 1:m
        push!(η, slack[i] / sqrt(A[i,:]' * H_inv_AT[:,i]))
        if η[end] >= m+2
            push!(redcons, i)
        end
    end
    keepers = m - length(redcons) <= maxconabs ? setdiff(1:m, redcons) :
    sort!(sortperm(η)[1:min(maxconabs, m)])
    A = A[keepers,:]
    b = b[keepers]
    for i in 1:m
        if i ∉ keepers
            delete(oa.model, oa.cutrefs[i])
        end
    end
    oa.cutrefs = oa.cutrefs[keepers]
    return (A,b)
end

"""
    ellipsoid_method(obj::AbstractVector, oracle::Function, r::Number)
Run an ellipsoid method to minimize `dot(obj, x)` such that `norm(x) ≤ r`
and `oracle(x) == true`. `oracle(x)` should either return `true` or a `Halfspace`.
"""
function ellipsoid_method(obj::AbstractVector, oracle::Function, r::Number;
opttol=1e-6, verbose=true, trackcalls=false)
    n = length(obj)
    if n <= 1
        error("The number of variables $n should be 2 or more.")
    end
    nobj = obj / norm(obj)
    trackcalls && (oraclecalls = 0)
    x = zeros(n)
    D = Matrix(r*I,n,n)
    # The ellipsoid we consider is {y: (y-x)' D^-1 (y-x) ≤ 1}
    while true
        # Add the constraint y: a'y >= a'x, i.e. -a'y <= -a'x
        if norm(x) > r
            a = -x / norm(x)
        else
            result = oracle(x)
            trackcalls && (oraclecalls += 1)
            if result isa Halfspace # x is infeasible
                a = -result.slope
            else # x is feasible
                # Using z = D^(-1/2)(y-x), min{obj' y: (y-x)' D^-1 (y-x) ≤ 1} =
                # min{obj' (x + D^(1/2) z): z'z ≤ 1}
                lb = obj' * x - sqrt(obj' * D * obj)
                relgap = (obj' * x - lb) / (1 + min(abs(obj' * x), abs(lb)))
                verbose && @printf("\rAbs. gap: %.3e\tRel. gap: %.3e", obj' * x - lb, relgap)
                if relgap <= opttol
                    verbose && println()
                    return trackcalls ? (x, oraclecalls) : x
                else
                    a = -nobj
                end
            end
        end
        # See Theorem 8.1 in "Introduction to Linear Optimization" by Bertsimas
        # and Tsitsiklis.
        try
            anorm = a / sqrt(a' * D * a)
            Danorm = D*anorm / (n+1)
            x .+= Danorm
            D .*= n^2/(n^2-1)
            D .+= - 2 * n^2/(n-1) * Danorm*Danorm'
            # x = x + (D*a)/(sqrt(a'*D*a) * (n+1))
            # D = n^2/(n^2-1) * (D - 2 * D*a*a'*D/(a'*D*a * (n+1)))
        catch e
            println("Numerical error: cannot compute the next ellipsoid. "*
            "The spectrum of the ellipsoid matrix follows.")
            display(eigen(D).values)
            rethrow(e)
        end
        if !issymmetric(D)
            D = 0.5*(D + D')
        end
    end
end


using FEASTSolver
using LinearAlgebra
using MatrixMarket
using Random
import FEASTSolver: iter_debug_print, update_R!, residuals, normalize!

function info(Λ, X, residuals, c, r)
    inside(x) = in_contour(x, c, r)
    in_eig = Λ[inside.(Λ)]
    in_res = residuals[inside.(Λ)]
    print("number inside : ")
    println(size(Λ[inside.(Λ)])[1])
    if sum(inside.(Λ)) > 0
        in_res_conv = in_res[in_res .<= 1e-5]
        in_eig_conv = in_eig[in_res .<= 1e-5]
        print("number inside converged : ")
        println(size(in_eig_conv)[1])
        print("max res inside: ")
        println(maximum(residuals[inside.(Λ)]))
        if size(in_res_conv, 1) > 0
            print("max res inside non spurious: ")
            println(maximum(in_res_conv))
        end
    end
    # display(in_eig_conv)
end

function moments_expand!(T, X::AbstractMatrix{ComplexF64}, nodes::Integer, iter::Integer;
    c=complex(0.0,0.0), r=1.0, debug=false, ϵ=10e-12, moments=2, spurious=1e-5)

    N, m₀ = size(X)
    Λ = zeros(ComplexF64, moments*m₀)
    θ = LinRange(π/nodes, 2*π-π/nodes, nodes)
    Q = zeros(ComplexF64, 2*moments, N, m₀)
    Q₀, Q₁ = zeros(ComplexF64, moments*N, moments*m₀), zeros(ComplexF64, moments*N, moments*m₀)
	Y = zeros(ComplexF64, N, m₀*moments)
    R = similar(Y, ComplexF64)

    for i=1:nodes
        z = (r*exp(θ[i]*im)+c)
		Temp = (T(z)\X) .* (r*exp(θ[i]*im)/nodes)
        Q[1,:,:] .+= Temp
        for j=1:2*moments
            Q[j,:,:] .+= Temp .* z^(j-1)
        end
    end
    for i=1:moments, j=1:moments
        Q₀[(i-1)*N+1:i*N, (j-1)*m₀+1:j*m₀] .= Q[i+j-1,:,:]
        Q₁[(i-1)*N+1:i*N, (j-1)*m₀+1:j*m₀] .= Q[i+j,:,:]
    end

    S = svd!(Q₀)
    F = eigen!(S.U' * Q₁ * S.V * Diagonal(1 ./ S.S))
    Y = S.U[1:N,:] * F.vectors
    Λ = F.values

    update_R!(Y, R, Λ, T)
    res = residuals(R, Λ, T)

	iter_debug_print(0, Λ, res, c, r, spurious)

	m₁ = moments*m₀
    Q₀, Q₁ = zeros(ComplexF64, moments*N, moments*m₁), zeros(ComplexF64, moments*N, moments*m₁)
	Q = zeros(ComplexF64, 2*moments, N, m₁)

    for i=1:nodes
        z = (r*exp(θ[i]*im)+c)
        resolvent = (1 ./(z .- Λ)) .* (r*exp(θ[i]*im)/nodes)

		Temp = (Y - T(z)\R) * Diagonal(resolvent)

        Q[1,:,:] .+= Temp
        for j=1:2*moments
            Q[j,:,:] .+= Temp .* z^(j-1)
        end
    end

    for i=1:moments, j=1:moments
        Q₀[(i-1)*N+1:i*N, (j-1)*m₁+1:j*m₁] .= Q[i+j-1,:,:]
        Q₁[(i-1)*N+1:i*N, (j-1)*m₁+1:j*m₁] .= Q[i+j,:,:]
    end

	S = svd!(Q₀)
    F = eigen!(S.U' * Q₁ * S.V * Diagonal(1 ./ S.S))
    Y = S.U[1:N,:] * F.vectors
    Λ = F.values

	R = similar(Y, ComplexF64)

	update_R!(Y, R, Λ, T)
    res = residuals(R, Λ, T)

	iter_debug_print(1, Λ, res, c, r, spurious)

	p = sortperm(res)
	res .= res[p]
	X .= Y[:, p][:,1:m₀]
	Λ .= Λ[p]
	R .= R[:, p]

    normalize!(X)
    Λ, X, res
end

#
# A0 = Matrix(mmread("data/quadraticM0.mtx"))
# A1 = Matrix(mmread("data/quadraticM1.mtx"))

Random.seed!(12345)
A0 = complex.(rand(100,100))
A1 = complex.(rand(100,100))
A0[:,1] .= 0

function T(z::ComplexF64)
	return (z + 0.4)*(z - 0.5)*(z + 0.4im)*(z-0.5im).*A1 .+ A0
end

R = 0.525
# C = complex(-0.8, 0.8)
C = complex(0, 0.0)

# e, v, res = nlfeast!(T, rand(ComplexF64,15,8), 2^6, 0, c=C, r=R, spurious=1e-3,debug=true)

e, v, res = nlfeast_moments!(T, rand(ComplexF64,100,30), 2^3, 10, c=C, r=R, moments=5, debug=true, ϵ=10e-16, store=false, spurious=1e-5)

# e, v, res = moments_expand!(T, rand(ComplexF64,15,7), 2^6, 1, c=C, r=R, moments=2, ϵ=10e-16, spurious=1e-3)
# e, v, res = beyn(T, A0, rand(ComplexF64,1000,120), 2^9; c=complex(-1.55,0.0), r=0.05)
# display(e)
# display(res)
#
inside(x) = in_contour(x, C, R)
# # print("\nmax res inside: ")
# # println(maximum(res[inside.(e)]))
# # print("number inside : ")
# # println(size(e[inside.(e)])[1])
# # # display(e)
# # display(res[inside.(e)])
#
# in_eig = e[inside.(e)]
# in_res = res[inside.(e)]
# in_res_conv = in_res[in_res .<= 1e-3]
# in_eig_conv = in_eig[in_res .<= 1e-3]
# if size(in_eig, 1) > 0
#     print("\nmax res inside: ")
#     println(maximum(res[inside.(e)]))
# end
# print("max res inside non spurious: ")
# println(maximum(in_res_conv))
# print("\nnumber inside : ")
# println(size(e[inside.(e)])[1])
# print("number inside converged : ")
# println(size(in_eig_conv)[1])

# Z = zeros(100,100)
# e, v, res = companion([(A0 .- 0.04.*A1), ((0.02im-0.02).*A1), (0.01im.*A1), ((-0.1-0.1im).*A1), A1])

info(e, v, res, C, R)

# display(res)
# println()
display(e[inside.(e)])


struct BrentMin{T}
    division::T
    detect_flatness::Bool
end
BrentMin(; division = (3 - sqrt(5)) / 2, detect_flatness = false) =
    BrentMin(division, detect_flatness)

function solve(
    problem::OptimizationProblem,
    approach::BrentMin,
    options::OptimizationOptions,
)
    _solve(problem, approach, options)
end

function _solve(prob, bm::BrentMin, options)
    time0 = time()
    a, b = bounds(prob)
    T = typeof(a)
    t = 1e-8
    c = bm.division
    v = w = x = a + c * (b - a)

    e = d = 0 * x

    f0 = fv = fw = fx = value(prob, x)
    p = q = r = 0 * x
    brent_iter = 0
    while brent_iter < options.maxiter
        brent_iter += 1
        m = (a + b) / 2
        tol = eps(T) * abs(x) + t
        if abs(x - m) > 2 * tol - (b - a) / 2 # stopping crit
            # fit parabola
            if abs(e) > tol
                r = (x - w) * (fx - fv)
                q = (x - v) * (fx - fw)
                p = (x - v) * q - (x - w) * r
                q = 2 * (q - r)
                if q > 0
                    p = -p
                else
                    q = -q
                end
                r = e
                e = d
            end

            if abs(p) < abs(q * r / 2) && p < q * (a - x) && p < q * (b - x)
                # then do the parapolic interpolation
                d = p / q
                u = x + d
                if u - a < 2 * tol || b - u < 2 * tol
                    d = sign(m-x)*tol
                end
            else
                # do golden section
                if x < m
                    e = b - x
                else
                    e = a - x
                end
                d = c * e
            end

            if abs(d) >= tol
                u = x + d
            else
                u = x + sign(d) * tol
            end

            fu = value(prob, u)
            if fu <= fx
                if u < x
                    b = x
                else
                    a = x
                end
                v = w
                fv = fw
                w = x
                fw = fx
                x = u
                fx = fu
            else
                if u < x
                    a = u
                else
                    b = u
                end
                if fu <= fw || w == x
                    v = w
                    fv = fw
                    w = u
                    fw = fu
                elseif fu <= fv || v == x || v == w
                    v = u
                    fv = fu
                end
            end
        else
            break
        end
    end

    return ConvergenceInfo(bm, (;x, f0, minimum=fx, time=time()-time0, iter=brent_iter), options)
end


using Graphs
using LightGraphs
using Metis
using Compat
using Base.Test

const copter2 = Metis.testgraph("copter2");
perm, iperm = nodeND(copter2)
@test all(invperm(perm) .== iperm)

srand(12321)
const a = convert(SparseMatrixCSC{Float64,Int32},sprand(1000,100,0.01))
const ata = a'a
perm, iperm = nodeND(ata,3)
@test all(invperm(perm) .== iperm)
perm, iperm = nodeND(ata,0)             # reset verbosity level

function counts{T<:Integer}(v::Vector{T})
    ans = zeros(Int,maximum(v))
    for vv in v
        ans[vv] += 1
    end
    ans
end

objval, part = partGraphKway(copter2, 6)
@test extrema(part) == (1,6)            # subset values are in the correct range
@test all(i->findfirst(part,i) != 0,1:6) # the 6 subsets are non-empty

# the sizes of partition subsets are not consistent across OS's
#@test counts(part) == [9076,9374,9384,9523,8978,9141]

function interface(part::Vector, g::Graphs.GenericAdjacencyList)
    (nv = length(part)) == length(g.vertices) || error("partition length != # of vertices")
    conn = falses(nv)                   # vertex connected to another subset?
    for i in 1:nv
        pp = part[i]
        conn[i] = any(j -> part[j] != pp, g.adjlist[i])
    end
    countnz(conn)
end

@test interface(part,copter2) < 6000    # 5907 on an Ubuntu system

objval, part = partGraphRecursive(copter2,6)
@test extrema(part) == (1,6)
@test all(i->findfirst(part,i) != 0,1:6)
@test interface(part,copter2) < 6000

const mdual = Metis.testgraph("mdual")
objval, part = partGraphKway(mdual, 10)
@test extrema(part) == (1,10)
@test all(i->findfirst(part,i) != 0,1:10)
@test interface(part,mdual) < 19000     # 18263 on an Ubuntu system

sizes, copterPart = vertexSep(copter2)

function testGraphPart(g,part)
    validPart = true
    for i in 1:length(part)
        partVal = part[i]
        if partVal == 0
            for j in g.adjlist[i]
                if part[j] == 1
                    println("Edge ($i,$j) connects sets 0 and 1")
                    validPart = false
                end
            end
        elseif partVal == 1
            for j in g.adjlist[i]
                if part[j] == 0
                    println("Edge ($i,$j) connects sets 1 and 0")
                    validPart = false
                end
            end
        elseif partVal != 2
            println("Vertex $i assigned to set $partVal")
            validPart = false
        end
    end
    validPart
end

@test testGraphPart(copter2,copterPart)

function appendel(I,J,V,i,j,v)
    push!(I,i)
    push!(J,j)
    push!(V,v)
end

function laplacian2d{T<:Integer}(nx::T,ny::T)
    n = nx*ny
    nzest = 5n
    I = @compat sizehint!(Int32[],nzest)
    J = @compat sizehint!(Int32[],nzest)
    V = @compat sizehint!(Float64[],nzest)
    for x in 1:nx
        for y in 1:ny
            s = x + (y-1)*nx
            appendel(I,J,V,s,s,2)
            x > 1 && appendel(I,J,V,s,s-1,-1)
            y > 1 && appendel(I,J,V,s,s-nx,-1)
        end
    end
    A = sparse(I,J,V,n,n)
    A + A'
end

const A = laplacian2d(100,110)

sizes, matPart = vertexSep(A)

function testMatPart(m,part)
    validPart = true
    for i in 1:length(part)
        partVal = part[i]
        if partVal == 0
            for k in m.colptr[i]:m.colptr[i+1]-1
                j = m.rowval[k]
                if part[j] == 1
                    println("Nonzero ($i,$j) connects sets 0 and 1")
                    validPart = false
                end
            end
        elseif partVal == 1
            for k in m.colptr[i]:m.colptr[i+1]-1
                j = m.rowval[k]
                if part[j] == 0
                    println("Nonzero ($i,$j) connects sets 1 and 0")
                    validPart = false
                end
            end
        elseif partVal != 2
            println("Vertex $i assigned to set $partVal")
            validPart = false
        end
    end
    validPart
end

@test testMatPart(A,matPart)

g = LightGraphs.TutteGraph()
x, y = partGraphKway(g, 6)
@test extrema(y) == (1,6)
@test length(y) == LightGraphs.nv(g)
@test all(i-> findfirst(y,i) != 0,1:6)

function interface(part::Vector, g::LightGraphs.Graph)
    (n = nv(g)) == length(part) || error("partition length != # of vertices")
    conn = falses(n)             # vertex connected to another subset?
    for i in 1:n
        pp = part[i]
        conn[i] = any(j -> part[j] != pp, map(dst,g.finclist[i]))
    end
    countnz(conn)
end


const ARDRegressor_ = sklm(:ARDRegression)
@sk_reg mutable struct ARDRegressor <: MMI.Deterministic
    n_iter::Int               = 300::(_ > 0)
    tol::Float64              = 1e-3::(_ > 0)
    alpha_1::Float64          = 1e-6::(_ > 0)
    alpha_2::Float64          = 1e-6::(_ > 0)
    lambda_1::Float64         = 1e-6::(_ > 0)
    lambda_2::Float64         = 1e-6::(_ > 0)
    compute_score::Bool       = false
    threshold_lambda::Float64 = 1e4::(_ > 0)
    fit_intercept::Bool       = true
    normalize::Bool           = false
    copy_X::Bool              = true
    verbose::Bool             = false
end
MMI.fitted_params(model::ARDRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    alpha     = fitresult.alpha_,
    lambda    = fitresult.lambda_,
    sigma     = fitresult.sigma_,
    scores    = fitresult.scores_
    )
descr(ARDRegressor, "Bayesian ARD regression.")

# =============================================================================
const BayesianRidgeRegressor_ = sklm(:BayesianRidge)
@sk_reg mutable struct BayesianRidgeRegressor <: MMI.Deterministic
    n_iter::Int         = 300::(_ ≥ 1)
    tol::Float64        = 1e-3::(_ > 0)
    alpha_1::Float64    = 1e-6::(_ > 0)
    alpha_2::Float64    = 1e-6::(_ > 0)
    lambda_1::Float64   = 1e-6::(_ > 0)
    lambda_2::Float64   = 1e-6::(_ > 0)
    compute_score::Bool = false
    fit_intercept::Bool = true
    normalize::Bool     = false
    copy_X::Bool        = true
    verbose::Bool       = false
end
MMI.fitted_params(model::BayesianRidgeRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    alpha     = fitresult.alpha_,
    lambda    = fitresult.lambda_,
    sigma     = fitresult.sigma_,
    scores    = fitresult.scores_
    )
descr(BayesianRidgeRegressor, "Bayesian ridge regression.")

# =============================================================================
const ElasticNetRegressor_ = sklm(:ElasticNet)
@sk_reg mutable struct ElasticNetRegressor <: MMI.Deterministic
    alpha::Float64      = 1.0::(_ ≥ 0)   # 0 is OLS
    l1_ratio::Float64   = 0.5::(0 ≤ _ ≤ 1)
    fit_intercept::Bool = true
    normalize::Bool     = false
    precompute::Union{Bool,AbstractMatrix} = false
    max_iter::Int       = 1_000::(_ ≥ 1)
    copy_X::Bool        = true
    tol::Float64        = 1e-4::(_ > 0)
    warm_start::Bool    = false
    positive::Bool      = false
    random_state::Any   = nothing  # Int, random state, or nothing
    selection::String   = "cyclic"::(_ in ("cyclic","random"))
end
MMI.fitted_params(model::ElasticNetRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    )
descr(ElasticNetRegressor, "Elastic net regression.")

# =============================================================================
const ElasticNetCVRegressor_ = sklm(:ElasticNetCV)
@sk_reg mutable struct ElasticNetCVRegressor <: MMI.Deterministic
    l1_ratio::Union{Float64,Vector{Float64}} = 0.5::(all(0 .≤ _ .≤ 1))
    eps::Float64        = 1e-3::(_ > 0)
    n_alphas::Int       = 100::(_ > 0)
    alphas::Any         = nothing::(_ === nothing || all(0 .≤ _))
    fit_intercept::Bool = true
    normalize::Bool     = false
    precompute::Union{Bool,String,AbstractMatrix} = "auto"
    max_iter::Int       = 1_000::(_ > 0)
    tol::Float64        = 1e-4::(_ > 0)
    cv::Any             = 5 # can be Int, Nothing or an iterable / cv splitter
    copy_X::Bool        = true
    verbose::Union{Bool, Int}  = 0
    n_jobs::Option{Int} = nothing
    positive::Bool      = false
    random_state::Any   = nothing
    selection::String   = "cyclic"::(_ in ("cyclic","random"))
end
MMI.fitted_params(model::ElasticNetCVRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    l1_ratio  = fitresult.l1_ratio_,
    mse_path  = fitresult.mse_path_,
    alphas    = fitresult.alphas_
    )
descr(ElasticNetCVRegressor, "Elastic net regression $CV.")

# =============================================================================
const HuberRegressor_ = sklm(:HuberRegressor)
@sk_reg mutable struct HuberRegressor <: MMI.Deterministic
    epsilon::Float64    = 1.35::(_ > 1.0)
    max_iter::Int       = 100::(_ > 0)
    alpha::Float64      = 1e-4::(_ > 0)
    warm_start::Bool    = false
    fit_intercept::Bool = true
    tol::Float64        = 1e-5::(_ > 0)
end
MMI.fitted_params(model::HuberRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    scale     = fitresult.scale_,
    outliers  = fitresult.outliers_
    )
descr(HuberRegressor, "Huber regression.")

# =============================================================================
const LarsRegressor_ = sklm(:Lars)
@sk_reg mutable struct LarsRegressor <: MMI.Deterministic
    fit_intercept::Bool      = true
    verbose::Union{Bool,Int} = false
    normalize::Bool          = true
    precompute::Union{Bool,String,AbstractMatrix} = "auto"
    n_nonzero_coefs::Int     = 500::(_ > 0)
    eps::Float64   = eps(Float64)::(_ > 0)
    copy_X::Bool   = true
    fit_path::Bool = true
end
MMI.fitted_params(model::LarsRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    alphas    = fitresult.alphas_,
    active    = fitresult.active_,
    coef_path = fitresult.coef_path_
    )
descr(LarsRegressor, "Lars regression.")

# =============================================================================
const LarsCVRegressor_ = sklm(:LarsCV)
@sk_reg mutable struct LarsCVRegressor <: MMI.Deterministic
    fit_intercept::Bool      = true
    verbose::Union{Bool,Int} = false
    max_iter::Int     = 500::(_ > 0)
    normalize::Bool   = true
    precompute::Union{Bool,String,AbstractMatrix} = "auto"
    cv::Any           = 5
    max_n_alphas::Int = 1_000::(_ > 0)
    n_jobs::Option{Int} = nothing
    eps::Float64      = eps(Float64)::(_ > 0)
    copy_X::Bool      = true
end
MMI.fitted_params(model::LarsCVRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    alpha     = fitresult.alpha_,
    alphas    = fitresult.alphas_,
    cv_alphas = fitresult.cv_alphas_,
    mse_path  = fitresult.mse_path_,
    coef_path = fitresult.coef_path_
    )
descr(LarsCVRegressor, "Lars regression $CV.")

# =============================================================================
const LassoRegressor_ = sklm(:Lasso)
@sk_reg mutable struct LassoRegressor <: MMI.Deterministic
    alpha::Float64      = 1.0::(_ ≥ 0) # should use alpha > 0 (or OLS)
    fit_intercept::Bool = true
    normalize::Bool     = false
    precompute::Union{Bool,AbstractMatrix} = false
    copy_X::Bool        = true
    max_iter::Int       = 1_000::(_ > 0)
    tol::Float64        = 1e-4::(_ > 0)
    warm_start::Bool    = false
    positive::Bool      = false
    random_state::Any   = nothing
    selection::String   = "cyclic"::(_ in ("cyclic","random"))
end
MMI.fitted_params(model::LassoRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    )
descr(LassoRegressor, "Lasso regression.")

# =============================================================================
const LassoCVRegressor_ = sklm(:LassoCV)
@sk_reg mutable struct LassoCVRegressor <: MMI.Deterministic
    eps::Float64        = 1e-3::(_ > 0)
    n_alphas::Int       = 100::(_ > 0)
    alphas::Any         = nothing::(_ === nothing || all(0 .≤ _))
    fit_intercept::Bool = true
    normalize::Bool     = false
    precompute::Union{Bool,String,AbstractMatrix} = "auto"
    max_iter::Int       = 1_000::(_ > 0)
    tol::Float64        = 1e-4::(_ > 0)
    copy_X::Bool        = true
    cv::Any             = 5
    verbose::Union{Bool,Int} = false
    n_jobs::Option{Int} = nothing
    positive::Bool      = false
    random_state::Any   = nothing
    selection::String   = "cyclic"::(_ in ("cyclic","random"))
end
MMI.fitted_params(model::LassoCVRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    alpha     = fitresult.alpha_,
    alphas    = fitresult.alphas_,
    mse_path  = fitresult.mse_path_,
    dual_gap  = fitresult.dual_gap_
    )
descr(LassoCVRegressor, "Lasso regression $CV.")

# =============================================================================
const LassoLarsRegressor_ = sklm(:LassoLars)
@sk_reg mutable struct LassoLarsRegressor <: MMI.Deterministic
    alpha::Float64      = 1.0::(_ ≥ 0) # 0 should be OLS
    fit_intercept::Bool = true
    verbose::Union{Bool, Int} = false
    normalize::Bool     = true
    precompute::Union{Bool,String,AbstractMatrix} = "auto"
    max_iter::Int       = 500::(_ > 0)
    eps::Float64        = eps(Float64)::(_ > 0)
    copy_X::Bool        = true
    fit_path::Bool      = true
    positive::Any       = false
end
MMI.fitted_params(model::LassoLarsRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    alphas    = fitresult.alphas_,
    active    = fitresult.active_,
    coef_path = fitresult.coef_path_
    )
descr(LassoLarsRegressor, "Lasso model fit with least angle regression (LARS).")

# =============================================================================
const LassoLarsCVRegressor_ = sklm(:LassoLarsCV)
@sk_reg mutable struct LassoLarsCVRegressor <: MMI.Deterministic
    fit_intercept::Bool = true
    verbose::Union{Bool, Int} = false
    max_iter::Int       = 500::(_ > 0)
    normalize::Bool     = true
    precompute::Union{Bool,String,AbstractMatrix} = "auto"
    cv::Any             = 5
    max_n_alphas::Int   = 1_000::(_ > 0)
    n_jobs::Option{Int} = nothing
    eps::Float64        = eps(Float64)::(_ > 0.0)
    copy_X::Bool        = true
    positive::Any       = false
end
MMI.fitted_params(model::LassoLarsCVRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    coef_path = fitresult.coef_path_,
    alpha     = fitresult.alpha_,
    alphas    = fitresult.alphas_,
    cv_alphas = fitresult.cv_alphas_,
    mse_path  = fitresult.mse_path_
    )
descr(LassoLarsCVRegressor, "Lasso model fit with least angle regression (LARS) $CV.")

# =============================================================================
const LassoLarsICRegressor_ = sklm(:LassoLarsIC)
@sk_reg mutable struct LassoLarsICRegressor <: MMI.Deterministic
    criterion::String   = "aic"::(_ in ("aic","bic"))
    fit_intercept::Bool = true
    verbose::Union{Bool, Int} = false
    normalize::Bool     = true
    precompute::Union{Bool,String,AbstractMatrix} = "auto"
    max_iter::Int       = 500::(_ > 0)
    eps::Float64        = eps(Float64)::(_ > 0.0)
    copy_X::Bool        = true
    positive::Any       = false
end
MMI.fitted_params(model::LassoLarsICRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    alpha     = fitresult.alpha_
    )
descr(LassoLarsICRegressor, "Lasso model with Lars using BIC or AIC for model selection.")

# =============================================================================
const LinearRegressor_ = sklm(:LinearRegression)
@sk_reg mutable struct LinearRegressor <: MMI.Deterministic
    fit_intercept::Bool = true
    normalize::Bool     = false
    copy_X::Bool        = true
    n_jobs::Option{Int} = nothing
end
MMI.fitted_params(model::LinearRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing)
    )
descr(LinearRegressor, "Ordinary least-square regression (OLS).")

# =============================================================================
const OrthogonalMatchingPursuitRegressor_ = sklm(:OrthogonalMatchingPursuit)
@sk_reg mutable struct OrthogonalMatchingPursuitRegressor <: MMI.Deterministic
    n_nonzero_coefs::Option{Int} = nothing
    tol::Option{Float64} = nothing
    fit_intercept::Bool  = true
    normalize::Bool      = true
    precompute::Union{Bool,String,AbstractMatrix} = "auto"
end
MMI.fitted_params(model::OrthogonalMatchingPursuitRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing)
    )
descr(OrthogonalMatchingPursuitRegressor, "Orthogonal Matching Pursuit (OMP) model.")

# =============================================================================
const OrthogonalMatchingPursuitCVRegressor_ = sklm(:OrthogonalMatchingPursuitCV)
@sk_reg mutable struct OrthogonalMatchingPursuitCVRegressor <: MMI.Deterministic
    copy::Bool            = true
    fit_intercept::Bool   = true
    normalize::Bool       = false
    max_iter::Option{Int} = nothing::(_ === nothing||_ > 0)
    cv::Any               = 5
    n_jobs::Option{Int}   = 1
    verbose::Union{Bool,Int} = false
end
MMI.fitted_params(model::OrthogonalMatchingPursuitCVRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    n_nonzero_coefs = fitresult.n_nonzero_coefs_
    )
descr(OrthogonalMatchingPursuitCVRegressor, "Orthogonal Matching Pursuit (OMP) model $CV.")

# =============================================================================
const PassiveAggressiveRegressor_ = sklm(:PassiveAggressiveRegressor)
@sk_reg mutable struct PassiveAggressiveRegressor <: MMI.Deterministic
    C::Float64                   = 1.0::(_ > 0)
    fit_intercept::Bool          = true
    max_iter::Int                = 1_000::(_ > 0)
    tol::Float64                 = 1e-4::(_ > 0)
    early_stopping::Bool         = false
    validation_fraction::Float64 = 0.1::(0 < _ < 1)
    n_iter_no_change::Int        = 5::(_ > 0)
    shuffle::Bool                = true
    verbose::Union{Bool,Int}     = 0
    loss::String                 = "epsilon_insensitive"::(_ in ("epsilon_insensitive","squared_epsilon_insensitive"))
    epsilon::Float64             = 0.1::(_ > 0)
    random_state::Any            = nothing
    warm_start::Bool             = false
    average::Union{Bool,Int}     = false
end
MMI.fitted_params(model::PassiveAggressiveRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing)
    )
descr(PassiveAggressiveRegressor, "Passive Aggressive Regressor.")

# =============================================================================
const RANSACRegressor_ = sklm(:RANSACRegressor)
@sk_reg mutable struct RANSACRegressor <: MMI.Deterministic
    base_estimator::Any         = nothing
    min_samples::Union{Int,Float64}     = 5::(_ isa Int ? _ ≥ 1 : (0 ≤ _ ≤ 1))
    residual_threshold::Option{Float64} = nothing
    is_data_valid::Any          = nothing
    is_model_valid::Any         = nothing
    max_trials::Int             = 100::(_ > 0)
    max_skips::Int              = typemax(Int)::(_ > 0)
    stop_n_inliers::Int         = typemax(Int)::(_ > 0)
    stop_score::Float64         = Inf::(_ > 0)
    stop_probability::Float64   = 0.99::(0 ≤ _ ≤ 1.0)
    loss::Union{Function,String}= "absolute_loss"::((_ isa Function) || _ in ("absolute_loss","squared_loss"))
    random_state::Any           = nothing
end
MMI.fitted_params(m::RANSACRegressor, (f, _, _)) = (
    estimator             = f.estimator_,
    n_trials              = f.n_trials_,
    inlier_mask           = f.inlier_mask_,
    n_skips_no_inliers    = f.n_skips_no_inliers_,
    n_skips_invalid_data  = f.n_skips_invalid_data_,
    n_skips_invalid_model = f.n_skips_invalid_model_
    )
descr(RANSACRegressor, "RANSAC regressor.")

# =============================================================================
const RidgeRegressor_ = sklm(:Ridge)
@sk_reg mutable struct RidgeRegressor <: MMI.Deterministic
    alpha::Union{Float64,Vector{Float64}} = 1.0::(all(_ .> 0))
    fit_intercept::Bool = true
    normalize::Bool     = false
    copy_X::Bool        = true
    max_iter::Int       = 1_000::(_ > 0)
    tol::Float64        = 1e-4::(_ > 0)
    solver::String      = "auto"::(_ in ("auto","svd","cholesky","lsqr","sparse_cg","sag","saga"))
    random_state::Any   = nothing
end
MMI.fitted_params(model::RidgeRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing)
    )
descr(RidgeRegressor, "Ridge regression.")

# =============================================================================
const RidgeCVRegressor_ = sklm(:RidgeCV)
@sk_reg mutable struct RidgeCVRegressor <: MMI.Deterministic
    alphas::Any              = (0.1, 1.0, 10.0)::(all(_ .> 0))
    fit_intercept::Bool      = true
    normalize::Bool          = false
    scoring::Any             = nothing
    cv::Any                  = 5
    gcv_mode::Option{String} = nothing::(_ === nothing || _ in ("auto","svd","eigen"))
    store_cv_values::Bool    = false
end
MMI.fitted_params(model::RidgeCVRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    alpha     = fitresult.alpha_,
    cv_values = model.store_cv_values ? fitresult.cv_values_ : nothing
    )
descr(RidgeCVRegressor, "Ridge regression $CV.")

# =============================================================================
const SGDRegressor_ = sklm(:SGDRegressor)
@sk_reg mutable struct SGDRegressor <: MMI.Deterministic
    loss::String             = "squared_loss"::(_ in ("squared_loss","huber","epsilon_insensitive","squared_epsilon_insensitive"))
    penalty::String          = "l2"::(_ in ("none","l2","l1","elasticnet"))
    alpha::Float64           = 1e-4::(_ > 0)
    l1_ratio::Float64        = 0.15::(_ > 0)
    fit_intercept::Bool      = true
    max_iter::Int            = 1_000::(_ > 0)
    tol::Float64             = 1e-3::(_ > 0)
    shuffle::Bool            = true
    verbose::Union{Int,Bool} = 0
    epsilon::Float64         = 0.1
    random_state::Any        = nothing
    learning_rate::String    = "invscaling"::(_ in ("constant","optimal","invscaling","adaptive"))
    eta0::Float64            = 0.01::(_ > 0)
    power_t::Float64         = 0.25::(_ > 0)
    early_stopping::Bool     = false
    validation_fraction::Float64 = 0.1::(0 < _ < 1)
    n_iter_no_change::Int    = 5::(_ > 0)
    warm_start::Bool         = false
    average::Union{Int,Bool} = false
end
MMI.fitted_params(model::SGDRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    average_coef      = model.average ? fitresult.average_coef_ : nothing,
    average_intercept = model.average ? ifelse(model.fit_intercept, fitresult.average_intercept_, nothing) : nothing
    )
descr(SGDRegressor, "Stochastic gradient descent-based regressor.")

# =============================================================================
const TheilSenRegressor_ = sklm(:TheilSenRegressor)
@sk_reg mutable struct TheilSenRegressor <: MMI.Deterministic
    fit_intercept::Bool = true
    copy_X::Bool        = true
    max_subpopulation::Int    = 10_000::(_ > 0)
    n_subsamples::Option{Int} = nothing::(_ === nothing||_ > 0)
    max_iter::Int       = 300::(_ > 0)
    tol::Float64        = 1e-3::(_ > 0)
    random_state::Any   = nothing
    n_jobs::Option{Int} = nothing
    verbose::Bool       = false
end
MMI.fitted_params(model::TheilSenRegressor, (fitresult, _, _)) = (
    coef      = fitresult.coef_,
    intercept = ifelse(model.fit_intercept, fitresult.intercept_, nothing),
    breakdown       = fitresult.breakdown_,
    n_subpopulation = fitresult.n_subpopulation_
    )
descr(TheilSenRegressor, "Theil-Sen regressor.")

# Metadata for Continuous -> Vector{Continuous}
const SKL_REGS_SINGLE = Union{
        Type{<:ARDRegressor},
        Type{<:BayesianRidgeRegressor},
        Type{<:ElasticNetRegressor},
        Type{<:ElasticNetCVRegressor},
        Type{<:HuberRegressor},
        Type{<:LarsRegressor},
        Type{<:LarsCVRegressor},
        Type{<:LassoRegressor},
        Type{<:LassoCVRegressor},
        Type{<:LassoLarsRegressor},
        Type{<:LassoLarsCVRegressor},
        Type{<:LassoLarsICRegressor},
        Type{<:LinearRegressor},
        Type{<:OrthogonalMatchingPursuitRegressor},
        Type{<:OrthogonalMatchingPursuitCVRegressor},
        Type{<:PassiveAggressiveRegressor},
        Type{<:RANSACRegressor},
        Type{<:RidgeRegressor},
        Type{<:RidgeCVRegressor},
        Type{<:SGDRegressor},
        Type{<:TheilSenRegressor}
    }

MMI.input_scitype(::SKL_REGS_SINGLE)  = Table(Continuous)
MMI.target_scitype(::SKL_REGS_SINGLE) = AbstractVector{Continuous}


using ApproxFun
using LinearAlgebra
using Plots; gr()
using Distributions
using OrdinaryDiffEq

println("Fourier Stuff")
S = Fourier()
N = 512
x = points(S, N)
y = points(S, N)
D2x = Derivative(S,2)[1:n,1:n]
D2 = kron(D2x, D2x)
T = ApproxFun.plan_transform(S, N)
Ti = ApproxFun.plan_itransform(S, N)

println("Caches")
u = zeros(N,N)
v = similar(u)
u_n = similar(u)
v_n = similar(u)
û_n = similar(u)
v̂_n = similar(u)
û_1 = similar(u)
v̂_1 = similar(u)
u_d2 = zeros(N*N)
v_d2 = zeros(N*N)

println("Setting up problem parameters")
Mᵤ =  1e3
Mᵥ =  1.0
γᵤ =  1.0
γᵥ =  1.0
a = 0.5
c = 0.00125

println("Initial Conditions")
r01 = 0.01 .+ rand(Uniform(-0.01/100, +0.01/100), N, N)
r02 = 0.25 .+ rand(Uniform(-0.25/100, +0.25/100), N, N)
r0 = cat(r01, r02, dims=3)

cache_ops = (D2,T,Ti, u, v, u_n, v_n, û_n, v̂_n, û_1, v̂_1, u_d2, v_d2)
p = (Mᵤ, Mᵥ, γᵤ, γᵥ, a, c, cache_ops)
function NTF_nl!(dr̂,r̂,p,t)
    Mᵤ, Mᵥ, γᵤ, γᵥ, a, c, cache_ops = p
    D2,T,Ti, u, v, u_n, v_n, û_n, v̂_n, û_1, v̂_1 = cache_ops
    û = @view r̂[:,:,1]
    v̂ = @view r̂[:,:,2]
    dû = @view dr̂[:,:,1]
    dv̂ = @view dr̂[:,:,2]
    mul!(u, Ti, û)
    mul!(v, Ti, v̂)

    @. u_n = 2u*(u-1)*(2u-1)
    mul!(û_n, T, u_n)
    mul!(u_d2, D2, vec(û))
    @. û_1 = (û_n-γᵤ*reshape(u_d2,N,N)) # Figure out how to do this so you don't have to reshape?
    mul!(û_1, D2, û_1)
    @. u_n = u*v
    mul!(u_n, T, u_n)
    dû = Mᵤ*û_1 - a*u_n + c

    @. v_n = 2v*(v-1)*(2v-1)
    mul!(v̂_n, T, v_n)
    mul!(v_d2, D2, vec(v̂))
    @. û_1 = (v̂_n-γᵥ*reshape(v_d2, N, N))
    mul!(v̂_1, D2, v̂_1)
    @. v_n = u*v #Make efficient, need only once
    mul!(v_n, T, v_n)
    dv̂ = Mᵥ*v̂_1 - a*v_n + c
end

using Symbolics
#println("Jacobian sparsity")
#dr0 = copy(r0)
#@time jac_sparsity = Symbolics.jacobian_sparsity((dr,r)->NTF_nl!(dr,r,p,0.0),dr0,r0)


using AMRVW
using Test
const A = AMRVW

include("test-amrvw-basics.jl")
include("test-passthrough.jl")
include("test-roots.jl")
include("test-twisted.jl")


using Test
using PyCall

haslib = pyimport("haslib")

mel = 1.6021766208000002e-22
A = 1e-10
a = 4.3*A
Ei = 4.196*mel
theta = 45.75

D = 5*mel
h = 0.1
xi = 0.38/A

crystal=haslib.Crystal(a, a, 120)
crystal.calc_reciprocals()
grid = HASlib.HexGrid(a)

@test grid.l.a1 ≈ crystal.a1[1:2]
@test grid.l.a2 ≈ crystal.a2[1:2]
@test grid.l.b1 ≈ crystal.b1_orig[1:2]
@test grid.l.b2 ≈ crystal.b2_orig[1:2]

function test_rotation(dir)
	crystal=haslib.Crystal(a, a, 120)
	grid = HASlib.HexGrid(a)
	crystal.rotate_crystal(dir...)
	b1,b2 = HASlib.rotate_reciprocal(grid.l,dir)
	(b1 ≈ crystal.b1[1:2]) && (b2 ≈ crystal.b2[1:2])
end

@test test_rotation([2,0])
@test test_rotation([0,2])
@test test_rotation([2,1])
@test test_rotation([2,-1])

b1,b2 = HASlib.rotate_reciprocal(grid.l,[2,-1])
mb1,mb2 = HASlib.rotate_reciprocal(grid.l,[-2,1])
@test b1 ≈ -mb1
@test b2 ≈ -mb2

b1,b2 = HASlib.rotate_reciprocal(grid.l,[1,0])
ERG = haslib.icc_2d_2.iCC_2D_2(h,1.0,a,D,xi,Ei,theta,b1,b2,floq=0,Gope=-1,Gclo=50,maxrad=20,sta_z=-1*A,end_z=6*A,st_p_wl=150)

m_He3 = 5.0082343984338715e-27
ki = sqrt(Ei/HASlib.h2m(m_He3))

channels = HASlib.get_channels(b1,b2,ki,theta*pi/180,5)
channels = HASlib.sort_and_cut(channels,-1,50)
result = cc(channels,a,h,D,xi,theta,m_He3,(-1*A,6*A),st_p_wl=150)

@test sort(result[:,1]) ≈ sort(ERG[:,7])

using BenchmarkTools
display(@benchmark cc(channels,a,h,D,xi,theta,m_He3,(-1*A,6*A),st_p_wl=150))
display(@benchmark haslib.CC_2D_old(h,a,D,xi,Ei,theta,b1,b2,G_ope=-1,G_clo=50,maxrad=20,sta_z=-1*A,end_z=6*A,st_p_wl=150))
display(@benchmark haslib.icc_2d_2.iCC_2D_2(h,1,a,D,xi,Ei,theta,b1,b2,Gope=-1,Gclo=50,floq=1,maxrad=20,sta_z=-1*A,end_z=6*A,st_p_wl=150))


using PyPlot, Bridge
λ(t) = 6/(t + 2)
T = 10.0
λmax = 3.0

P = InhomogPoisson(λ)
tt = sample(ThinningAlg(λmax), T, P)

plot(Bridge.piecewise(SamplePath(tt, collect(1:length(tt))))..., label="N(t)")
legend()