using Random, Simplices, Test
using Simplices: Vertex, Simplex, windingnumber, windingangle
using Simplices: centre, assessconvergence, position, value
using Simplices: bestvertex, issortedbyangle, hypervolume

@testset "Simplices tests" begin

  Random.seed!(0)

  @testset "Vertices" begin
    v1 = Vertex([0.0, 1.0], 2.0 - 3.0im)
    v2 = Vertex([0.0, 1.0], 2.0 - 3.0im)
    @test !(v1 == v2)
    @test isequal(v1, v2)
  end

  @testset "Simplex winding number" begin

    irrelevant = [0.0, 0.0]
    @testset "Simplex encloses zero" begin
      v1 = Vertex(irrelevant, 1.0 - im)
      v2 = Vertex(irrelevant, 0.0 + im)
      v3 = Vertex(irrelevant, -1.0 - im)
      encloseszero = Simplex([v1, v2, v3])
      @test isapprox(1, abs(windingangle(encloseszero)) / (2π))
      @test windingnumber(encloseszero) == 1
    end

    @testset "Simplex doesn't enclose zero" begin
      v1 = Vertex(irrelevant, 1.0 + im)
      v2 = Vertex(irrelevant, 2.0 + im)
      v3 = Vertex(irrelevant, 1.0 + im * 2)
      doesntenclosezero = Simplex([v1, v2, v3])
      @test abs(windingangle(doesntenclosezero)) / (2π) < 1.0e-3
      @test windingnumber(doesntenclosezero) == 0
      v1 = Vertex(irrelevant, 1.0 + im)
      v2 = Vertex(irrelevant, 0.0 + 2*im)
      v3 = Vertex(irrelevant, -1.0 + im)
      doesntenclosezero = Simplex([v1, v2, v3])
      @test abs(windingangle(doesntenclosezero)) / (2π) < 1.0e-3
      @test windingnumber(doesntenclosezero) == 0
      v1 = Vertex(irrelevant, 1.0 + 0im)
      v2 = Vertex(irrelevant, 0.0 + im)
      v3 = Vertex(irrelevant, 1.0 + 0im)
      doesntenclosezero = Simplex([v1, v2, v3])
      @test abs(windingangle(doesntenclosezero)) / (2π) < 1.0e-3
      @test windingnumber(doesntenclosezero) == 0
    end
  end

  @testset "Simplex convergence" begin
    @testset "Simplices with identical vertices is converged" begin
      dim = 2
      T = Float64
      U = ComplexF64
      pos = T[1.0, 1.0]
      val = rand(U)
      v1 = Vertex(pos, val)
      v2 = Vertex(pos, val)
      v3 = Vertex(pos, val)
      s = Simplex([v1, v2, v3])

      defaults = Simplices.convergenceconfig(dim, T)
      returncode = assessconvergence(s, defaults)
      @test returncode == :XTOL_REACHED
    end

    @testset "Simplices in a simplex eps apart are converged" begin
      dim = 2
      T = Float64
      U = ComplexF64
      v1 = Vertex([one(T), one(T)], rand(U))
      v2 = Vertex([one(T), one(T) + eps(T)], rand(U))
      v3 = Vertex([one(T) + eps(T), one(T)], rand(U))
      s = Simplex([v1, v2, v3])
      defaults = Simplices.convergenceconfig(dim, T)
      returncode = assessconvergence(s, defaults)
      @test returncode == :XTOL_REACHED
    end

    @testset "Simplices in a chain eps apart are converged" begin
      dim = 2
      T = Float64
      U = ComplexF64
      v1 = Vertex([one(T), one(T)], rand(U))
      v2 = Vertex([one(T), one(T) + eps(T)], rand(U))
      v3 = Vertex([one(T), one(T) + eps(T) + eps(T)], rand(U))
      s = Simplex([v1, v2, v3])
      defaults = Simplices.convergenceconfig(dim, T)
      returncode = assessconvergence(s, defaults)
      @test returncode == :XTOL_REACHED
      @assert all(isapprox(position(v1)[d], position(v2)[d],
                        atol=defaults[:xtol_abs][d],
                        rtol=defaults[:xtol_rel][d]) for d ∈ 1:dim)
      @assert all(isapprox(position(v2)[d], position(v3)[d],
                        atol=defaults[:xtol_abs][d],
                        rtol=defaults[:xtol_rel][d]) for d ∈ 1:dim)
      @assert !all(isapprox(position(v1)[d], position(v3)[d],
                        atol=defaults[:xtol_abs][d],
                        rtol=defaults[:xtol_rel][d]) for d ∈ 1:dim)
      @test returncode == :XTOL_REACHED
    end

    @testset "Simplices vertices are in order" begin
      for _ ∈ 1:10
        dim = 2
        T = Float64
        U = ComplexF64
        v1 = Vertex(rand(dim), rand(U))
        v2 = Vertex(rand(dim), rand(U))
        v3 = Vertex(rand(dim), rand(U))
        s = Simplex([v1, v2, v3])
        p = s.permabs
        @test abs(value(s[p[1]])) < abs(value(s[p[2]]))
        @test abs(value(s[p[2]])) < abs(value(s[p[3]]))
        @test issortedbyangle(s)
      end
    end
  end

  @testset "Shapes" begin
    @testset "centres" begin
      v1 = Vertex([0.0, 0.0], one(ComplexF64))
      v2 = Vertex([1.0, 0.0], one(ComplexF64))
      v3 = Vertex([0.0, 1.0], one(ComplexF64))
      s = Simplex([v1, v2, v3])
      @test all(centre(s) .== [1/3, 1/3])
    end

    @testset "extremas" begin
      s = Simplex(x->im, [1.0, 3.0], 1.0)
      exs = Simplices.extrema(s)
      @test exs[1] == (1.0, 2.0)
      @test exs[2] == (3.0, 4.0)
    end

    @testset "hypervolumes" begin
      x0 = rand(2)
      a, b = rand(2)
      v1 = Vertex(x0 .+ [0.0, 0.0], one(ComplexF64))
      v2 = Vertex(x0 .+ [a, 0.0], one(ComplexF64))
      v3 = Vertex(x0 .+ [0.0, b], one(ComplexF64))
      s = Simplex([v1, v2, v3])
      @test hypervolume(s) ≈ (a * b) / 2
      x0 = rand(3)
      a, b, c = rand(3)
      v1 = Vertex(x0 .+ [0, 0, 0], one(ComplexF64))
      v2 = Vertex(x0 .+ [a, 0, 0], one(ComplexF64))
      v3 = Vertex(x0 .+ [0, b, 0], one(ComplexF64))
      v4 = Vertex(x0 .+ [0, 0, c], one(ComplexF64))
      s = Simplex([v1, v2, v3, v4])
      @test hypervolume(s) ≈ (a * b * c) / 6
    end
  end

  @testset "Partition of hypercubes" begin
    @testset "unity" begin
      for dims in 1:6
        positions = Simplices.partitionunitypositions(dims, Float64)
        ss = [Simplex(x->one(ComplexF64), [p for p in ps]) for ps in positions]
        for i ∈ 1:10
          x = rand(dims)
          xins = [x ∈ s for s in ss]
          @test sum(xins) == 1 # can only exist in 1 simplex
        end
      end
    end
  @testset "random hypercube" begin
      for dims in 1:6
        a = rand(dims) .- 0.5
        b = a .+ rand(dims)
        ss = Simplices.partitionhypercube(x->1, a, b)
        for i ∈ 1:10
          x = a .+ rand(dims) .* (b .- a)
          xins = [x ∈ s for s in ss]
          @test sum(xins) == 1 # can only exist in 1 simplex
        end
      end
    end
   @testset "partition simplices don't overlap" begin
      for dims in 1:4 # anything higher is too expensive
        ss = Simplices.partitionhypercube(x->1, zeros(dims), ones(dims))
        for a in ss, b in ss
          if isequal(a, b)
            @test a ∈ b
          else
            @test !(a ∈ b)
          end
        end
      end
    end
  end

end


include("../examples/rijke.jl")
#include("../src/clvs.jl")
using PyPlot
using OrdinaryDiffEq
using LinearAlgebra
using JLD
using SharedArrays
using Distributed
function compute_rayleigh_criterion(n_samples)
    d_u = 3
    Jray = SharedArray{Float64}(n_samples)
    Jray .= 0.
    beta = LinRange(6.0,9.0,n_samples)
    nSteps = 20000
    nRunup = 400000
    answer = @distributed for i = 1:n_samples
    	println("beta = ", beta[i])
    	s = [beta[i], 0.2]
    	u = Rijke_ODE(rand(N), s, nRunup)
    	for j = 1:nSteps
        	u = Rijke_ODE(u, s, 1)
			Jray[i] += 0.5*sum(u[Ng+1:2*Ng].*zetaj.*u[Ng+1:2*Ng])/nSteps
    	end
    	@show Jray[i]
    end
    wait(answer)
    save("../data/attractor/Jray.jld", "beta", beta,
     "Jray", Jray)
end
function plot_rayleigh()
	filename = string("../data/attractor/",
					  "Jray_actual.jld")
	X = load(filename)
	Jray = X["Jray"]
	beta = X["beta"]
	fig, ax = subplots(1,1)
	ax.plot(beta, Jray, "b.", ms=4.0)
	ax_in = fig.add_axes([0.2,0.6,0.2,0.2])
	ax_in.plot(beta[6.7 .<= beta .<= 7.3],
			Jray[6.7 .<= beta .<= 7.3], "b.",ms=4.0)
	ax.xaxis.set_tick_params(labelsize=25)
	ax.yaxis.set_tick_params(labelsize=25)
	ax_in.xaxis.set_tick_params(labelsize=20)
	ax_in.yaxis.set_tick_params(labelsize=20)

	ax.set_xlabel(L"$\beta$", fontsize=25)
	ax.set_ylabel(L"$\langle J_{\rm ray}\rangle$", fontsize=25)
	ax.grid(true)
	ax_in.grid(true)
end


using TriangleMesh
using Plots
using Triplot

poly = polygon_regular(5)
poly = polygon_Lshape()
mesh = create_mesh(poly, set_area_max=true, quality_meshing=true)

function plotTriMesh(mesh)
    x,y = mesh.point[1,:], mesh.point[2,:]

    plot() # init plot
    for vertex_ids in eachcol(mesh.cell)
        ids = vcat(vertex_ids, vertex_ids[1])
        plot!(x[ids],y[ids],linecolor=:black)
    end
    display(plot!(legend=false,ratio=1))
end
plotTriMesh(mesh)

function to_rgba(x::UInt32)
    a = ((x & 0xff000000)>>24)/255
    b = ((x & 0x00ff0000)>>16)/255
    g = ((x & 0x0000ff00)>>8)/255
    r = (x & 0x000000ff)/255
    RGBA(r, g, b, a)
end

