# using Revise
using Test, BifurcationKit, LinearAlgebra, Setfield, ForwardDiff
const BK = BifurcationKit

norminf = (x)->norm(x, Inf)
source_term(x; a = 0.5, b = 0.01) = 1 + (x + a*x^2)/(1 + b*x^2)
limits(x,i,N, b) = (i<1||i>N) ? b : x[i]

function F_chan(x, p)
	α, β = p
	N = length(x)
	f = similar(x)
	N = length(x)
	ind(x,i) = limits(x,i,N,β)
	for i=1:N
		f[i] = (ind(x,i-1) - 2 * x[i] + ind(x,i+1)) * (N-1)^2 + α * source_term(x[i], b = β)
	end
	return f
end

J_chan(u0, p) = ForwardDiff.jacobian(u -> F_chan(u,p), u0)
jet = BK.getJet(F_chan, J_chan)
par_chan = (α = 3.3, β = 0.01)

n = 101
	sol = [(i-1)*(n-i)/n^2+0.1 for i=1:n]
	opt_newton = NewtonPar(tol = 1e-9, verbose = false)
	out, hist, flag = newton(jet[1], jet[2], sol, par_chan,
							opt_newton, normN = norminf)

# _J = J_chan(sol, par_chan)
# heatmap(_J - _J', yflip=true)

# test with secant continuation
opts_br0 = ContinuationPar(dsmin = 0.01, dsmax = 0.15, ds= 0.01, pMax = 4.1, maxSteps = 150, newtonOptions = opt_newton, detectBifurcation = 3)
	br, = continuation(jet[1], jet[2],
		out, par_chan, (@lens _.α), opts_br0; )
####################################################################################################
# deflation newton solver, test of jacobian expression
deflationOp = DeflationOperator(2.0, dot, 1.0, [out])

# for testing
show(deflationOp)
length(deflationOp)
deflationOp(2out, out)
push!(deflationOp, rand(n))
deleteat!(deflationOp, 2)

chanDefPb = DeflatedProblem(jet[1], jet[2], deflationOp)

opt_def = setproperties(opt_newton; tol = 1e-10, maxIter = 1000)
outdef1, = newton(chanDefPb, out .* (1 .+0.01*rand(n)), par_chan, opt_def)

# we now compare the jacobians for the deflated problem either using finite differences or the explicit jacobian
rhs = rand(n)
J_def_fd = BK.finiteDifferences(u->chanDefPb(u, par_chan),1.5*out)
res_fd =  J_def_fd \ rhs

Jacdf = (u0, pb::DeflatedProblem, ls = opt_def.linsolve ) -> (return (u0, par_chan, pb, ls))
Jacdfsolver = DeflatedLinearSolver()

res_explicit = Jacdfsolver(Jacdf(1.5out, chanDefPb, opt_def.linsolver),rhs)[1]

# Test jacobian expression for deflated problem
@test norm(res_fd - res_explicit,Inf64) < 1e-4

opt_def = setproperties(opt_newton; tol = 1e-10, maxIter = 1000)
outdef1, = newton(
		jet[1], jet[2],
		out.*(1 .+ 0.1*rand(n)), par_chan,
		opt_def, deflationOp)

####################################################################################################
# Fold continuation, test of Jacobian expression
outfold = newtonFold(jet[1], jet[2], br, 2; startWithEigen = true)
@test  outfold[3] && outfold[4] == 2
outfold = newtonFold(jet[1], jet[2], br, 2; startWithEigen = true, issymmetric = true)
@test  outfold[3] && outfold[4] == 2
outfold = newton(jet[1], jet[2], br, 2; startWithEigen = true, issymmetric = true)
@test  outfold[3] && outfold[4] == 2

optcontfold = ContinuationPar(dsmin = 0.001, dsmax = 0.15, ds= 0.01, pMax = 4.1, pMin = 0., newtonOptions = NewtonPar(verbose=false, tol = 1e-8), maxSteps = 50, detectBifurcation = 2)
outfoldco, = continuationFold(jet[1], jet[2], br, 2, (@lens _.β), optcontfold; startWithEigen = true, updateMinAugEveryStep = 1, plot = false)
outfoldco, = continuationFold(jet[1], jet[2], br, 2, (@lens _.β), optcontfold; startWithEigen = true, updateMinAugEveryStep = 1, issymmetric = true, plot = false)
outfoldco, = continuationFold(jet[1], jet[2], br, 2, (@lens _.β), optcontfold; startWithEigen = true, updateMinAugEveryStep = 1, issymmetric = true, d2F = jet[3], plot = false)

# manual handling
indfold = 1
foldpt = FoldPoint(br, indfold)
foldpb = FoldProblemMinimallyAugmented(
		jet[1], jet[2],
		nothing, nothing,
		(@lens _.α),
		br.specialpoint[indfold].x,
		br.specialpoint[indfold].x,
		opts_br0.newtonOptions.linsolver)
foldpb(foldpt, par_chan) |> norm

outfold, = newtonFold(jet[1], jet[2], foldpt, par_chan, (@lens _.α), br.specialpoint[indfold].x, br.specialpoint[indfold].x, NewtonPar(verbose=false))
	# println("--> Fold found at α = ", outfold.p, " from ", br.specialpoint[indfold].param)

# user defined Fold Problem
indfold = 1

# we define the following wrappers to be able to use ForwardDiff
Bd2Vec(x) = vcat(x.u, x.p)
Vec2Bd(x) = BorderedArray(x[1:end-1], x[end])
foldpbVec(x,p) = Bd2Vec(foldpb(Vec2Bd(x),p))

rhs = rand(n+1)
Jac_fold_fdMA(u0) = ForwardDiff.jacobian( u -> foldpbVec(u, par_chan), u0)
J_fold_fd = Jac_fold_fdMA(Bd2Vec(foldpt))
res_fd =  J_fold_fd \ rhs

Jac_fold_MA(u0, p, pb::FoldProblemMinimallyAugmented) = (return (x=u0, params=p, fldpb = pb))
jacFoldSolver = BK.FoldLinearSolverMinAug()
debugTmpForσ = zeros(n+1,n+1) # temporary array for debugging σ
res_explicit = jacFoldSolver(Jac_fold_MA(foldpt, par_chan, foldpb), Vec2Bd(rhs); debugArray = debugTmpForσ)

Jac_fold_MA(foldpt, 0.01, foldpb)[2]

# test whether the Jacobian Matrix for the Fold problem is correct
@test debugTmpForσ[1:end-1,1:end-1] - J_fold_fd[1:end-1,1:end-1] |> norminf == 0

# we test the expression for σp
σp_fd = J_fold_fd[end,end]
σp_fd_ana = debugTmpForσ[end,end]
@test σp_fd ≈ σp_fd_ana rtol = 1e-5

# we test the expression for σx
σx_fd = J_fold_fd[end,1:end-1]
σx_ana = debugTmpForσ[end,1:end-1]
@test σx_fd ≈ σx_ana rtol = 1e-2

σx_fd - σx_ana |> norminf
debugTmpForσ[1:end-1,1:end-1] - J_fold_fd[1:end-1,1:end-1] |> norminf

# check our solution of the bordered problem
res_exp = debugTmpForσ \ rhs
@test norm(res_exp - Bd2Vec(res_explicit[1]), Inf64) < 1e-10
#############################################
# same with foldpb.issymmetric = true
@set! foldpb.issymmetric = true
rhs = rand(n+1)
Jac_fold_fdMA(u0) = ForwardDiff.jacobian( u -> foldpbVec(u, par_chan), u0)
J_fold_fd = Jac_fold_fdMA(Bd2Vec(foldpt))
res_fd =  J_fold_fd \ rhs

Jac_fold_MA(u0, p, pb::FoldProblemMinimallyAugmented) = (return (x=u0, params=p, fldpb = pb))
jacFoldSolver = BK.FoldLinearSolverMinAug()
debugTmpForσ = zeros(n+1,n+1) # temporary array for debugging σ
res_explicit = jacFoldSolver(Jac_fold_MA(foldpt, par_chan, foldpb), Vec2Bd(rhs); debugArray = debugTmpForσ)

Jac_fold_MA(foldpt, 0.01, foldpb)[2]

# test whether the Jacobian Matrix for the Fold problem is correct
@test debugTmpForσ[1:end-1,1:end-1] - J_fold_fd[1:end-1,1:end-1] |> norminf == 0

# we test the expression for σp
σp_fd = J_fold_fd[end,end]
σp_fd_ana = debugTmpForσ[end,end]
@test σp_fd ≈ σp_fd_ana rtol = 1e-5

# we test the expression for σx
σx_fd = J_fold_fd[end,1:end-1]
σx_ana = debugTmpForσ[end,1:end-1]
@test σx_fd ≈ σx_ana rtol = 1e-2

σx_fd - σx_ana |> norminf
debugTmpForσ[1:end-1,1:end-1] - J_fold_fd[1:end-1,1:end-1] |> norminf

# check our solution of the bordered problem
res_exp = debugTmpForσ \ rhs
@test norm(res_exp - Bd2Vec(res_explicit[1]), Inf64) < 1e-10
####################################################################################################
# Use of different eigensolvers
opt_newton = NewtonPar(tol = 1e-8, verbose = false, eigsolver = EigKrylovKit())
opts_br0 = ContinuationPar(dsmin = 0.01, dsmax = 0.15, ds= 0.01, pMax = 4.1, maxSteps = 250, newtonOptions = opt_newton, detectFold = true, detectBifurcation = 1, nev = 15)

br, = continuation(jet[1], jet[2], out, par_chan, (@lens _.α), opts_br0, recordFromSolution = (x,p)->norm(x,Inf64), plot = false, verbosity = 0)

opts_br0 = ContinuationPar(dsmin = 0.01, dsmax = 0.15, ds= 0.01, pMax = 4.1, maxSteps = 250, newtonOptions = NewtonPar(tol =1e-8), detectFold = true, detectBifurcation = 1, nev = 15)

br, = continuation(jet[1], jet[2], out, par_chan, (@lens _.α),opts_br0,plot = false, verbosity = 0)


#RSDeltaSigma: exampleHBF() function
#-------------------------------------------------------------------------------

function HBFex1()
	q1 = [
		 0 -4 -5;
		 1  1  1;
		 0 -4 -5;
		-1 -1 -1;
		-1 -3 -5;
		 1  1  1;
		-3 -5  0;
		-1 -1  0;
	]
	q2 = [
		-1 -3 -7;
		 1  1  1;
		-2 -5 -7;
		-1  1  1;
		-3 -8  0;
		 1 -1  0;
		-4 -6 -8;
		-1 -1 -1;
		-4 -8  0;
		 1 -1  0;
		-5 -7 -8;
		-1 -1 -1;
		-5  0  0;
		 1  0  0;
		-6 -7  0;
		-1 -1  0;
		-7 -8  0;
		-1 -1  0;
		-6  0  0;
		 1  0  0;
	]
	return (q1,q2)
end

"""`HBFex2()`

Coefficients from toolbox v2.0 documentation.
"""
function HBFex2()
	q1 = [
		 0 -4 -7;
		 1 -1  1;
		-1 -3 -6;
		-1 -1 -1;
		-2 -4 -7;
		 1 -1  1;
	]
	q2 = [
		-1 -3 -8;
		 1  1 -1;
		-2 -4 -9;
		-1  1 -1;
		-3 -5 -9;
		 1 -1  1;
		-4 -7 -8;
		-1  1  1;
		-5 -8 -11;
		 1 -1 -1;
		-6 -9 -11;
		-1  1 -1;
	]
	return (q1,q2)
end

"""`Saramaki88()`

Coefficients from "Efficient VLSI-Realizable Decimators for Sigma-Delta
Analog-to-Digital Converters," T. Saramaki and H. Tenhunen, ISCAS 1988,
pp 1525-2528.
"""
function Saramaki88()
	q1 = [
		 0 -2;
		 1 -1;
		-2 -9;
		-1  1;
	]
	q2 = [
		-1 -3  0;
		 1  1  0;
		-3 -5 -6;
		-1 -1 -1;
		-4 -6  0;
		 1  1  0;
	]
	return (q1,q2)
end

"""`Saramaki90()`

Coefficients from "Multiplier-Free Decimation Algorithms for Superresolution
Oversampled Converters," T. Saramaki, T. Karema, T. Ritoniemi, and H. Tenhunen
ISCAS 1990, vol. 4. pp 1525-2528.
"""
function Saramaki90()
	q1 = [
		 0 -4 -5;
		 1  1  1;
		 0 -4 -5;
		-1 -1 -1;
		-1 -3 -5;
		 1  1  1;
		-3 -5 -20;
		-1 -1  1;
	]
	q2 = [
		-1 -3 -7;
		 1  1  1;
		-2 -5 -7;
		-1  1  1;
		-3 -8  0;
		 1 -1  0;
		-4 -6 -8;
		-1 -1 -1;
		-4 -8  0;
		 1 -1  0;
		-5 -7 -8;
		-1 -1 -1;
		-5  0  0;
		 1  0  0;
		-6 -7  0;
		-1 -1  0;
		-6  0  0;
		 1  0  0;
		-7 -8  0;
		-1 -1  0;
		-6  0  0;
		 1  0  0;
	]
	return (q1,q2)
end

function exampleHBF(n::Int = 0)
	fnlist = [HBFex1, HBFex2, Saramaki88, Saramaki90]
	if !in(n, keys(fnlist))
		throw("exampleHBF: No code for example $n.")
	end
	fn = fnlist[n]
	(q1, q2) = fn()
	F1 = convertForm(bunquantize(q1),q1)
	F2 = convertForm(bunquantize(q2),q2)
	return (F1, F2)
end

#Last line


#=
ndgrid-range
ndgrid for a set of AbstractRange inputs.
=#

export ndgrid

"""
    GridAR{T,d,D} <: AbstractGrid{T,d,D}
The `d`th component of `D`-dimensional `ndgrid(x, y ...)`
where `1 ≤ d ≤ D` and `x, y, ...` are each an `AbstractRange`.
"""
struct GridAR{T,d,D} <: AbstractGrid{T,d,D}
    dims::Dims{D}
    first0::T # first - step
    step::T

    function GridAR(dims::Dims{D}, v::AbstractRange{T}, d::Int) where {D, T}
        1 ≤ d ≤ D || throw(ArgumentError("$d for $dims"))
        new{T,d,D}(dims, first(v) - step(v), step(v))
    end
end

Base.size(a::GridAR) = a.dims
Base.eltype(::GridAR{T}) where T = T

@inline Base.@propagate_inbounds function Base.getindex(
    a::GridAR{T,d,D},
    i::Vararg{Int,D},
) where {T,d,D}
     @boundscheck checkbounds(a, i...)
     return a.first0 + (@inbounds i[d]) * a.step
end

#=
"""
    (xg, yg, ...) = ndgrid(v1, v2,, ...)
Construct `ndgrid` tuple for `AbstractRange` inputs.

# Example

```jldoctest
julia> ndgrid(1:3, 5:2:7)
([1 1; 2 2; 3 3], [5 7; 5 7; 5 7])
```
"""
function ndgrid(vs::Vararg{AbstractRange})
    D = length(vs)
    dims = ntuple(i -> length(vs[i]), D)
    T = Tuple{ntuple(i -> GridAR{eltype(vs[i]), i, D}, D)...} # return type
    return ntuple(i -> GridAR(dims, vs[i], i), D)::T