# plotting
x,y = mesh.point[1,:], mesh.point[2,:]
z,t = (@. sin(pi*x)*sin(pi*y)), mesh.cell
img = to_rgba.(Triplot.rasterize(x,y,z,t)')
plot([extrema(x)...], [extrema(y)...], img, yflip=false)


############# Controlled Gates and 2 qubit Gates#############
function CU!(C::Circuit, U::AbstractArray{ComplexF64,2},control::Int,target::Int)
    #Implentation of CU gate using Kronnecker products.
    Uf = sparse(U)

    a = length(C.qreg)

    if (target==control)
        error("Target qubit and control qubit must be Different!")
    elseif ((target>a) || (control>a)) || ((target<=0) || (control<=0))
        error("The target and control qubit must be within the size of the register!")
    end

    M1 = spzeros(ComplexF64,2,2)
    M2 = copy(M1)
    I2 = sparse(I*(1.0+0.0*im),2,2)

    conket1 = sparse([0.0+0.0*im 0.0+0.0*im;
                      0.0+0.0*im 1.0+0.0*im])
    conket0 = sparse([1.0+0.0*im 0.0+0.0*im
                      0.0+0.0*im 0.0+0.0*im])

    for i ∈ 1:a

        if (i==1)
            if (i==control)
                M1 = conket1
            elseif (i==target)
                M1 = Uf
            else
                M1 = I2
            end
        else
            if (i==control)
                M1 = kron(conket1,M1)
            elseif (i==target)
                M1 = kron(Uf,M1)
            else
                M1 = kron(I2,M1)
            end
        end
    end

    for i ∈ 1:a

        if (i==1==control)
            M2 = conket0
        elseif (i==1!==control)
            M2 = I2
        else
            if (i==control)
                M2 = kron(conket0,M2)
            else
                M2 = kron(I2,M2)
            end
        end
    end
    return push!(C.gates,M1+M2)
end

function CU!(C::Circuit,U::AbstractSparseArray{ComplexF64,Int},control::Int,target::Int)
    #Implentation of CU gate using Kronnecker products.
    a = length(C.qreg)

    if (target==control)
        error("Target qubit and control qubit must be Different!")
    elseif ((target>a) || (control>a)) || ((target<=0) || (control<=0))
        error("The target and control qubit must be within the size of the register!")
    end

    M1 = spzeros(ComplexF64,2,2)
    M2 = copy(M1)
    I2 = sparse(I*(1.0+0.0*im),2,2)

    conket1 = sparse([0.0+0.0*im 0.0+0.0*im;
                      0.0+0.0*im 1.0+0.0*im])
    conket0 = sparse([1.0+0.0*im 0.0+0.0*im
                      0.0+0.0*im 0.0+0.0*im])

    for i ∈ 1:a

        if (i==1)
            if (i==control)
                M1 = conket1
            elseif (i==target)
                M1 = U
            else
                M1 = I2
            end
        else
            if (i==control)
                M1 = kron(conket1,M1)
            elseif (i==target)
                M1 = kron(U,M1)
            else
                M1 = kron(I2,M1)
            end
        end
    end

    for i ∈ 1:a

        if (i==1==control)
            M2 = conket0
        elseif (i==1!==control)
            M2 = I2
        else
            if (i==control)
                M2 = kron(conket0,M2)
            else
                M2 = kron(I2,M2)
            end
        end
    end
    return push!(C.gates,M1+M2)
end

function CU!(C::Circuit,λ::Float64,control::Int,target::Int)
    #Implentation of CU gate using Kronnecker products.
    a = length(C.qreg)
    U = sparse([1.0  0.0
                0.0  exp(im*λ)])

    if (target==control)
        error("Target qubit and control qubit must be Different!")
    elseif ((target>a) || (control>a)) || ((target<=0) || (control<=0))
        error("The target and control qubit must be within the size of the register!")
    end

    M1 = spzeros(ComplexF64,2,2)
    M2 = copy(M1)
    I2 = sparse(I*(1.0+0.0*im),2,2)

    conket1 = sparse([0.0+0.0*im 0.0+0.0*im;
                      0.0+0.0*im 1.0+0.0*im])
    conket0 = sparse([1.0+0.0*im 0.0+0.0*im
                      0.0+0.0*im 0.0+0.0*im])

    for i ∈ 1:a

        if (i==1)
            if (i==control)
                M1 = conket1
            elseif (i==target)
                M1 = U
            else
                M1 = I2
            end
        else
            if (i==control)
                M1 = kron(conket1,M1)
            elseif (i==target)
                M1 = kron(U,M1)
            else
                M1 = kron(I2,M1)
            end
        end
    end

    for i ∈ 1:a

        if (i==1==control)
            M2 = conket0
        elseif (i==1!==control)
            M2 = I2
        else
            if (i==control)
                M2 = kron(conket0,M2)
            else
                M2 = kron(I2,M2)
            end
        end
    end
    return push!(C.gates,M1+M2)
end

###Methods for doing a number of gates in a row
function CU!(C::Circuit,U::AbstractArray{ComplexF64,2},controls::AbstractVector{Int},targets::AbstractVector{Int})
    #convenience method for doing mutliple CU operations at a time.
    a = length(C.qreg)
    b = length(controls)
    if length(targets)!==b
        error("The set of controls and targets must be the same length or the operation is ambiguous.")
    elseif ((maximum(controls)>a) || (maximum(targets)>a)) || ((minimum(controls)<=0) || (minimum(targets)<=0))
        error("The target and control qubits must all lie within the register!")
    end

    Uf = sparse(U)
    for i in 1:b
        CU!(C,Uf,controls[i],targets[i])
    end
end

function CU!(C::Circuit,U::AbstractArray{ComplexF64,2},control::Int,targets::AbstractVector{Int})
    #convenience method for doing mutliple CU operations at a time.
    a = length(C.qreg)
    if ((control>a) || (maximum(targets)>a)) || ((control<=0) || (minimum(targets)<=0))
        error("The target and control qubits must all lie within the register!")
    end

    Uf = sparse(U)
    for i in targets
        CU!(C,Uf,control,i)
    end
end

function CU!(C::Circuit,U::AbstractArray{ComplexF64,2},controls::AbstractVector{Int},target::Int)
    #convenience method for doing mutliple CU operations at a time.
    a = length(C.qreg)
    if ((maximum(controls)>a) || (target>a)) || ((minimum(controls)<=0) || (target<=0))
        error("The target and control qubits must all lie within the register!")
    end

    Uf = sparse(U)
    for i in controls
        CU!(C,Uf,i,target)
    end
end

function CU!(C::Circuit,U::AbstractSparseArray{ComplexF64,Int},controls::AbstractVector{Int},targets::AbstractVector{Int})
    #convenience method for doing mutliple CU operations at a time.
    a = length(C.qreg)
    b = length(controls)
    if length(targets)!==b
        error("The set of controls and targets must be the same length or the operation is ambiguous.")
    elseif ((maximum(controls)>a) || (maximum(targets)>a)) || ((minimum(controls)<=0) || (minimum(targets)<=0))
        error("The target and control qubits must all lie within the register!")
    end

    for i in 1:b
        CU!(C,U,controls[i],targets[i])
    end
end

function CU!(C::Circuit,U::AbstractSparseArray{ComplexF64,Int},control::Int,targets::AbstractVector{Int})
    #convenience method for doing mutliple CU operations at a time.
    a = length(C.qreg)
    if ((control>a) || (maximum(targets)>a)) || ((control<=0) || (minimum(targets)<=0))
        error("The target and control qubits must all lie within the register!")
    end

    for i in targets
        CU!(C,U,control,i)
    end
end

function CU!(C::Circuit,U::AbstractSparseArray{ComplexF64,Int},controls::AbstractVector{Int},target::Int)
    #convenience method for doing mutliple CU operations at a time.
    a = length(C.qreg)
    if ((maximum(controls)>a) || (target>a)) || ((minimum(controls)<=0) || (target<=0))
        error("The target and control qubits must all lie within the register!")
    end

    for i in controls
        CU!(C,U,i,target)
    end
end

function CU!(C::Circuit,λ::Float64,controls::AbstractVector{Int},targets::AbstractVector{Int})
    #convenience method for doing mutliple CU operations at a time.
    a = length(C.qreg)
    b = length(controls)
    if length(targets)!==b
        error("The set of controls and targets must be the same length or the operation is ambiguous.")
    elseif ((maximum(controls)>a) || (maximum(targets)>a)) || ((minimum(controls)<=0) || (minimum(targets)<=0))
        error("The target and control qubits must all lie within the register!")
    end

    for i in 1:b
        CU!(C,λ,controls[i],targets[i])
    end
end

function CU!(C::Circuit,λ::Float64,control::Int,targets::AbstractVector{Int})
    #convenience method for doing mutliple CU operations at a time.
    a = length(C.qreg)
    if ((control>a) || (maximum(targets)>a)) || ((control<=0) || (minimum(targets)<=0))
        error("The target and control qubits must all lie within the register!")
    end

    for i in targets
        CU!(C,λ,control,i)
    end
end

function CU!(C::Circuit,λ::Float64,controls::AbstractVector{Int},target::Int)
    #convenience method for doing mutliple CU operations at a time.
    a = length(C.qreg)
    if ((maximum(controls)>a) || (target>a)) || ((minimum(controls)<=0) || (target<=0))
        error("The target and control qubits must all lie within the register!")
    end

    for i in controls
        CU!(C,λ,i,target)
    end
end

function CX!(C::Circuit,control::Int,target::Int)
    #Implentation of CNOT gate using Kronnecker products.
    X = sparse([0.0+0.0*im 1.0+0.0*im;
                1.0+0.0*im 0.0+0.0*im])

    CU!(C,X,control,target)
end

function CX!(C::Circuit,controls::AbstractVector{Int},targets::AbstractVector{Int})
    #convenience method for doing mutliple CX operations at a time.
    X = sparse([0.0+0.0*im 1.0+0.0*im;
                1.0+0.0*im 0.0+0.0*im])

    CU!(C,X,controls,targets)
end

function CX!(C::Circuit,control::Int,targets::AbstractVector{Int})
    #convenience method for doing mutliple CX operations at a time.
    X = sparse([0.0+0.0*im 1.0+0.0*im;
                1.0+0.0*im 0.0+0.0*im])

    CU!(C,X,control,targets)
end

function CX!(C::Circuit,controls::AbstractVector{Int},target::Int)
    #convenience method for doing mutliple CX operations at a time.
    X = sparse([0.0+0.0*im 1.0+0.0*im;
                1.0+0.0*im 0.0+0.0*im])

    CU!(C,X,controls,target)
end

function CZ!(C::Circuit,control::Int,target::Int)
    #Implentation of CZ gate using Kronnecker products.
    Z = sparse([1.0+0.0*im  0.0+0.0*im;
                0.0+0.0*im -1.0+0.0*im])

    CU!(C,Z,control,target)
end

function CZ!(C::Circuit,controls::AbstractVector{Int},targets::AbstractVector{Int})
    #convenience method for doing mutliple CZ operations at a time.
    Z = sparse([1.0+0.0*im  0.0+0.0*im;
                0.0+0.0*im -1.0+0.0*im])

    CU!(C,Z,controls,targets)
end

function CZ!(C::Circuit,control::Int,targets::AbstractVector{Int})
    #convenience method for doing mutliple CZ operations at a time.
    Z = sparse([1.0+0.0*im  0.0+0.0*im;
                0.0+0.0*im -1.0+0.0*im])

    CU!(C,Z,control,targets)
end

function CZ!(C::Circuit,controls::AbstractVector{Int},target::Int)
    #convenience method for doing mutliple CZ operations at a time.
    Z = sparse([1.0+0.0*im  0.0+0.0*im;
                0.0+0.0*im -1.0+0.0*im])

    CU!(C,Z,controls,target)
end

function Swap!(C::Circuit,b1::Int,b2::Int)
    #quick and dirty implmentation of Swap Gate
    a = length(C.qreg)
    if b1==b2
        error("This ins't really a swap, you just tried to swap a qubit with itself!")
    elseif ((b1>a) || (b2>a)) || ((b1<=0) || (b2<=0))
        error("Both bits must be within the size of the register!")
    end
    CX!(C,b2,b1)
    CX!(C,b1,b2)
    CX!(C,b2,b1)
end

function CCX!(C::Circuit,c1::Int,c2::Int,target::Int)
    #Implementation of the Toffoli gate using Kronnecker products.
    X = sparse([0.0+0.0*im 1.0+0.0*im;
                1.0+0.0*im 0.0+0.0*im])

    a = length(C.qreg)

    if (target==c1) || (target==c2)
        error("Target qubit and control qubits must be Different!")
    elseif (c1==c2)
        error("The control qubits must be different from one another!")
    elseif ((target>a) || (c1>a) || (c2>a)) || ((target<=0) || (c1<=0) ||(c2<=0) )
        error("The target and control qubits must be within the size of the register!")
    end

    M1 = spzeros(ComplexF64,2,2)
    M2 = copy(M1)
    M3 = copy(M1)
    M4 = copy(M1)
    I2 = sparse(I*(1.0+0.0*im),2,2)

    conket1 = sparse([0.0+0.0*im 0.0+0.0*im;
                      0.0+0.0*im 1.0+0.0*im])
    conket0 = sparse([1.0+0.0*im 0.0+0.0*im
                      0.0+0.0*im 0.0+0.0*im])

    for i ∈ 1:a

        if (i==1)
            if (i==c1)||(i==c2)
                M1 = conket1
            elseif (i==target)
                M1 = X
            else
                M1 = I2
            end
        else
            if (i==c1)||(i==c2)
                M1 = kron(conket1,M1)
            elseif (i==target)
                M1 = kron(X,M1)
            else
                M1 = kron(I2,M1)
            end
        end
    end

    for i ∈ 1:a

        if (i==1)
            if (i==c1)||(i==c2)
                M2 = conket0
            else
                M2 = I2
            end
        else
            if (i==c1)||(i==c2)
                M2 = kron(conket0,M2)
            else
                M2 = kron(I2,M2)
            end
        end
    end

    for i ∈ 1:a

        if (i==1)
            if (i==c1)
                M3 = conket0
            elseif (i==c2)
                M3 = conket1
            else
                M3 = I2
            end
        else
            if (i==c1)
                M3 = kron(conket0,M3)
            elseif (i==c2)
                M3 = kron(conket1,M3)
            else
                M3 = kron(I2,M3)
            end
        end
    end

    for i ∈ 1:a

        if (i==1)
            if (i==c1)
                M4 = conket1
            elseif (i==c2)
                M4 = conket0
            else
                M4 = I2
            end
        else
            if (i==c1)
                M4 = kron(conket1,M4)
            elseif (i==c2)
                M4 = kron(conket0,M4)
            else
                M4 = kron(I2,M4)
            end
        end
    end
    return push!(C.gates,M1+M2+M3+M4)

end

function CSwap!(C::Circuit,control,bit1,bit2)
    #Controlled swap for order finding algorithm
    CX!(C,bit2,bit1)
    CCX!(C,control,bit1,bit2)
    CX!(C,bit2,bit1)
end


using QRupdate
using LinearOperators

macro fprintf(args...)
    :( f()=@printf $(args...); f() )
end

type OptimResults
    # Statistics
    itns::Int
    time::Float64
    stat::Symbol
    exitmsg::AbstractString
    nprodA::Int
    nprodAt::Int
    numtrim::Int

    # State
    active::Vector{Int}
    state::Vector{Int}
    y::Vector{Float64}
    S::Matrix{Float64}
    R::Matrix{Float64}
end

# Constructor
function OptimResults(m::Int, n::Int)
    active = zeros(Int,0)
    state = zeros(Int,n)
    y = zeros(m)
    S = zeros(m,0)
    R = zeros(0,0)
    OptimResults(0,0,:exit_unknown,"Exit Unknown",0,0,0,
                 active, state, y, S, R)
end

"""
Solve the optimization problem

   DP:       minimize_y   - b'y  +  1/2 lambda y'y
             subject to   bl <= A'y <= bu
using given A,b,bl,bu,lambda.  When bl = -e, bu = e = ones(n,1),
DP is the dual Basis Pursuit problem

   BPdual:   maximize_y   b'y  -  1/2 lambda y'y
             subject to   ||A'y||_inf  <=  1.

General call to BPdual:
(active,state,x,y,S,R,info) = BPdual(A,b,bl,bu,lambda,active,state,y,S,R)

EXAMPLE
=======
Generate data:
[A,b,lambda] = gendata(1);

Solve BPdual:
[active,state,x,y,S,info] = BPdual(A,b,-1,1,lambda);

Resolve with a different lambda
lambda = lambda / 10;
[active,state,x,y,S,R,info] = BPdual(A,b,-1,1,lambda,active,state,y,S,R);

INPUT
=====
A          is an m-by-n linear operator.
b          is an m-vector.
lambda     is a nonnegative scalar.
bl, bu     are n-vectors.
active,state,y,S,R  may be empty (i.e., equal to [])
           or output from BPdual with some previous value of lambda.

OUTPUT
======
active     is an nact-vector of indices j with nact = length(active),
           listing which of the constraints bl(j) <= A(:,j)'y <= bu(j)
           is active.
state      is an n-vector giving the state of each dual constraint.
           state(j) = -2  if         Aj'*y < bl(j)  infeas below
                    = -1  if         Aj'*y = bl(j)  active below
                    = +1  if         Aj'*y = bu(j)  active above
                    = +2  if         Aj'*y > bu(j)  infeas above
                    =  0  if bl(j) < Aj'*y < bu(j)  inactive
x          is an nact-vector of primal solution values.
           The n-vector xx = zeros(n,1), xx(active) = x
           solves the primal BP problem (see below).
y          is an m-vector of BPdual solution values.
S          is the submatrix A(:,active).
R          is the Cholesky factorization of S.
info       is a structure with the following components:
           .time    total solution time (seconds)
           .stat    = 0  solution is optimal
                    = 1  too many iterations
                    = 2  current S is rank-deficient
                    = 3  dual infeasible point
           .itn     number of iterations
           .exitmsg exit message
           .nprodA  number of products with A
           .nprodAt number of products with A'
           .rNorm   2-norm of the final residual r = b - A(:,active)*x
           .xNorm   1-norm of the primal solution
           .numtrim number of constraints trimmed from final working set

BPdual Toolbox
Copyright 2008, Michael P. Friedlander and Michael A. Saunders
http://www.cs.ubc.ca/labs/scl/bpdual

The primal BP problem may be written

BP:        min_{x,y} ||x||_1 +  1/2 lambda||y||_2^2
              st         Ax     +      lambda  y  =  b.

If lambda > 0, we have the residual vector r = lambda y = b - Ax,
and BP is equivalent to the original BP-denoising problem

   BPDN:      min_{x,y} lambda ||x||_1 +  1/2 ||r||_2^2
              st         Ax     +      r  =  b.

09 Jun 2007: First version of BPdual.
             Michael Friedlander and Michael Saunders.
             A can be dense or sparse (but not an operator).
             S = A(:,active) holds active columns.
             Least-squares subproblems solved by x = S\g.
29 Aug 2007: "A" is now an operator.
09 Nov 2007: Generalize the l1 dual constraints from -e <= A'y <= e
             to bl <= A'y <= bu.
"""
function bpdual(# Positional arguments
                A::Union{AbstractMatrix,AbstractLinearOperator},
                b::Vector,
                λin::Real,
                bl::Vector,
                bu::Vector;

                # Keyword arguments
                inform::OptimResults = OptimResults(size(A)...),
                coldstart::Bool = true,
                trim::Int = 0,
                itnMax::Int = 10*maximum(size(A)),
                feaTol::Real = 5e-05,
                λmin::Real = √eps(1.),
                optTol::Real = 1e-05,
                gapTol::Real = 1e-06,
                pivTol::Real = 1e-12,
                actMax::Real = Inf,
                loglevel::Int = 1)

    time0 = time()

    # ------------------------------------------------------------------
    # Check input arguments.
    # ------------------------------------------------------------------
    m, n = size(A)

    active = inform.active
    state = inform.state
    y = inform.y
    S = inform.S
    R = inform.R

    coldstart = isempty(S)

    # ------------------------------------------------------------------
    # Grab input options and set defaults where needed.
    # ------------------------------------------------------------------
    tieTol = feaTol + 1e-8        # Perturbation to break ties
    λ = max(λin, λmin)

    # ------------------------------------------------------------------
    # Print log header.
    # ------------------------------------------------------------------
    if loglevel > 0
        @fprintf("\n")
        @fprintf(" %s\n",repeat("=",80))
        @fprintf(" %s\n","BPdual")
        @fprintf(" %s\n",repeat("=",80))
        @fprintf(" %-20s: %8i %5s"    ,"No. rows"          ,m       ,"")
        @fprintf(" %-20s: %8.2e\n"    ,"λ"                 ,λ          )
        @fprintf(" %-20s: %8i %5s"    ,"No. columns"       ,n       ,"")
        @fprintf(" %-20s: %8.2e\n"    ,"Optimality tol"    ,optTol     )
        @fprintf(" %-20s: %8i %5s"    ,"Maximum iterations",itnMax  ,"")
        @fprintf(" %-20s: %8.2e\n"    ,"Duality tol"       ,gapTol     )
        @fprintf(" %-20s: %8i %5s"    ,"Support trimming"  ,trim,    "")
        @fprintf(" %-20s: %8.2e\n"    ,"Pivot tol"         ,pivTol     )
    end

    # ------------------------------------------------------------------
    # Initialize local variables.
    # ------------------------------------------------------------------
    EXIT_STATUS =
        Dict(:exit_optimal=>"optimal solution found -- full Newton step",
             :exit_singular_ls=>"singular least-squares subproblem",
             :exit_iterations=>"too many iterations",
             :exit_infeasible=>"dual infeasible point",
             :exit_requested=>"user requested exit",
             :exit_actmax=>"max no. of active constraints reached",
             :exit_small_gap=>"optimal solution found -- small duality gap",
             :exit_unknown=>"unknown exit")
    eFlag = :exit_unknown
    itn       = 0
    step      = 0.0
    p         = 0                         # index of added constraint
    q         = 0                         # index of deleted constraint
    svar      = ""                        # ...and their string value.
    r         = zeros(m)
    zerovec   = zeros(n)
    nBrks     = 0
    numtrim   = 0
    nprodA    = 0
    nprodAt   = 0

    # ------------------------------------------------------------------
    # Cold/warm-start initialization.
    # ------------------------------------------------------------------
    if coldstart
        x = zeros(0)
        z = zeros(n)

    else
        # Trim active-set components that correspond to huge bounds,
        # and then find the min-norm change to y that puts the active
        # constraints exactly on their bounds.
        g = b - λ*y              # Compute steepest-descent dir
        triminf!(active,state,S,R,bl,bu,g)
        nact = length(active)
        x = zeros(nact)
        y = restorefeas(y, active, state, S, R, bl, bu)
        z = A'*y; nprodAt = nprodAt+1

    end

    # Make sure initial point is feasible.
    # bl ≤ z ≤ bu  -->  bl - z ≤ 0
    #              -->  z - bu ≤ 0
    for j in 1:n
        infeasible = (bl[j] - z[j] > feaTol) || (z[j] - bu[j] > feaTol)
        if infeasible
            eFlag = :exit_infeasible
        end
    end

    # ------------------------------------------------------------------
    # Main loop.
    # ------------------------------------------------------------------
    while true

        g = b - λ*y            # Steepest-descent direction

        condS = condest(R)

        if condS > 1e+10
            eFlag = :exit_singular_ls
            # Pad x with enough zeros to make it compatible with S.
            resize!(x, size(S,2))

        else
            x, dy = newtonstep(S, R, g, x, λ)

        end
        r = b - S*x

        # ------------------------------------------------------------------
        # Print to log.
        # ------------------------------------------------------------------
        yNorm = norm(y,2)
        rNorm = norm(r,2)
        xNorm = norm(x,1)

        pObj, dObj, rGap = objectives(x,y,active,b,bl,bu,λ,rNorm,yNorm)
        nact = length(active)   # size of the working set

        if loglevel > 0
            if q != 0
                svar = @sprintf("%8i%1s", q, "-")
            elseif p != 0
                svar = @sprintf("%8i%1s", p ," ")
            else
                svar = @sprintf("%8s%1s"," "," ")
            end
            if mod(itn,50)==0
                @fprintf("\n %4s  %8s %9s %7s %11s %11s %11s %7s %7s\n",
                            "Itn","step","Add/Drop","Active","rNorm2",
                            "xNorm1","Objective","RelGap","condS")
            end
            @fprintf(
            " %4i  %8.1e %8s %7i %11.4e %11.4e %11.4e %7.1e %7.1e\n",
                    itn, step, svar, nact, rNorm, xNorm,
                    dObj, rGap, condS)
        end

        # ------------------------------------------------------------------
        # Check exit conditions.
        # ------------------------------------------------------------------
        if eFlag == :exit_unknown && rGap < gapTol && nact > 0
            eFlag = :exit_small_gap
        end
        if eFlag == :exit_unknown && itn >= itnMax
            eFlag = :exit_iterations
        end
        if eFlag == :exit_unknown && nact >= actMax
            eFlag = :exit_actmax
        end

        # If this is an optimal solution, trim multipliers before exiting.
        if eFlag == :exit_optimal || eFlag == :exit_small_gap
            if trim == 1
                # Optimal trimming. λ0 may be different from λ.
                # Recompute gradient just in case.
                g = b - λin*y
                trimx!(x,S,R,active,state,g,b,λ,feaTol,optTol,loglevel)
                numtrim = nact - length(active)
                nact = length(active)

            elseif trim == 2
                # Threshold trimming.
                # Not yet implemented.
            end
        end

        # Act on any live exit conditions.
        if eFlag != :exit_unknown
            break
        end

        # --------------------------------------------------------------
        # New iteration starts here.
        # --------------------------------------------------------------
        itn += 1
        p = q = 0

        # --------------------------------------------------------------
        # Compute search direction dz.
        # --------------------------------------------------------------
        if norm(dy,Inf) < eps(1.0)       # save a mat-vec prod if dy is 0
            dz = zeros(n)
        else
            dz = A'*dy
            nprodAt += 1
        end

        # ---------------------------------------------
        # Find step to the nearest inactive constraint
        # ---------------------------------------------
        pL, stepL, pU, stepU = step_to_bnd(z, dz, bl, bu, state, tieTol, pivTol)

        step = min(1.0, stepL, stepU)
        hitBnd = step < 1.0
        if hitBnd                    # We bumped into a new constraint
            if step == stepL
                p = pL
                state[p] = -1
            else
                p = pU
                state[p] = +1
            end
        end

        y += step*dy                 # Update dual variables.
        if mod(itn,50)==0
            z = A'*y                 # Occasionally set z = A'y directly.
            nprodAt += 1
        else
            z += step*dz
        end

        if hitBnd
            zerovec[p] = 1            # Extract a = A(:,p)
            a = A*zerovec
            nprodA += 1
            zerovec[p] = 0
            R = qraddcol(S,R,a)       # Update R
            S = [S a]                 # Expand S, active
            push!(active, p)
            push!(x, 0)

        else
            #------------------------------------------------------
            # step = 1.  We moved to a minimizer for the current S.
            # See if we can delete an active constraint.
            #------------------------------------------------------
            drop = false
            if length(active) > 0
                dropl = (state[active].==-1) & (x .> +optTol)
                dropu = (state[active].==+1) & (x .< -optTol)
                dropa = dropl | dropu
                drop = any(dropa)
            end

            # Delete an active constraint.
            if drop
                nact = length(active)
                qa = indmax(abs(x.*dropa))
                q = active[qa]
                state[q] = 0
                S = S[:, 1:nact .!= qa]
                deleteat!(active, qa)
                deleteat!(x, qa)
                R = qrdelcol(R,qa)

            else
                eFlag = :exit_optimal

            end

        end # if step < 1

  end # while true

  tottime = time() - time0

  if loglevel > 0
      @fprintf("\n EXIT BPdual -- %s\n\n",EXIT_STATUS[eFlag])
      @fprintf(" %-20s: %8i %5s","No. significant nnz",sparsity(x),"")
      @fprintf(" %-20s: %8i\n","Products with A",nprodA)
      @fprintf(" %-20s: %8i %5s","No. trimmed nnz",numtrim,"")
      @fprintf(" %-20s: %8i\n","Products with At",nprodAt)
      @fprintf(" %-20s: %8.1e %5s","Solution time (sec)",tottime,"")
      @fprintf("\n")
  end

   # Gather exit data.
   return (x, r, OptimResults(itn,
                              tottime,
                              eFlag,
                              EXIT_STATUS[eFlag],
                              nprodA,
                              nprodAt,
                              numtrim,
                              active, state, y, S, R)
           )
end # function bpdual

# ----------------------------------------------------------------------
# Find step to first bound.
# ----------------------------------------------------------------------
function step_to_bnd(z, dz, bl, bu, state, tieTol, pivTol)

    n = length(dz)
    stepL = stepLtie = stepU = stepUtie = Inf
    pL = pU = 0

    for i in 1:n

        sL = bl[i] -  z[i]
        sU =  z[i] - bu[i]

        # skip variables that aren't free.
        state[i] != 0 && continue

        # variable moves to LOWER bound.
        if dz[i] < -pivTol
            tmp = (sL-tieTol)/dz[i]
            if tmp < stepLtie
                stepLtie = tmp
                stepL = max(0, sL/dz[i])
                pL = i
            end
        end

        # variable moves to UPPER bound.
        if dz[i] > +pivTol
            tmp = -(sU-tieTol)/dz[i]
            if tmp < stepUtie
                stepUtie = tmp
                stepU = max(0, -sU/dz[i])
                pU = i
            end
        end

    end # if

    return (pL, stepL, pU, stepU)

end

# ----------------------------------------------------------------------
# Compute the primal and dual objective values, and the duality gap:
#
#    DP:  minimize_y   - b'y  +  1/2 λ y'y
#         subject to   bl <= A'y <= bu
#
#    PP:  minimize_x   bl'neg(x) + bu'pos(x) + 1/2 λ y'y
#         subject to   Ax + λ y = b.
# ----------------------------------------------------------------------
function objectives(x, y, active,
                    b, bl, bu,
                    λ, rNorm, yNorm)

    bigNum = 1e20

    if isempty(x)
        blx = 0.
        bux = 0.
    else
        blx = bl[active]
        blx[blx .< -bigNum] = 0
        blx = dot(blx, min(x,0))

        bux = bu[active]
        bux[bux .>  bigNum] = 0
        bux = dot(bux, max(x,0))
    end

    if λ > eps(1.0)
        pObj = blx + bux + rNorm^2/2    # primal objective
    else
        pObj = blx + bux
    end
    dObj  = λ*yNorm^2/2 - dot(b,y)         # dual   objective
    maxpd = max(1, pObj, dObj)
    rGap  = abs(pObj+dObj)/maxpd               # relative duality gap

    return (pObj, dObj, rGap)

end # function objectives

# ----------------------------------------------------------------------
# Compute a Newton step.  This is a step to a minimizer of the EQP
#
#   min   g'dy + 1/2 λ dy'dy   subj to   S'dy = 0.
#
# The optimality conditions are given by
#
#   [ -λ I   S ] [dy] = [ h ], with  h = b - lam y - S x,
#   [   S'     ] [dx]   [ 0 ]
#
# where x is an estimate of the Lagrange multiplier. Thus, dx solves
# min ||S dx - h||. On input, g = b - lam y. Alternatively, solve the LS problem
# min ||S  x - g||.
# ----------------------------------------------------------------------
function newtonstep(S, R, g, x, λ)
    m, n = size(S)
    if m==0 || n==0
        dy = g/λ            # Steepest descent
        return (x, dy)
    end
    x, dr = csne(R,S,g)     # LS problem
    if m > n                # Overdetermined system
        dy = dr/λ           # dy is the scaled residual
    else                    # System is square or underdetermined
        dy = zeros(m)       # Anticipate that the residual is 0
    end
    return (x, dy)
end # function newtonstep

# ----------------------------------------------------------------------
# Compute the infeasibility of z relative to bounds bl/bu.
#
#  (bl - z) <= 0  implies  z is   feasible wrt to lower bound
#           >  0  ...           infeasible ...
#  (z - bu) <= 0  implies  z is   feasible wrt to lower bound.
#           >  0  ...           infeasible ...
function infeasibilities(bl::Vector, bu::Vector, z::Vector)
    sL = bl - z
    sU = z - bu
    return (sL, sU)
end

# ----------------------------------------------------------------------
# Trim working constraints with "infinite" bounds.
function triminf!(active::Vector, state::Vector, S::Matrix, R::Matrix,
                  bl::Vector, bu::Vector, g::Vector, b::Vector)

    bigbnd = 1e10
    nact = length(active)

    # Generate a list of constraints to trim.
    tlistbl = find( state[active] .== -1 & bl[active] .< -bigbnd )
    tlistbu = find( state[active] .== +1 & bu[active] .> +bigbnd )
    tlist   = [tlistbl; tlistbu]

    if isempty(tlist)
        return
    end

    for q in tlist

        qa = active[q]          # Index of active constraint to delete
        nact = nact - 1
        S = S[:,1:size(S,2) .!= qa] # Delete column from S
        deleteat!(active, qa)   # Delete index from active set
        R = qrdelcol(R, qa)     # Recompute new QR factorization
        state[q] = 0            # Mark constraint as free
        x = csne(R, S, g)       # Recompute multipliers

        rNorm = norm(b - S*x)
        xNorm = norm(x, 1)
        @fprintf("%4i %4i %4i %10.3e %10.e %10.e %10.e\n",
                k,q,nact,bl[q],bu[q],rNorm,xNorm)
    end

end

# ----------------------------------------------------------------------
# n = sparsity(x,threshold) returns the number of elements in x that
# contain most of the weight of the vector.
function sparsity(x::Vector, threshold::Real = 0.9995)
    threshold = clamp(threshold, 0, 1)
    x = sort(abs(x), rev=true)
    tsum = sum(x)
    csum = 0.0
    j = 0
    while j < length(x)
        csum += x[j+1] / tsum
        if csum >= threshold
            j += 1
        end
    end
    return j
end

# ----------------------------------------------------------------------
# Estimate condition number of R'R.
function condest(R::Matrix)
    if isempty(R)
        condS = 1.0
    else
        n = size(R,1)
        rmin = +Inf
        rmax = -Inf
        for i in 1:n
            rii = R[i,i]
            rmin = rii < rmin ? rii : rmin
            rmax = rii > rmax ? rii : rmax
        end
        condS = rmax / rmin
    end
    return condS
end # function condest

# ----------------------------------------------------------------------
# Trim unneeded constraints from the active set.
#
# `trimx` assumes that the current active set is optimal, ie, 1. x has
# the correct sign pattern, and 2. z := A'y is feasible.
#
# Keep trimming the active set until one of these conditions is
# violated. Condition 2 isn't checked directly. Instead, we simply check
# if the resulting change in y, dy, is small.  It would be more
# appropriate to check that dz := A'dy is small, but we want to avoid
# incurring additional products with A. (Implicitly we're assuming that
# A has a small norm.)
function trimx!(x,S,R,active,state,g,b,λ,featol,opttol,loglevel)

    k = 0
    nact = length(active)
    xabs = abs(x)
    xmin, qa = findmin(xabs)
    gNorm = norm(g,Inf)

    while xmin < opttol

        e = sign(x.*(xabs .> opttol)) # Signs of significant multipliers
        q = active[qa] # Index of the corresponding constraint
        a = S[:,qa]    # Save the col from S in case we need to add it back.
        xsmall = x[qa] # Value of candidate multiplier

        # Trim quantities related to the small multiplier.
        deleteat!(e, qa)
        deleteat!(active, qa)
        S = S[:,1:nact .!= qa]
        R = qrdelcol(R, qa)

        # Recompute the remaining multipliers and their signs.
        (x,dy) = csne(R,S,g)           # min ||g - Sx||_2
        xabs = abs(x)
        et = sign(x.*(xabs .> opttol))
        dyNorm = norm(dy,Inf)/λ        # dy = (g - Sx) / lambda

        # Check if the trimmed active set is still optimal
        if any( et .!= e ) || (dyNorm / max(1,gNorm) > featol)
            R = qraddcol(S,R,a)
            S = [S a]
            push!(active, q)
            x = csne(R, S, g)
            break
        end

        if loglevel > 0 && mod(k,50)==0
            @fprintf("\n %4s  %8s %8s  %7s  %10s  %10s  %10s\n",
                    "Itn","xSmall","Add/Drop","Active","rNorm2",
                    "xNorm1","dyNorm")
        end

        # The trimmed x is still optimal.
        k += 1
        nact -= 1
        state[q] = 0               # Mark this constraint as free.
        rNorm = norm(b - S*x)
        xNorm = norm(x,1)

        # Grab the next canddate multiplier.
        xmin, qa = findmin(xabs)

        # Logging.
        if loglevel > 0
            @fprintf(" %4iT %8.1e %8i- %7i  %10.4e  %10.4e  %10.4e\n",
                    k,xsmall,q,nact,rNorm,xNorm,dyNorm)
        end

    end # while

    return (x,S,R,active,state)

end # function trimx


"""
An infinite length 2 feature data stream generator based on the sine function
as described in [1].

# Arguments:
* `stream::Stream`: A container of data stream options
* `classification_function`: One of the 4 functions `sine_func_[0, 1, 2, 3]`.
    Determines how the label is generated from features.

[1] Gama, Joao, et al.'s 'Learning with drift
detection.' Advances in artificial intelligence–SBIA 2004. Springer Berlin
Heidelberg, 2004. 286-295."

"""
struct SineGenerator{F<:Function} <: AbstractDataStream
    stream::Stream
    classification_function::F
    classification_function_idx::Int
end

function sine_func_0(a, b)
    return a >= sin(b) ? 0 : 1
end
function sine_func_1(a, b)
    return a < sin(b) ? 0 : 1
end
function sine_func_2(a, b)
    return a >= 0.5 + 0.3 * sin(3 * pi * b) ? 0 : 1
end
function sine_func_3(a, b)
    return a < 0.5 + 0.3 * sin(3 * pi * b) ? 0 : 1
end

function select_sine_func(classification_function_idx::Int64)
    classification_function = sine_func_0
    if classification_function_idx == 1
        classification_function = sine_func_1
    end
    if classification_function_idx == 2
        classification_function = sine_func_2
    end
    if classification_function_idx == 3
        classification_function = sine_func_3
    end
    return classification_function
end

function classify_instance(s::SineGenerator, X)
    s.classification_function(X...)
end

"""
Initialize a SineGenerator using an Int to select a classification_function.
"""
function SineGenerator(classification_function_idx::Int64)
    classification_function = select_sine_func(classification_function_idx)
    return SineGenerator(Stream("SineGenerator", 2), classification_function, classification_function_idx)
end

"""
Each feature is uniform random in range [0, 1]
"""
function next_num_feature(s::SineGenerator, f_id::Int)
    rand(s.stream.random_state)
end


module PolynomialBasesTests

using Test
using Gridap

p1 = Point(1,2)
p2 = Point(0,3)
p = [p1,p2]

T = VectorValue{3,Float64}
orders = (1,1)
basis = MonomialBasis(T,orders)

v = evaluate(basis,p)
g = evaluate(∇(basis),p)

test_basis(basis,p,v,g)

T = Float64
orders = (1,1)
basis = MonomialBasis(T,orders)

v = evaluate(basis,p)
g = evaluate(∇(basis),p)

test_basis(basis,p,v,g)

T = VectorValue{2,Float64}
filter(e,order) = sum(e) <= order
order = 1
basis = MonomialBasis(2,T,filter,order)

@test length(basis) == 6

T = VectorValue{2,Float64}
basis = GradMonomialBasis(T,2)
v = evaluate(basis,p)
g = evaluate(∇(basis),p)

@test length(basis) == 12

p = Point{3,Int}[(1,2,3),(3,2,4)]

T = VectorValue{3,Float64}
order = 1
basis = CurlGradMonomialBasis(T,order)
v = evaluate(basis,p)
g = evaluate(∇(basis),p)

@test length(basis) == 6

end # module


function DiscretizedAffineEstimator(mhist::MCEBHistogram, kernel::ContinuousUnivariateDistribution)
    DiscretizedAffineEstimator(mhist, x->pdf(kernel, x))
end

"""
    SincKernel(h)

Implements the `SincKernel` with bandwidth `h` to be used for kernel density estimation
through the `KernelDensity.jl` package. The sinc kernel is defined as follows:
```math
K_{\\text{sinc}}(x) = \\frac{\\sin(x)}{\\pi x}
```
It is not typically used for kernel density estimation, because this kernel is not
a density itself. However, it is particularly well suited to deconvolution problems
and estimation of very smooth densities because its Fourier transform is the following:
```math
K^*_{\\text{sinc}}(t) = \\mathbf 1( t \\in [-1,1])
```
"""
struct SincKernel <: ContinuousUnivariateDistribution
   h::Float64 #Bandwidth
end

_default_bandwidth(::Type{SincKernel}, m::Integer) = 1/sqrt(log(m))
SincKernel(; m = 1_000) = SincKernel(default_bandwidth(SincKernel, m))

cf(a::SincKernel, t) = one(Float64)*(-1/a.h <= t <= 1/a.h)

function pdf(a::SincKernel, t)
   if t==zero(Float64)
       return(one(Float64)/pi/a.h)
   else
       return(sin(t/a.h)/pi/t)
   end
end

"""
    DeLaValleePoussinKernel(h)

Implements the `DeLaValleePoussinKernel` with bandwidth `h` to be used for kernel density estimation
through the `KernelDensity.jl` package. The De La Vallée-Poussin kernel is defined as follows:
```math
K_V(x) = \\frac{\\cos(x)-\\cos(2x)}{\\pi x^2}
```
Its use case is similar to the [`SincKernel`](@ref), however it has the advantage of being integrable
(in the Lebesgue sense). Its Fourier transform is the following:
```math
K^*_V(t) = \\begin{cases}
 1, & \\text{ if } x\\in[-1,1] \\\\
 0, &\\text{ if } |t| \\geq 2 \\\\
 2-|t|,& \\text{ if } |t| \\in [1,2]
 \\end{cases}
```
"""
struct DeLaValleePoussinKernel <: ContinuousUnivariateDistribution
   h::Float64 #Bandwidth
end

_default_bandwidth(a::Type{DeLaValleePoussinKernel}, m::Integer) = 1/sqrt(log(m))

function DeLaValleePoussinKernel(; m = 1_000)
     DeLaValleePoussinKernel(default_bandwidth(DeLaValleePoussinKernel, m))
end

function cf(a::DeLaValleePoussinKernel, t)
   if abs(t * a.h) <= 1
       return(one(Float64))
   elseif abs(t * a.h) <= 2
       return(2*one(Float64) - abs(t * a.h))
   else
       return(zero(Float64))
   end
end

function pdf(a::DeLaValleePoussinKernel, t)
   if t==zero(Float64)
       return(3*one(Float64)/2/pi/a.h)
   else
       return(a.h*(cos(t/a.h)-cos(2*t/a.h))/pi/t^2)
   end
end

"""
    KDEInfinityBandOptions(; kernel=DeLaValleePoussinKernel
                             bandwidth = nothing,
                             a_min,
                             a_max,
                             nboot = 1000)

This struct contains hyperparameters that will be used for constructing a neighborhood
of the marginal density. The steps of the method (and corresponding hyperparameter meanings)
are as follows
* First a kernel density estimate ``\\bar{f}`` of the data is fit with `kernel` as the
kernel and `bandwidth` (the default `bandwidth = nothing` corresponds to automatic
bandwidth selection).
* Second, a Poisson bootstrap with `nboot` replication will be used to estimate a ``L_{\\infty}``
neighborhood ``c_m`` of the true density ``f`` which is such that with probability tending to 1:
```math
\\sup_{x \\in [a_{\\text{min}} , a_{\\text{max}}]} | \\bar{f}(x) - f(x)| \\leq c_m
```
Note that the bound is valid from `a_min` to `a_max`.

## Reference:
  > Paul Deheuvels and Gérard Derzko. Asymptotic certainty bands for kernel density
  > estimators based upon a bootstrap resampling scheme. In Statistical models and methods
  > for biomedical and technical systems, pages 171–186. Springer, 2008
"""
Base.@kwdef struct KDEInfinityBandOptions{T<:Real,
                                         S<:Union{T, Nothing}}
   a_min::T
   a_max::T
   npoints::Integer = 4096
   bandwidth::S = nothing
   kernel = DeLaValleePoussinKernel
   nboot::Integer = 1000
   η_infl::T = 1.1
   n_interp::Integer = 25
end

"""
    KDEInfinityBand

The result of running `StatsBase.fit(opt::KDEInfinityBandOptions, Xs)`. Here `opt` is an instance
of [`KDEInfinityBandOptions`](@ref) and `Xs` is a vector of samples distributed according to
a density ```f``.

## Fields:
* `a_min`,`a_max`, `kernel`: These are the same as the fields in `opt:KDEInfinityBandOptions`.
* `C∞`: This is the Poisson Bootstrap point estimate of ``\\sup_{x \\in [a_{\\text{min}} , a_{\\text{max}}]} | \\bar{f}(x) - f(x)|``
* `fitted_kde`: The fitted `KernelDensity` object.
"""
Base.@kwdef struct KDEInfinityBand{T<:Real}
    C∞::T
    a_min::T
    a_max::T
    fitted_kde
    interp_kde
    kernel
    η_infl::T = 1.1
    n_interp::Integer = 25
end

function _default_bandwidth(kernel, Xs::AbstractVector, ::Nothing)
    _default_bandwidth(kernel, Xs)
end

function _default_bandwidth(kernel, Xs::AbstractVector)
    _default_bandwidth(kernel, length(Xs))
end

function _default_bandwidth(kernel, Xs::AbstractVector, h::Number)
    h
end

function StatsBase.fit(opt::KDEInfinityBandOptions, Xs)
    a_min = opt.a_min
    a_max = opt.a_max
    npoints = opt.npoints
    kernel = opt.kernel
    bandwidth = _default_bandwidth(kernel, Xs, opt.bandwidth)
    nboot = opt.nboot
    tmp_res = certainty_banded_KDE(Xs, a_min, a_max; npoints = npoints,
                         kernel = kernel, bandwidth = bandwidth,
                         nboot = nboot)
    tmp_res = @set tmp_res.η_infl = opt.η_infl
    tmp_res = @set tmp_res.n_interp = opt.n_interp
    tmp_res
end

function certainty_banded_KDE(Xs, a_min, a_max;
                        npoints = 4096,
                        kernel=DeLaValleePoussinKernel,
                        bandwidth = _default_bandwidth(kernel, length(Xs)),
                        nboot = 1_000)

    kernel = kernel(bandwidth)
    h = bandwidth
    m = length(Xs)

    lo, hi = extrema(Xs)
    lo_kde, hi_kde = min(a_min, lo - 6*h), max(a_max, hi + 6*h)
    midpts = range(lo_kde, hi_kde; length = npoints)

    fitted_kde = kde(Xs, KernelDensity.UniformWeights(m), midpts, kernel)
    interp_kde = InterpKDE(fitted_kde)
    midpts_idx = findall( (midpts .>= a_min) .& (midpts .<= a_max) )
    C∞_boot = Vector{Float64}(undef, nboot)

    for k =1:nboot
        # Poisson bootstrap to estimate certainty band
        Z_pois = rand(Poisson(1), m)
        ws =  Weights(Z_pois/sum(Z_pois))
        f_kde_pois =  kde(Xs, ws, midpts, kernel)
        C∞_boot[k] = maximum(abs.(fitted_kde.density[midpts_idx] .- f_kde_pois.density[midpts_idx]))
    end

    C∞ = median(C∞_boot)

    KDEInfinityBand(C∞=C∞,
                    a_min=a_min,
                    a_max=a_max,
                    fitted_kde=fitted_kde,
                    interp_kde=interp_kde,
                    kernel=kernel)
end

function certainty_banded_KDE(Xs, Zs_discr::DiscretizedStandardNormalSamples; kwargs...)
    a_min, a_max = extrema(Zs_discr.mhist.grid)
    certainty_banded_KDE(Xs, a_min, a_max; kwargs...)
end

# plotting of infinite band

@recipe function f(ctband::KDEInfinityBand;
                        show_bands=true)
    y = ctband.fitted_kde.density
    x = ctband.fitted_kde.x
    xlim --> (ctband.a_min, ctband.a_max)
    ylab --> "Density"

    seriestype  -->  :path
    linewidth --> 2
    color --> :purple

    infty_bound = ctband.C∞ * ctband.η_infl

    if show_bands
        _ys_lower, _ys_upper = density_bands_to_ribbons(y, infty_bound)
        ribbons --> (_ys_lower, _ys_lower)
        fillalpha --> 0.2
    end

    xguide --> L"x"
    x,y
end

"""
    set_neighborhood(Zs_discr::DiscretizedStandardNormalSamples,
                          fkde::KDEInfinityBand)

Return a [`DiscretizedStandardNormalSamples`](@ref) instance identical to `Zs_discr` but with
`f_min`,`f_max`,`var_proxy` set based on the fit from `fkde`
(cf.[`KDEInfinityBand`](@ref)).

"""
function set_neighborhood(Zs_discr::DiscretizedStandardNormalSamples,
                          fkde::KDEInfinityBand)

    @unpack n_interp, η_infl = fkde
    grid = Zs_discr.mhist.grid
    fkde_interp = fkde.interp_kde.itp
    n = nobs(Zs_discr)

    f_max = zeros(length(grid) + 1)
    f_min = zeros(length(grid) + 1)
    var_proxy =  zeros(length(grid) + 1)

    C∞ = η_infl * fkde.C∞
    # fill in intermediate points first
    for i in 2:length(grid)
        x_left = grid[i-1]
        x_right = grid[i]
        h = x_right - x_left  # step(grid)...
        interp_res = fkde_interp(range(x_left, x_right, length=n_interp))
        f_max[i] = maximum(interp_res)*h + C∞*h
        f_min[i] = max( minimum(interp_res)*h - C∞*h, 0.0)
        var_proxy[i] = max(maximum(interp_res)*h, C∞*h)
    end

    #fill in first and last element

    # DKW constant with log(n) instead of log(2/alpha)
    C_dkw = sqrt(log(n)/2n)
    f_max[1] = pdf(Zs_discr.mhist)[1] + C_dkw
    f_min[1] = max(pdf(Zs_discr.mhist)[1] - C_dkw, 0.0)
    var_proxy[1] = f_max[1]
    f_max[end] = pdf(Zs_discr.mhist)[end] + C_dkw
    f_min[end] = max(pdf(Zs_discr.mhist)[end] - C_dkw, 0.0)
    var_proxy[end] = f_max[end]

    Zs_discr = @set Zs_discr.var_proxy = var_proxy
    Zs_discr = @set Zs_discr.f_min = f_min
    Zs_discr = @set Zs_discr.f_max = f_max
    Zs_discr
end


using DiffEqBayes
using OrdinaryDiffEq, ParameterizedFunctions, RecursiveArrayTools
using Stan, Mamba, Test

ProjDir = dirname(@__FILE__)

f1 = @ode_def_nohes LotkaVolterraTest4 begin
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a=>1.5 b=>1.0 c=>3.0 d=>1.0

cd(ProjDir) do

  isdir("tmp") && rm("tmp", recursive=true)
  isfile("LotkaVolterra2-1.pdf") &&
  rm("LotkaVolterra2-1.pdf")

  println("Four parameter case")
  u0 = [1.0,1.0]
  tspan = (0.0,10.0)
  prob1 = ODEProblem(f1,u0,tspan)
  sol = solve(prob1,Tsit5())
  t = collect(range(1, stop=10, length=10))
  randomized = VectorOfArray([(sol(t[i]) + .01randn(2)) for i in 1:length(t)])
  data = convert(Array,randomized)
  priors = [Truncated(Normal(1.5,0.01),0,2),Truncated(Normal(1.0,0.01),0,1.5),Truncated(Normal(3.0,0.01),0,4),Truncated(Normal(1.0,0.01),0,2)]

  bayesian_result = stan_inference(prob1,t,data,priors;num_samples=100,num_warmup=500,vars =("StanODEData",InverseGamma(4,1)))
  theta1 = bayesian_result.chain_results[:,["theta.1"],:]
  theta2 = bayesian_result.chain_results[:,["theta.2"],:]
  theta3 = bayesian_result.chain_results[:,["theta.3"],:]
  theta4 = bayesian_result.chain_results[:,["theta.4"],:]

  global sim2 = bayesian_result.chain_results[1:end,
    ["theta.1", "theta.2", "theta.3", "theta.4"], :]
  ## Plotting
  p = plot(sim2, [:trace, :mean, :density, :autocor], legend=true);
  draw(p, ncol=4, filename="LotkaVolterra2", fmt=:pdf)

  @test mean(theta1.value[:,:,1]) ≈ 1.5 atol=1e-1
  @test mean(theta2.value[:,:,1]) ≈ 1.0 atol=1e-1
  @test mean(theta3.value[:,:,1]) ≈ 3.0 atol=1e-1
  @test mean(theta4.value[:,:,1]) ≈ 1.0 atol=1e-1
end


struct NewtonTrustRegion
end

struct NewtonTrustRegionCache{Tx} <: AbstractSolverCache
    x::Tx
    xold::Tx
    r::Tx
    r_predict::Tx
    p::Tx
    p_c::Tx
    pi::Tx
    d::Tx
end
function NewtonTrustRegionCache(df)
    x = similar(df.x_f) # Current point
    xold = similar(x) # Old point
    r = similar(df.F)       # Current residual
    r_predict = similar(x)  # predicted residual
    p = similar(x)          # Step
    p_c = similar(x)          # Cauchy point
    pi = similar(x)
    d = similar(x)          # Scaling vector
    NewtonTrustRegionCache(x, xold, r, r_predict, p, p_c, pi, d)
end
macro trustregiontrace(stepnorm)
    esc(quote
        if tracing
            dt = Dict()
            if extended_trace
                dt["x"] = copy(cache.x)
                dt["f(x)"] = copy(value(df))
                dt["g(x)"] = copy(jacobian(df))
                dt["delta"] = delta
                dt["rho"] = rho
            end
            update!(tr,
                    it,
                    maximum(abs, value(df)),
                    $stepnorm,
                    dt,
                    store_trace,
                    show_trace)
        end
    end)
end

function dogleg!(p::AbstractArray{T}, p_c::AbstractArray{T}, p_i,
                 r::AbstractArray{T}, d::AbstractArray{T}, J::AbstractMatrix{T},
                 delta::Real) where T
    try
        copyto!(p_i, J \ vec(r)) # Gauss-Newton step
    catch e
        if isa(e, LAPACKException) || isa(e, SingularException)
            # If Jacobian is singular, compute a least-squares solution to J*x+r=0
            U, S, V = svd(convert(Matrix{T}, J)) # Convert to full matrix because sparse SVD not implemented as of Julia 0.3
            k = sum(S .> eps())
            mrinv = V * Matrix(Diagonal([1 ./ S[1:k]; zeros(eltype(S), length(S)-k)])) * U' # Moore-Penrose generalized inverse of J
            vecpi = vec(p_i)
            mul!(vecpi,mrinv,vec(r))
        else
            throw(e)
        end
    end
    rmul!(p_i, -one(T))

    # Test if Gauss-Newton step is within the region
    if wnorm(d, p_i) <= delta
        copyto!(p, p_i)   # accepts equation 4.13 from N&W for this step
    else
        # For intermediate we will use the output array `p` as a buffer to hold
        # the gradient. To make it easy to remember which variable that array
        # is representing we make g an alias to p and use g when we want the
        # gradient

        # compute g = J'r ./ (d .^ 2)
        g = p
        mul!(vec(g), transpose(J), vec(r))
        g .= g ./ d.^2

        # compute Cauchy point
        p_c .= -wnorm(d, g)^2 / sum(abs2, J*vec(g)) .* g

        if wnorm(d, p_c) >= delta
            # Cauchy point is out of the region, take the largest step along
            # gradient direction
            rmul!(g, -delta/wnorm(d, g))

            # now we want to set p = g, but that is already true, so we're done

        else
            # from this point on we will only need p_i in the term p_i-p_c.
            # so we reuse the vector p_i by computing p_i = p_i - p_c and then
            # just so we aren't confused we name that p_diff
            p_i .-= p_c
            p_diff = p_i

            # Compute the optimal point on dogleg path
            b = 2 * wdot(d, p_c, d, p_diff)
            a = wnorm(d, p_diff)^2
            tau = (-b + sqrt(b^2 - 4a*(wnorm(d, p_c)^2 - delta^2)))/(2a)
            p_c .+= tau .* p_diff
            copyto!(p, p_c)
        end
    end
end

function trust_region_(df::OnceDifferentiable,
                       initial_x::AbstractArray{T},
                       xtol::Real,
                       ftol::Real,
                       iterations::Integer,
                       store_trace::Bool,
                       show_trace::Bool,
                       extended_trace::Bool,
                       factor::Real,
                       autoscale::Bool,
                       cache = NewtonTrustRegionCache(df)) where T
    copyto!(cache.x, initial_x)
    value_jacobian!!(df, cache.x)
    cache.r .= value(df)
    check_isfinite(cache.r)

    it = 0
    x_converged, f_converged, converged = assess_convergence(value(df), ftol)
    delta = convert(real(T), NaN)
    rho = convert(real(T), NaN)
    if converged
        tr = SolverTrace()
        name = "Trust-region with dogleg"
        if autoscale
            name *= " and autoscaling"
        end

        return SolverResults(name,
        #initial_x, reshape(cache.x, size(initial_x)...), norm(cache.r, Inf),
        initial_x, copy(cache.x), norm(cache.r, Inf),
        it, x_converged, xtol, f_converged, ftol, tr,
        first(df.f_calls), first(df.df_calls))
    end

    tr = SolverTrace()
    tracing = store_trace || show_trace || extended_trace
    @trustregiontrace convert(real(T), NaN)
    nn = length(cache.x)
    if autoscale
        for j = 1:nn
            cache.d[j] = norm(view(jacobian(df), :, j))
            if cache.d[j] == zero(cache.d[j])
                cache.d[j] = one(cache.d[j])
            end
        end
    else
        fill!(cache.d, one(real(T)))
    end

    delta = factor * wnorm(cache.d, cache.x)
    if delta == zero(delta)
        delta = factor
    end

    eta = convert(real(T), 1e-4)

    while !converged && it < iterations
        it += 1

        # Compute proposed iteration step
        dogleg!(cache.p, cache.p_c, cache.pi, cache.r, cache.d, jacobian(df), delta)

        copyto!(cache.xold, cache.x)
        cache.x .+= cache.p
        value!(df, cache.x)

        # Ratio of actual to predicted reduction (equation 11.47 in N&W)
        mul!(vec(cache.r_predict), jacobian(df), vec(cache.p))
        cache.r_predict .+= cache.r
        rho = (sum(abs2, cache.r) - sum(abs2, value(df))) / (sum(abs2, cache.r) - sum(abs2, cache.r_predict))

        if rho > eta
            # Successful iteration
            cache.r .= value(df)
            jacobian!(df, cache.x)

            # Update scaling vector
            if autoscale
                for j = 1:nn
                    cache.d[j] = max(convert(real(T), 0.1) * real(cache.d[j]), norm(view(jacobian(df), :, j)))
                end
            end

            x_converged, f_converged, converged = assess_convergence(cache.x, cache.xold, cache.r, xtol, ftol)
        else
            cache.x .-= cache.p
            x_converged, converged = false, false
        end

        @trustregiontrace euclidean(cache.x, cache.xold)

        # Update size of trust region
        if rho < 0.1
            delta = delta/2
        elseif rho >= 0.9
            delta = 2 * wnorm(cache.d, cache.p)
        elseif rho >= 0.5
            delta = max(delta, 2 * wnorm(cache.d, cache.p))
        end
    end

    name = "Trust-region with dogleg"
    if autoscale
        name *= " and autoscaling"
    end
    return SolverResults(name,
                         initial_x, copy(cache.x), norm(cache.r, Inf),
                         it, x_converged, xtol, f_converged, ftol, tr,
                         first(df.f_calls), first(df.df_calls))
end

function trust_region(df::OnceDifferentiable,
                      initial_x::AbstractArray{T},
                      xtol::Real,
                      ftol::Real,
                      iterations::Integer,
                      store_trace::Bool,
                      show_trace::Bool,
                      extended_trace::Bool,
                      factor::Real,
                      autoscale::Bool,
                      cache = NewtonTrustRegionCache(df)) where T
    trust_region_(df, initial_x, convert(real(T), xtol), convert(real(T), ftol), iterations, store_trace, show_trace, extended_trace, convert(real(T), factor), autoscale, cache)
end


## Preamble
using Base.Test

using CompScienceMeshes
using BEAST

## The actual tests

κ = ω = 1.0

Γ = meshsegment(1.0, 0.5)
X = lagrangec0d1(Γ)
@test numvertices(Γ)-2 == numfunctions(X)

hypersingular = HyperSingular(κ)
identityop    = Identity()
doublelayer   = DoubleLayer(κ)

@time N = assemble(hypersingular, X, X)
@time I = assemble(identityop, X, X)

@test size(N) == (numfunctions(X), numfunctions(X))
@test size(I) == (numfunctions(X), numfunctions(X))
@test rank(I) == numfunctions(X)

@time e = assemble(PlaneWaveNeumann(κ, point(0.0, 1.0)), X)
@test length(e) == numfunctions(X)

x1 = N \ e;

# Testing duallagrangec0d1
Γ1 = meshcircle(1.0, 2.5,3)          # creating a triangle
Γ2 = barycentric_refinement(Γ1)      # creating the refined mesh

X1 =duallagrangec0d1(Γ1,Γ2)          # creating the basis functions
X1 = duallagrangec0d1(Γ1, Γ2, x->false, Val{2})
@test numcells(Γ1) == numfunctions(X1)  # making sure it is assigned according to the coarse mesh segments
@test length(X1.fns[1])== 6           # making sure each segment represent 6 shapes inside it
@test length(X1.fns[numfunctions(X1)])== 6  # making sure the last segment functions contains 6 shapes as well

## Test linear Lagrange elements on triangles and the computation of their curl
T = Float64
Degr = 1
Dim1 = 3
NumF = 3
f = BEAST.LagrangeRefSpace{T,Degr,Dim1,NumF}()
sphere = readmesh(joinpath(dirname(@__FILE__),"assets","sphere5.in"))
s = chart(sphere, first(cells(sphere)))
t = neighborhood(s, [1,1]/3)
v = f(t, Val{:withcurl})

A = volume(s)
@test v[1][2] == (s[3]-s[2])/2A
@test v[2][2] == (s[1]-s[3])/2A
@test v[3][2] == (s[2]-s[1])/2A

@test v[1][1] ≈ 1/3
@test v[2][1] ≈ 1/3
@test v[3][1] ≈ 1/3

## Test the construction of continuous linear Lagrange elements on 2D surfaces
using CompScienceMeshes
using BEAST
using Base.Test

m = meshrectangle(1.0, 1.0, 0.5, 3)
X = lagrangec0d1(m)
x = refspace(X)

@test numfunctions(x) == 3

@test numfunctions(X) == 1
@test length(X.fns[1]) == 6

## test the scalar trace for Lagrange functions
using CompScienceMeshes
using BEAST
using Base.Test

p1 = point(0,0,0)
p2 = point(1,0,0)
p3 = point(0,1,0)
p4 = point(1,1,1)
c1 = index(1,2,3)

m = Mesh([p1,p2,p3,p4],[c1])
b = Mesh([p1,p2], [index(1,2)])
X = lagrangec0d1(m, boundary(m))
@test numfunctions(X) == 3

Y = BEAST.strace(X, b)

@test numfunctions(X) == 3
@test numfunctions(Y) == 3

@test length(Y.fns[1]) == 1
@test length(Y.fns[2]) == 1
@test length(Y.fns[3]) == 0

sh = Y.fns[1][1]; @test (sh.cellid, sh.refid, sh.coeff) == (1, 1, 1.0)
sh = Y.fns[2][1]; @test (sh.cellid, sh.refid, sh.coeff) == (1, 2, 1.0)

x = refspace(X)

cell = chart(m, first(cells(m)))
face = chart(b, first(cells(b)))
Q = BEAST.strace(x, cell, 3, face)
@test Q == [1 0 0; 0 1 0]

## test Lagrange construction on Junctions
using CompScienceMeshes
using BEAST
using Base.Test

m1 = meshrectangle(1.0, 0.5, 0.5)
m2 = rotate(m1, 0.5π*[1,0,0])
m3 = rotate(m1, 1.0π*[1,0,0])
m = weld(m1, m2, m3)

X = lagrangec0d1(m)
x = refspace(X)
@test numfunctions(X) == 1
@test length(X.fns[1]) == 9

p = point(0.5, 0.0, 0.0)
for s in X.fns[1]
    cell = m.faces[s.cellid]
    patch = chart(m, cell)
    bary = carttobary(patch, p)
    mp = neighborhood(patch, bary)

end

## Test the dual pieweise constant lagrange elemetns
using CompScienceMeshes
using BEAST
using Base.Test

width, height = 1.0, 1.0
h = 0.5

m = meshrectangle(width, height, h)
b = meshsegment(width, width, 3)

X = duallagrangecxd0(m, b)
@test numfunctions(X) == 2
@test isa(refspace(X), BEAST.LagrangeRefSpace)

@test length(X.fns[1]) == 6
@test length(X.fns[2]) == 12

fine = geometry(X)
for fn in X.fns
    n = length(fn)
    for sh in fn
        @test sh.refid == 1
        cellid = sh.cellid
        cell = fine.faces[cellid]
        ptch = chart(fine, cell)
        @test sh.coeff * volume(ptch) ≈ 1/n
    end
end

## Test the construction of dual piecewise linear, globally continuous elements
using CompScienceMeshes
using BEAST
using Base.Test

m = meshrectangle(1.0, 1.0, 0.25)
j = meshsegment(1.0, 1.0, 3)

X = duallagrangec0d1(m,j)
x = refspace(X)

@test numfunctions(X) == numcells(m)

isonjunction = inclosure_gpredicate(j)
els, ad = BEAST.assemblydata(X)
for p in 1:numcells(m)
    el = els[p]
    for r in 1:numfunctions(x)
        vert = el[r]
        isonjunction(vert) || continue
        for (i,w) in ad[p,r]
            @test w == 0
        end
    end
end

## end of file


# Canonical form of multivariate normal

### Generic types

immutable MvNormalCanon{P<:AbstractPDMat,V<:Union(Vector{Float64},ZeroVector{Float64})} <: AbstractMvNormal
    μ::V    # the mean vector
    h::V    # potential vector, i.e. inv(Σ) * μ
    J::P    # precision matrix, i.e. inv(Σ)
end

typealias FullNormalCanon MvNormalCanon{PDMat,Vector{Float64}}
typealias DiagNormalCanon MvNormalCanon{PDiagMat,Vector{Float64}}
typealias IsoNormalCanon  MvNormalCanon{ScalMat,Vector{Float64}}

typealias ZeroMeanFullNormalCanon MvNormalCanon{PDMat,ZeroVector{Float64}}
typealias ZeroMeanDiagNormalCanon MvNormalCanon{PDiagMat,ZeroVector{Float64}}
typealias ZeroMeanIsoNormalCanon  MvNormalCanon{ScalMat,ZeroVector{Float64}}

### Constructors

function MvNormalCanon{P<:AbstractPDMat}(μ::Vector{Float64}, h::Vector{Float64}, J::P)
    length(μ) == length(h) == dim(J) || throw(DimensionMismatch("Inconsistent argument dimensions"))
    MvNormalCanon{P,Vector{Float64}}(μ, h, J)
end

function MvNormalCanon{P<:AbstractPDMat}(J::P)
    z = ZeroVector(Float64, dim(J))
    MvNormalCanon{P,ZeroVector{Float64}}(z, z, J)
end

function MvNormalCanon{P<:AbstractPDMat}(h::Vector{Float64}, J::P)
    length(h) == dim(J) || throw(DimensionMismatch("Inconsistent argument dimensions"))
    MvNormalCanon{P,Vector{Float64}}(J \ h, h, J)
end

MvNormalCanon(h::Vector{Float64}, J::Matrix{Float64}) = MvNormalCanon(h, PDMat(J))
MvNormalCanon(h::Vector{Float64}, prec::Vector{Float64}) = MvNormalCanon(h, PDiagMat(prec))
MvNormalCanon(h::Vector{Float64}, prec::Float64) = MvNormalCanon(h, ScalMat(length(h), prec))

MvNormalCanon(J::Matrix{Float64}) = MvNormalCanon(PDMat(J))
MvNormalCanon(prec::Vector{Float64}) = MvNormalCanon(PDiagMat(prec))
MvNormalCanon(d::Int, prec::Float64) = MvNormalCanon(ScalMat(d, prec))

### Show

distrname(d::IsoNormalCanon) = "IsoNormalCanon"
distrname(d::DiagNormalCanon) = "DiagNormalCanon"
distrname(d::FullNormalCanon) = "FullNormalCanon"

distrname(d::ZeroMeanIsoNormalCanon) = "ZeroMeanIsoNormalCanon"
distrname(d::ZeroMeanDiagNormalCanon) = "ZeroMeanDiagormalCanon"
distrname(d::ZeroMeanFullNormalCanon) = "ZeroMeanFullNormalCanon"

### conversion between conventional form and canonical form

meanform{C,V}(d::MvNormalCanon{C,V}) = MvNormal{C,V}(d.μ, inv(d.J))

canonform{C}(d::MvNormal{C,Vector{Float64}}) = (J = inv(d.Σ); MvNormalCanon(d.μ, J * d.μ, J))
canonform{C}(d::MvNormal{C,ZeroVector{Float64}}) = MvNormalCanon(inv(d.Σ))

### Basic statistics

length(d::MvNormalCanon) = length(d.μ)
mean(d::MvNormalCanon) = convert(Vector{Float64}, d.μ)

var(d::MvNormalCanon) = diag(inv(d.J))
cov(d::MvNormalCanon) = full(inv(d.J))
invcov(d::MvNormalCanon) = full(d.J)
logdetcov(d::MvNormalCanon) = -logdet(d.J)

### Evaluation

sqmahal(d::MvNormalCanon, x::DenseVector{Float64}) = quad(d.J, x - d.μ)
sqmahal!(r::DenseVector{Float64}, d::MvNormalCanon, x::DenseMatrix{Float64}) = quad!(r, d.J, x .- d.μ)

# Sampling (for GenericMvNormal)

unwhiten_winv!(J::AbstractPDMat, x::DenseVecOrMat) = unwhiten!(inv(J), x)
unwhiten_winv!(J::PDiagMat, x::DenseVecOrMat) = whiten!(J, x)
unwhiten_winv!(J::ScalMat, x::DenseVecOrMat) = whiten!(J, x)

_rand!(d::MvNormalCanon, x::DenseVecOrMat) = add!(unwhiten_winv!(d.J, randn!(x)), d.μ)


using LinearAlgebra: getindex
using Distributions: eltype
push!(LOAD_PATH, abspath(@__DIR__,".."))
using Pkg
pkg"activate ."

using Test
using LinearAlgebra,Statistics,StatsBase,Distributions
using Plots,NamedColors ; theme(:dark)
using SparseArrays
using SpikingRNNs; const global S = SpikingRNNs
using BenchmarkTools
using FFTW

function onesparsemat(w::Real)
  mat=Matrix{Float64}(undef,1,1) ; mat[1,1]=w
  return sparse(mat)
end

##

n = 6
myβ = 1.1
mywdistr=Uniform(0.2,0.3)
mywmat = sprand(n,n,0.5,n->rand(mywdistr,n))
mywmat_nonsp = convert(Matrix{Float64},mywmat)
myin = rand(Uniform(0.1,0.5),n)
tfake = NaN
p1 = S.PopulationHawkesExp(n,myβ)
ps1 = S.PSHawkes(p1)
conn1 = S.ConnectionHawkes(ps1,mywmat,ps1)
# rates to test
p1_in = S.PopInputStatic(ps1,myin)
myntw = S.RecurrentNetwork(tfake,(ps1,),(p1_in,),(conn1,) )

#
nspikes = 500_000
# initialize
S.hawkes_initialize!(myntw)

my_act = Vector{Tuple{Int64,Float64}}(undef,nspikes)
for k in 1:nspikes
  S.dynamics_step!(tfake,myntw)
  idx_fire = findfirst(ps1.isfiring)
  t_now = ps1.time_now[1]
  my_act[k] = (idx_fire,t_now)
end

myspikes_all = [S.hawkes_get_spiketimes(i,my_act) for i in 1:n]
myTmax = maximum(x->x[end],myspikes_all)+eps()
@info "Total duration $(round(myTmax;digits=1)) s"
rates_all =  map(x->length(x)/myTmax, myspikes_all )
@info "Rates are $(round.(rates_all;digits=2)) "

# rates with Fourier
# analytic rate from fourier Eq between 6 and  7 in Hawkes
ratefou = let G0 =  mywmat_nonsp .* S.fou_interaction_kernel.(0,p1)
  inv(I-G0)*myin |> real
end

@info "With Fourier, rates are $(round.(ratefou;digits=2)) Hz"

##
# covariance numerical

mydt = 0.2
myτmax = 80.0
mytaus = S.get_times(mydt,myτmax)
ntaus = length(mytaus)
cov_num = S.covariance_density_numerical(myspikes_all,mydt,myτmax)

##

# numerical cov densities
plot(mytaus[2:end],cov_num[1,1,2:end] ; linewidth = 2)
plot!(mytaus[2:end],cov_num[2,2,2:end] ; linewidth = 2)
plot!(mytaus[2:end],cov_num[2,5,2:end] ; linewidth = 2)

## now with Fourier!

# Pick eq 12 from Hawkes
function four_high_res(dt::Real,Tmax::Real)
  k1 = 2
  k2 = 0.2
  n = size(mywmat_nonsp,1)
  myτmax,mydt = Tmax * k1, dt*k2
  mytaus = S.get_times(mydt,myτmax)
  nkeep = div(length(mytaus),k1)
  myfreq = S.get_frequencies_centerzero(mydt,myτmax)
  G_omega = map(mywmat_nonsp) do w
    ifftshift( w .* S.fou_interaction_kernel.(myfreq,p1))
  end
  D = Diagonal(ratefou)
  M = Array{ComplexF64}(undef,n,n,length(myfreq))
  Mt = similar(M,Float64)
  for i in eachindex(myfreq)
    G = getindex.(G_omega,i)
    # M[:,:,i] = (I-G)\(G*D+D*G'-G*D*G')/(I-G')
    M[:,:,i] = (I-G)\D*(G+G'-G*G')/(I-G')
  end
  for i in 1:n,j in 1:n
    Mt[i,j,:] = real.(ifft(M[i,j,:])) ./ mydt
  end
  return mytaus[1:nkeep],Mt[:,:,1:nkeep]
end

taush,Cfou=four_high_res(mydt,myτmax)

# numerical cov densities
plot(mytaus[2:end],cov_num[1,1,2:end] ; linewidth = 2)
plot!(taush[1:end],Cfou[1,1,1:end] ; linewidth=3, linestyle=:dash)

plot(mytaus[2:end],cov_num[2,2,2:end] ; linewidth = 2)
plot!(taush[1:end],Cfou[2,2,1:end] ; linewidth=3, linestyle=:dash)

plot(mytaus[2:end],cov_num[2,5,2:end] ; linewidth = 2)
plot!(taush[3:end],Cfou[2,5,3:end] ; linewidth=3, linestyle=:dash)

plot(mytaus[2:end],cov_num[1,6,2:end] ; linewidth = 2)
plot!(taush[3:end],Cfou[1,6,3:end] ; linewidth=3, linestyle=:dash)


using jInv.Mesh
using Multigrid

println("************************************************* Example 2D ******************************************************");

domain = [0.0, 1.0, 0.0, 1.0];
n = [512,512];
Mr = getRegularMesh(domain,n)
G = getNodalGradientMatrix(Mr);
Ar = G'*G;
Ar = Ar + 1e-4*norm(Ar,1)*speye(size(Ar,2));

levels      = 4;
numCores 	= 8;
maxIter     = 10;
relativeTol = 1e-10;
relaxType   = "Jac";
relaxParam  = 0.75;
relaxPre 	= 1;
relaxPost   = 1;
cycleType   ='V';
coarseSolveType = "NoMUMPS";

MG = getMGparam(levels,numCores,maxIter,relativeTol,relaxType,relaxParam,relaxPre,relaxPost,cycleType,coarseSolveType,0.5,0.0);

N = size(Ar,2);

b = Ar*rand(N,2);
x = zeros(N,2);

MGsetup(Ar,Mr,MG,Float64,size(b,2),true);

println("****************************** Stand-alone GMG: ******************************")
solveMG(MG,b,x,true);

println("****************************** Stand-alone GMG: iterative coarsest ***********************")
coarseSolveType = "BiCGSTAB"
MG = getMGparam(levels,numCores,maxIter,relativeTol,relaxType,relaxParam,relaxPre,relaxPost,cycleType,coarseSolveType,0.5,0.0);
MGsetup(Ar,Mr,MG,Float64,size(b,2),true);
solveMG(MG,b,x,true);
println("****************************** GMRES preconditioned with GMG: (only one rhs...) ******************************")
x[:] = 0.0
b = vec(b[:,1]);
x = vec(x[:,1]);
relaxType   = "Jac-GMRES";
coarseSolveType = "NoMUMPS"
MG = getMGparam(levels,numCores,maxIter,relativeTol,relaxType,relaxParam,relaxPre,relaxPost,cycleType,coarseSolveType,0.5,0.0);
MGsetup(Ar,Mr,MG,Float64,size(b,2),true);
solveGMRES_MG(Ar,MG,b,x,true,2)

Ar = 0;
b = 0;
x = 0;
Mr = 0;

copySolver(MG);
destroyCoarsestLU(MG);
clear!(MG)


#File: efficfron.jl
#Author:
#Date: 29-June-2020

#import XL data
[names,~,returns] = getXLData('hfma_matlab_data','10 Hedge Funds');

N = size(returns,2); ## data points
am1 = fMean(returns,1); #monthly mean
covar = cov(returns); #covariance matrix

#plot mean-variance efficient frontier
portopt(am1,covar,N)


velocity(r::AbstractArray, t::Int, δt::Int) = (r[t+δt, :] .- r[t, :]) ./ δt
c(v1::AbstractVector, v2::AbstractVector) = v1 ⋅ v2

function vacf(
    df::DataFrame,
    id::Symbol,
    x::Symbol,
    y::Symbol,
    max_τ::Int,
    δt::Int,
)
    result = DataFrame(n = Int[], τ = Int[], value = Float64[])
    @inbounds for n = 1:maximum(df[!, id])
        data = extract(df, n, id, [x, y])
        T = size(data, 1)
        T < max_τ ? maxt = T : maxt = max_τ
        @inbounds for τ = 0:maxt-δt-1
            s = 0.0
            @simd for t = 1:(T-δt-τ)
                vₜ = velocity(data, t, δt)
                v2 = velocity(data, t + τ, δt)
                s += c(vₜ, v2)
            end
            s /= (T - δt - τ)
            push!(result, [n, τ, s])
        end
    end
    result
end


# # Hyperelasticity
#
# **Keywords**: *hyperelasticity*, *finite strain*, *large deformations*, *Newton's method*,
# *conjugate gradient*, *automatic differentiation*
#
# ![hyperelasticity.png](hyperelasticity.png)
#
# *Figure 1*: Cube loaded in torsion modeled with a hyperelastic material model and
# finite strain.
#
#-
#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`hyperelasticity.ipynb`](@__NBVIEWER_ROOT_URL__/examples/hyperelasticity.ipynb).
#-
# ## Introduction
#
# In this example we will solve a problem in a finite strain setting using an
# hyperelastic material model. In order to compute the stress we will use automatic
# differentiation, to solve the non-linear system we use Newton's
# method, and for solving the Newton increment we use conjugate gradients.
#
# The weak form is expressed in terms of the first Piola-Kirchoff stress ``\mathbf{P}``
# as follows: Find ``u \in \mathbb{U}`` such that
#
# ```math
# \int_{\Omega} \nabla_{\mathbf{X}} \delta \mathbf{u} : \mathbf{P}(\mathbf{u})\ \mathrm{d}\Omega =
# \int_{\Omega} \delta \mathbf{u} \cdot \mathbf{b}\ \mathrm{d}\Omega + \int_{\Gamma^\mathrm{N}}
# \delta \mathbf{u} \cdot \mathbf{t}\ \mathrm{d}\Gamma
# \quad \forall \delta \mathbf{u} \in \mathbb{U}^0,
# ```
#
# where ``\mathbf{u}`` is the unknown displacement field, ``\mathbf{b}`` is the body force acting
# on the reference domain, ``\mathbf{t}`` is the traction acting on the Neumann part of the reference
# domain's boundary, and where ``\mathbb{U}`` and ``\mathbb{U}^0`` are suitable trial and test sets.
# ``\Omega`` denotes the reference domain, which is also called reference or material domain.
# Gradients are defined with respect to the reference domain, here denoted with  an ``X``.
# Formally this is expressed as ``(\nabla_{\mathbf{X}} \bullet)_{ij} := \frac{\partial(\bullet)_i}{\partial X_j}``.
# Note that for large deformation problems it is also possibile that gradients and integrals
# are defined on the deformed domain, which is also called the current or spatial domain, depending
# on the specific formulation.
#

using Ferrite, Tensors, TimerOutputs, ProgressMeter
import KrylovMethods, IterativeSolvers

# ## Hyperelastic material model
#
# The stress can be derived from an energy potential, defined in
# terms of the right Cauchy-Green tensor ``\mathbf{C} = \mathbf{F}^{\mathrm{T}} \mathbf{F}``,
# where ``\mathbf{F} = \mathbf{I} + \nabla_{\mathbf{X}} \mathbf{u}`` is the deformation gradient.
# We shall use a neo-Hookean model, where the potential can be written as
#
# ```math
# \Psi(\mathbf{C}) = \frac{\mu}{2} (I_1 - 3) - \mu \ln(J) + \frac{\lambda}{2} \ln(J)^2,
# ```
#
# where ``I_1 = \mathrm{tr}(\mathbf{C})`` is the first invariant, ``J = \sqrt{\det(\mathbf{C})}``
# and ``\mu`` and ``\lambda`` material parameters.
# From the potential we obtain the second Piola-Kirchoff stress ``\mathbf{S}`` as
#
# ```math
# \mathbf{S} = 2 \frac{\partial \Psi}{\partial \mathbf{C}},
# ```
#
# and the tangent of ``\mathbf{S}`` as
#
# ```math
# \frac{\partial \mathbf{S}}{\partial \mathbf{C}} = 4 \frac{\partial \Psi}{\partial \mathbf{C}}.
# ```
#
# We can implement the material model as follows, where we utilize automatic differentiation
# for the stress and the tangent, and thus only define the potential:

struct NeoHooke
    μ::Float64
    λ::Float64
end

function Ψ(C, mp::NeoHooke)
    μ = mp.μ
    λ = mp.λ
    Ic = tr(C)
    J = sqrt(det(C))
    return μ / 2 * (Ic - 3) - μ * log(J) + λ / 2 * log(J)^2
end

function constitutive_driver(C, mp::NeoHooke)
    ## Compute all derivatives in one function call
    ∂²Ψ∂C², ∂Ψ∂C = Tensors.hessian(y -> Ψ(y, mp), C, :all)
    S = 2.0 * ∂Ψ∂C
    ∂S∂C = 4.0 * ∂²Ψ∂C²
    return S, ∂S∂C
end;

# Finally, for the finite element problem we need ``\mathbf{P}`` and
# ``\frac{\partial \mathbf{P}}{\partial \mathbf{F}}``, which can be
# obtained by using the following relations:
#
# ```math
# \begin{align*}
# \mathbf{P} &= \mathbf{F} \cdot \mathbf{S},\\
# \frac{\partial \mathbf{P}}{\partial \mathbf{F}} &= [\mathbf{F} \bar{\otimes} \mathbf{I}] :
# \frac{\partial \mathbf{S}}{\partial \mathbf{C}} : [\mathbf{F}^\mathrm{T} \bar{\otimes} \mathbf{I}]
# + \mathbf{I} \bar{\otimes} \mathbf{S}.
# \end{align*}
# ```
#
# ## Newton's Method
#
# As mentioned above, to deal with the non-linear weak form we first linearize
# the problem such that we can apply Newton's method, and then apply the FEM to
# discretize the problem. Skipping a detailed derivation, Newton's method can
# be expressed as:
# Given some initial guess ``\mathbf{u}^0``, find a sequence ``\mathbf{u}^{k}`` by iterating
#
# ```math
# \mathbf{u}^{k+1} = \mathbf{u}^{k} - \Delta \mathbf{u}^{k}
# ```
#
# until some termination condition has been met. Therein we determine ``\Delta \mathbf{u}^{k}``
#
# ```math
# \mathbf{K}(\mathbf{u}^{k}) \Delta \mathbf{u}^{k} = \mathbf{g}(\mathbf{u}^{k})
# ```
#
# where $\mathbf{g}$ is the global residual and $\mathbf{K} = \frac{\partial g}{\partial u}$
# the Jacobi matrix, such that
#
# ```math
# K_{ij} = \int_{\Omega} \nabla_{\mathbf{X}} \delta \mathbf{u}_{i} : \frac{\partial \mathbf{P}}{\partial \mathbf{F}} : \nabla_{\mathbf{X}} \delta \mathbf{u}_{j} \, \mathrm{d} \Omega
# ```
# and
#
# ```math
# g_{i} = \int_{\Omega} \nabla_{\mathbf{X}} \delta \mathbf{u}_{i} : \mathbf{P} - \delta \mathbf{u}_{i} \cdot \mathbf{b} \, \mathrm{d} \Omega
# ```
#
# A detailed derivation can be found in every continuum mechanics book, which has a
# chapter about finite elasticity theory. We used "Nonlinear solid mechanics: a continuum
# approach for engineering science." by Gerhard Holzapfel (chapter 8) as a reference.
#
# ## Finite element assembly
#
# The element routine for assembling the residual and tangent stiffness is implemented
# as usual, with loops over quadrature points and shape functions:

function assemble_element!(ke, ge, cell, cv, fv, mp, ue)
    ## Reinitialize cell values, and reset output arrays
    reinit!(cv, cell)
    fill!(ke, 0.0)
    fill!(ge, 0.0)

    b = Vec{3}((0.0, -0.5, 0.0)) # Body force
    t = Vec{3}((0.1, 0.0, 0.0)) # Traction
    ndofs = getnbasefunctions(cv)

    for qp in 1:getnquadpoints(cv)
        dΩ = getdetJdV(cv, qp)
        ## Compute deformation gradient F and right Cauchy-Green tensor C
        ∇u = function_gradient(cv, qp, ue)
        F = one(∇u) + ∇u
        C = tdot(F)
        ## Compute stress and tangent
        S, ∂S∂C = constitutive_driver(C, mp)
        P = F ⋅ S
        I = one(S)
        ∂P∂F = otimesu(F, I) ⊡ ∂S∂C ⊡ otimesu(F', I) + otimesu(I, S)

        ## Loop over test functions
        for i in 1:ndofs
            ## Test function + gradient
            δui = shape_value(cv, qp, i)
            ∇δui = shape_gradient(cv, qp, i)
            ## Add contribution to the residual from this test function
            ge[i] += ( ∇δui ⊡ P - δui ⋅ b ) * dΩ

            ∇δui∂P∂F = ∇δui ⊡ ∂P∂F # Hoisted computation
            for j in 1:ndofs
                ∇δuj = shape_gradient(cv, qp, j)
                ## Add contribution to the tangent
                ke[i, j] += ( ∇δui∂P∂F ⊡ ∇δuj ) * dΩ
            end
        end
    end

    ## Surface integral for the traction
    for face in 1:nfaces(cell)
        if onboundary(cell, face)
            reinit!(fv, cell, face)
            for q_point in 1:getnquadpoints(fv)
                dΓ = getdetJdV(fv, q_point)
                for i in 1:ndofs
                    δui = shape_value(fv, q_point, i)
                    ge[i] -= (δui ⋅ t) * dΓ
                end
            end
        end
    end
end;

# Assembling global residual and tangent
function assemble_global!(K, f, dh, cv, fv, mp, u)
    n = ndofs_per_cell(dh)
    ke = zeros(n, n)
    ge = zeros(n)

    ## start_assemble resets K and f
    assembler = start_assemble(K, f)

    ## Loop over all cells in the grid
    @timeit "assemble" for cell in CellIterator(dh)
        global_dofs = celldofs(cell)
        ue = u[global_dofs] # element dofs
        @timeit "element assemble" assemble_element!(ke, ge, cell, cv, fv, mp, ue)
        assemble!(assembler, global_dofs, ge, ke)
    end
end;

# Define a main function, with a loop for Newton iterations

function solve()
    reset_timer!()

    ## Generate a grid
    N = 10
    L = 1.0
    left = zero(Vec{3})
    right = L * ones(Vec{3})
    grid = generate_grid(Tetrahedron, (N, N, N), left, right)

    ## Material parameters
    E = 10.0
    ν = 0.3
    μ = E / (2(1 + ν))
    λ = (E * ν) / ((1 + ν) * (1 - 2ν))
    mp = NeoHooke(μ, λ)

    ## Finite element base
    ip = Lagrange{3, RefTetrahedron, 1}()
    qr = QuadratureRule{3, RefTetrahedron}(1)
    qr_face = QuadratureRule{2, RefTetrahedron}(1)
    cv = CellVectorValues(qr, ip)
    fv = FaceVectorValues(qr_face, ip)

    ## DofHandler
    dh = DofHandler(grid)
    push!(dh, :u, 3) # Add a displacement field
    close!(dh)

    function rotation(X, t, θ = deg2rad(60.0))
        x, y, z = X
        return t * Vec{3}(
            (0.0,
            L/2 - y + (y-L/2)*cos(θ) - (z-L/2)*sin(θ),
            L/2 - z + (y-L/2)*sin(θ) + (z-L/2)*cos(θ)
            ))
    end

    dbcs = ConstraintHandler(dh)
    ## Add a homogenous boundary condition on the "clamped" edge
    dbc = Dirichlet(:u, getfaceset(grid, "right"), (x,t) -> [0.0, 0.0, 0.0], [1, 2, 3])
    add!(dbcs, dbc)
    dbc = Dirichlet(:u, getfaceset(grid, "left"), (x,t) -> rotation(x, t), [1, 2, 3])
    add!(dbcs, dbc)
    close!(dbcs)
    t = 0.5
    Ferrite.update!(dbcs, t)

    ## Pre-allocation of vectors for the solution and Newton increments
    _ndofs = ndofs(dh)
    un = zeros(_ndofs) # previous solution vector
    u  = zeros(_ndofs)
    Δu = zeros(_ndofs)
    ΔΔu = zeros(_ndofs)
    apply!(un, dbcs)

    ## Create sparse matrix and residual vector
    K = create_sparsity_pattern(dh)
    g = zeros(_ndofs)

    ## Perform Newton iterations
    newton_itr = -1
    NEWTON_TOL = 1e-8
    prog = ProgressMeter.ProgressThresh(NEWTON_TOL, "Solving:")

    while true; newton_itr += 1
        u .= un .+ Δu # Current guess
        assemble_global!(K, g, dh, cv, fv, mp, u)
        normg = norm(g[Ferrite.free_dofs(dbcs)])
        apply_zero!(K, g, dbcs)
        ProgressMeter.update!(prog, normg; showvalues = [(:iter, newton_itr)])

        if normg < NEWTON_TOL
            break
        elseif newton_itr > 30
            error("Reached maximum Newton iterations, aborting")
        end

        ## Compute increment using cg! from IterativeSolvers.jl
        @timeit "linear solve (KrylovMethods.cg)" ΔΔu′, flag, relres, iter, resvec = KrylovMethods.cg(K, g; maxIter = 1000)
        @assert flag == 0
        @timeit "linear solve (IterativeSolvers.cg!)" IterativeSolvers.cg!(ΔΔu, K, g; maxiter=1000)

        apply_zero!(ΔΔu, dbcs)
        Δu .-= ΔΔu
    end

    ## Save the solution
    @timeit "export" begin
        vtk_grid("hyperelasticity", dh) do vtkfile
            vtk_point_data(vtkfile, dh, u)
        end
    end

    print_timer(title = "Analysis with $(getncells(grid)) elements", linechars = :ascii)
    return u
end

# Run the simulation

u = solve();

## test the result                #src
using Test                        #src
@test norm(u) ≈ 4.865189736192834 #src

#md # ## Plain program
#md #
#md # Here follows a version of the program without any comments.
#md # The file is also available here: [`hyperelasticity.jl`](hyperelasticity.jl).
#md #
#md # ```julia
#md # @__CODE__
#md # ```