end
=#


using SparseArrays
using MatrixNetworks

function hypergraph_largest_component(H::SparseMatrixCSC)
    m,n = size(H)
    A = [spzeros(m,m) H; sparse(H') spzeros(n,n)]
    lcc, pcc = largest_component(A)
    p_nodes = pcc[m+1:end]
    p_edges = pcc[1:m]
    H = H[p_edges,p_nodes]
    return H, p_edges, p_nodes
end

# These are some useful functions for converting back and forth between
# different ways to store hypergraphs.
function NeighborList(He2n::SparseArrays.SparseMatrixCSC{Float64,Int64},Hn2e::SparseArrays.SparseMatrixCSC{Float64,Int64})
    """
    NeighborList: given hypergraph (edge x node) indicence and (node x edge)
    incidence, return a list of neighbors for each node.
    """
    m,n = size(He2n)
    rp = He2n.rowval
    ci = He2n.colptr
    Neighbs = Vector{Vector{Int64}}()
    for i = 1:n
        # Get list of edges that node i is in
        Edges = rp[ci[i]:ci[i+1]-1]

        # build a set of neighboring nodes
        NeighbSet = Vector{Int64}()
        for j = 1:length(Edges)
            ej = Edges[j]               # For each node ej adjacent to I...
            nodes = getnodes(Hn2e,ej)   # ...get nodes in edge ej
            append!(NeighbSet,nodes)
        end
        push!(Neighbs, sort(setdiff(unique(NeighbSet),i)))
    end

    return Neighbs
end

function NeighborList(node2edge::Vector{Vector{Int64}},edge2node::Vector{Vector{Int64}})
    """
    NeighborList: given edge to node list and node to edge list, compute
    neighbors
    """
    n = length(node2edge)
    m = length(edge2node)
    Neighbs = Vector{Vector{Int64}}()
    for i = 1:n
        # Get list of edges that node i is in
        Edges = node2edge[i]

        # build a set of neighboring nodes
        NeighbSet = Vector{Int64}()
        for j = 1:length(Edges)
            ej = Edges[j]               # For each node ej adjacent to I...
            nodes = edge2node[ej]       # ...get nodes in edge ej
            append!(NeighbSet,nodes)
        end
        push!(Neighbs, sort(setdiff(unique(NeighbSet),i)))
    end

    return Neighbs
end

function getedges(He2n::SparseArrays.SparseMatrixCSC{Float64,Int64},I::Int64)
    """
    Given the edge-by-node incidence matrix He2n, return the set of hyperedges
    that node I is in.
    """
    first = He2n.colptr[I]
    last = He2n.colptr[I+1]-1
    edges = He2n.rowval[first:last]

    return edges
end

function getnodes(Hn2e::SparseArrays.SparseMatrixCSC{Float64,Int64},J::Int64)
    """
    Given the node-by-edge incidence matrix Hn2e, return the set of nodes
    contained in hyperedge J
    """
    first = Hn2e.colptr[J]
    last = Hn2e.colptr[J+1]-1
    nodes = Hn2e.rowval[first:last]
    mult = Hn2e.nzval[first:last]
    edge = Vector{Int64}()
    # need to adjust for multiplicities
    for t = 1:length(nodes)
        node = nodes[t]
        for k = 1:mult[t]
            push!(edge,node)
        end
    end

    return edge
end

function incidence2elist(H::SparseArrays.SparseMatrixCSC{Float64,Int64},nodelist::Bool=false)
    """
    This converts a hypergraph in incidence matrix form to hyperedge list form.
    Incidence matrix form:  H[e,u] = 1  iff node u is in hyperedge e
    Edgelist form: Hyperedges[j] = array of nodes in hyperedge j

    nodelist == true: means you actually want to get a map from node IDs to hyperedge ids
    """
    if ~nodelist
        # unless you want the node2edge map, transpose first
        H = SparseArrays.sparse(H')
    end
    rp = H.rowval
    ci = H.colptr
    nz = H.nzval
    Hyperedges = Vector{Vector{Int64}}()
    n,m = size(H)

    for i = 1:m
        startedge = ci[i]
        endedge = ci[i+1]-1
        nodes = rp[startedge:endedge]
        mult = nz[startedge:endedge]
        edge = Vector{Int64}()
        # need to adjust for multiplicities
        for t = 1:length(nodes)
            node = nodes[t]
            for k = 1:mult[t]
                push!(edge,node)
            end
        end

        push!(Hyperedges,edge)
    end
    return Hyperedges
end

function elist2incidence(Hyperedges::Vector{Vector{Int64}}, N::Int64)
    """
    Converts a hyperedge list into a binary incidence matrix for the hypergraph.

    This is the exact inverse of incidence2elist
    """
    U = Vector{Int64}()
    E = Vector{Int64}()
    M = length(Hyperedges)
    for enum = 1:length(Hyperedges)
        e = Hyperedges[enum]
        for node in e
            push!(U,node)
            push!(E,enum)
        end
    end

    H = SparseArrays.sparse(E,U,ones(length(U)),M,N)
    return H
end

function RandomSplittingFunctions(Edges,probs)
    ## Generate random splitting functions
    EdgesW = Vector{Vector{Float64}}()
    for edge in Edges
        k = length(edge)
        w = random_scb_function(k,probs[1],probs[2])
        while ~check_cb_submodular(w,0)
            w = random_scb_function(k,probs[1],probs[2])
        end
        push!(EdgesW,w)
    end
    return EdgesW
end

function CliqueExpansion(Edges::Vector{Vector{Int64}},n::Int64,weighted::Bool=true,binary::Bool=false)
    """
    Weighted clique expansion where a hyperedge e is expanded to a
    weighted clique with each edge having weight 1/(|e| - 1)
    """
    I = Int64[]
    J = Int64[]
    V = Float64[]
    for edge in Edges
        k = length(edge)
        for i = 1:k-1
            ei = edge[i]
            for j = i+1:k
                ej = edge[j]
                push!(I,ei)
                push!(J,ej)
                if weighted
                    push!(V, 1/k)
                else
                    push!(V, 1)
                end
            end
        end
    end
    A = SparseArrays.sparse(I,J,V,n,n)
    for i = 1:n
        A[i, i] = 0.0
    end
    SparseArrays.dropzeros!(A)
    A = SparseArrays.sparse(A+A')
    if binary
        I, J, V = SparseArrays.findnz(A)
        A = SparseArrays.sparse(I, J, 1, n, n)
    end
    return A
end


"""
## CircPlot

Plot circular data in a number of ways:

- `cplot_histogram`: Plot a polar histogram of angular data.
- `cplot_pts`: Plot angular data as points on the edge of a circle.
- `chistogram`: Get a set of `Plots.Shape`s which can be plotted as a polar
  histogram at any point on a map, for example.

By default, all input to routines is in radians, but all routines have a positional
argument to set `degrees` to `true` when not using radians.
"""
module CircPlot

import StatsBase

using Plots

export
    chistogram,
    cplot_histogram,
    cplot_pts

#=
Exported routines
=#
"""
    cplot_pts(a::Array, degrees=false; azimuth=false, hold=false, args...) -> ::Plots.Plot

Use Plots to plot a set of points on the circumference of the unit circle.

Use `azimuth=true` to plot angles clockwise from north.

Supply any additional arguments to `Plots.plot()` as keyword arguments; for example,
to set the size of points, do

    julia> cplot_pts(angles; markersize=20)
"""
function cplot_pts(a, degrees::Bool=false;
                   azimuth::Bool=false, axial::Bool=false, line::Bool=true,
                   padding=0.05, R=1, centre=(0,0), kwargs...)
    circum = 0:0.01:2π
    markerstyle = (0.1,:blue)
    degrees && (a = deg2rad.(a))
    azimuth && (a = π/2 .- a)
    R += padding
    x0, y0 = centre
    p = plot(aspect_ratio=:equal,
        xlim=(-R+x0,R+x0), ylim=(-R+y0,R+y0), legend=false)
    line && plot!(p, x0.+cos.(circum), y0.+sin.(circum), l=:black)
    scatter!(p, x0.+cos.(a), y0.+sin.(a), m=markerstyle; kwargs...)
    axial && scatter!(p, x0.+cos.(a+π), y0.+sin.(a+π), m=markerstyle; kwargs...)
    if azimuth
        xlabel!(p, "East")
        ylabel!(p, "North")
    else
        xlabel!(p, "x")
        ylabel!(p, "y")
    end
    p
end

"""
    cplot_histogram(a, binwidth, degrees=false; azimuth=false, axial=false, weights=ones(length(a)), circ=nothing; kwargs...) -> ::Plots.Plot

Plot a polar histogram of a set of angles `a` using bins `binwidth` wide.
`binwidth` should be given in the same convention (radians or °) as the data.
If π (or 180°) is not an integer multiple of `binwidth`, then the nearest value
for which this is true is selected.

Use `azimuth=true` to plot angles clockwise from north.

Use `axial=true` if data are π-ambiguous (i.e., are orientational rather than
directional data).

`weights` must have the same length as `a`, and if passed in will scale each
point `i` by `weights[i]`.

Set options for the surrounding circle using `circ=Dict(:opt=>val)`, where
the keys and values in the Dict will be passed to Plots for plotting the circle
shape around the plot.

Additional keyword arguments `kwargs...` are passed to `Plots.plot()` when
plotting the bars.
"""
function cplot_histogram(a, binwidth, degrees::Bool=false;
                         azimuth::Bool=false, axial::Bool=false,
                         weights=ones(eltype(a), length(a)), circ=nothing,
                         kwargs...)
    hist = fit_hist(a, binwidth, degrees, axial, weights=weights)
    n = length(hist.weights)
    edges = hist.edges[1]
    azimuth && (edges = π/2 .- collect(edges))
    R = maximum(hist.weights)
    p = plot(aspect_ratio=:equal, xlim=(-R,R), ylim=(-R,R), legend=false)
    circ != nothing && plot!(p, circle(R); circ...)
    plot!(p, sector.(edges[1:n], edges[2:n+1], hist.weights), fill=:black; kwargs...)
    p
end

"""
    chistogram(a, binwidth, degrees::Bool=true, x=0, y=0; axial=false, azimuth=false, scale=1, maxr=nothing) -> bars::Array{Shape}

Return a set of Shapes which show a polar histogram created from the set of
angles `a` with bins `binwidth` wide, centred at `x` and `y`.  These shapes can
be plotted with `plot(bars)`.

Set `axial` to `true` for data with π-ambiguity and `azimuth` to `true` if data are
clockwise from north.

`scale` determine what length to scale the radii by, or alternatively set `maxr` to
the radius (in plot units) of the maximum bin.  Set `minweight` to be the minimum
bar count required for the plot to fill the entire maximum radius from `maxr`.
"""
function chistogram(a, binwidth, degrees=false, x=0, y=0;
        weights=ones(a), axial=false, azimuth=false, scale=1, maxr=nothing,
        minweight=1, minobs=nothing)
    if minobs != nothing
        warning("`minobs` is deprecated: use `minweight` instead")
        minweight = minobs
    end
    hist = fit_hist(a, binwidth, degrees, axial, weights=weights)
    n = length(hist.weights)
    edges = hist.edges[1]
    R = maximum(hist.weights)
    maxr != nothing && (scale = maxr/max(minweight, R))
    azimuth && (edges = π/2 .- collect(edges))
    sector.(edges[1:n], edges[2:n+1], hist.weights*scale, x, y)
end

#=
Unexported routines
=#
"""    fit_hist(a, binwidth, degrees=false, axial=false; weights=ones(a))

Return a `StatsBase.Histogram` with the results of binning the angles `a`
into bins `binwidth` wide.  If `axial` is `true`, then directions have
π symmetry.  Optionally sypply a `weights` array containing the weighting
for each value in `a`."""
function fit_hist(a, binwidth, degrees::Bool, axial::Bool;
                  weights=ones(eltype(a), length(a)))
    binwidth isa Bool &&
        throw(ArgumentError("`binwidth` supplied as a `Bool`, but should be a real number"))
    length(weights) == length(a) ||
        throw(ArgumentError("`weights` must have same length as number of angles"))
    n = (degrees ? round(Int, 360/binwidth) : round(Int, 2pi/binwidth))
    n = max(1, n)
    binwidth = 2pi/n
    bins = range(0, stop=2pi, length=n+1)
    data = degrees ? deg2rad.(mod.(a, 360)) : mod.(a, 2pi)
    # FIXME: Many angles are given as whole numbers: this cludge helps
    # an issue where floating point values sometimes are binned one way,
    # sometimes another, giving polar histograms which depend on the range
    # in which angles are given.
    data .+= 10*eps(float(eltype(a)))
    axial && (data .= mod.(data, pi))
    h = StatsBase.fit(StatsBase.Histogram, data,
                      StatsBase.FrequencyWeights(weights), bins, closed=:left)
    axial && (h.weights[end÷2+1:end] .= h.weights[1:end÷2])
    h
end

"""Return a sensible number of segments between two angles in radians θ₁ and θ₂."""
number_of_segments(θ₁, θ₂) = max(2, round(Int, abs(θ₂ - θ₁)*90/π))

"""Return a circle shape with radius `r` centred at `x` and `y`."""
circle(r, x=0, y=0) = Shape(arc(0, 2π, r, x, y))

"""Return an arc (part of a circle's circumference) between `θ₁` and `θ₂` radians
anticlockwise from the +x axis, with radius `r`, centred at `x` and `y`."""
arc(θ₁, θ₂, r=1, x=0, y=0) = Tuple{Float64,Float64}[(x+r*cos(θ), y+r*sin(θ))
    for θ in range(θ₁, stop=θ₂, length=number_of_segments(θ₁, θ₂))]

"""Return a sector (pie wedge) between angles `θ₁` and `θ₂` radians anticlockwise
from the x-axis with radius `r`, centred at `x` and `y`."""
sector(θ₁, θ₂, r=1, x=0, y=0)::Shape =
    Shape(vcat(arc(θ₁, θ₂, r, x, y), Tuple{Float64,Float64}((x, y))))

"""Return a filled sector (arcshape) between two angles θ₁ and θ₂ radians
anticlockwise from the x-axis, and two radii, `r₁` and `r₂`, centred at `x` and `y`."""
arcshape(θ₁, θ₂, r₁, r₂, x=0, y=0) =
    Shape(vcat(arc(θ₁, θ₂, r₁, x, y), arc(θ₂, θ₁, r₂, x, y)))

end # module


using StatsModelComparisons, StanSample
using Test

println()
include("model_comparison_tests.jl")
include("test_demo_wells.jl")
include("test_cars.jl")


function weights!(w::Vector{Float64},c::Vector{Float64},ϵ::Float64,IS::Vector{Float64},
    s::Vector{Float64})
    for i = 1:length(w)
        w[i] = c[i]/((ϵ + IS[i])^2);
    end
    s[1] = sum(w);
    for i = 1:length(w)
        w[i] /= s[1];
    end
end


put_real_eigenvalue_first(λ, T) = put_real_eigenvalue_first(λ, T, 0)
function put_real_eigenvalue_first(λ::Vector{Complex{BigFloat}}, T::Matrix{Complex{BigFloat}}, n_attempt::Int64)
    i_real = findfirst(λ .== real.(λ))
    if i_real == nothing
        return λ, T
    elseif i_real == 1
        return λ, T
    else
        n = length(λ)
        perm_mat = diagm(1=>ones(n-1), (1-n)=>ones(1))
        λ_mat = diagm(0=>λ)
        λ_perm_order = Int64.(collect(1:n)' * perm_mat)[:]
        λ_perm = λ[λ_perm_order]
        T_perm = T * perm_mat
        (n < n_attempt) && error("too many attempts")
        return put_real_eigenvalue_first(λ_perm, T_perm, n_attempt + 1)
    end
end

function big_eigen(M::Matrix{BigFloat})
    n, n2 = size(M)
    (n == n2) || error("matrix must be square")
    M_ = deepcopy(M)
    λ = eigvals(M_)
    T = zeros(Complex{BigFloat}, n, n)
    for k = 1:n
        sss = svd(M - I*λ[k])  # I'm aware this is massivly ineffecient
        T[:, k] = sss.Vt[end, :]'   # this does conjugate transpose
    end
    λ, T = put_real_eigenvalue_first(λ, T)
    return λ, T
end


# Unit test for the Nonlinear Arnoldi method (in src/method_nlar.jl)
# The "Gun" problem form gun_native.jl

using NonlinearEigenproblemsTest
using NonlinearEigenproblems
using Test
using LinearAlgebra
using Random

@bench @testset "Nonlinear Arnoldi" begin
	TOL = 1e-10;

    @info "Testing gun problem"
    Random.seed!(0)

    nep = nep_gallery("nlevp_native_gun")
    n = size(nep,1);

    #The eigenvalues of the gun problem lie within a semi-circle centred approximately around the "shift" with an approximate radius "scale"
    shift = 250^2;
    scale = 330^2-220^2;

    #Run Quasi-Newton for initial guess obtained from the knowledge of the eigenvalue distribution
    λ_ref,v_ref = quasinewton(nep, λ = shift+scale*(-0.131403+0.00759532im), v = ones(n), logger = displaylevel, tol = TOL/50, maxit = 500)
    @info "Eigenvalue computed by quasi newton: $λ_ref"

    #Shift and scale the NEP(mainly to avoid round-off errors because of the large entries in the coefficient matrices)
    nep1_spmf = SPMF_NEP(get_Av(nep),get_fv(nep))
    nep1 = shift_and_scale(nep1_spmf,shift=shift,scale=scale);

    #Run NLAR on the shifted and scaled NEP (neigs set to 1 to save time. Method works for the case of finding multiple eigenvalues as well)
    λ,u = nlar(nep1, tol=TOL, λ=0, maxit=100, neigs=2, R=0.01, logger=displaylevel, v=ones(n),
        eigval_sorter=residual_eigval_sorter, inner_solver_method=IARInnerSolver(),
        qrfact_orth=false, num_restart_ritz_vecs=8, max_subspace=150)

    λ_shifted = λ[1];v=u[:,1];

    #Rescale and shift back to the eigenvalue of the original problem
    λ_orig = shift+scale*λ_shifted;

    #Test residual and distance from reference eigenvalue
    @test norm(compute_Mlincomb(nep, λ_orig, v)) < sqrt(TOL)*50
    @test norm(compute_Mder(nep, λ_orig) * v) < sqrt(TOL)*50
    @test norm(λ_ref - λ_orig) < sqrt(TOL)*500
    ###########################################################################################################3

    #Testing PEP
    @info "Testing a PEP"
    pepnep = nep_gallery("pep0",60)
    Dc,Vc = polyeig(pepnep,DefaultEigSolver)
    c = sortperm(abs.(Dc))
    @info " 6 smallest eigenvalues found by polyeig() according to the absolute values: $(Dc[c[1:6]])"

    λ,u = nlar(pepnep, tol=TOL, maxit=60, neigs=3, R=0.001, logger=displaylevel,
        λ=Dc[4]+1e-2, v=ones(size(pepnep,1)), eigval_sorter=residual_eigval_sorter,
        inner_solver_method=IARInnerSolver(), qrfact_orth=false, errmeasure=ResidualErrmeasure(pepnep))

    verify_lambdas(3, pepnep, λ, u, TOL)

    @info "Testing errors thrown"
    nep = nep_gallery("pep0",4)
    @test_throws NEPCore.NoConvergenceException λ,u= nlar(nep, tol=1e-20, maxit=2, neigs=3)

end


cep_prime(n::Order,q::Number,z::Number,N::Integer=matsize3(n,q)) = cep_prime(n,q,z:z,N)[1]
function cep_prime(n::Order,q::Number,z::AbstractVector,N::Integer=matsize3(n,q))
    An = cep_coeffs(n,q,N) # Calculate coefficients
    # Evaluate sine for all z values
    C = [-2k*sin(2k*z) for z in z, k in 0:N-1]
    # Calculate function at z, for all n
    return C * An
end

cea_prime(n::Order,q::Number,z::Number,N::Integer=matsize3(n,q)) = cea_prime(n,q,z:z,N)[1]
function cea_prime(n::Order,q::Number,z::AbstractVector,N::Integer=matsize3(n,q))
    An = cea_coeffs(n,q,N) # Calculate coefficients
    # Evaluate sine for all z values
    C = [-(2k+1)*sin((2k+1)*z) for z in z, k in 0:N-1]
    # Calculate function at z, for all n
    return C * An
end

sea_prime(n::Order,q::Number,z::Number,N::Integer=matsize3(n,q)) = sea_prime(n,q,z:z,N)[1]
function sea_prime(n::Order,q::Number,z::AbstractVector,N::Integer=matsize3(n,q))
    Bn = sea_coeffs(n,q,N) # Calculate coefficients
    # Evaluate cosine for all z values
    S = [(2k+1)*cos((2k+1)*z) for z in z, k in 0:N-1]
    # Calculate function at z, for all n
    return S * Bn
end

sep_prime(n::Order,q::Number,z::Number,N::Integer=matsize3(n,q)) = sep_prime(n,q,z:z,N)[1]
function sep_prime(n::Order,q::Number,z::AbstractVector,N::Integer=matsize3(n,q))
    Bn = sep_coeffs(n,q,N) # Calculate coefficients
    # Evaluate cosine for all z values
    S = [(2k+2)*cos((2k+2)*z) for z in z, k in 0:N-1]
    # Calculate function at z, for all n
    return S * Bn
end


using HOODESolver
using LinearAlgebra
using Test

function testcoeffloat()
    tab=[ (big"0.1", big"0.0001"), (big"0.00001", big"0.01")]
    @testset "test coefexp float vs bigfloat" begin
        for (bigeps, bigdt) in tab
            epsilon = convert(Float64,bigeps)
            dt = convert(Float64, bigeps)
            bigceab = CoefExpAB(10,bigeps,32,bigdt)
            ceab = CoefExpAB(10,epsilon,32,bigdt)
            @test isapprox(bigceab.tab_coef, ceab.tab_coef, atol=1e-15, rtol=1e-14)
    #       println("epsilon=$epsilon norm=$(norm(bigceab.tab_coef-ceab.tab_coef, Inf))")
        end
    end
 end
function testcomp()
    @testset "test solve float vs bigfloat" begin

        bigeps = big"0.0001"
        epsilon = convert(Float64,bigeps)
        bigu0=rand(BigFloat,4)
        u0 = convert.(Float64, bigu0)
        A =  [0 0 1 0; 0 0 0 0;-1 0 0 0; 0 0 0 0]
        fct = (u,p,t) -> [0, u[4], -2u[1]*u[2], -u[2]-u[1]^2+u[2]^2] # Henon-Heiles
        bigprob = HOODEProblem(fct, bigu0, (big"0.0",big"3.0"), missing, A, bigeps)
        prob = HOODEProblem(fct, u0, (0.0, 3.0), missing, A, epsilon)
        bigsol = solve(bigprob)
        sol = solve(prob)

        for i=1:101
            @test isapprox(
    bigsol.interp.u_caret[i],
    sol.interp.u_caret[i],
    atol=1e-13,
    rtol=1e-12
)
        end
    end
end

testcoeffloat()
testcomp()


"""
    function latticeCompactHexagonal(box,r)

Make a lattice with compact hexagonal structure in box with spheres of radius r.

# Returns
[posX,posY,posZ] Where pos are arrays with the positions of the cells for each component.
"""
function compactHexagonal(box::Array{<:Real,2},r::Number)

    d = size(box)
    if d != (3,2)
        error("Incompatible number of dimensions ", d)
    end

    #Make first dimension
    lineX = Array(box[1,1]:2*r:box[1,2])

    #Make second dimension
    lineY = fill(box[2,1]+r,length(lineX))
    nY = ceil(INT,(box[2,2]-box[2,1])/(4*sin(pi/3)*r))
    areaX = zeros(FLOAT,1)
    areaY = zeros(FLOAT,1)
    dx = 2*r*cos(pi/3)
    dy = 2*r*sin(pi/3)
    for i in 0:nY-1
        append!(areaX,lineX)
        append!(areaX,lineX.+dx)

        append!(areaY,lineY.+dy*2*i)
        append!(areaY,lineY.+dy*2*i.+dy)
    end

    #Make third dimension
    areaZ = fill(box[2,1]+r,length(areaX))
    nZ = ceil(INT,(box[3,2]-box[3,1])/(4*sin(pi/3)*r))
    volumeX = zeros(FLOAT,1)
    volumeY = zeros(FLOAT,1)
    volumeZ = zeros(FLOAT,1)
    dy = 2*r/2/cos(pi/6)
    dz = 2*sqrt(2/3)*r
    for i in 0:nZ-1
        append!(volumeX,areaX)
        append!(volumeX,areaX)

        append!(volumeY,areaY)
        append!(volumeY,areaY.+dy)

        append!(volumeZ,areaZ.+dz*2*i)
        append!(volumeZ,areaZ.+dz*2*i.+dz)
    end

    #Return
    return [volumeX volumeY volumeZ]

end


# This file is a part of AstroLib.jl. License is MIT "Expat".
# Copyright (C) 2016 Mosè Giordano.

function _kepler_solver{T<:AbstractFloat}(M::T, e::T)
    if 0 <= e <= 1
        # M is assumed to be in the range [-pi, pi], see Markley (1995), page 2.
        # First restrict it to [0, 2pi], then move values above pi to [-pi, 0].
        M = cirrange(M, 2.0*pi)
        M > pi && (M -= 2.0*pi)
        if M == 0
            return 0.0
        else
            # equation (20)
            α = (3.0*pi*pi + 1.6*pi*(pi - abs(M))/(1.0 + e))/(pi*pi - 6.0)
            # equation (5)
            d = 3.0*(1.0 - e) + α*e
            # equation (9)
            q = 2.0*α*d*(1.0 - e) - M*M
            # equation (10)
            r = 3.0*α*d*(d - 1.0 + e)*M + M*M*M
            # equation (14)
            w = (abs(r) + sqrt(q*q*q + r*r))^(2.0/3.0)
            # equation (15)
            E1 = (2.0*r*w/(w*w + w*q + q*q) + M)/d
            # equation (21)
            f0 = E1 - e*sin(E1) - M
            # equation (25)
            f1 = 1.0 - e*cos(E1)
            # equation (26)
            f2 = e*sin(E1)
            # equation (27)
            f3 = 1.0 - f1
            # equation (28)
            f4 = -f2
            # equation (22)
            δ3 = -f0/(f1 - f0*f2/(2.0*f1))
            # equation (23)
            δ4 = -f0/(f1 + δ3*f2/2.0 + δ3*δ3*f3/6.0)
            # equation (24)
            δ5 = -f0/(f1 + δ4*f2/2.0 + δ4*δ4*f3/6.0 + δ4*δ4*δ4*f4/24.0)
            return E1 + δ5 # equation 29
        end
    else
        error("eccentricity must be in the range [0, 1]")
    end
end

"""
    kepler_solver(M, e) -> E

### Purpose ###

Solve Kepler's equation in the elliptic motion regime (\$0 \\leq e \\leq 1\$)
and return eccentric anomaly \$E\$.

### Explanation ###

In order to find the position of a body in elliptic motion (e.g., in the
two-body problem) at a given time \$t\$, one has to solve the [Kepler's
equation](https://en.wikipedia.org/wiki/Kepler%27s_equation)

\$\$ M(t) = E(t) - e\\sin E(t)\$\$

where \$M(t) = (t - t_{0})/P\$ is the mean anomaly, \$E(t)\$ the eccentric
anomaly, \$e\$ the eccentricity of the orbit, \$t_0\$ is the time of periapsis
passage, and \$P\$ is the period of the orbit.  Usually the eccentricity is
given and one wants to find the eccentric anomaly \$E(t)\$ at a specific time
\$t\$, so that also the mean anomaly \$M(t)\$ is known.

### Arguments ###

* `M`: mean anomaly.  This can be either a scalar or an array
* `e`: eccentricity, in the elliptic motion regime (\$0 \\leq e \\leq 1\$)

### Output ###

The eccentric anomaly \$E\$, restricted to the range \$[-\\pi, \\pi]\$.  If an
array of mean anomalies is provided in input, an array of the same length as `M`
is returned.

### Method ###

Many different numerical methods exist to solve Kepler's equation.  This
function implements the algorithm proposed in Markley (1995) Celestial Mechanics
and Dynamical Astronomy, 63, 101
(DOI:[10.1007/BF00691917](http://dx.doi.org/10.1007/BF00691917)).  This method
is not iterative, requires only four transcendental function evaluations, and
has been proved to be fast and efficient over the entire range of elliptic
motion \$0 \\leq e \\leq 1\$.

### Example ###

(1) Find the eccentric anomaly for an orbit with eccentricity \$e = 0.7\$ and
for \$M(t) = 8\\pi/3\$.

``` julia
ecc = 0.7;
E = kepler_solver(8pi/3, ecc)
# => 2.5085279492864223
```

(2) Plot the eccentric anomaly as a function of mean anomaly for eccentricity
\$e = 0\$, \$0.5\$, \$0.9\$.  Recall that `kepler_solver` gives \$E \\in [-\\pi,
\\pi]\$, use `cirrange` to have it in \$[0, 2\\pi]\$.  Use
[PyPlot.jl](https://github.com/stevengj/PyPlot.jl) for plotting.

``` julia
using PyPlot
M=linspace(0, 2pi, 1001)[1:end-1];
for ecc in (0, 0.5, 0.9); plot(M, cirrange(kepler_solver(M, ecc), 2pi)); end
```

### Notes ###

The true anomaly can be calculated with `trueanom` function.
"""
kepler_solver(M::Real, e::Real) =
    _kepler_solver(promote(float(M), float(e))...)

function kepler_solver{R1<:Real,R2<:Real}(M::AbstractArray{R1}, e::R2)
    typee = promote_type(float(R1), float(R2))
    E = similar(M, typee)
    for i in eachindex(M)
        E[i] = kepler_solver(M[i], e)
    end
    return E
end


using JuMP, Gurobi

mutable struct Node
    Level::Int
    Model::JuMP.Model
    Zbound::Float64
    Xrelax::Array{Float64}
    Status::Symbol
end

mutable struct Best
    Zstar::Float64
    Zbound::Float64
    Xstar::Array{Float64}
    Intsols::Int
    Visited::Int
end

function IndeciseVariable(node::Node)
    sizeX = length(node.Xrelax)
    nonInteger = Array{Float64}(sizeX)
    for i=1:sizeX
        if node.Model.colCat[i] == :Bin
            nonInteger[i] = abs(node.Xrelax[i] - 0.5)
        else
            nonInteger[i] = 2 #Variable is not binary
        end
    end
    mostInDoubt = indmin(nonInteger) #get index of the most fractionary variable
    if nonInteger[mostInDoubt] == 0.5  #all variables are 0 or 1
        return false
    end
    return mostInDoubt
end

function Bound(node::Node, best::Best, sense::Symbol)
    errortolerance = 1e-4 #if 0.01% away from optimal -> current answer is optimal
    if node.Status != :Optimal #Bound by Infeasabilty
        return false
    end
    if sense == :Max #Bound by Limit
        if node.Zbound < best.Zbound
            return false
        end
    else
        if node.Zbound > best.Zbound
            return false
        end
    end
    if abs(node.Zbound - best.Zstar) <= errortolerance #Bound by Optimality
        return false
    end
end

function Branch(node::Node)
    indeciseVariable = IndeciseVariable(node)
    if indeciseVariable == false #if all binary variables are integer -> no branching
        return false, false
    end

    leftChild = deepcopy(node.Model)
    # Forces the most indecise variable to be 0
    leftChild.colUpper[indeciseVariable] = 0
    leftChild.colLower[indeciseVariable] = 0

    rightChild = deepcopy(node.Model)
    # Forces the most indecise variable to be 0
    rightChild.colUpper[indeciseVariable] = 1
    rightChild.colLower[indeciseVariable] = 1

    return leftChild, rightChild
end

function InitializeHeadNode(model::JuMP.Model, sense::Symbol)
            #Level,    Model,   Zbound,             Xrelax,                 Status
    node = Node(0, deepcopy(model), Inf, Array{Float64}(length(model.colUpper)), :None)
    if sense == :Min
        node.Zbound = -Inf
    end
    return node
end

function InitializeBest(model::JuMP.Model)
    sense = getobjectivesense(model)
               #Zstar, Zbound        Xstar,                    Intsols    Visited
    best = Best( -Inf, Inf, Array{Float64}(length(model.colUpper)), 0 ,0) #initialize Best as a Max problem
    if sense == :Min
        best.Zstar = Inf
        best.Zbound = -Inf
    end
    return best
end

function SolveRelax(model::JuMP.Model, solver::MathProgBase.AbstractMathProgSolver)
    setsolver(model, solver)
    status = solve(model, relaxation = true) #solve model relaxation
    return status, getobjectivevalue(model), model.colVal
end

function UpdateBest(best::Best, node::Node, sense::Symbol, binaryVariables::Array{Int})
    if sense == :Max
        if (node.Zbound > best.Zstar && all(isinteger, node.Xrelax[binaryVariables])) #if the bound is better and if all binary variables are integer
            best.Zstar = node.Zbound
            best.Zbound = node.Zbound
            best.Xstar = node.Xrelax
            best.Intsols = best.Intsols + 1
        end
        if (node.Zbound < best.Zbound && node.Zbound > best.Zstar)
            best.Zbound = node.Zbound
        end
    else
        if (node.Zbound < best.Zstar && all(isinteger, node.Xrelax[binaryVariables]))
            best.Zstar = node.Zbound
            best.Zbound = node.Zbound
            best.Xstar = node.Xrelax
            best.Intsols = best.Intsols + 1
        end
        if (node.Zbound > best.Zbound && node.Zbound < best.Zstar)
            best.Zbound = node.Zbound
        end
    end
end

function BinaryVariables(model::JuMP.Model)
    binaryVariableIndexes = Array{Int}(0)
    for i=1:length(model.colCat)
        if model.colCat[i] == :Bin
            push!(binaryVariableIndexes,i) #gets index of all binary variables
        end
    end
    return binaryVariableIndexes
end

function SolveMIP(model::JuMP.Model)
    tic()
    t1 = time_ns()
    t2 = 0.01
    iter = 0
    maxiter = 1000 #max iterations
    tmax = 180 #max time in nanoseconds
    level = 0
    solver = GurobiSolver(OutputFlag=0)
    sense = getobjectivesense(model)
    best = InitializeBest(model)
    nodes = Array{Node}(1)
    binaryVariableIndexes = BinaryVariables(model) #gets index of all binary variables
    nodes[1] = InitializeHeadNode(model, sense) #root node
    nodes[1].Status, nodes[1].Zbound, nodes[1].Xrelax = SolveRelax(nodes[1].Model,solver)
    model.ext[:status] = nodes[1].Status
    UpdateBest(best, nodes[1], sense, binaryVariableIndexes) #updates best answer
    best.Visited = best.Visited + 1

    while (!isempty(nodes) && iter <= maxiter && t2 <= tmax)
        level = level + 1
        leftChild, rightChild = Branch(nodes[end])
        pop!(nodes)

        if leftChild != false
            statusLeftChild, boundLeftChild, xRelaxLeftChild = SolveRelax(leftChild, solver)
            best.Visited = best.Visited + 1
            nodeLeftChild = Node(level, leftChild, boundLeftChild, xRelaxLeftChild, statusLeftChild)
            UpdateBest(best, nodeLeftChild, sense, binaryVariableIndexes)
            if Bound(nodeLeftChild, best, sense) != false
                push!(nodes, nodeLeftChild)
            end
        end

        if rightChild != false
            statusRightChild, boundRightChild, xRelaxRightChild = SolveRelax(rightChild, solver)
            best.Visited = best.Visited + 1
            nodeRightChild = Node(level, rightChild, boundRightChild, xRelaxRightChild, statusRightChild)
            UpdateBest(best,nodeRightChild, sense, binaryVariableIndexes)
            if Bound(nodeRightChild, best, sense) != false
                push!(nodes, nodeRightChild)
            end
        end
        iter = iter + 1
        if iter%50 == 0
            t2 = (time_ns() - t1)/1e9
        end
    end

    if ((iter >= maxiter || t2 >= tmax) && all(isinteger, best.Xstar[binaryVariableIndexes])) #have found not optimal integer solutions
        model.ext[:status] = :SubOptimal
    elseif ((iter >= maxiter || t2 >= tmax) && ~(all(isinteger, best.Xstar[binaryVariableIndexes]))) #have not found any solutions
        model.ext[:status] = :NotSolved
    end

    model.ext[:nodes] = best.Visited
    model.ext[:intsols] = best.Intsols
    model.colVal = best.Xstar
    model.objVal = best.Zstar
    model.objBound = best.Zbound
    time = toc()
    model.ext[:time] = time
end


module TestDiffMap
	using ManifoldLearning
	using Base.Test

	d = 2
	t = 1
	X, L = swiss_roll()
	Y = transform(DiffMap, X; d=d, t=t)

	@test outdim(Y) == d
	@test size(projection(Y), 2) == size(X, 2)
end


@testset "integrate_gauss_legendre" begin
    # TODO: Test on randomly generated polynomials
    # TODO: Test with alternating endpoints
    # TODO: Test with wide endpoints
    # TODO: Test with non-finite endpoints

    @test Arblib.overlaps(
        ArbExtras.integrate_gauss_legendre(sin, Arb(0), Arb(1)),
        1 - cos(Arb(1)),
    )

    @test Arblib.overlaps(
        ArbExtras.integrate_gauss_legendre(x -> x^3, Arb(0), Arb(1)),
        Arb(1 // 4),
    )

    @test Arblib.overlaps(
        ArbExtras.integrate_gauss_legendre(x -> 1 + x + x^2 + x^3, Arb(1), Arb(3)),
        Arb(104 // 3),
    )

    @test Arblib.overlaps(
        ArbExtras.integrate_gauss_legendre(sin ∘ exp, Arb(0), Arb(4)),
        real(Arblib.integrate(sin ∘ exp, 0, 4)),
    )

    @test isnan(ArbExtras.integrate_gauss_legendre(inv, Arb(-1), Arb(1)))
end

@testset "integrate" begin
    # TODO: Add more tests

    @test Arblib.overlaps(ArbExtras.integrate(sin, Arb(0), Arb(1)), 1 - cos(Arb(1)))

    @test Arblib.overlaps(ArbExtras.integrate(x -> x^3, Arb(0), Arb(1)), Arb(1 // 4))

    @test Arblib.overlaps(
        ArbExtras.integrate(sin ∘ exp, Arb(0), Arb(4)),
        real(Arblib.integrate(sin ∘ exp, 0, 4)),
    )

    @test isnan(ArbExtras.integrate(inv, Arb(-1), Arb(1)))

    # depth_start > 0
    @test Arblib.overlaps(
        ArbExtras.integrate(sin, Arb(0), Arb(1), depth_start = 4),
        1 - cos(Arb(1)),
    )

    @test Arblib.overlaps(
        ArbExtras.integrate(x -> x^3, Arb(0), Arb(1), depth_start = 5),
        Arb(1 // 4),
    )
end


"""
ChebyshevWeight{kind,T}()

is a quasi-vector representing the Chebyshev weight of the specified kind on -1..1.
That is, `ChebyshevWeight{1}()` represents `1/sqrt(1-x^2)`, and
`ChebyshevWeight{2}()` represents `sqrt(1-x^2)`.
"""
struct ChebyshevWeight{kind,T} <: AbstractJacobiWeight{T} end
ChebyshevWeight{kind}() where kind = ChebyshevWeight{kind,Float64}()
ChebyshevWeight() = ChebyshevWeight{1,Float64}()

getproperty(w::ChebyshevWeight{1,T}, ::Symbol) where T = -one(T)/2
getproperty(w::ChebyshevWeight{2,T}, ::Symbol) where T = one(T)/2

"""
Chebyshev{kind,T}()

is a quasi-matrix representing Chebyshev polynomials of the specified kind (1, 2, 3, or 4)
on -1..1.
"""
struct Chebyshev{kind,T} <: AbstractJacobi{T} end
Chebyshev{kind}() where kind = Chebyshev{kind,Float64}()

AbstractQuasiArray{T}(::Chebyshev{kind}) where {T,kind} = Chebyshev{kind,T}()
AbstractQuasiMatrix{T}(::Chebyshev{kind}) where {T,kind} = Chebyshev{kind,T}()

const WeightedChebyshev{kind,T} = WeightedBasis{T,<:ChebyshevWeight{kind},<:Chebyshev{kind}}

WeightedChebyshev{kind}() where kind = ChebyshevWeight{kind}() .* Chebyshev{kind}()
WeightedChebyshev{kind,T}() where {kind,T} = ChebyshevWeight{kind,T}(λ) .* Chebyshev{kind,T}(λ)

const ChebyshevTWeight = ChebyshevWeight{1}
const ChebyshevUWeight = ChebyshevWeight{2}
const ChebyshevT = Chebyshev{1}
const ChebyshevU = Chebyshev{2}
const WeightedChebyshevT = WeightedChebyshev{1}
const WeightedChebyshevU = WeightedChebyshev{2}

# conveniences...perhaps too convenient
Chebyshev() = Chebyshev{1}()
WeightedChebyshev() = WeightedChebyshevT()

chebyshevt() = ChebyshevT()
chebyshevt(d::AbstractInterval{T}) where T = ChebyshevT{float(T)}()[affine(d, ChebyshevInterval{T}()), :]
chebyshevt(d::Inclusion) = chebyshevt(d.domain)
chebyshevt(S::AbstractQuasiMatrix) = chebyshevt(axes(S,1))
chebyshevu() = ChebyshevU()
chebyshevu(d::AbstractInterval{T}) where T = ChebyshevU{float(T)}()[affine(d, ChebyshevInterval{T}()), :]
chebyshevu(d::Inclusion) = chebyshevu(d.domain)
chebyshevu(S::AbstractQuasiMatrix) = chebyshevu(axes(S,1))

"""
     chebyshevt(n, z)

computes the `n`-th Chebyshev polynomial of the first kind at `z`.
"""
chebyshevt(n::Integer, z::Number) = Base.unsafe_getindex(ChebyshevT{typeof(z)}(), z, n+1)
"""
     chebyshevt(n, z)

computes the `n`-th Chebyshev polynomial of the second kind at `z`.
"""
chebyshevu(n::Integer, z::Number) = Base.unsafe_getindex(ChebyshevU{typeof(z)}(), z, n+1)

chebysevtweight(d::AbstractInterval{T}) where T = ChebyshevTWeight{float(T)}[affine(d,ChebyshevInterval{T}())]
chebysevuweight(d::AbstractInterval{T}) where T = ChebyshevUWeight{float(T)}[affine(d,ChebyshevInterval{T}())]

==(a::Chebyshev{kind}, b::Chebyshev{kind}) where kind = true
==(a::Chebyshev, b::Chebyshev) = false
==(::Chebyshev, ::Jacobi) = false
==(::Jacobi, ::Chebyshev) = false
==(::Chebyshev, ::Legendre) = false
==(::Legendre, ::Chebyshev) = false

OrthogonalPolynomial(w::ChebyshevWeight{kind,T}) where {kind,T} = Chebyshev{kind,T}()
orthogonalityweight(P::Chebyshev{kind,T}) where {kind,T} = ChebyshevWeight{kind,T}()

function getindex(w::ChebyshevTWeight, x::Number)
    x ∈ axes(w,1) || throw(BoundsError())
    inv(sqrt(1-x^2))
end

function getindex(w::ChebyshevUWeight, x::Number)
    x ∈ axes(w,1) || throw(BoundsError())
    sqrt(1-x^2)
end

sum(::ChebyshevWeight{1,T}) where T = convert(T,π)
sum(::ChebyshevWeight{2,T}) where T = convert(T,π)/2

normalizationconstant(::ChebyshevT{T}) where T = Vcat(sqrt(inv(convert(T,π))), Fill(sqrt(2/convert(T,π)),∞))

Jacobi(C::ChebyshevT{T}) where T = Jacobi(-one(T)/2,-one(T)/2)
Jacobi(C::ChebyshevU{T}) where T = Jacobi(one(T)/2,one(T)/2)

#######
# transform
#######

factorize(L::SubQuasiArray{T,2,<:ChebyshevT,<:Tuple{<:Inclusion,<:OneTo}}) where T =
    TransformFactorization(grid(L), plan_chebyshevtransform(Array{T}(undef, size(L,2))))

# TODO: extend plan_chebyshevutransform
factorize(L::SubQuasiArray{T,2,<:ChebyshevU,<:Tuple{<:Inclusion,<:OneTo}}) where T<:FastTransforms.fftwNumber =
    TransformFactorization(grid(L), plan_chebyshevutransform(Array{T}(undef, size(L,2))))

########
# Jacobi Matrix
########

jacobimatrix(C::ChebyshevT{T}) where T =
    Tridiagonal(Vcat(one(T), Fill(one(T)/2,∞)), Zeros{T}(∞), Fill(one(T)/2,∞))

jacobimatrix(C::ChebyshevU{T}) where T =
    SymTridiagonal(Zeros{T}(∞), Fill(one(T)/2,∞))

# These return vectors A[k], B[k], C[k] are from DLMF.
recurrencecoefficients(C::ChebyshevT) = (Vcat(1, Fill(2,∞)), Zeros{Int}(∞), Ones{Int}(∞))
recurrencecoefficients(C::ChebyshevU) = (Fill(2,∞), Zeros{Int}(∞), Ones{Int}(∞))

# special clenshaw!
function copyto!(dest::AbstractVector{T}, v::SubArray{<:Any,1,<:Expansion{<:Any,<:ChebyshevT}, <:Tuple{AbstractVector{<:Number}}}) where T
    f = parent(v)
    (x,) = parentindices(v)
    P,c = arguments(f)
    clenshaw!(paddeddata(c), x, dest)
end

###
# Mass matrix
###

@simplify function *(Tc::QuasiAdjoint{<:Any,<:ChebyshevT}, wT::WeightedChebyshevT)
    V = promote_type(eltype(Tc), eltype(wT))
    Diagonal([convert(V,π); Fill(convert(V,π)/2,∞)])
end

@simplify function *(Tc::QuasiAdjoint{<:Any,<:ChebyshevU}, wT::WeightedChebyshevU)
    V = promote_type(eltype(Tc), eltype(wT))
    Diagonal(Fill(convert(V,π)/2,∞))
end

##########
# Derivatives
##########

# Ultraspherical(1)\(D*Chebyshev())
@simplify function *(D::Derivative{<:Any,<:ChebyshevInterval}, S::ChebyshevT)
    T = promote_type(eltype(D),eltype(S))
    A = _BandedMatrix((zero(T):∞)', ℵ₀, -1,1)
    ApplyQuasiMatrix(*, ChebyshevU{T}(), A)
end

#####
# Conversion
#####

\(::Chebyshev{kind,T}, ::Chebyshev{kind,V}) where {kind,T,V} = SquareEye{promote_type(T,V)}(ℵ₀)

function \(U::ChebyshevU, C::ChebyshevT)
    T = promote_type(eltype(U), eltype(C))
    _BandedMatrix(Vcat(-Ones{T}(1,∞)/2,
                        Zeros{T}(1,∞),
                        Hcat(Ones{T}(1,1),Ones{T}(1,∞)/2)), ℵ₀, 0,2)
end

function \(w_A::WeightedChebyshevT, w_B::WeightedChebyshevU)
    wA,A = w_A.args
    wB,B = w_B.args
    T = promote_type(eltype(w_A), eltype(w_B))
    _BandedMatrix(Vcat(Fill(one(T)/2, 1, ∞), Zeros{T}(1, ∞), Fill(-one(T)/2, 1, ∞)), ℵ₀, 2, 0)
end

\(w_A::WeightedChebyshevU, w_B::WeightedChebyshevT) = inv(w_B \ w_A)
\(T::ChebyshevT, U::ChebyshevU) = inv(U \ T)

####
# interrelationships
####

# (18.7.3)

function \(A::ChebyshevT, B::Jacobi)
    J = Jacobi(A)
    Diagonal(J[1,:]) * (J \ B)
end

function \(A::Jacobi, B::ChebyshevT)
    J = Jacobi(B)
    (A \ J) * Diagonal(inv.(J[1,:]))
end

function \(A::Chebyshev, B::Jacobi)
    J = Jacobi(A)
    Diagonal(A[1,:] .\ J[1,:]) * (J \ B)
end

function \(A::Jacobi, B::Chebyshev)
    J = Jacobi(B)
    (A \ J) * Diagonal(J[1,:] .\ B[1,:])
end

function \(A::Jacobi, B::ChebyshevU)
    T = promote_type(eltype(A), eltype(B))
    (A.a == A.b == one(T)/2) || throw(ArgumentError())
    Diagonal(B[1,:] ./ A[1,:])
end

# TODO: Toeplitz dot Hankel will be faster to generate
function \(A::ChebyshevT, B::Legendre)
    T = promote_type(eltype(A), eltype(B))
   UpperTriangular( BroadcastMatrix{T}((k,j) -> begin
            (iseven(k) == iseven(j) && j ≥ k) || return zero(T)
            k == 1 && return Λ(convert(T,j-1)/2)^2/π
            2/π * Λ(convert(T,j-k)/2) * Λ(convert(T,k+j-2)/2)
        end, 1:∞, (1:∞)'))
end

\(A::AbstractJacobi, B::Chebyshev) = ApplyArray(inv,B \ A)

function \(A::Jacobi, B::WeightedBasis{<:Any,<:JacobiWeight,<:Chebyshev})
    w, T = B.args
    J = Jacobi(T)
    wJ = w .* J
    (A \ wJ) * (J \ T)
end

function \(A::Chebyshev, B::WeightedBasis{<:Any,<:JacobiWeight,<:Jacobi})
    J = Jacobi(A)
    (A \ J) * (J \ B)
end

function \(A::Chebyshev, B::WeightedBasis{<:Any,<:JacobiWeight,<:Chebyshev})
    J = Jacobi(A)
    (A \ J) * (J \ B)
end

####
# sum
####

function _sum(A::WeightedBasis{T,<:ChebyshevUWeight,<:ChebyshevU}, dims) where T
    w, U = A.args
    @assert dims == 1
    Hcat(convert(T, π)/2, Zeros{T}(1,∞))
end

function _sum(A::WeightedBasis{T,<:ChebyshevWeight,<:Chebyshev}, dims) where T
    @assert dims == 1
    Hcat(convert(T, π), Zeros{T}(1,∞))
end

function cumsum(T::ChebyshevT{V}; dims::Integer) where V
    @assert dims == 1
    Σ = _BandedMatrix(Vcat(-one(V) ./ (-2:2:∞)', Zeros{V}(1,∞), Hcat(one(V), one(V) ./ (4:2:∞)')), ℵ₀, 0, 2)
    ApplyQuasiArray(*, T, Vcat((-1).^(0:∞)'* Σ, Σ))
end

cumsum(f::Expansion{<:Any,<:ChebyshevT}) = cumsum(f.args[1]; dims=1) * f.args[2]

####
# algebra
####

broadcastbasis(::typeof(+), ::ChebyshevT, U::ChebyshevU) = U
broadcastbasis(::typeof(+), U::ChebyshevU, ::ChebyshevT) = U


using IntUtils

function main()
    if length(ARGS) > 0
        n = parse(Int64, ARGS[1])
        a = Dict{Int64,Int64}[]
        for i in 2:n
            push!(a, prime_factors_table(i))
        end
        prod = 1
        for i in 2:n
            maxp = 0
            for j in 1:(n-1)
                p = get(a[j], i, 0)
                if p > maxp
                    maxp = p
                end
            end
            if maxp > 0
                prod *= i ^ maxp
            end
        end

        @printf("The smallest number that can be divided by each of the numbers from 1 to %d without any remainder is %d\n", n, prod)
    else
        error("this program needs a single numeric argument")
    end

end

main()


@testset "Utilities" begin
  a, b, c = P2(0,0), P2(1,0), P2(0,1)
  @test signarea(a, b, c) == T(0.5)
  a, b, c = P2(0,0), P2(0,1), P2(1,0)
  @test signarea(a, b, c) == T(-0.5)

  p1, p2, p3 = P2(0,0), P2(1,1), P2(0.25,0.5)
  s = Segment(P2(0.5,0.0), P2(0.0,1.0))
  @test sideof(p1, s) == :LEFT
  @test sideof(p2, s) == :RIGHT
  @test sideof(p3, s) == :ON

  p1, p2, p3 = P2(0.5,0.5), P2(1.5,0.5), P2(1,1)
  c = Chain(P2[(0,0),(1,0),(1,1),(0,1),(0,0)])
  @test sideof(p1, c) == :INSIDE
  @test sideof(p2, c) == :OUTSIDE
  @test sideof(p3, c) == :INSIDE
end


using Revise
using JuMP,Gurobi,BilevelJuMP
includet("./bilevelJuMPExtension.jl")
# min -4sum(x) -3sum(y)
# s.t.
# y = argmin_y sum(y)
#      2x + y <= 4
#       x +2y <= 4
#       x     >= 0
#           y >= 0
#
# sol: x = 2, y = 0

# opt is the orginal problem
opt =  quote
    function some_model()
        blm = BilevelModel()
        sz = 10
        um = Upper(blm)
        lm = Lower(blm)
        @variable(um, x[1:sz])
        @variable(lm, y[1:sz])

        @objective(um, Min, -4sum(x) -3sum(y))

        @objective(lm, Min, sum(y))

        for i in 1:sz
            @constraint(lm, 2x[i]+y[i] <= 4)
            @constraint(lm, x[i]+2y[i] <= 4)
            @constraint(lm, x[i] >= 0)
            @constraint(lm, y[i] >= 0)
        end
        @constraint(lm,con[i in 1:sz],y[i]==1) #sign
        return blm
    end
end

# test.1 simple removal and addition of constarint named "con"
printstyled("****REMOVAL OF CONSTRAINT CON****\n";color=:blue)
opt = @delete_cons(opt,con)
mdl = eval(opt)()
optimize!(mdl,Gurobi.Optimizer())
@assert value(mdl[:y][1]) ≈ 0
printstyled("****ADDITION OF CONSTRAINT CON****\n";color=:blue)
opt = @add_lower_cons(opt,con,opt[:y][1] == 1)
mdl = eval(opt)()
optimize!(mdl,Gurobi.Optimizer())
@assert value(mdl[:y][1]) ≈ 1
printstyled("****SIMPLE TEST PASSED****\n";color=:green)

#test.2 complex removal and addition of constarint named "con" within loops
opt = [opt,opt]
n = 10
@time begin
for ii = 1:2
    for i = 1:2
        for j = 1:2
            printstyled("****REMOVAL OF CONSTRAINT CON****\n";color=:blue)
            opt[ii] = @delete_cons(opt[ii],con)
            f = eval(opt[ii])
            optimize!(f(),Gurobi.Optimizer(OutputFlag=0))
            sth = [0.5,0.3]
            idx = 2
            printstyled("****ADDITION OF CONSTRAINT CON****\n";color=:blue)
            opt[ii] = @add_lower_cons(opt[ii],con,sum(opt[ii][:y][i] + sth[idx]^2 for i in 1:n) == 10 + 0.1*(i+j))
            # println(opt[ii])
            f = eval(opt[ii])
            optimize!(f(),Gurobi.Optimizer(OutputFlag=0))
        end
    end
end
end
printstyled("****COMPLEX ADDITION TEST WITH LOOPS PASSED****\n";color=:green)

#test.3 batch addtion of constraints
@time begin
    for t in 1:2
        printstyled("****REMOVAL OF CONSTRAINT CON****\n";color=:blue)
        opt[t] = @delete_cons(opt[t],con)
        f = eval(opt[t])
        optimize!(f(),Gurobi.Optimizer(OutputFlag=0))
        sth = [1 for i in 1:n]
        printstyled("****ADDITION OF CONSTRAINT CON****\n";color=:blue)
        opt[t] = @add_lower_cons(opt[t],con[i = 1:n],opt[t][:y][i] == sth[i])
        # println(opt[t])
        f = eval(opt[t])
        optimize!(f(),Gurobi.Optimizer(OutputFlag=0))
    end
end
printstyled("****BATCH CONSTRAINTS ADDITION PASSED****\n";color=:green)


using DifferentialEquations: ODEProblem, DifferentialEquations

@inline function __integrate(u::Function, domain::Tuple{T,T}; kwargs...) where T <: AbstractFloat
	prob = ODEProblem( (f, p, τ) -> u(τ), zero(T), domain)
	return DifferentialEquations.solve(prob; kwargs...)
end

@inline function __getNorm(u::Function, c::Compression; kwargs...)
	σ, μ = c.σ, c.μ
	sol = __integrate(
		τ -> abs(τ) ≈ 1.0 ? 0.0 : σ * abs2( u(σ * c.decomp(τ) + μ) ) * c.inv_jacobian(τ),
		(-1.0, 1.0);
		kwargs...
	)
	return t -> abs(sol(c.comp((t - μ) / σ))), sol(1.0)
end

"""
    generateCouplings(u; compression, kwargs...)

inputs
- u:            mode function
- compression:  which compression of ℝ → (-1, 1) will be used for the differential equation
								Options are the keys of the `Compressions` dictionary
								or any `Compression` Object
- kwargs:				all other kwargs get's passed to the numerical integrator

output:
- (gᵢ, gₒ, norm): the coupling constant calculated from the mode function and its norm function.
"""
# TODO test with (df, f, p, τ) lambda expression

function generateCoupling(u::Function; compression::Union{Symbol, Compression} = :exponential, kwargs...)
	if isa(compression, Symbol)
		@assert compression ∈ keys(Compressions)

		return _generateCoupling(u, Compressions[compression]; kwargs...)
	end

	return _generateCoupling(u, compression; kwargs...)
end

function _generateCoupling(u::Function, c::Compression; kwargs...)
	norm, Norm = __getNorm(u, c; kwargs...)
	gᵢ(t) =   u(t) / √(abs(Norm - norm(t)) + (u(t) == 0.0))
	gₒ(t) = - u(t) / √(norm(t) + (u(t) == 0.0))

	return gᵢ, gₒ, norm
end

# """
#     modeContribution(u, v; compression)
#
# inputs
# - u:            mode function of the cavity
# - v:						mode function falsely contributing to the cavity
# - compression:  which compression of ℝ → (-1, 1) will be used for the differential equation
# 								Options are the keys of the `Compressions` dictionary
# 								or any `Compression` Object
#
# output:
# - α:   the contribution of mode u to the cavity
# """
# # TODO test with (df, f, p, τ) lambda expression
# function modeContribution(u::Mode, v::Mode; compression::Union{Symbol, Compression} = :exponential)
# 	if isa(compression, Symbol)
# 		@assert compression ∈ keys(Compressions)
#
# 		return _modeContribution(u, v, Compressions[compression])
# 	end
#
# 	return _modeContribution(u, v, compression)
# end
#
# function _modeContribution(uM::Mode, vM::Mode, c::Compression)
# 	u, v = uM.modeFunction, vM.modeFunction
# 	dfdτ(f, p, τ) = begin
# 		abs(τ) == 1.0 && return 0.0
#
# 		t = c.decomp(τ)
# 		df = u(t) * v(t) * c.inv_jacobian(τ)
# 	end
#
# 	prob = ODEProblem(dfdτ, 0., (-1.0, 1.0))
#   sol = solve(prob, reltol = 1e-5, abstol = 1e-8)
#
# 	return t -> sol(c.comp(t)) / √(uM.norm(t))
# end
#
# """
#     modifiedMode(u, v; compression)
#
# inputs
# - u:            `Mode` of the first cavity
# - v:						target `Mode` for the second cavity
# - [α]:					predetermined modeContribution of v to the first cavity. Optional.
# 								Useful, if α should be calculated with a different compression than the
# 								modified mode.
# - compression:  which compression of ℝ → (-1, 1) will be used for the differential equation
# 								Options are the keys of the `Compressions` dictionary
# 								or any `Compression` Object
#
# output:
# - ṽ:   modified `Mode` for the second cavity, accounting for the scattering at the first cavity
# """
# # TODO test with (df, f, p, τ) lambda expression
# function modifiedMode(u::Mode, v::Mode, α::Union{Function, Nothing} = nothing;
# 		compression::Union{Symbol, Compression} = :exponential
# 	)
#
# 	if isa(compression, Symbol)
# 		@assert compression ∈ keys(Compressions)
# 		compression = Compressions[compression]
# 	end
#
# 	isnothing(α) && (α = _modeContribution(u, v, compression))
# 	return _modifiedMode(u, v, α, compression)
# end
#
# function _modifiedMode(u::Mode, v::Mode, α::Function, compression::Compression)
# 	vMod(t) = v.modeFunction(t) + u.gₒ(t) * α(t)
#
# 	return Mode(vMod, compression = compression)
# end


using LinearAlgebra
using SparseArrays

# Calculate the objective of logistic regression:
# \frac{1}{n} \sum_{i=1}^n \log( 1 + \exp( -w^T x_i y_i ) ) + \lambda/2 \|w\|^2.
function obj(
    X::SparseMatrixCSC{Float64, Int64},
    y::Vector{Int64},
    W::Matrix{Float64},
    λ::Float64
)
    n, _ = size(X)
    XW = X*W
    objval = 0.0
    for i = 1:n
        prob = softmax(XW[i,:])
        objval += -log(prob[ y[i] ])
    end
    objval /= n
    objval += λ/2*norm(W)^2
    return objval
end

# Calculate the objective of neural network model:
function obj(
    Xt::SparseMatrixCSC{Float64, Int64},
    Y::Flux.OneHotArray,
    W::Flux.Chain,
    λ::Float64
)
    loss(x,y) = Flux.Losses.crossentropy(W(x), y)
    l = 0.0
    num_data = size(Xt, 2)
    for i = 1:num_data
        l += loss(Xt[:,i], Y[:,i])
    end
    sqnorm(w) = sum(abs2, w)
    return l/num_data + (λ/2) * sum(sqnorm, params(W))
end

# Line-search for softmax
function lineSearch(
    X::SparseMatrixCSC{Float64, Int64},
    y::Vector{Int64},
    D::Matrix{Float64},
    W::Matrix{Float64},
    G::Matrix{Float64},
    λ::Float64
)
    maxLineSearchSteps = 100
    GTD = dot(G, D)
    # @show(GTD)
    Wnorm = norm(W)
    XW = X*W
    XD = X*D

    n = size(X, 1)
    objval = 0.0
    @inbounds for i = 1:n
        prob = softmax(XW[i,:])
        objval += -log(prob[ y[i] ])
    end
    objval /= n
    objval += λ/2*Wnorm^2

    Wnew = zeros(Float64, size(W))
    XWnew = zeros(Float64, size(XW))

    η = 1.0
    β = 1e-2
    for t = 1:maxLineSearchSteps
        Wnew .= W - η*D
        XWnew .= XW - η*XD
        objNew = 0.0
        @inbounds for i = 1:n
            prob = softmax(XWnew[i,:])
            objNew += -log(prob[ y[i] ])
        end
        objNew /= n
        objNew += λ/2*norm(Wnew)^2
        if objNew > objval - β*η*GTD
            η *= 0.5
        else
            break
        end
        if t == maxLineSearchSteps
            @warn("Reached maximum linesearch steps.")
        end
    end
    return η
end

# Line-search for FedDCD
function lineSearch2(
    X::SparseMatrixCSC{Float64, Int64},
    Y::Vector{Int64},
    y::Matrix{Float64},  # dual variable
    D::Matrix{Float64},
    W::Matrix{Float64},
    G::Matrix{Float64},
    λ::Float64
)
    maxLineSearchSteps = 100
    GTD = dot(G, D)
    # @show(GTD)
    Wnorm = norm(W)
    XW = X*W
    XD = X*D
    WTy = dot(W, y)
    DTy = dot(D, y)

    n = size(X, 1)
    objval = 0.0
    @inbounds for i = 1:n
        prob = softmax(XW[i,:])
        objval += -log(prob[ Y[i] ])
    end
    objval /= n
    objval -= WTy   # objective have an additional term -W^Ty
    objval += λ/2*Wnorm^2

    Wnew = zeros(Float64, size(W))
    XWnew = zeros(Float64, size(XW))

    η = 1.0
    β = 1e-2
    for t = 1:maxLineSearchSteps
        Wnew .= W - η*D
        XWnew .= XW - η*XD
        objNew = 0.0
        @inbounds for i = 1:n
            prob = softmax(XWnew[i,:])
            objNew += -log(prob[ Y[i] ])
        end
        objNew /= n
        objNew -= ( WTy - η*DTy )
        objNew += λ/2*norm(Wnew)^2
        if objNew > objval - β*η*GTD
            η *= 0.5
        else
            break
        end
        if t == maxLineSearchSteps
            @warn("Reached maximum linesearch steps.")
        end
    end
    return η
end

# Calculate the accuracy
function accuracy(
    X::SparseMatrixCSC{Float64, Int64},
    y::Vector{Int64},
    W::Matrix{Float64},
)
    n = size(X, 1)
    XW = X*W
    ret = 0
    for i = 1:n
        prob = softmax( XW[i,:] )
        if argmax(prob) == y[i]
            ret += 1
        end
    end
    return ret/n
end

# Calculate the accuracy
function accuracy(
    X::SparseMatrixCSC{Float64, Int64},
    Y::Flux.OneHotArray,
    W::Flux.Chain,
)
    Xt = copy(X')
    acc(x,y) = 1.0* ( Flux.onecold(W(x)) == Flux.onecold(y) )
    num_data = size(Xt, 2)
    out = 0.0
    for i = 1:num_data
        out += acc(Xt[:,i], Y[:,i])
    end
    return out/num_data
end

# Calculate the stochastic gradient
function getStochasticGrad(
    Xt::SparseMatrixCSC{Float64, Int64},
    y::Vector{Int64},
    W::Matrix{Float64},
    idx::Int64
)
    x = Xt[:,idx]
    Wtx = W'*x
    s = softmax( Wtx )
    s[ y[idx] ] -= 1
    g = x * s'
    return g
end

# Calculate gradient
function getGradient(
    X::SparseMatrixCSC{Float64, Int64},
    Xt::SparseMatrixCSC{Float64, Int64},
    y::Vector{Int64},
    W::Matrix{Float64},
    λ::Float64
)
    n, d = size(X)
    _, K = size(W)
    g = zeros(Float64, d, K)
    XW = X*W
    @inbounds for i = 1:n
        s = softmax( XW[i,:] )
        s[ y[i] ] -= 1
        x = Xt[:,i]
        I, V = findnz(x)
        @inbounds for col = 1:K
            @inbounds for j = 1:length(I)
                row = I[j]
                val = V[j]
                g[row, col] += val*s[col]
            end
        end
    end
    g ./= n
    g += λ.*W
    return g
end

# Add a dense matrix with a sparse matrix in place, A = A + x * b', x sparse, b dense.
function FastHvMatrixUpdate!(
    A::Matrix{Float64},
    x::SparseVector{Float64,Int64},
    b::Vector{Float64}
)
    I, V = findnz(x)
    K = length(b)
    @inbounds for col = 1:K
        @inbounds for idx = 1:length(I)
            row = I[idx]
            v = V[idx]
            A[row, col] += v*b[col]
        end
    end
    return nothing
end

# Fast Hessian-vector product for softmax classification
function FastHv(
    X::SparseMatrixCSC{Float64, Int64},
    Xt::SparseMatrixCSC{Float64, Int64},
    W::Matrix{Float64},
    λ::Float64,
    V::Matrix{Float64}
)
    # Time complexity is O( nnz(X)*K )
    n, d = size(X)
    _, K = size(W)
    XW = X*W
    XV = X*V
    P = zeros(Float64, n, K)
    ret = zeros(Float64, d, K)
    for i = 1:n
        P[i,:] = softmax( XW[i,:] )
    end
    @inbounds for i = 1:n
        # Complexity of the loop: nnz(xi)*K
        # Calculate (x_i x_i^T V ( Λ - p_i p_i^T ) )
        xi = Xt[:, i]
        ppi = P[i, :]
        # First calculate x_i^T V
        xtV = XV[i, :]
        # # Calculate x_i^T V ( Λ - p_i p_i^T )
        xtVppt = xtV .* ppi
        xtVppt -= dot(xtV, ppi) .* ppi
        # # Calculate x_i x_i^T V ( Λ - p_i p_i^T )
        FastHvMatrixUpdate!( ret, xi, xtVppt)
        # ret += xi*xtVppt'
    end
    ret ./= n
    ret += λ.*V
    return ret
end

# Compute Newton direction. Use cg to solve the linear system H D = g.
function ComputeNewtonDirection(
    X::SparseMatrixCSC{Float64, Int64},
    Xt::SparseMatrixCSC{Float64, Int64},
    W::Matrix{Float64},
    λ::Float64,
    g::Matrix{Float64}
)
    n, d = size(X)
    _, K = size(W)

    H = LinearMap(v->vec(FastHv(X, Xt, W, λ, reshape(v,d,K))), d*K, issymmetric=true, isposdef=true)
    D = cg(H, vec(g), abstol=1e-8, reltol=1e-4, maxiter=1000)

    return reshape(D, d, K)
end

# Using Newton's method to solve the softmax classification problem.
function SoftmaxNewtonMethod(
    X::SparseMatrixCSC{Float64, Int64},
    Xt::SparseMatrixCSC{Float64, Int64},
    y::Vector{Int64},
    W::Matrix{Float64},
    λ::Float64,
    maxIter::Int64=20,
    tol::Float64=1e-4
)
    startTime = time()
    for iter = 1:maxIter
        objval = obj(X, y, W, λ)
        g = getGradient(X, Xt, y, W, λ)
        gnorm = norm(g)
        @printf("Iter %3d, obj: %4.5e, gnorm: %4.5e, time: %4.2f\n", iter, objval, gnorm, time()-startTime)
        if gnorm < tol
            break
        end
        # Compute Newton direction.
        D = ComputeNewtonDirection( X, Xt, W, λ, g)
        η = lineSearch(X, Y, D, W, g, λ)
        W .-= η*D
    end
    return W
end


A=[1. 2. 3.; 4. 5. 6.; 7. 8. 9.]

    function cost(X::ProductRepr)
        return cost([submanifold_components(X)...])
    end
    function cost(X::Array{Matrix{Float64},1})
        return -0.5 * norm(transpose(X[1]) * A * X[2])^2
    end

    function egrad(X::Array{Matrix{Float64},1})
        U = X[1]
        V = X[2]
        AV = A*V
        AtU = transpose(A)*U
        return [ -AV*(transpose(AV)*U), -AtU*(transpose(AtU)*V) ];
    end

    function rgrad(M::ProductManifold, X::ProductRepr)
        eG = egrad([submanifold_components(M,X)...])
        x = [submanifold_components(M,X)...]
        return Manifolds.ProductRepr(project.(M.manifolds, x, eG)...)
    end

    function e2rHess(M::Grassmann, x, ξ, eGrad::Matrix{T},Hess::Matrix{T}) where T<:Union{U, Complex{U}} where U<:AbstractFloat
        pxHess = project(M,x,Hess)
        xtGrad = x'*eGrad
        ξxtGrad = ξ*xtGrad
        return pxHess - ξxtGrad
    end

    function eHess(X::Array{Matrix{Float64},1}, H::Array{Matrix{Float64},1})
        U = X[1]
        V = X[2]
        Udot = H[1]
        Vdot = H[2]
        AV = A*V
        AtU = transpose(A)*U
        AVdot = A*Vdot
        AtUdot = transpose(A)*Udot
        return [ -(AVdot*transpose(AV)*U + AV*transpose(AVdot)*U + AV*transpose(AV)*Udot),
                 -(AtUdot*transpose(AtU)*V + AtU*transpose(AtUdot)*V + AtU*transpose(AtU)*Vdot)
            ]
    end

    function rhess(M::ProductManifold, X::ProductRepr, H::ProductRepr)
        x = [submanifold_components(M,X)...]
        h = [submanifold_components(M,H)...]
        eG = egrad(x)
        eH = eHess(x,h)
        return Manifolds.ProductRepr(e2rHess.(M.manifolds, x, h, eG, eH)...)
    end

@testset "Manopt Trust-Region" begin
    import Random: seed!
    seed!(42);

    N = Grassmann(3,2)
    M = N × N

    x = random_point(M)

    @test_throws ErrorException trust_regions(M, cost, rgrad, x, rhess; ρ_prime = 0.3)
    @test_throws ErrorException trust_regions(M, cost, rgrad, x, rhess; Δ_bar = -0.1)
    @test_throws ErrorException trust_regions(M, cost, rgrad, x, rhess; Δ = -0.1)
    @test_throws ErrorException trust_regions(M, cost, rgrad, x, rhess; Δ_bar = 0.1, Δ = 0.11)

    X = trust_regions(M, cost, rgrad, x, rhess; Δ_bar=4*sqrt(2*2) )
    opt = trust_regions(M, cost, rgrad, x, rhess; Δ_bar=4*sqrt(2*2), return_options=true )
    @test isapprox(M,X,get_solver_result(opt))

    @test cost(X) + 142.5 ≈ 0 atol=10.0^(-13)

    XuR = trust_regions(M, cost, rgrad, x, rhess;
        Δ_bar=4*sqrt(2*2),
        useRandom = true
    )

    @test cost(XuR) + 142.5 ≈ 0 atol=10.0^(-12)

    XaH = trust_regions(
        M,
        cost,
        rgrad,
        x,
        (p,x,ξ) -> approxHessianFD(
            p,
            x,
            x -> rgrad(p,x),
            ξ;
            stepsize=2^(-9),
            transport=ProductVectorTransport(ProjectionTransport(),ProjectionTransport())
        );
        stopping_criterion = StopWhenAny(
            StopAfterIteration(2000),
            StopWhenGradientNormLess(10^(-6))
        ),
        Δ_bar=4*sqrt(2*2),
    )
    @test cost(XaH) + 142.5 ≈ 0 atol=10.0^(-10)

    ξ = random_tangent(M,x)
    @test_throws ErrorException getHessian(SubGradientProblem(M,cost,rgrad),x, ξ)

    # Test the random step trust region
    p = HessianProblem(M, cost, rgrad, rhess, (M,x,ξ) -> ξ)
    o = TrustRegionsOptions(x, StopAfterIteration(2000), 10.0^(-8),
        sqrt(manifold_dimension(M)), retract, true, 0.1, 1000.)
    @test step_solver!(p,o,0) == nothing

    η = truncated_conjugate_gradient_descent(M, cost, rgrad, x, ξ, rhess, 0.5)
    ηOpt = truncated_conjugate_gradient_descent(M, cost, rgrad, x, ξ, rhess, 0.5; return_options=true)
    @test submanifold_components(get_solver_result(ηOpt)) == submanifold_components(η)
end


using FMMLIB2D
using SpecialFunctions
using Compat.LinearAlgebra
using Compat.Random
using Compat.Test

if VERSION < v"0.7"
    srand(0)
else
    Random.seed!(0)
end

function direct_self(source, dipstr)
    pot = zero(source)
    grad = zero(source)
    hess = zero(source)
    N = length(source)
    for i=1:N
        for j=1:N
            if i == j
                continue
            end
            pot[i] += dipstr[j] / (source[i] - source[j])
            grad[i] += -dipstr[j] / (source[i] - source[j])^2
            hess[i] += 2*dipstr[j] / (source[i] - source[j])^3
        end
    end
    return pot, grad, hess
end

function direct_targ(source, dipstr, target)
    M = length(target)
    N = length(source)
    pot = zero(target)
    grad = zero(target)
    hess = zero(target)
    for i=1:M
        for j=1:N
            pot[i] += dipstr[j] / (target[i] - source[j])
            grad[i] += -dipstr[j] / (target[i] - source[j])^2
            hess[i] += 2*dipstr[j] / (target[i] - source[j])^3
        end
    end
    return pot, grad, hess
end

function runtest()
    iprec = 5
    N = 30
    M = 40
    source = rand(N) + 1im*rand(N)
    dipstr = rand(N) + 1im*rand(N)
    target = rand(M) + 1im*rand(M)
    ifpot = 1
    ifgrad = 1
    ifhess = 1
    ifpottarg = 1
    ifgradtarg = 1
    ifhesstarg = 1

    refpottarg, refgradtarg, refhesstarg = direct_targ(source, dipstr, target)
    refpot, refgrad, refhess = direct_self(source, dipstr)

    @testset "Direct interface" begin
        pot, grad, hess, pottarg, gradtarg, hesstarg =
            zfmm2dparttarg(iprec, source, dipstr, ifpot, ifgrad, ifhess,
                           target, ifpottarg, ifgradtarg, ifhesstarg)

        pot_relerr = norm(refpot-pot, Inf) / norm(refpot, Inf)
        grad_relerr = norm(refgrad-grad, Inf) / norm(refgrad, Inf)
        hess_relerr = norm(refhess-hess, Inf) / norm(refhess, Inf)

        pottarg_relerr = norm(refpottarg-pottarg, Inf) / norm(refpottarg, Inf)
        gradtarg_relerr = norm(refgradtarg-gradtarg, Inf) / norm(refgradtarg, Inf)
        hesstarg_relerr = norm(refhesstarg-hesstarg, Inf) / norm(refhesstarg, Inf)

        @test pot_relerr < 1e-15
        @test grad_relerr < 1e-15
        @test hess_relerr < 1e-15
        @test pottarg_relerr < 1e-15
        @test gradtarg_relerr < 1e-15
        @test hesstarg_relerr < 1e-15
    end

    @testset "Keyword interface" begin
        U = zfmm2d(source=source, dipstr=dipstr)
        @test norm(U.pot-refpot, Inf) / norm(refpot, Inf) < 1e-14
        U = zfmm2d(source=source, dipstr=dipstr, target=target)
        @test norm(U.pottarg-refpottarg, Inf) / norm(refpot, Inf) < 1e-14
    end

    @testset "Tolerance test" begin
        N = 1000
        source = rand(N) + 1im*rand(N)
        dipstr = rand(N) + 1im*rand(N)
        refpot, _, _ = direct_self(source, dipstr)
        for tolexp=-14:-1
            TOL = 0.49*10.0^tolexp
            U = zfmm2d(source=source, tol=TOL, dipstr=dipstr)
            relerr = norm(U.pot-refpot) / norm(refpot)
            @test relerr < TOL
        end
    end
end

runtest()


## This file has methods to stably divide the jacobi singularity

function dirichlet_divide_singularity{T<:Number}(b::Bool,v::Vector{T})
    n=length(v)
    w=zeros(T,n-1)
    s=b?1:-1

    if n==1
        return w
    elseif n==2
        w[1]=-s*v[2]
        return w
    end

    w[n-1]=-2s*v[n]

    for k=n-2:-1:2
        @inbounds w[k]=-2s*(v[k+1] - .5w[k+1])
    end

    @inbounds w[1]=-s*(v[2] - .5w[2])

    w
end

# dirichletrange_divide_singularity divides by (1-x^2)
# where the coefficients are ChebyshevDirichlet{1,1}
# assume that v[1] == v[2] ==0.
# the output coefficients are Chebyshev

function dirichlet_divide_singularity{T<:Number}(v::Vector{T})
    n=length(v)
    if n ≤ 2    # assumes v[1]==v[2]==0 which is deleted
        return zeros(T,1)
    end

    w=zeros(T,n-2)

    w[n-2]=-4v[n]
    if n>3
        w[n-3]=-4v[n-1]
    end

    for k=n-4:-1:2
        @inbounds w[k]=-4*(v[k+2] - .25w[k+2])
    end

    w[1]=-2v[3]
    if n≥5
        w[1]+=0.5w[3]
    end

    w
end

divide_singularity(s::Bool,f::Fun)=Fun(dirichlet_divide_singularity(s,Fun(f,s?ChebyshevDirichlet{0,1}:ChebyshevDirichlet{1,0}).coefficients),
                                                                    Chebyshev(domain(f)))

divide_singularity(f::Fun)=Fun(dirichlet_divide_singularity(Fun(f,ChebyshevDirichlet{1,1}).coefficients),Chebyshev(domain(f)))

function divide_singularity(s::Tuple{Int,Int},f::Fun)
    if s[1]>0 && s[2]>0
        divide_singularity((s[1]-1,s[2]-1),divide_singularity(f))
    elseif s[1]>0
        divide_singularity((s[1]-1,s[2]),divide_singularity(false,f))
    elseif s[2]>0
        divide_singularity((s[1],s[2]-1),divide_singularity(true,f))
    else
        @assert s[1]==s[2]==0
        f
    end
end


# Rational Quadratic ARD Covariance Function

"""
    RQArd <: StationaryARD{WeightedSqEuclidean}

ARD Rational Quadratic kernel (covariance)
```math
k(x,x') = σ²(1 + (x - x')ᵀL⁻²(x - x')/(2α))^{-α}
```
with length scale ``ℓ = (ℓ₁, ℓ₂, …)``, signal standard deviation ``σ``, and shape parameter
``α`` where ``L = diag(ℓ₁, ℓ₂, …)``.
"""
mutable struct RQArd <: StationaryARD{WeightedSqEuclidean}
    "Inverse squared length scale"
    iℓ2::Vector{Float64}
    "Signal variance"
    σ2::Float64
    "Shape parameter"
    α::Float64
    "Priors for kernel parameters"
    priors::Array

    """
        RQArd(ll::Vector{Float64}, lσ::Float64, lα::Float64)

    Create `RQArd` with length scale `exp.(ll)`, signal standard deviation `exp(lσ)`, and
    shape parameter `exp(lα)`.
    """
    RQArd(ll::Vector{Float64}, lσ::Float64, lα::Float64) =
        new(exp.(-2 .* ll), exp(2 * lσ), exp(lα), [])
end

function set_params!(rq::RQArd, hyp::VecF64)
    length(hyp) == num_params(rq) || throw(ArgumentError("RQArd kernel has $(num_params(rq_ard)) parameters"))
    @views @. rq.iℓ2 = exp(-2 * hyp[1:(end-2)])
    rq.σ2 = exp(2 * hyp[end-1])
    rq.α = exp(hyp[end])
end

get_params(rq::RQArd) = [-log.(rq.iℓ2) / 2; log(rq.σ2) / 2; log(rq.α)]
get_param_names(rq::RQArd) = [get_param_names(rq.iℓ2, :ll); :lσ; :lα]
num_params(rq::RQArd) = length(rq.iℓ2) + 2

Statistics.cov(rq::RQArd,r::Float64) = rq.σ2*(1+0.5*r/rq.α)^(-rq.α)

@inline dk_dll(rq::RQArd, r::Float64, wdiffp::Float64) =
    rq.σ2 * wdiffp * (1 + r / (2 * rq.α))^(-rq.α - 1)
@inline function dk_dlα(rq::RQArd, r::Float64)
    part = (1 + r / (2 * rq.α))
    return rq.σ2 * part^(-rq.α) * (r / (2 * part) - rq.α * log(part))
end
@inline function dKij_dθp(rq::RQArd, X::MatF64, i::Int, j::Int, p::Int, dim::Int)
    if p <= dim
        return dk_dll(rq, distij(metric(rq),X,i,j,dim), distijk(metric(rq),X,i,j,p))
    elseif p==dim+1
        return dk_dlσ(rq, distij(metric(rq),X,i,j,dim))
    else
        return dk_dlα(rq, distij(metric(rq),X,i,j,dim))
    end
end
@inline function dKij_dθp(rq::RQArd, X::MatF64, data::StationaryARDData, i::Int, j::Int, p::Int, dim::Int)
    return dKij_dθp(rq,X,i,j,p,dim)
end


module TestTensorNetworks
using Test

using LinearAlgebra: tr

using Catlab.CategoricalAlgebra.CSets, Catlab.WiringDiagrams
using Catlab.LinearAlgebra.TensorNetworks

# Parsing and code generation
#############################

# Parsing
#--------

parsed = @tensor_network (i,j,k,ℓ) D[i,j,k] = A[i,ℓ] * B[j,ℓ] * C[k,ℓ]
d = RelationDiagram{Symbol}(3)
add_box!(d, 2, name=:A); add_box!(d, 2, name=:B); add_box!(d, 2, name=:C)
add_junctions!(d, 4, variable=[:i,:j,:k,:ℓ])
set_junction!(d, [1,4,2,4,3,4])
set_junction!(d, 1:3, outer=true)
@test parsed == d

parsed = @tensor_network D[i,j,k] = A[i,ℓ] * B[j,ℓ] * C[k,ℓ]
@test parsed == d
parsed = @tensor_network D[i,j,k] := A[i,ℓ] * B[j,ℓ] * C[k,ℓ]
@test parsed == d

# Degenerate case: single term.
parsed = @tensor_network B[i,j,k] = A[i,j,k]
d = singleton_diagram(RelationDiagram{Symbol}, 3, name=:A)
set_subpart!(d, :variable, [:i,:j,:k])
@test parsed == d

# Degenerate case: no terms.
parsed = @tensor_network A[i,j] = 1
d = RelationDiagram{Symbol}(2)
add_junctions!(d, 2, variable=[:i,:j])
set_junction!(d, 1:2, outer=true)
@test parsed == d

# Round trip: parse then compile
#-------------------------------

macro roundtrip_tensor(tensor)
  quote
    compiled = gen_tensor_notation(@tensor_network($tensor),
                                   assign_op=:(=), assign_name=:out)
    @test compiled == $(QuoteNode(tensor))
  end
end

@roundtrip_tensor out[i,k] = A[i,j] * B[j,k]
@roundtrip_tensor out[i,j,k] = A[i,ℓ] * B[j,ℓ] * C[k,ℓ]
@roundtrip_tensor out = u[i] * A[i,j] * v[j]
@roundtrip_tensor out[j] = α * a[i] * B[i,j]

# Evaluation
############

# Binary matrix multiply.
A, B = randn((3,4)), randn((4,5))
network = @tensor_network C[i,k] = A[i,j] * B[j,k]
@test oapply(network, [A,B]) ≈ A*B

# Ternary matrix multiply.
C = randn((5,6))
network = @tensor_network D[i,ℓ] = A[i,j] * B[j,k] * C[k,ℓ]
@test oapply(network, [A,B,C]) ≈ A*B*C

# Trace.
A = randn((5,5))
network = @tensor_network out[] = A[i,i]
@test oapply(network, [A])[] ≈ tr(A)

# Tensor product.
A, B = randn((3,4)), randn((5,6))
network = @tensor_network C[i,j,k,ℓ] = A[i,j] * B[k,ℓ]
@test oapply(network, [A,B]) ≈ (reshape(A, (3,4,1,1)) .* reshape(B, (1,1,5,6)))

end


#  Compute r = A * x + s - b in-place
function primal_kkt_condition!(r::AbstractVector{T}, A::AbstractMatrix{T}, x::AbstractVector{T}, s::SplitVector{T}, b::AbstractVector{T}) where {T}
	# r = A * x
	mul!(r, A, x)
	@. r += s
	@. r -= b
	return nothing
end

# Compute r = P * x + q - A' * μ in-place
function dual_kkt_condition!(r::AbstractVector{T}, r_temp::AbstractVector{T}, P::AbstractMatrix{T}, x::AbstractVector{T}, q::AbstractVector{T}, A::AbstractMatrix{T}, μ::AbstractVector{T}) where {T}
	mul!(r, P, x)
	@. r += q

	mul!(r_temp, A', μ)
	@. r -= r_temp
	return nothing
end

# Einv * x
unscale_primal!(x::AbstractVector{T}, Einv::AbstractMatrix{T}) where {T} = lmul!(Einv, x)

# cinv * Dinv * x
function unscale_dual!(x::AbstractVector{T}, Dinv::AbstractMatrix{T}, cinv) where {T}
		lmul!(Dinv, x)
		@. x *= cinv
		return nothing
end

function calculate_residuals!(ws::COSMO.Workspace,  IGNORESCALING_FLAG::Bool = false)
	# reuse vectors from temporary workspace
	r_prim = ws.utility_vars.vec_m
	r_dual = ws.utility_vars.vec_n
	r_temp = ws.utility_vars.vec_n2

	# r_prim = A * x + s - b
	primal_kkt_condition!(r_prim, ws.p.A, ws.vars.x, ws.vars.s, ws.p.b)

	# ∇f0 + ∑ νi ∇hi(x*) == 0 condition
	# r_dual = P * x + q - A' * μ
	dual_kkt_condition!(r_dual, r_temp, ws.p.P, ws.vars.x, ws.p.q, ws.p.A, ws.vars.μ)

	if (ws.settings.scaling != 0 && !IGNORESCALING_FLAG)
		# unscale primal residual: r_prim = Einv * ( A * x + s - b)
		unscale_primal!(r_prim, ws.sm.Einv)

		# unscale dual residual: r_dual = cinv * Dinv * (P * x + q - A' * μ)
		unscale_dual!(r_dual, ws.sm.Dinv, ws.sm.cinv)
	end

	return norm(r_prim, Inf), norm(r_dual, Inf)

end

function max_res_component_norm(ws::COSMO.Workspace, IGNORESCALING_FLAG::Bool = false)
		r_prim = ws.utility_vars.vec_m
		r_dual = ws.utility_vars.vec_n

		# Determine if the components inside max{ } have to be unscaled
		unscale = (ws.settings.scaling != 0 && !IGNORESCALING_FLAG)

		# Calculate: max { ||Einv * Ax ||_Inf, ||Einv * s||_Inf, ||Einv * b||_Inf }
		# Einv * A * x
		mul!(r_prim, ws.p.A, ws.vars.x)
		unscale && unscale_primal!(r_prim, ws.sm.Einv)
		max_norm_prim = norm(r_prim, Inf)

		# Einv * s
		@. r_prim = ws.vars.s
		unscale && unscale_primal!(r_prim, ws.sm.Einv)
		max_norm_prim = max.(max_norm_prim, norm(r_prim, Inf))

		# Einv * b
		@. r_prim = ws.p.b
		unscale && unscale_primal!(r_prim, ws.sm.Einv)
		max_norm_prim = max.(max_norm_prim, norm(r_prim, Inf))

		# Calculate: max { ||cinv * Dinv * P * x ||_Inf, ||cinv * Dinv * q||_Inf, ||cinv * Dinv * A' * μ||_Inf }
		# ||cinv * Dinv * P * x ||_Inf
		mul!(r_dual, ws.p.P, ws.vars.x)
		unscale && unscale_dual!(r_dual, ws.sm.Dinv, ws.sm.cinv)
		max_norm_dual = norm(r_dual, Inf)

		# ||cinv * Dinv * q||_Inf
		@. r_dual = ws.p.q
		unscale && unscale_dual!(r_dual, ws.sm.Dinv, ws.sm.cinv)
		max_norm_dual = max.(max_norm_dual, norm(r_dual, Inf))

		# ||cinv * Dinv * A' * μ||_Inf
		mul!(r_dual, ws.p.A', ws.vars.μ)
		unscale && unscale_dual!(r_dual, ws.sm.Dinv, ws.sm.cinv)
		max_norm_dual = max.(max_norm_dual, norm(r_dual, Inf))

	return max_norm_prim, max_norm_dual
end

function has_converged(ws::COSMO.Workspace, r_prim::Float64, r_dual::Float64)
	max_norm_prim, max_norm_dual = max_res_component_norm(ws)
	settings = ws.settings
	ϵ_prim = settings.eps_abs + settings.eps_rel * max_norm_prim
	ϵ_dual = settings.eps_abs + settings.eps_rel * max_norm_dual

	# if an optimal objective value was specified for the problem check if current solution is within specified accuracy
	obj_true_FLAG = true
	if !isnan(settings.obj_true)
		current_cost = ws.sm.cinv[] * (1/2 * ws.vars.x' * ws.p.P * ws.vars.x + ws.p.q' * ws.vars.x)[1]
		obj_true_FLAG = abs(settings.obj_true - current_cost) <= settings.obj_true_tol
	end

	return (r_prim < ϵ_prim  && r_dual < ϵ_dual && obj_true_FLAG)
end


using OrdinaryDiffEq

function SinCosDiffEqToSolve!(ddu,du,u,p,t)
    ddu[1] = -u[1]
end

prob = SecondOrderODEProblem(SinCosDiffEqToSolve!,[2.0],[3.0],(0.0,1.0))
sol = solve(prob, Tsit5())
sol = solve(prob, Vern9())
sol = solve(prob, ROCK4())
sol = solve(prob, Rodas4())
sol = solve(prob, AutoVern9(Rodas4()))
sol = solve(prob, TRBDF2())


# A positive fraction whose numerator is less than its denominator is called a
# proper fraction. For any denominator, d, there will be d−1 proper fractions;
# for example, with d = 12: 1/12, 2/12, 3/12, 4/12, 5/12, 6/12, 7/12, 8/12,
# 9/12, 10/12, 11/12.
#
# We shall call a fraction that cannot be cancelled down a resilient fraction.
# Furthermore we shall define the resilience of a denominator, R(d), to be the
# ratio of its proper fractions that are resilient; for example, R(12) = 4/11.
# In fact, d = 12 is the smallest denominator having a resilience R(d) < 4/10.
#
# Find the smallest denominator d, having a resilience R(d) < 15499/94744.

using ProjectEulerSolutions

# Slower, explicitly calculate the gcd for each fraction, skipping some when
# the number is a multiple of 2.  This is quite slow when n gets large, so we
# will not use it in the benchmarking.
function resilience_slow(n::Integer)::Float64
    count = 1
    step = 1
    if rem(n,2) == 0
        step = 2
    end
    for i in 1:step:n-2
        if gcd(i,n) == 1
            count += 1
        end
    end
    return count / (n-1)
end

# We are actually finding the solution to Euler's totient function divided by
# n-1.
function resilience_fast(n::Integer)::Float64
    return totient(n) / (n - 1)
end

# Brute force approach, with some logic.  The minimum resilience values will be
# when the denominator is a multiple of primes.  We find which multiple
# satisfies the initial constraint and then iterate through smaller
# multiplicative values.
function p243solution(thresh::Float64=0.4)::Integer
    primes = sieve_eratosthenes(30)
    index = 2
    while resilience_fast(prod(primes[1:index])) > thresh
        index += 1
    end
    starting_val = prod(primes[1:index-1])
    mult = 2
    while resilience_fast(mult * starting_val) > thresh
        mult += 1
    end
    return mult * starting_val
end

p243 = Problems.Problem(p243solution)

Problems.benchmark(p243, 15499/94744)


using DataFrames
using CSV
using GLM
using Lathe

#df_gas = DataFrame(CSV.File("../resources/data1.csv"))
df_gas = DataFrame(CSV.File("../resources/data1_v1.csv"))

colnames = Symbol[]
for i in string.(names(df_gas))
    i = replace(strip(i), " " => "_")
    i = replace(i, "-" => "")
    i = replace(i, "/" => "")
    i = replace(i, "(" => "")
    i = replace(i, ")" => "")
    i = replace(i, "%" => "")
    i = replace(i, "," => "")
    i = replace(i, "*" => "")
    println(i)
    push!(colnames, Symbol(i))
    #push!(colnames, Symbol(replace(replace(replace(strip(i), " " => "_"), "-" => "_"), "/" => "_")))
end

rename!(df_gas, colnames)
println(names(df_gas))
formulasGasT = FormulaTerm[]

fm0 = @formula(Electrical_Power_MW ~ Percent_of_Gas_Turbine_Electrical_Load)
push!(formulasGasT, fm0)
#fm1 = @formula(Natural_Gas_Flow_lbhr ~ Percent_of_Gas_Turbine_Electrical_Load)
fm1 = @formula(HHV_Heat_Input_MMBTUhr ~ Percent_of_Gas_Turbine_Electrical_Load)
push!(formulasGasT, fm1)
fm2 = @formula(CO2_Emissions_lbhr ~ Percent_of_Gas_Turbine_Electrical_Load)
push!(formulasGasT, fm2)
fm3 = @formula(HP_Turbine_Power_kW ~ Percent_of_Gas_Turbine_Electrical_Load)
push!(formulasGasT, fm3)
fm4 = @formula(IP_Turbine_Power_kW ~ Percent_of_Gas_Turbine_Electrical_Load)
push!(formulasGasT, fm4)
fm5 = @formula(Total_kW ~ Percent_of_Gas_Turbine_Electrical_Load) # auxiliary
push!(formulasGasT, fm5)

coefReg = Vector{Float64}[]
coef_gas_df = DataFrame(y = String[], b = Float64[], a = Float64[])
for i in formulasGasT
    linReg = lm(i, df_gas)
    y = string(terms(i)[1])
    b = coef(linReg)[1]
    a = coef(linReg)[2]
    push!(coef_gas_df, (y, b, a))
    println(linReg)
    println(deviance(linReg))
    println(coef(linReg))
end

CSV.write("gas_coeffs.csv", coef_gas_df)

df_steam_t = DataFrame(CSV.File("../resources/data2.csv"))
coln_steam = Symbol[]
for i in string.(names(df_steam_t))
    i = replace(strip(i), " " => "_")
    i = replace(i, "-" => "")
    i = replace(i, "/" => "")
    i = replace(i, "(" => "")
    i = replace(i, ")" => "")
    i = replace(i, "%" => "")
    i = replace(i, "," => "")
    i = replace(i, "*" => "")
    i = replace(i, "+" => "")
    println(i)
    push!(coln_steam, Symbol(i))
end

rename!(df_steam_t, coln_steam)
#LP_Turbine_Power_kW
#Condenser_Duty_MMBTUhr
#DAC_Steam_Duty_MMBTUhr
#Total_ST_and_BOP_Auxiliary_Loads_kW
formulasSteamT = FormulaTerm[]
#
fm0st = @formula(LP_Turbine_Power_kW ~ Percent_of_Gas_Turbine_Electrical_Load)
push!(formulasSteamT, fm0st)
fm1st = @formula(Duty_of_CC_Reboiler_MMBTUhr ~ Percent_of_Gas_Turbine_Electrical_Load)
push!(formulasSteamT, fm1st)
fm2st = @formula(DAC_Steam_Duty_MMBTUhr ~ Percent_of_Gas_Turbine_Electrical_Load)
push!(formulasSteamT, fm2st)
fm3st = @formula(Total_ST_and_BOP_Auxiliary_Loads_kW ~ Percent_of_Gas_Turbine_Electrical_Load)
push!(formulasSteamT, fm3st)
coef_steam_df = DataFrame(y = String[], b = Float64[], a = Float64[])
for i in formulasSteamT
    linReg = lm(i, df_steam_t)
    y = string(terms(i)[1])
    b = coef(linReg)[1]
    a = coef(linReg)[2]
    push!(coef_steam_df, (y, b, a))
    println(linReg)
    println(deviance(linReg))
    println(coef(linReg))
end

CSV.write("steam_coeffs.csv", coef_steam_df)


lines = readlines("inputs/05.txt")
function day_5(lines)
  grid = zeros(Int64, 1000,1000)
  pairs = []

  for line in lines
    x1 = parse(Int64, split(line,',')[1])
    y2 = parse(Int64, split(line,',')[3])
    mid = split(line,',')[2]
    y1 = parse(Int64, split(mid, ' ')[1])
    x2 = parse(Int64, split(mid, ' ')[3])
    push!(pairs, [x1, y1, x2, y2])
  end

  for p in pairs
    x1, y1, x2, y2 = p
    if x1 == x2
      for y in min(y1, y2):max(y1, y2)
        grid[x1, y] += 1
      end
    elseif y1 == y2
      for x in min(x1, x2):max(x1, x2)
        grid[x, y1] += 1
      end
    end
  end
  count(i -> i > 1, grid)
end

println(day_5(lines))


# Used "Conjugate Bayesian analysis of the Gaussian distribution" by Murphy as
# a reference.  Note that there were some typos in that document so the code
# here may not correspond exactly.

struct NormalInverseWishart{T<:Real,V<:AbstractVector{T},M<:AbstractMatrix{T}} <: ContinuousUnivariateDistribution
    dim::Int
    zeromean::Bool
    mu::V
    kappa::T              # This scales precision (inverse covariance)
    Lamchol::Cholesky{T,M}  # Covariance matrix (well, sqrt of one)
    nu::T

    function NormalInverseWishart{T}(mu::AbstractVector{T}, kappa::T,
                                     Lamchol::Cholesky{T,M}, nu::T) where {T<:Real, M<:AbstractMatrix{T}}
        # Probably should put some error checking in here
        d = length(mu)
        zmean::Bool = true
        for i = 1:d
            if !iszero(mu[i])
                zmean = false
                break
            end
        end
        new{T,typeof(mu),M}(d, zmean, mu, T(kappa), Lamchol, T(nu))
    end
end

function NormalInverseWishart(mu::AbstractVector{U}, kappa::Real,
                                Lamchol::Cholesky{S}, nu::Real) where {S<:Real, U<:Real}
    T = promote_type(eltype(mu), typeof(kappa), typeof(nu), S)
    return NormalInverseWishart{T}(convert(AbstractVector{T}, mu),
                                   T(kappa),
                                   Cholesky{T}(Lamchol),
                                   T(nu))
end

function NormalInverseWishart(mu::AbstractVector{U}, kappa::Real,
                              Lambda::AbstractMatrix{S}, nu::Real) where {S<:Real, U<:Real}
    T = promote_type(eltype(mu), typeof(kappa), typeof(nu), S)
    return NormalInverseWishart{T}(convert(AbstractVector{T}, mu),
                                   T(kappa),
                                   Cholesky{T}(cholesky(Lambda)),
                                   T(nu))
end

function insupport(::Type{NormalInverseWishart}, x::AbstractVector{T}, Sig::AbstractMatrix{T}) where T<:Real
    return (all(isfinite, x) &&
           size(Sig, 1) == size(Sig, 2) &&
           isApproxSymmmetric(Sig) &&
           size(Sig, 1) == length(x) &&
           hasCholesky(Sig))
end

"""
    params(niw::NormalInverseWishart)

The parameters are
* μ::AbstractVector{T<:Real} the expected mean vector
* Λchol::Cholesky{T<:Real} the Cholesky decomposition of the scale matrix
* κ::T<:Real prior pseudocount for the mean
* ν::T<:Real prior pseudocount for the covariance
"""
params(niw::NormalInverseWishart) = (niw.mu, niw.Lamchol, niw.kappa, niw.nu)

pdf(niw::NormalInverseWishart, x::AbstractVector{T}, Sig::AbstractMatrix{T}) where T<:Real =
        exp(logpdf(niw, x, Sig))

function logpdf(niw::NormalInverseWishart, x::AbstractVector{T}, Sig::AbstractMatrix{T}) where T<:Real
    if !insupport(NormalInverseWishart, x, Sig)
        return -Inf
    else
        p = size(x, 1)

        nu = niw.nu
        kappa = niw.kappa
        mu = niw.mu
        Lamchol = niw.Lamchol
        hnu = 0.5 * nu
        hp = 0.5 * p

        # Normalization
        logp::T = hnu * logdet(Lamchol)
        logp -= hnu * p * log(2.)
        logp -= logmvgamma(p, hnu)
        logp -= hp * (log(2.0*pi) - log(kappa))

        # Inverse-Wishart
        logp -= (hnu + hp + 1.) * logdet(Sig)
        logp -= 0.5 * tr(Sig \ convert(AbstractMatrix, Lamchol))

        # Normal
        z = niw.zeromean ? x : x - mu
        logp -= 0.5 * kappa * invquad(PDMat(Sig), z)

        return logp

    end
end

function rand(niw::NormalInverseWishart)
    Sig = rand(InverseWishart(niw.nu, niw.Lamchol))
    mu = rand(MvNormal(niw.mu, Sig ./ niw.kappa))
    return (mu, Sig)
end


@testset "ChebyshevPoints" begin
    @test_throws Exception ChebyshevPoints{Int}(2)
    @test eltype(@inferred(ChebyshevPoints(2))) == Float64

    for T in Floats
        @test eltype(@inferred(ChebyshevPoints{T}(3))) == T
        @test typeof(@inferred(ChebyshevPoints{T}(3)[1])) == T

        @test ChebyshevPoints{T}(0) ≈ T[]
        @test ChebyshevPoints{T}(1) ≈ T[0]  atol = eps(real(T))
        @test ChebyshevPoints{T}(2) ≈ T[-1,1]
        @test ChebyshevPoints{T}(3) ≈ T[-1,0,1]
        @test ChebyshevPoints{T}(4) ≈ T[-1,-0.5,0.5,1]
        @test ChebyshevPoints{T}(5) ≈ T[-1,-1/sqrt(T(2)),0,1/sqrt(T(2)),1]
    end
end


using Polynomials
using Jacobi
using Test

@testset "Jacobi.jl" begin
    include("testjacobi.jl")
    include("testgauss.jl")
end