using OrdinaryDiffEq, DiffEqDevTools, Sundials, ModelingToolkit, ODEInterfaceDiffEq,
      Plots, DASSL, DASKR
using LinearAlgebra

@variables t y1(t)=1.0 y2(t)=2.0 y3(t)=3.0
@parameters p1=77.27 p2=8.375e-6 p3=0.161
D = Differential(t)
eqs = [
  D(y1) ~ p1*(y2+y1*(1-p2*y1-y2))
  D(y2) ~ (y3-(1+y1)*y2)/p1
  D(y3) ~ p3*(y1-y3)
]
@named sys = ODESystem(eqs)
simpsys = structural_simplify(sys)
mmprob = ODEProblem(sys,[],(0.0,30.0))
daeprob = DAEProblem(sys,[D(y1)=>77.26935286375,
                              D(y2)=>-0.012941633234114146,
                              D(y3)=>-0.322],[],(0.0,30.0))
odaeprob = ODAEProblem(simpsys,[],(0.0,30.0))

ref_sol = solve(daeprob,IDA(),abstol=1/10^14,reltol=1/10^14);
ode_ref_sol = solve(odaeprob,CVODE_BDF(),abstol=1/10^14,reltol=1/10^14);

probs = [mmprob,daeprob,odaeprob]
refs = [ref_sol,ref_sol,ode_ref_sol];

plot(ref_sol)

abstols = 1.0 ./ 10.0 .^ (6:9)
reltols = 1.0 ./ 10.0 .^ (2:5);
setups = [Dict(:prob_choice => 1, :alg=>Rosenbrock23()),
          Dict(:prob_choice => 1, :alg=>Rodas4()),
          #Dict(:prob_choice => 1, :alg=>FBDF()),
          Dict(:prob_choice => 1, :alg=>QNDF()),
          Dict(:prob_choice => 1, :alg=>rodas()),
          Dict(:prob_choice => 1, :alg=>radau()),
          Dict(:prob_choice => 1, :alg=>RadauIIA5()),
          Dict(:prob_choice => 2, :alg=>DFBDF()),
          Dict(:prob_choice => 2, :alg=>IDA()),
]

wp = WorkPrecisionSet(probs,abstols,reltols,setups;print_names=true,
                      save_everystep=false,appxsol=refs,maxiters=Int(1e5),numruns=10)
plot(wp)

setups = [Dict(:prob_choice => 1, :alg=>Rosenbrock23()),
          Dict(:prob_choice => 1, :alg=>Rodas4()),
          Dict(:prob_choice => 2, :alg=>IDA()),
          Dict(:prob_choice => 3, :alg=>Rosenbrock23()),
          Dict(:prob_choice => 3, :alg=>Rodas4()),
          Dict(:prob_choice => 3, :alg=>CVODE_BDF()),
          Dict(:prob_choice => 3, :alg=>TRBDF2()),
          Dict(:prob_choice => 3, :alg=>KenCarp4()),
          ]
wp = WorkPrecisionSet(probs,abstols,reltols,setups;
                      save_everystep=false,appxsol=refs,maxiters=Int(1e5),numruns=10)
plot(wp)

abstols = 1.0 ./ 10.0 .^ (6:8)
reltols = 1.0 ./ 10.0 .^ (2:4);
setups = [Dict(:prob_choice => 3, :alg=>Rosenbrock23()),
          Dict(:prob_choice => 3, :alg=>Rodas4()),
          Dict(:prob_choice => 2, :alg=>IDA()),
          Dict(:prob_choice => 2, :alg=>DASSL.dassl()),
          Dict(:prob_choice => 2, :alg=>DASKR.daskr()),
          ]
wp = WorkPrecisionSet(probs,abstols,reltols,setups;
                      save_everystep=false,appxsol=refs,maxiters=Int(1e5),numruns=10)
plot(wp)

abstols = 1.0 ./ 10.0 .^ (6:9)
reltols = 1.0 ./ 10.0 .^ (2:5);
setups = [Dict(:prob_choice => 1, :alg=>Rosenbrock23()),
          Dict(:prob_choice => 1, :alg=>Rodas4()),
          #Dict(:prob_choice => 1, :alg=>FBDF()), # too slow
          Dict(:prob_choice => 1, :alg=>QNDF()),
          Dict(:prob_choice => 1, :alg=>rodas()),
          Dict(:prob_choice => 1, :alg=>radau()),
          Dict(:prob_choice => 1, :alg=>RadauIIA5()),
          Dict(:prob_choice => 2, :alg=>DFBDF()),
          Dict(:prob_choice => 2, :alg=>IDA()),
          ]
gr()
wp = WorkPrecisionSet(probs,abstols,reltols,setups;error_estimate = :l2,
                      save_everystep=false,appxsol=refs,maxiters=Int(1e5),numruns=10)
plot(wp)

abstols = 1.0 ./ 10.0 .^ (6:9)
reltols = 1.0 ./ 10.0 .^ (2:5);
setups = [Dict(:prob_choice => 1, :alg=>Rosenbrock23()),
          Dict(:prob_choice => 1, :alg=>Rodas4()),
          Dict(:prob_choice => 2, :alg=>IDA()),
          Dict(:prob_choice => 3, :alg=>Rosenbrock23()),
          Dict(:prob_choice => 3, :alg=>Rodas4()),
          Dict(:prob_choice => 3, :alg=>CVODE_BDF()),
          Dict(:prob_choice => 3, :alg=>TRBDF2()),
          Dict(:prob_choice => 3, :alg=>KenCarp4()),
          ]
wp = WorkPrecisionSet(probs,abstols,reltols,setups;error_estimate = :l2,
                      save_everystep=false,appxsol=refs,maxiters=Int(1e5),numruns=10)
plot(wp)

abstols = 1.0 ./ 10.0 .^ (7:12)
reltols = 1.0 ./ 10.0 .^ (4:9)

setups = [Dict(:prob_choice => 1, :alg=>Rodas5()),
          Dict(:prob_choice => 3, :alg=>Rodas5()),
          Dict(:prob_choice => 1, :alg=>Rodas4()),
          Dict(:prob_choice => 3, :alg=>Rodas4()),
          #Dict(:prob_choice => 1, :alg=>FBDF()),
          Dict(:prob_choice => 1, :alg=>QNDF()),
          Dict(:prob_choice => 1, :alg=>rodas()),
          Dict(:prob_choice => 1, :alg=>radau()),
          Dict(:prob_choice => 1, :alg=>RadauIIA5()),
          Dict(:prob_choice => 2, :alg=>DFBDF()),
          Dict(:prob_choice => 2, :alg=>IDA()),
          ]
gr()
wp = WorkPrecisionSet(probs,abstols,reltols,setups;
                      save_everystep=false,appxsol=refs,maxiters=Int(1e5),numruns=10)
plot(wp)

wp = WorkPrecisionSet(probs,abstols,reltols,setups;error_estimate = :l2,
                      save_everystep=false,appxsol=refs,maxiters=Int(1e5),numruns=10)
plot(wp)

using SciMLBenchmarks
SciMLBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


using StaticArrays
@testset "Correlated Wiener Process" begin

using DiffEqNoiseProcess, DiffEqBase, Test, Statistics

# oop
ρ = 0.3
Γ = [1 ρ;ρ 1]
W = CorrelatedWienerProcess(Γ, 0., zeros(2), zeros(2))

dt = 0.1
calculate_step!(W,dt,nothing,nothing)
for i in 1:10
  accept_step!(W,dt,nothing,nothing)
end

prob = NoiseProblem(W,(0.0,1.0))
sol = solve(prob;dt=0.1)

# inplace
W = CorrelatedWienerProcess!(Γ, 0., zeros(2), zeros(2))

dt = 0.1
calculate_step!(W,dt,nothing,nothing)
for i in 1:10
  accept_step!(W,dt,nothing,nothing)
end

prob = NoiseProblem(W,(0.0,1.0))
sol = solve(prob;dt=0.01)

output_func = (sol,i) -> (sol.dW,false)
ensemble_prob = EnsembleProblem(prob,output_func=output_func)
@time sol = Array(solve(ensemble_prob,dt=dt,trajectories=1_000_000))

@test zero(mean(sol,dims=2)[:]) ≈ mean(sol,dims=2)[:] atol=1e-2
@test Γ ≈ cov(sol,dims=2)/dt rtol=1e-2

# with StaticArrays
Γ = @SMatrix [1.0 ρ
              ρ 1.0]
W = CorrelatedWienerProcess(Γ, 0., @SVector(zeros(2)), @SVector(zeros(2)))

dt = 0.1
calculate_step!(W,dt,nothing,nothing)
for i in 1:10
  accept_step!(W,dt,nothing,nothing)
end

prob = NoiseProblem(W,(0.0,1.0))
sol = solve(prob;dt=0.1)

@test typeof(sol[1]) <: SArray
@test typeof(sol[end]) <: SArray

end


using Base.Threads
using LoopVectorization
import TriangleMesh
using Printf
using CairoMakie
import AbstractPlotting.GeometryBasics

#--------------------------------------------------------------------#

function tplot(Mesh, v)
    p   = [AbstractPlotting.GeometryBasics.Polygon( Point2f0[ (Mesh.point[1, Mesh.cell[j,i]],Mesh.point[2, Mesh.cell[j,i]]) for j=1:3]    )
    for i in 1:Mesh.n_cell]
    display(poly(p, color = v ))
    return
end

function Tanalytic!( mesh, xc , yc, T, nvert_ele, a, b, c, d, alp, bet, vec )
    # Evaluate T analytic on barycentres: vec 0 or 1
    if vec==0 # Loop version vectorized at compilation using @avx
        @avx for i=1:mesh.n_cell
            xc[i] = 0;
            yc[i] = 0;
            for j=1:nvert_ele
                xc[i] += 1.0/nvert_ele*( mesh.point[1,mesh.cell[j,i]] )
                yc[i] += 1.0/nvert_ele*( mesh.point[2,mesh.cell[j,i]] )
            end
        end
        @avx for i=1:mesh.n_cell
            T[i] = exp(alp.sin(a*xc[i] + c*yc[i]) + bet*cos(b*xc[i] + d*yc[i]));
        end
    else # Standard vectorization
        # X    = mesh.point[1,mesh.cell[:,:]]'  # X coordinates of vertices of each element
        # Y    = mesh.point[2,mesh.cell[:,:]]'  # Y coordinates of vertices of each element
        # sumX = sum(X, dims=2)
        # sumY = sum(Y, dims=2)
        # @. xc = 1.0/nvert_ele * sumX[:]
        # @. yc = 1.0/nvert_ele * sumY[:]
    for j=1:nvert_ele # as fast as above and less cryptic
        @. xc = 0.0
        @. yc = 0.0
        @. xc = xc + 1.0/nvert_ele * mesh.point[1,mesh.cell[j,:]]
        @. yc = yc + 1.0/nvert_ele * mesh.point[2,mesh.cell[j,:]]
    end
    @. T = exp(alp*sin(a*xc + c*yc) + bet*cos(b*xc + d*yc));
end
return
end

#--------------------------------------------------------------------#

function main()

# Create sides of mesh
xmin, xmax = 0, 1
ymin, ymax = 0, 1
nx, ny     = 20, 20

# Four corners of the domain
px = [0.0 1.0 1.0 0.0]
py = [0.0 0.0 1.0 1.0]
sx = [ 1 2 3 4 ]
sy = [ 2 3 4 1 ]
st = [ 1 1 1 1 ]          # segment markers
p  = vcat(px, py)         # points
s  = vcat(sx, sy)         # segments

# TriangleMesh.Polygon_pslg:
# n_point           :: Int32
# point             :: Array{Float64,2}
# n_point_marker    :: Int32
# point_marker      :: Array{Int32,2}
# n_point_attribute :: Int32
# point_attribute   :: Array{Float64,2}
# n_segment         :: Int32
# segment           :: Array{Int32,2}
# segment_marker    :: Array{Int32,1}
# n_hole            :: Int32
# hole              :: Array{Float64,2}

# Triangulation
holes    = Array{Float64}(undef,2,0)
domain   = TriangleMesh.Polygon_pslg(size(p,2), p, 0, Array{Int64}(undef,2,0), 0, Array{Float64}(undef,2,0),  size(s,2), s, st[:], 0, holes)
switches = "Dpenq33IAa0.01"#"Dpenq33o2IAa"
mesh     = TriangleMesh.create_mesh(domain, switches)
nvert_el = 3 # vertices per element
println(size(mesh.cell))

# Generate function to be visualised on mesh
alp = 0.1; bet = 0.3; a = 5.1; b = 4.3; c = -6.2; d = 3.4;
xc = zeros(mesh.n_cell)
yc = zeros(mesh.n_cell)
T  = zeros(mesh.n_cell)
ncalls = 4
# A) Loop version with @avx
vec = 0
@printf("Looped, %d times:\n", ncalls)
for icall=1:ncalls
    @time Tanalytic!(mesh, xc ,yc , T, nvert_el, a, b, c, d, alp, bet, vec)
end
# B) Standard vectorisation
vec = 1
@printf("Vectorized, %d times:\n", ncalls)
for icall=1:ncalls
    @time Tanalytic!(mesh, xc ,yc , T, nvert_el, a, b, c, d, alp, bet, vec)
end
# # Visualise
tplot(mesh, T)

return
end

main()


using FastGaussQuadrature
using SpecialFunctions
using ApproxFun
using Interpolations

export TracyWidom

struct TracyWidom <: ContinuousUnivariateDistribution
    beta::Integer
end

function cdf(d::TracyWidom, x::Real)
    if d.beta == 1
        _fredholm_det(_V_airy_kernel, x, 100)
    elseif d.beta == 2
        _fredholm_det(_airy_kernel, x, 100)
    elseif d.beta == 4
        (_fredholm_det(_V_airy_kernel, x * sqrt(2), 100) + _fredholm_det(_V_airy_kernel, x * sqrt(2), 100, -1))/2
    else
        throw(DomainError(d.beta))
    end
end

_TWfs = Dict()

function _TWf(x, beta)
    if !(beta in keys(_TWfs))
        dom = Chebyshev(-5..5)
        D = Derivative(dom)
        _TWfs[beta] = D * Fun(x -> cdf(TracyWidom(beta), x), dom)
    end
    _TWfs[beta](x)
end

function pdf(d::TracyWidom, x::Real)
    _TWf(x, d.beta)
end

_TWinvs = Dict()

function _TWinv(t, beta)
    (0 < t < 1) || throw(DomainError(t))

    if !(beta in keys(_TWinvs))
        # bounds outside which we use limiting representation
        u = 4
        l = -5

        # TODO: make a more principled grid using tail behaviour
        n = 1000
        y = range(l, u, length=n)
        x = cdf.(TracyWidom(beta), y)

        # Transform limiting ends to maintain differentiability
        righta = pdf(TracyWidom(beta), u)/(beta * sqrt(u) * exp(-beta * u^(3/2)*2/3))
        rightb = x[end] - righta * (1 - exp(-beta * u^(3/2) * 2/3))

        lefta = pdf(TracyWidom(beta), l)/(beta * l^2/8 * exp(beta * l^3/24))
        leftb = x[1] - lefta * exp(-beta * abs(l)^3/24)

        itp = interpolate((x,), y, Gridded(Linear()))

        _TWinvs[beta] = [righta, rightb, x[end], lefta, leftb, x[1], itp]
    end

    # TODO: make the indexing less terrible
    if t > _TWinvs[beta][3]
        s = (t - _TWinvs[beta][2])/_TWinvs[beta][1]
        return (log(1/(1 - s)) * 3/(2 * beta))^(2/3)
    elseif t < _TWinvs[beta][6]
        s = (t - _TWinvs[beta][5])/_TWinvs[beta][4]
        return -(log(1/s) * 24/beta)^(1/3)
    else
        return _TWinvs[beta][end](t)
    end

end

function quantile(d::TracyWidom, q::Real)
    _TWinv(q, d.beta)
end

function _phi(s, xi)
    s + 10 * tan(pi * (xi + 1)/4)
end

function _Dphi(s, xi)
    pi * 5/2 * sec(pi * (xi + 1)/4)^2
end

"""
Transform a kernel on L_2(s, infinity) to one on L_2(-1, 1)
"""
function _transform_kernel(kernel, s)
    (x, y) -> sqrt(_Dphi(s, x) * _Dphi(s, y)) * kernel(_phi(s, x), _phi(s, y))
end

"""
Approximate fredholm determinant of a kernel on (s, infinity)
"""
function _fredholm_det(kernel, s, N, z=1)

    (nodes, weights) = gausslegendre(N)
    trans_kernel = _transform_kernel(kernel, s)

    M = I - z * sqrt.(weights * weights') .* trans_kernel.(nodes', nodes)
    det(M)
end

function _airy_kernel(x, y)
    if x == y
        airyaiprime(x)^2 - x * airyai(x)^2
    else
        (airyai(x) * airyaiprime(y) - airyaiprime(x) * airyai(y))/(x - y)
    end
end

function _V_airy_kernel(x, y)
    airyai((x + y)/2)/2
end


function gyration(df::DataFrame, id::Symbol, x::Symbol, y::Symbol)
    spt.add_dr!(df)
    N = maximum(df.TrackID)
    result = DataFrame(TrackID = Int[], value = Float64[])
    dr = mean(filter(!isnan, df.dr))
    @inbounds for n = 1:maximum(df[!, id])
        data = extract(df, n, id, x, y)
        T = size(data, 1)
        x̄ = mean(data[:, 1])
        ȳ = mean(data[:, 2])
        c = 0.0
        @inbounds for t = 1:T
            x = abs2(data[t, 1] - x̄)
            y = abs2(data[t, 2] - ȳ)
            c += (x + y)
        end
        c *= ((sqrt(π / 2) / dr) * 1 / T)
        push!(result, [n, c])
    end
    result
end


struct MaxSafeAccel
    dt::Float64
end

"""
Calculate the maximum safe acceleration that will allow the car to avoid a collision if the car in front slams on its brakes
"""
function max_safe_acc(safety::MaxSafeAccel, c::Context, idx::Int)
    dt = safety.dt

    fidx = fore_index(c, idx)
    if fidx == 0
        return Inf
    else
        n_brake_acc = nullable_max_safe_acc(headway(c, idx),
                                            vel(state(c, idx)),
                                            vel(state(c, fidx)),
                                            brake_limit(vehicle(c, idx)), dt)
        return get(n_brake_acc, -brake_limit(vehicle(c, idx)))
    end
end

"""
Return max_safe_acc or an empty Nullable if the discriminant is negative.
"""
function nullable_max_safe_acc(gap, v_behind, v_ahead, braking_limit, dt)
    bp = braking_limit
    v = v_behind
    vo = v_ahead
    g = gap
    # VVV see mathematica notebook
    discriminant = 8.*g*bp + bp^2*dt^2 - 4.*bp*dt*v + 4.*vo^2
    if discriminant >= 0.0
        return Nullable{Float64}(- (bp*dt + 2.*v - sqrt(discriminant)) / (2.*dt))
    else
        return Nullable{Float64}()
    end
end

function is_safe(safety::MaxSafeAccel, c::Context, idx::Int, a::AccelSpeed)
    return a.long_accel <= max_safe_acc(safety, c, idx)
end

struct RoadEdges
    dt::Float64
end

function is_safe(safety::RoadEdges, c::Context, idx::Int, a::AccelSpeed)
    y = pos(state(c, idx))[2]
    w = width(vehicle(c, idx))
    rw = c.roadway
    next_y = y + a.lat_speed*safety.dt
    return next_y - w/2 >= -rw.lanewidth/2 && next_y + w/2 <= (rw.n_lanes-0.5)*rw.lanewidth
end


module TestStructuredGraphicalLinearAlgebra
using Test

using Catlab, Catlab.Doctrines
using Catlab.LinearAlgebra.StructuredGraphicalLinearAlgebra

R = ℝ(FreeStructuredLinearFunctions.Ob)
A, B = Ob(FreeStructuredLinearFunctions, :A, :B)
f, g = Hom(:f, A, B), Hom(:g, A, B)
u, v = Hom(:u, R, A), Hom(:g, R, B)

# Domains and codomains
@test dom(plus(f,g)) == A
@test codom(plus(f,g)) == B

# Unicode syntax
@test f+g == plus(f,g)

# Banded Matrices
Dᵤ = diag(u)
@test dom(Dᵤ) == A
@test codom(Dᵤ) == A

Dᵥ = diag(v)
@test dom(Dᵥ) == B
@test codom(Dᵥ) == B

@test dom(Dᵥ⋅Dᵥ) == B
@test codom(Dᵥ⋅Dᵥ) == B

@test dom(Dᵥ⊕Dᵥ) == B⊕B
@test codom(Dᵥ⊕Dᵥ) == B⊕B
@test dom(Dᵤ⊕Dᵥ) == A⊕B
@test codom(Dᵤ⊕Dᵥ) == A⊕B

end


# ==============================================================================
# Modele JuMP pour calculer la relaxation linéaire du 2SPA sur un objectif donne, avec eventuellement une ϵ-contrainte
function computeLinearRelax2SPA(nbvar::Int, nbcontraintes::Int, L::Array{Int,2}, c1::Array{Int,1}, c2::Array{Int,1}, epsilon, obj::Int)

    model = Model(with_optimizer(GLPK.Optimizer))
#    model = Model(with_optimizer(Gurobi.Optimizer, GUROBI_ENV))
#    set_optimizer_attributes(model, "OutputFlag" => 0)

    @variable(model, 0.0 <= x[1:nbvar] <= 1.0 )
    @constraint(model, [i=1:nbcontraintes],(sum((x[j]*L[i,j]) for j in 1:nbvar)) == 1)
    if obj == 1
        @objective(model, Min, sum((c1[i])*x[i] for i in 1:nbvar))
        @constraint(model, sum((c2[i])*x[i] for i in 1:nbvar) <= epsilon)
    else
        @objective(model, Min, sum((c2[i])*x[i] for i in 1:nbvar))
        @constraint(model, sum((c1[i])*x[i] for i in 1:nbvar) <= epsilon)
    end
    optimize!(model)
    return objective_value(model), value.(x)
end


immutable DiscreteUniform <: DiscreteUnivariateDistribution
    a::Int
    b::Int
    function DiscreteUniform(a::Real, b::Real)
        if a < b
            new(int(a), int(b))
        else
            error("a must be less than b")
        end
    end
end

DiscreteUniform(b::Integer) = DiscreteUniform(0, b)
DiscreteUniform() = DiscreteUniform(0, 1)

function cdf(d::DiscreteUniform, k::Real)
    if k < d.a
        return 0.0
    elseif <= d.b
        return (ifloor(k) - d.a + 1.0) / (d.b - d.a + 1.0)
    else
        return 1.0
    end
end

entropy(d::DiscreteUniform) = log(d.b - d.a + 1.0)

insupport(d::DiscreteUniform, x::Number) = isinteger(x) && d.a <= x <= d.b

function kurtosis(d::DiscreteUniform)
    n = d.b - d.a + 1.0
    return -(6.0 / 5.0) * (n^2 + 1.0) / (n^2 - 1.0)
end

mean(d::DiscreteUniform) = (d.a + d.b) / 2.0

median(d::DiscreteUniform) = (d.a + d.b) / 2.0

function mgf(d::DiscreteUniform, t::Real)
    a, b = d.a, d.b
    return (exp(t * a) - exp(t * (b + 1))) /
           ((b - a + 1.0) * (1.0 - exp(t)))
end

function cf(d::DiscreteUniform, t::Real)
    a, b = d.a, d.b
    return (exp(im * t * a) - exp(im * t * (b + 1))) /
           ((b - a + 1.0) * (1.0 - exp(im * t)))
end

modes(d::DiscreteUniform) = [d.a:d.b]

function pdf(d::DiscreteUniform, x::Real)
    if insupport(d, x)
        return (1.0 / (d.b - d.a + 1))
    else
        return 0.0
    end
end

function quantile(d::DiscreteUniform, p::Real)
    n = d.b - d.a + 1
    return d.a + ifloor(p * n)
end

rand(d::DiscreteUniform) = d.a + rand(0:(d.b - d.a))

skewness(d::DiscreteUniform) = 0.0

var(d::DiscreteUniform) = ((d.b - d.a + 1.0)^2 - 1.0) / 12.0


using DataStructures
using Test

# Function to check to make sure we are on the canvas
function inbounds(canvas_size, loc)

    # Make sure we are not beneath or to the left of the canvas
    if minimum(Tuple(loc)) < 1
        return false

    # Make sure we are not to the right of the canvas
    elseif loc[2] > canvas_size[2]
        return false

    # Make sure we are not above the canvas
    elseif loc[1] > canvas_size[1]
        return false
    else
        return true
    end
end

function find_neighbors(canvas, loc::CartesianIndex, old_val, new_val)

    # Finding north, south, east, west neighbors
    possible_neighbors = [loc + CartesianIndex(0, 1),
                          loc + CartesianIndex(1, 0),
                          loc + CartesianIndex(0, -1),
                          loc + CartesianIndex(-1, 0)]

    # Exclusing neighbors that should not be colored
    neighbors =  []
    for possible_neighbor in possible_neighbors
        if inbounds(size(canvas), possible_neighbor) &&
           canvas[possible_neighbor] == old_val
            push!(neighbors, possible_neighbor)
        end
    end

    return neighbors
end

function stack_fill!(canvas, loc::CartesianIndex, old_val, new_val)
    if new_val == old_val
        return
    end

    s = Stack{CartesianIndex}()
    push!(s, loc)

    while length(s) > 0
        current_loc = pop!(s)
        if canvas[current_loc] == old_val
            canvas[current_loc] = new_val
            possible_neighbors = find_neighbors(canvas, current_loc,
                                                old_val, new_val)
            for neighbor in possible_neighbors
                push!(s,neighbor)
            end
        end

    end
end

function queue_fill!(canvas, loc::CartesianIndex, old_val, new_val)
    if new_val == old_val
        return
    end

    q = Queue{CartesianIndex}()
    enqueue!(q, loc)

    # Coloring the initial location
    canvas[loc] = new_val

    while length(q) > 0
        current_loc = dequeue!(q)

        possible_neighbors = find_neighbors(canvas, current_loc,
                                            old_val, new_val)

        # Coloring as we are enqueuing neighbors
        for neighbor in possible_neighbors
            canvas[neighbor] = new_val
            enqueue!(q,neighbor)
        end

    end
end

function recursive_fill!(canvas, loc::CartesianIndex, old_val, new_val)

    if (old_val == new_val)
        return
    end

    canvas[loc] = new_val

    possible_neighbors = find_neighbors(canvas, loc, old_val, new_val)
    for possible_neighbor in possible_neighbors
        recursive_fill!(canvas, possible_neighbor, old_val, new_val)
    end
end

function main()

    # Creation of a 5x5 grid with a single row of 1.0 elements
    grid = zeros(5,5)
    grid[3,:] .= 1

    # Create solution grid
    answer_grid = zeros(5,5)
    answer_grid[1:3, :] .= 1

    # Start filling at 1,1
    start_loc = CartesianIndex(1,1)

    @testset "Fill Methods" begin
        # Use recursive method and reinitialize grid
        recursive_fill!(grid, start_loc, 0.0, 1.0)
        @test grid == answer_grid

        grid[1:2,:] .= 0

        # Use queue method and reinitialize grid
        queue_fill!(grid, start_loc, 0.0, 1.0)
        @test grid == answer_grid

        grid[1:2,:] .= 0

        # Use stack method and reinitialize grid
        stack_fill!(grid, start_loc, 0.0, 1.0)
        @test grid == answer_grid
    end

end

main()


# You can see here what is Agent-based Model and its pros and cons.
# Throughout this post, I would like to display how we can model the COVID-19 spread.

# You can build ABM(Agent-Based Model) pretty easily with Agents.jl.
# There are five steps to build a ABM.

# but before we proceeding it, we need to think about our model.

# ## The SIR model
# SIR stands for Succeptable, Infected and Recovered.
# It is just modelling the general cases that once you got COVID-19 and recovered, you are not infected anymore.
# With the three states(S, I, R), there are one more states of agents which is dead.
# Therefore, we need to know what is going on here.

# ## Set a Space
# First of all, you need to set space the agent will move around.
# Agents.jl provide several options(you can check that [here](https://juliadynamics.github.io/Agents.jl/stable/api/#Available-spaces))
# Here for your simple model.
using Agents


using InvariantMeasures
using ValidatedNumerics

# Testing whether the "shooting method" beats taking successive preimages

# f = x -> 2*x + 0.1*InvariantMeasures.sinpi(2*x)
# preim1 = root(x ->f(x)-1., 0..4, 0)

f = x-> 2.5x*(1-x) + 0.01*x*x + 0.3*x*x*x + 0.003*x*x*x*x
preim1 = @interval(0.5)

mypreim(y) = root(x -> f(x)-y, hull(0..0,preim1), 0)

n = 15

y0 = @interval(0.21)

# method 1
ys = fill(y0, n+1)
ys[1] = y0
for i = 1:n
    ys[i+1] = mypreim(ys[i])
end

# method 2

zs = fill(hull(0..0,preim1), n)
fs = fill(f, n)

nthpreimage!(y0, fs, zs)

z = zs[1]

[zs[end:-1:1] ys[2:end]]
[diam.(zs[end:-1:1]) diam.(ys[2:end])]

diam.(zs[end:-1:1]) ./ diam.(ys[2:end])


using Images

function cubicbezier!(xy::Matrix,
                      img::Matrix = fill(RGB(255.0, 255.0, 255.0), 17, 17),
                      col::ColorTypes.Color = convert(eltype(img), Gray(0.0)),
                      n::Int = 20)
    t = collect(0:n) ./ n
    M = hcat((1 .- t) .^ 3, # a
             3t .* (1 .- t) .^ 2, # b
             3t .^ 2 .* (1 .- t), # c
             t .^ 3) # d
    p = floor.(Int, M * xy)
    for i in 1:n
        drawline!(img, p[i, :]..., p[i+1, :]..., col)
    end
    return img
end

xy = [16 1; 1 4; 3 16; 15 11]
cubicbezier!(xy)


using Test, Clines, LinearFractionalTransformations

@testset "Construction" begin
    C = Circle(1 + 0im, im, -im)
    @test radius(C) == 1
    @test center(C) == 0

    D = Circle(0, 0, 1)
    @test C == D

    E = Circle(0im, 1)
    @test C == E

    P = three_points(E)
    @test C == Cline(P...)

    L = Line(1, 2)
    @test L == Line(-1, 3)
    @test L == Cline(1, 5, Inf)

    LL = dilate(L)
    @test LL == L
end

@testset "Properties" begin
    C = Circle(2, 3, 2)
    @test area(C) == 4π
    @test circumference(C) == 4π

    L = Cline(0, 1 + im, 2 + 2im)
    @test slope(L) == 1
end

@testset "Set Ops" begin
    C = Circle(0im, 6)
    D = Circle(4 + 0im, 2)
    @test D ⊆ C
    S = C ∩ D
    @test length(S) == 1

    D = Circle(4 + 0im, 3)
    S = C ∩ D
    @test length(S) == 2

    D = Circle(4 + 0im, 1)
    S = C ∩ D
    @test isempty(S)
end

@testset "Intersection" begin
    C = Circle(0, 0, 1)
    L = Line(0, 3 + 4im)
    @test length(C ∩ L) == 2
    D = Circle(2, 0, 1)
    @test length(C ∩ D) == 1

    D = Circle(5, 5, 1)
    @test C ∩ D == Set{Complex{Float64}}()

    X = Line(0, 1)
    Y = Line(im, 2 + im)
    @test length(X ∩ Y) == 0

    Y = Cline(-5, 2, Inf)
    @test X == Y
end

@testset "LFT" begin
    C = Circle(1 + im, 3)
    F = LFT()
    D = F(C)
    @test abs(center(C) - center(D)) <= get_tolerance()
    @test abs(radius(C) - radius(D)) <= get_tolerance()

    D = Circle(2 - im, 5)
    F = LFT(C, D)
    @test F(C) == D
    @test (inv(F))(D) == C

    F = LFT(C)
    X = Line(0, 1)
    @test (inv(F))(X) == C

end

@testset "Inversion" begin
    X = Line(0, 1)
    z = 3 - 2im
    @test inv(X, z) == conj(z)

    C = Circle(2, 3, 5)
    L = Line(0, 1 - im)

    LL = inv(C, L)
    @test inv(C, LL) == L

    F = inv(L)
    @test F(F(C)) == C

end

@testset "Collinearity" begin
    pts = 2 * collect(1:3) - im * collect(1:3)
    @test collinear(pts...)
    @test isa(Cline(pts...), Line)

    C = Circle(4, 6, 10)
    PP = three_points(C)
    @test !collinear(PP...)
end

@testset "Kiss" begin
    pts = [0, 2im, 3 + im]
    CC = kiss(pts...)
    @test length(CC[1] ∩ CC[2]) == 1
    S = soddy(CC...)
    @test length(CC[3] ∩ S) == 1
end


using CompScienceMeshes
using Test

for T in [Float32, Float64]
    local p1 = point(T,1,0,0)
    local p2 = point(T,0,1,0)
    local p3 = point(T,0,0,1)
    local p4 = point(T,0,0,0)

    local tet = simplex(p1,p2,p3,p4)
    c = neighborhood(tet, point(T,1,1,1)/4)

    @test cartesian(c) ≈ point(T, 0.25, 0.25, 0.25)
    @test parametric(c) ≈ point(T, 0.25, 0.25, 0.25)

    @test length(c) == 3
    @test c[1] ≈ 0.25
    @test c[2] ≈ 0.25
    @test c[3] ≈ 0.25

    @test jacobian(1.234) ≈ 1

    @test tangents(c,1) ≈ point(T,1,0,0)
    @test tangents(c,2) ≈ point(T,0,1,0)
    @test tangents(c,3) ≈ point(T,0,0,1)

    tri = simplex(p1,p2,p4)
    c_tri = neighborhood(tri, point(T,1,1)/3)

    @test normal(c_tri) ≈ point(T,0,0,1)
end


function assess(α::T, β::T, x::T, tol::T) where T <: AbstractFloat
    Nopt = ceil(Integer, x^(1/α))
    term = NaN
    Nmin = -1
    for n = 1:Nopt
        term = exp( loggamma(n*α-β+1) - n*log(x) )
        if term < tol
            Nmin = n
            break
        end
    end
    return Nmin, term
end


export FPSystem

"""
    FPSystem(polynomials, vars) <: AbstractSystem

Create a polynomial system using the [`FixedPolynomials`](https://github.com/JuliaAlgebra/FixedPolynomials.jl) package.
"""
struct FPSystem{T} <: AbstractSystem
    system::FP.System{T}
end

FPSystem(polys::Vector{<:MP.AbstractPolynomial}, vars) = FPSystem(FP.System(polys, vars))
function FPSystem(polys::Vector{<:MP.AbstractPolynomial};
    variables=nothing, parameters=nothing)

    parameters === nothing || error("Parameters are not supported by FPSystem")

    if variables === nothing
        FPSystem(FP.System(polys))
    else
        FPSystem(FP.System(polys, variables))
    end
 end

struct FPSystemCache{JC<:FP.JacobianConfig} <: AbstractSystemCache
    config::JC
end

cache(F::FPSystem, x) = FPSystemCache(FP.config(F.system, x))

Base.size(F::FPSystem) = (length(F.system), FP.nvariables(F.system))

evaluate!(u, F::FPSystem, x, c::FPSystemCache) = FP.evaluate!(u, F.system, x, c.config)
evaluate(F::FPSystem, x, c::FPSystemCache) = FP.evaluate(F.system, x, c.config)
jacobian!(U, F::FPSystem, x, c::FPSystemCache) = FP.jacobian!(U, F.system, x, c.config)
jacobian(F::FPSystem, x, c::FPSystemCache) = FP.jacobian(F.system, x, c.config)
function evaluate_and_jacobian!(u, U, F::FPSystem, x, c::FPSystemCache)
    FP.evaluate_and_jacobian!(u, U, F.system, x, c.config)
end
function evaluate_and_jacobian(F::FPSystem, x, c::FPSystemCache)
    FP.evaluate_and_jacobian(F.system, x, c.config)
end

weylnorm2(F::FPSystem) = FP.weyldot(F.system.polys, F.system.polys)


using transD_GP.MT1D, PyPlot
##
ntimesperdecade = 10
T = 10 .^(-3:1/ntimesperdecade:6)
f = 1 ./T
## model 1 conductor above
ρ1 = [10, 100]
z = [0, 5000]
h = diff(z)
Z = MT1D.Z_f(f, ρ1, h)
fig = MT1D.plotcurve(T, Z)
## model 2 conductor below
ρ2 = [100, 10]
Z = MT1D.Z_f(f, ρ2, h)
MT1D.plotcurve(T, Z, fig)
## and to do this with the models plotted
fig = MT1D.plotmodelcurve(T, ρ1, z)
MT1D.plotmodelcurve(T, ρ2, z, fig)


using Hungarian
using LinearAlgebra
using Random
using Distributions

"""
Run the algorithm from "Maximum Quadratic Assignment Problem: Reduction
from Maximum Label Cover and LP-based Approximation Algorithm" by
Makarychev, Manokaran, and Sviridenko (MMS) to approximately solve the QAP.
"""
function qap_mms(A, B, obj)
	@assert obj == :max
	x, y = adams_johnson_linearization(A, B, :max, integer=false)
	perm = rounding_mms(A, B, x, y)
	P = Matrix(I, size(A)...)[perm,:]
	return P, perm
end

"""
LP rounding step of the MMS algorithm.
"""
function rounding_mms(A, B, x, y)
	N = size(A, 1)
	n2 = Int(floor(N/2))

	La = randperm(N)[1:n2]
	Lb = randperm(N)[1:n2]

	Ra = setdiff(1:N, La)
	Rb = setdiff(1:N, Lb)

	ϕ̃ = fill(-1, N)
	for u in La
		probs = x[u,Lb]
		push!(probs, 1 - sum(probs))
		probvec!(probs)
		p′ = rand(Categorical(probs))
		p = (p′ == length(probs)) ? -1 : p′
		ϕ̃[u] = p
	end

	ϕ = fill(-1, N)
	for p in Lb
		us = findall(==(p), ϕ̃)
		if isempty(us) continue end
		u = rand(us)
		ϕ[u] = p
	end

	La = findall(!=(-1), ϕ)
	if isempty(La)
		return rounding_mms(A, B, x, y)
	end

	W = [sum(A[u,v] * B[ϕ[u],q] for u in La) for v in Ra, q in Rb]
	W = maximum(W) .- W
	matching = Hungarian.munkres(W)
	ψ = fill(-1, N)
	for i in 1:Int(ceil(N/2))
		j = findfirst(matching[i,:] .== Hungarian.STAR)
		ψ[Ra[i]] = Rb[j]
	end

	perm = [(ϕ[i] > 0) ? ϕ[i] : (ψ[i] > 0) ? ψ[i] : -1 for i in 1:N]

	missing_ind = findall(==(-1), perm)
	missing_vals = shuffle(setdiff(1:N, perm))
	perm[missing_ind] = missing_vals

	return perm
end

function probvec!(xs::Array{<:Real,1})
	z = zero(xs[1])
	s = z
	for i in 1:length(xs)
		if xs[i] < z
			xs[i] = z
		end
		s += xs[i]
	end
	if s != 1
		for i in 1:length(xs)
			xs[i] = xs[i] / s
		end
	end
	return xs
end


include("splinesSerial.jl")

controlpoints=[
   [[ 0,0,0],[0 ,3  ,4],[0,6,3],[0,10,0]],
   [[ 3,0,2],[2 ,2.5,5],[3,6,5],[4,8,2]],
   [[ 6,0,2],[8 ,3 , 5],[7,6,4.5],[6,10,2.5]],
   [[10,0,0],[11,3  ,4],[11,6,3],[10,9,0]]]

bicubic = larBicubicSurface(controlpoints)
point = [0.7,0.2]
result = bicubic(point)


"""
    Morse{T<:AbstractFloat} <: EmpiricalPotential{NamedTuple{(:D, :α, :σ)},NamedTuple{(:rcutoff,)}}

The Morse potential is a simple two-body intermolecular potential with three typical parameters, primarily describing the interaction of neutral atoms with more complex bond interactions. Formally, the interaction between two atoms at a distance, ``r``, is given by

```math
\\begin{equation}
V_{M}(r; D, \\alpha, \\sigma, rcutoff, species) =
    \\begin{cases}
    0 & r > rcutoff \\\\
    D \\left( 1 - e^{\\alpha(r - \\sigma)}\right)^2 & r < rcutoff.
    \\end{cases}
\\end{equation}
```

Users must supply three parameters, ``D`` (units of energy), ``\\alpha`` (units of inverse distance), ''\\sigma'' (units of distance), and radial cutoff (units of distance).
"""
struct Morse{T<:AbstractFloat} <: EmpiricalPotential{NamedTuple{(:D, :α, :σ)},NamedTuple{(:rcutoff,)}}
    D::T
    α::T
    σ::T
    rcutoff::T
    species::Tuple
    function Morse(D, α, σ, rcutoff, species)
        D, α, σ, rcutoff = promote(austrip(D), α, austrip(σ), austrip(rcutoff))
        new{typeof(rcutoff)}(D, α, σ, rcutoff, Tuple(species))
    end
    function Morse{T}(; D::T, α::T, σ::T, rcutoff::T, species::Tuple) where {T}
        new{T}(D, α, σ, rcutoff, species)
    end
end

get_rcutoff(m::Morse) = m.rcutoff
get_species(m::Morse) = m.species

_morse_exp(R::AbstractFloat, m::Morse) = exp(-m.α * (R - m.σ))

potential_energy(R::AbstractFloat, m::Morse) = m.D * (1 - _morse_exp(R, m))^2
force(R::AbstractFloat, r::SVector{3}, m::Morse) = (2 * m.D * m.α * (1 - _morse_exp(R, m)) * _morse_exp(R, m) / R)r


plasma_frequency(n, q, m) = sqrt(4*pi*n*q^2 / m)
gyro_frequency(q, B, m) = q*B / (m * 3e10)
electron_collision_time(n, T, q, m, l) = 3*sqrt(m)*T^(3/2) / (4*sqrt(2*pi)*l*n*q^4)
ion_collision_time(n, T, q, m, l) = 3*sqrt(m)*T^(3/2) / (4*sqrt(pi)*l*n*q^4)
sound_speed(T, m) = sqrt(T / m)
plasma_beta(n, T, B) = 4*pi*n*T / B^2

#NOTE: input cgs units, except T0, input eV for T0
function dimensionless_parameters(a, R0, n0, T0, B0)

    # constants
    c = 3.00e10
    q = 4.8032e-10
    me = 9.1094e-28
    mp = 1.6726219e-24

    # convert T0 from eV to erg
    T0 = T0 * 1.60218e-12

    # assume deuterium ions
    mi = 2 * mp

    # important parameters
    wpe = plasma_frequency(n0, q, me)
    wci = gyro_frequency(q, B0, mi)
    tau_i = ion_collision_time(n0, T0, q, mi, 1)
    tau_e = electron_collision_time(n0, T0, q, me, 1)
    beta0 = plasma_beta(n0, T0, B0)
    cs0 = sound_speed(T0, mi)
    t0 = sqrt(a*R0/2) / cs0

    # parameters for gdb model
    am = R0*beta0/a
    ad = cs0^2 * t0 / (wci*a^2)
    ki = 2*3.9*t0*tau_i*T0 / (R0^2 * mi)
    ke = 2*3.2*t0*tau_e*T0 / (R0^2 * me)
    er = 2*a/R0
    eg = 0.08*tau_i / t0
    ev = cs0*t0/R0
    de2 = c / (a*wpe)
    eta = 0.51*t0*de2 / tau_e

    return Float64[am, ad, ki, ke, er, eg, ev, de2, eta]
end

function partial_s(f, f_x, f_y, psi_x, psi_y, Ds, am)

    return Ds*f + am*(psi_x .* f_y - psi_y .* f_x)
end

function partial_ss(f, f_s, f_x, f_y, f_xy, f_xx, f_yy, psi_x, psi_y, psi_xy, psi_xx, psi_yy, Dss, Ds, Dx, Dy, am)

    _a = Dss*f + am*Ds*(psi_x.*f_y - psi_y.*f_x) + am*(psi_x.*Dy*f_s - psi_y.*Dx*f_s)
    _b = (f_x.*psi_xy + f_xx.*psi_y - f_y.*psi_xx - psi_x.*f_xy).*psi_y
    _c = (f_x.*psi_yy - f_y.*psi_xy - f_yy.*psi_x + psi_y.*f_xy).*psi_x

    return _a + am^2 * (_b - _c)
end

function lnn_total_ion_derivative(phi_c, Pe_c, n, j_s, u_s, er, ev, ad)

    _a = er*phi_c - er*ad*Pe_c ./ n
    _b = er*ad*j_s ./ n - ev*u_s
    return _a + _b
end

function lnTe_total_electron_derivative_adiabatic(lnn_Dt, j, lnn_s, n, Te, lnTe_c, er, ad)

    _a = (2/3)*(lnn_Dt + er*ad*j.*lnn_s./n)
    _b = -(5/3)*er*ad*(Te .* lnTe_c)

    return _a + _b
end

function lnTi_total_ion_derivative_adiabatic(lnn_Dt, Ti, lnTi_c, er, ad)

    _a = (2/3)*lnn_Dt
    _b = (5/3)*er*ad*(Ti .* lnTi_c)

    return _a + _b
end

# NOTE: this one is written to be species-ambiguous since it applies to Te and Ti
function lnT_partial_thermal(lnT_s, lnT_ss, T, n, k)

    _a = (7/2)*lnT_s.^2 + lnT_ss
    return k*T.^(5/2) .* _a ./ n
end

function u_total_ion_derivative(Pe_s, Pi_s, G_s, n, Ti, u_c, ev, eg, er, ad)

    _a = -ev*(Pe_s + Pi_s) ./ n
    _b = -4*ev*eg*G_s ./ n
    _c = er*ad*(Ti .* u_c)

    return +(_a, _b, _c)
end

function general_partial_derivative(f_Dt, f_x, f_y, f_s, phi_x, phi_y, v)

    return f_Dt - phi_x.*f_y + phi_y.*f_x - v.*f_s
end

function w_partial_t(n, lnn_x, lnn_y, lnn_xy, lnn_xx, lnn_yy,
                     Pe_x, Pe_y, Pe_xy, Pe_xx, Pe_yy, Pe_xxx, Pe_yyy, Pe_xxy, Pe_xyy, Pe_c,
                     phi_x, phi_y, phi_xy, phi_xx, phi_yy, phi_xxx, phi_yyy, phi_xxy, phi_xyy,
                     G_c, Pi_c, j_s, ad, eg)

    _a = (lnn_x.*lnn_y + lnn_xy).*(Pe_y.*phi_y - Pe_x.*phi_x)
    _b = (lnn_x.^2 + lnn_xx).*Pe_x.*phi_y - (lnn_y.^2 + lnn_yy).*Pe_y.*phi_x
    _c = -Pe_x.*lnn_x.^2 .*phi_y + Pe_x.*lnn_x.*lnn_y.*phi_x
    _d = Pe_x.*lnn_x.*phi_xy - Pe_x.*lnn_y.*phi_xx + Pe_xx.*lnn_x.*phi_y
    _e = -Pe_xx.*lnn_y.*phi_x - Pe_xx.*phi_xy - Pe_xxx.*phi_y
    _f = -Pe_y.*lnn_x.*lnn_y.*phi_y + Pe_y.*lnn_x.*phi_yy
    _g = Pe_y.*lnn_y.^2 .*phi_x - Pe_y.*lnn_y.*phi_xy + Pe_yy.*lnn_x.*phi_y
    _h = -Pe_yy.*lnn_y.*phi_x + Pe_yy.*phi_xy + Pe_yyy.*phi_x + phi_x.*Pe_xxy
    _i = phi_xx.*Pe_xy - ad.*phi_y.*Pe_xyy - phi_yy.*Pe_xy
    _j = lnn_x.*phi_x.*phi_xy - lnn_x.*phi_xx.*phi_y + lnn_y.*phi_x.*phi_yy - lnn_y.*phi_y.*phi_xy
    _k = -lnn_x.*phi_xx.*phi_y + lnn_y.*phi_x.*phi_yy - lnn_y.*phi_y.*phi_xy
    _l = phi_x.*phi_yyy + phi_x.*phi_xxy - phi_xxx.*phi_y - phi_y.*phi_xyy

    _A = ad * +(_a, _b, _c, _d, _e, _f, _g, _h, _i)
    _B = n .* +(_j, _k, _l)

    return _A + _B - eg*G_c - Pe_c - Pi_c + j_s
end

function A_partial_t(phi_x, phi_y, phi_s, j, jn, jn_x, jn_y, jn_s, Pe_s, u, n, Te, de2, ev, ad, am, eta)

    _a = de2 * (phi_x .* jn_y - phi_y .* jn_x)
    _b = de2 * (u - jn / ev).*jn_s / ev
    _c = (phi_s - ad*Pe_s ./ n) / am
    _d = eta*(j ./ Te.^(3/2))

    return +(_a, _b, _c, _d)
end

function Abohm_def(Te, Ti, n, phi, er, ev, ad, de2)

    cs = sqrt.(Te + Ti)
    lambda = 2.695
    jbohm = ev * n.*cs.*(1 .- exp.(lambda .- phi ./(ad*Te))) / (ad*er)

    return -de2*jbohm
end

function G_def(Ti, phi_c, Pi_c, u_s, er, ev, ad)

    return Ti.^(5/2) .* (er*(phi_c + ad*Pi_c) - 4*ev*u_s)
end


using DiffEqBase, DiffEqOperators, LinearAlgebra, Zygote, Test
const A = rand(Float32, 300, 300)
f(du, u, p, t) = mul!(du, A, u)
f(u, p, t) = A * u
x = rand(Float32, 300)
v = rand(Float32, 300)
du = similar(x)

J = VecJacOperator(f, x)
actual_vjp = Zygote.jacobian(x -> f(x, nothing, nothing), x)[1]' * v
@test J * v ≈ actual_vjp
J = VecJacOperator(f, x; autodiff = false)
@test J * v ≈ actual_vjp


module IGA

# import dependent packages
using FastGaussQuadrature, PyCall, PyPlot

# export types
export Degree, KnotVector, BezierPatch, NURBS, QuadRule

# export functionality from other libs
export ndgrid, allcomb, linearspace, boundary
export buildvector, uniqueknots, globalrefine, findspan
export dimsplinespace, grevillepoints, onebasisfuneval, bsplinebasiseval, dersbsplinebasisfuns, dersbsplinebasisfunsinterpolationmatrix
export knotinsertionoperator!, knotinsertionoperator, degreeelevationoperator, bezierdecompoperator
export dim

# export functionality from nurbs.jl
export nsdim, nbezierpatches, ienarray
export refine!, globalrefine!, degreeelevate!
export boundary, ∂

# export functionality from this module
export dersbezierbasisfuns, assembly, plot, plot_stress_contours, L2residual

include("functions/baseextension.jl")
include("functions/knotvector.jl")
include("functions/bsplinebasisfuns.jl")
include("functions/knotrefinement.jl")
include("functions/nurbs.jl")
include("functions/quadrature.jl")

######################################################################################################
####################################### Shape function routines ######################################
######################################################################################################

# bernstein basis functions on [-1,1] and their derivatives
dersbezierbasisfuns(p::Int, u::Float64, nout::Int) = dersbsplinebasisfuns(p, buildvector([-1.0;1.0],[p+1;p+1]), p+1, u, nout)'

# bernstein basis functions on [-1,1] and their derivatives at a set of points u
function dersbezierbasisfuns(p::Int, u::Vector{Float64}, nout::Int)
    ϕ = zeros(p+1,length(u),nout)
    for k in 1:length(u)
        ϕ[:,k,:] = dersbezierbasisfuns(p, u[k], nout)
    end
    return ϕ
end

# Compute 1D shape function and derivatives in a single point
function dersbezierbasisfuns(S::BezierPatch{1}, ϕ::Vector{Float64}, dϕ::Vector{Float64})

    # initialize
    C = S.dec
    W = S.wpts
    P = S.cpts

    # NURBS are B-splines in projective three-space
    Bʷ  = broadcast(*, W, C[1]'* ϕ)
    ∇Bʷ = broadcast(*, W, C[1]'* dϕ)

    # compute weightfunction and derivative
    w  = sum(Bʷ)
    ∇w = sum(∇Bʷ)

    # compute NURBS basisfunction and derivative
    R  = Bʷ / w
    ∇R = (∇Bʷ * w - Bʷ * ∇w) / w^2

    # return quadtities in physical space
    return R, ∇R
end

# Compute 1D shape function and derivatives in an array of points
function dersbezierbasisfuns(S::BezierPatch{1}, ϕ::Matrix{Float64}, dϕ::Matrix{Float64})
    R, ∇R = zeros(ϕ), zeros(dϕ)
    for k in 1:size(R,2)
        R[:,k], ∇R[:,k] = dersbezierbasisfuns(S, ϕ[:,k], dϕ[:,k])
    end
    return R, ∇R
end

# Compute 2D shape function, derivatives and determinant in a single point
function dersbezierbasisfuns(S::BezierPatch{2}, ϕ::NTuple{2,Vector{Float64}}, dϕ::NTuple{2,Vector{Float64}})

    # initialize
    C = S.dec
    W = S.wpts
    P = S.cpts

    # compute B-spline basis functions
    ϕ  = ntuple(k-> C[k]'* ϕ[k], 2)
    dϕ = ntuple(k-> C[k]'*dϕ[k], 2)

    # NURBS are B-splines in projective 3-space
    Bʷ  = broadcast(*, W,  kron(ϕ[2], ϕ[1]))
    ∇Bʷ = broadcast(*, W, [kron(ϕ[2], dϕ[1]) kron(dϕ[2], ϕ[1])])

    # compute weightfunction
    w  = sum(Bʷ)
    ∇w = sum(∇Bʷ, 1)

    # compute NURBS functions and derivatives
    R  = Bʷ / w
    ∇R = (∇Bʷ * w - Bʷ * ∇w) / w^2

    # return quadtities in physical space
    return R, ∇R
end

# Compute 2D shape function and derivatives in an array of points
function dersbezierbasisfuns(S::BezierPatch{2}, ϕ::NTuple{2,Matrix{Float64}}, dϕ::NTuple{2,Matrix{Float64}})
    n = prod([size(ϕ[k],1) for k in 1:2]); m = ntuple(k -> size(ϕ[k],2), 2)
    R, ∇R = zeros(n,prod(m)), zeros(n,prod(m),2)
    i = 1
    for k in 1:m[2]
        for l in 1:m[1]
            R[:,i], ∇R[:,i,:] = dersbezierbasisfuns(S, (ϕ[1][:,k], ϕ[2][:,l]), (dϕ[1][:,k], dϕ[2][:,l]))
            i+=1
        end
    end
    return R, ∇R
end

######################################################################################################
########################################## Assembly routines #########################################
######################################################################################################

typealias Forcing  Function
typealias Traction Function
typealias BoundaryDisplacement Function
typealias Material Function

function assembly(S::NURBS, D::Material, f::Forcing, g::BoundaryDisplacement, h::Traction)

    #########################  Initialize #############################
    p     = S.deg                            # polynomial degree
    nsd   = nsdim(S)                         # number of space dimensions
    ndofs = dimsplinespace(S)                # dimension of the spline space
    dims  = nbezierpatches(S)                # number of elements in each direction

    IEN   = S.ien                            # set IEN-array
    ID    = S.id                             # set ID-array

    neq   = maximum(ID)                      # number of equations
    K     = spzeros(neq,neq)                 # allocate space for the stiffness matrix
    F     = zeros(neq)                       # allocate space for the righthandside forcing

    ########################  Precompute data #########################

    # precompute Gauss-Legendre nodes and weights
    qr = ntuple(k -> QuadRule(p[k]+1,"legendre"), nsd)

    # pre-compute univariate Bernstein polynomials and their derivatives
    ϕ = ntuple(k -> dersbezierbasisfuns(p[k], qr[k].points, 2), nsd)

    ########################  Assembly loop over elements #########################

    # integration over domain Ω
    for e in 1:size(IEN,2)

        # initialize
        A = IEN[:,e]                       # global node numbers
        P = ID[:,A]                          # global equation numbers
        i,j = ind2sub(dims,e)                # element subscripts

        # construct Bezier Patch
        Sb = BezierPatch(tuple(p...), S.cpts[A,:], S.wpts[A], (S.dec[1][:,:,i], S.dec[2][:,:,j]))

        # compute element contribution of the stiffness matrix and the forcing vector
        Kₑ, Fₑ = element_stiffness_and_forcing(Sb, ϕ, qr, D, f)

        # add contributions of element stiffness matrix and right hand side forcing
        I = P.!=0                  # index internal displacements
        K[P[I],P[I]] += Kₑ[I[:],I[:]]
        F[P[I]]      += Fₑ[I[:]]

        # subtract Dirichlet boundary displacements
        for i in 1:2
            for a in 1:size(P,2)
                if P[i,a]==0
                    F[P[I]] -= Kₑ[I,a] * g(vec(S.cpts[A[a],:]))[i]
                end
            end
        end
    end

    # integration over boundary ∂Ω
    for dir = 1:2
        I = 1:nsd.!=dir
        for k in 1:2
            C = ∂(S,dir,k) # boundary k, direction dir, of S
            for e in 1:size(C.ien,2)

                # initialize
                A = C.ien[:,e]                       # global node numbers
                P = ID[:,A][:]                       # global equation numbers
                J = P.!=0                            # dofs with Neumann data

                # construct Bezier Patch
                Cb = BezierPatch(tuple(C.deg...), S.cpts[A,:], S.wpts[A], (C.dec[1][:,:,e],))

                # compute element contribution of the stiffness matrix and the forcing vector
                F[P[J]] += neumann_boundary_condition(Cb, ϕ[I], qr[I], h)[J]

            end
        end
    end

    return K, F
end

# computation of element stiffness and forcing
function element_stiffness_and_forcing(Sb::BezierPatch{2}, ϕ::NTuple{2,Array{Float64,3}}, qr::NTuple{2,QuadRule}, D::Material, f::Forcing)

    # initialize
    nen = prod(Sb.deg+1)                 # number of element nodes
    ned = 2                              # number of element dofs per node
    Kₑ = zeros(nen*ned, nen*ned)         # allocate space for element stiffness matrix
    Fₑ = zeros(nen*ned)                  # allocate space for element forcing vector

    # loop over quadrature points
    for j in 1:dim(qr[2])
        for i in 1:dim(qr[1])

            # shape function routine
            R, ∇R = dersbezierbasisfuns(Sb,(ϕ[1][:,i,1],ϕ[2][:,j,1]), (ϕ[1][:,i,2],ϕ[2][:,j,2]))

            # compute geometrical properties
            J = ∇R' * Sb.cpts                                # Jacobian
            x = vec(R'  * Sb.cpts)                           # quadrature point in physical space
            α = qr[1].weights[i]*qr[2].weights[j]*det(J)     # quadrature weight in physical space

            # compute strain components
            N = (inv(J) * ∇R')'
            B = Matrix{Float64}[[N[a,1]  0.0;
                                 0.0    N[a,2];
                                 N[a,2] N[a,1]] for a in 1:nen]

            # compute element stiffness contribution of quadrature point i,j
            for a in 1:nen
                p = ned*(a-1)+(1:ned)
                for b in a:nen
                    q = ned*(b-1)+(1:ned)
                    Kₑ[p,q] += (B[a]' * D(x) * B[b]) * α
                end
            end

            # compute element forcing contribution of quadrature point i,j
            Fₑ += ((f(x) * α) * R')[:]
        end
    end

    # Use symmetry to determine the remainder of the element stiffness matrix
    for a in 1:nen
        p = ned*(a-1)+(1:ned)
        for b in a+1:nen
            q = ned*(b-1)+(1:ned)
            Kₑ[q,p] = Kₑ[p,q]'
        end
    end
    return Kₑ, Fₑ
end

# computation of element tractions
function neumann_boundary_condition(Cb::BezierPatch{1}, ϕ::Tuple{Array{Float64,3}}, qr::Tuple{QuadRule}, h::Function)

    # initialize
    nen = prod(Cb.deg+1)                 # number of element nodes
    ned = 2                              # number of element dofs per node
    Fₑ = zeros(nen*ned)                  # allocate space for element forcing vector

    # loop over quadrature points
    for i in 1:dim(qr[1])

        # shape function routine
        R, ∇R = dersbezierbasisfuns(Cb, ϕ[1][:,i,1], ϕ[1][:,i,2])

        # compute geometrical properties
        t = vec(∇R' * Cb.cpts)                           # Tangent vector
        x = vec(R'  * Cb.cpts)                           # quadrature point in physical space
        α = qr[1].weights[i]*norm(t)                     # quadrature weight in physical space

        # compute element forcing contribution of quadrature point i,j
        Fₑ += ((h(x) * α) * R')[:]
    end

    return Fₑ
end

######################################################################################################
########################################## Postprocessing routines ###################################
######################################################################################################

import PyPlot.plot

# plot the NURBS geometry
function plot(S::NURBS, gran=(3,3))
    # initialize
    if isempty(S.ien)==true
        S.ien = ienarray(S)
    end
    if isempty(S.dec[1])==true
        S.dec = bezierdecompoperator(S)
    end
    dims = ntuple(i -> size(S.dec[i],3), 2)
    u    = ntuple(i -> linearspace(-1.0,1.0,gran[i]), 2)
    ϕ    = ntuple(i -> dersbezierbasisfuns(S.deg[i], u[i], 2), 2)

    # loop over elements
    for e in 1:prod(dims)
        i,j = ind2sub(dims,e)          # element subscripts
        A   = S.ien[:,e]                 # global node numbers

        # construct Bezier Patch
        Sb  = BezierPatch(tuple(S.deg...), S.cpts[A,:], S.wpts[A], (S.dec[1][:,:,i], S.dec[2][:,:,j]))

        # compute 2-dimensional rational basis functions
        R, ∇R = dersbezierbasisfuns(Sb, (ϕ[1][:,:,1], ϕ[2][:,:,1]), (ϕ[1][:,:,2], ϕ[2][:,:,2]))

        # plot surface
        X = ntuple(i -> reshape(R' * Sb.cpts[:,i], gran), size(Sb.cpts,2))
        contourf(X[1], X[2], zeros(X[1]))

        # plot boundary
        plot(vec(X[1][:,1]),   vec(X[2][:,1]),"k")
        plot(vec(X[1][:,end]), vec(X[2][:,end]),"k")
        plot(vec(X[1][1,:]),   vec(X[2][1,:]),"k")
        plot(vec(X[1][end,:]), vec(X[2][end,:]),"k")
    end
end

# plot the contourlevels of stress
function plot_stress_contours(G::NURBS{2}, S::NURBS{2}, ii::Int, D::Material, levels::Vector{Float64}, gran=(4,4))

    # initialize
    p   = G.deg
    nen = prod(p+1)
    dims = ntuple(i -> size(S.dec[i],3), 2)

    # sample points
    u = ntuple(k -> linearspace(-1.0,1.0,gran[k]), 2)

    # pre-compute univariate Bernstein polynomials and their derivatives
    ϕ = ntuple(k -> dersbezierbasisfuns(p[k], u[k], 2), 2)

    # loop over elements
    for e in 1:size(S.ien,2)
        i,j = ind2sub(dims,e)          # element subscripts
        A   = G.ien[:,e]               # global node numbers

        # construct Bezier Patch
        Sb = BezierPatch(tuple(p...), G.cpts[A,:], G.wpts[A], (G.dec[1][:,:,i], G.dec[2][:,:,j]))

        # compute the stress at the sample points
        σd = zeros(gran[1],gran[2],3)
        X = zeros(gran); Y = zeros(gran)
        for j in 1:gran[2]
            for i in 1:gran[1]

                # shape function routine
                R, ∇R = dersbezierbasisfuns(Sb,(ϕ[1][:,i,1],ϕ[2][:,j,1]), (ϕ[1][:,i,2],ϕ[2][:,j,2]))

                # compute geometrical properties
                J = ∇R' * Sb.cpts                                # Jacobian
                x = vec(R'  * Sb.cpts)                           # quadrature point in physical space

                if det(J)==0.0
                    J = eye(2)
                end

                # compute strain basis functions
                N = (inv(J) * ∇R')'
                B = Matrix{Float64}[[N[a,1]  0.0;
                                     0.0    N[a,2];
                                     N[a,2] N[a,1]] for a in 1:nen]

                # compute the stress in point x
                for a in 1:nen
                    save = D(x)*B[a]*S.cpts[A[a],:]'
                    for k in 1:3
                        σd[i,j,k] += save[k]
                    end
                end

                # save the coordinates
                X[i,j] = x[1]; Y[i,j] = x[2]
            end
        end

        cp = contourf(X, Y, σd[:,:,ii],levels)
        if e==1
            colorbar(cp)
        end
    end
end

typealias Geometry NURBS{2}
typealias Solution NURBS{2}
typealias ExactSolution Function

# Compute the L2 residial in stress
function L2residual(G::Geometry, S::Solution, σ::ExactSolution, D::Material)

    # initialize
    p     = S.deg                            # polynomial degree
    nsd   = nsdim(S)                         # number of space dimensions
    nen   = prod(p+1)
    ndofs = dimsplinespace(S)                # dimension of the spline space
    dims  = nbezierpatches(S)                # number of elements in each direction
    IEN   = S.ien                            # set IEN-array

    # precompute Gauss-Legendre nodes and weights
    qr = ntuple(k -> QuadRule(p[k]+1,"legendre"), nsd)

    # pre-compute univariate Bernstein polynomials and their derivatives
    ϕ = ntuple(k -> dersbezierbasisfuns(p[k], qr[k].points, 2), nsd)

    # integration over domain Ω
    L²res = zeros(3)
    L21 = zeros(3)
    L22 = zeros(3)
    for e in 1:size(IEN,2)

        # initialize
        A = IEN[:,e]                         # global node numbers
        i,j = ind2sub(dims,e)                # element subscripts

        # construct Bezier Patch
        Sb = BezierPatch(tuple(p...), G.cpts[A,:], G.wpts[A], (G.dec[1][:,:,i], G.dec[2][:,:,j]))

        # loop over quadrature points
        for j in 1:dim(qr[2])
            for i in 1:dim(qr[1])

                # shape function routine
                R, ∇R = dersbezierbasisfuns(Sb,(ϕ[1][:,i,1],ϕ[2][:,j,1]), (ϕ[1][:,i,2],ϕ[2][:,j,2]))

                # compute geometrical properties
                J = ∇R' * Sb.cpts                                # Jacobian
                x = vec(R'  * Sb.cpts)                           # quadrature point in physical space
                α = qr[1].weights[i]*qr[2].weights[j]*det(J)     # quadrature weight in physical space

                # compute strain components
                N = (inv(J) * ∇R')'

                # compute the stress in point x
                σd = zeros(3)
                for a in 1:nen
                    # strain basis function `a'
                    B = [N[a,1]    0.0;
                            0.0 N[a,2];
                         N[a,2] N[a,1]]

                    # compute the stress in quadrature point i,j
                    σd += D(x) * B * S.cpts[A[a],:]'
                end
                # compute element forcing contribution of quadrature point i,j
                L²res += (σd - σ(x)).^2 * α
            end
        end

    end
    return sqrt(L²res)
end

end


module CompressedSensing

using LinearAlgebra
using SparseArrays

using LazyInverses
using WoodburyFactorizations
using UpdatableQRFactorizations

const AbstractMatOrFac{T} = Union{AbstractMatrix{T}, Factorization{T}}

using Base.Threads
using StatsBase: sample, mean

using JuMP
using JuMP: Model, optimizer_with_attributes, @variable, @objective, @constraint
using JuMP: SecondOrderCone
using Clp # for linear program
using ECOS # ECOS is interior point method (more accurate but slower on large problems)
# using SCS # operator splitting method, fast but not high accuracy

abstract type Update{T} end
(U::Update)(x) = update!(U, x)

# Compressed Sensing algorithms
include("util.jl") # sparse data generator
include("oblivious.jl")
include("matchingpursuit.jl")
include("forward.jl")
include("backward.jl")
include("twostage.jl")
include("stepwise.jl")
include("sbl.jl")
include("basispursuit.jl")

end


"""
    population_data()

Returns a dictionary `D` with the US population data used in section 9.2.
The items in the dictionary are three vectors of length 100.

* `D["population"]`:  2010 population in millions for ages 0,...,99
* `D["birth_rate"]`:  birth rate
* `D["death_rate"]`:  death rate.
"""
function population_data()
population = [
   3.94415;
   3.97807;
   4.09693;
   4.11904;
   4.06317;
   4.05686;
   4.06638;
   4.03058;
   4.04649;
   4.14835;
   4.17254;
   4.11442;
   4.10624;
   4.11801;
   4.16598;
   4.24282;
   4.31614;
   4.39529;
   4.50085;
   4.58523;
   4.51913;
   4.35429;
   4.26464;
   4.19857;
   4.24936;
   4.26235;
   4.15231;
   4.24887;
   4.21525;
   4.22308;
   4.28567;
   3.97022;
   3.98685;
   3.88015;
   3.83922;
   3.95643;
   3.80209;
   3.93445;
   4.12188;
   4.36480;
   4.38327;
   4.11498;
   4.07610;
   4.10511;
   4.21150;
   4.50887;
   4.51976;
   4.53526;
   4.53880;
   4.60590;
   4.66029;
   4.46463;
   4.50085;
   4.38035;
   4.29200;
   4.25471;
   4.03751;
   3.93639;
   3.79493;
   3.64127;
   3.62113;
   3.49260;
   3.56318;
   3.48388;
   2.65713;
   2.68076;
   2.63914;
   2.64936;
   2.32367;
   2.14232;
   2.04312;
   1.94932;
   1.86427;
   1.73696;
   1.68449;
   1.62008;
   1.47107;
   1.45533;
   1.40012;
   1.37119;
   1.30851;
   1.21287;
   1.16142;
   1.07481;
   0.98572;
   0.91472;
   0.81421;
   0.71291;
   0.64062;
   0.53800;
   0.43556;
   0.34499;
   0.28139;
   0.21698;
   0.16944;
   0.12972;
   0.09522;
   0.06814;
   0.04590;
   0.03227;
];
d = [
   0.00623;
   0.00044;
   0.00027;
   0.00020;
   0.00016;
   0.00012;
   0.00011;
   0.00011;
   0.00012;
   0.00011;
   0.00010;
   0.00013;
   0.00013;
   0.00015;
   0.00020;
   0.00025;
   0.00037;
   0.00047;
   0.00064;
   0.00071;
   0.00076;
   0.00087;
   0.00087;
   0.00088;
   0.00094;
   0.00092;
   0.00095;
   0.00093;
   0.00099;
   0.00101;
   0.00103;
   0.00109;
   0.00110;
   0.00114;
   0.00115;
   0.00120;
   0.00131;
   0.00137;
   0.00146;
   0.00156;
   0.00162;
   0.00185;
   0.00201;
   0.00216;
   0.00243;
   0.00258;
   0.00298;
   0.00325;
   0.00351;
   0.00387;
   0.00413;
   0.00454;
   0.00494;
   0.00533;
   0.00571;
   0.00602;
   0.00670;
   0.00710;
   0.00769;
   0.00828;
   0.00860;
   0.00932;
   0.00998;
   0.01101;
   0.01250;
   0.01282;
   0.01404;
   0.01515;
   0.01687;
   0.01830;
   0.01967;
   0.02133;
   0.02347;
   0.02562;
   0.02800;
   0.03083;
   0.03441;
   0.03711;
   0.04126;
   0.04448;
   0.04964;
   0.05539;
   0.06149;
   0.06803;
   0.07673;
   0.08561;
   0.09540;
   0.10636;
   0.11802;
   0.13385;
   0.15250;
   0.16491;
   0.18738;
   0.20757;
   0.22688;
   0.25196;
   0.27422;
   0.29239;
   0.32560;
   0.34157;
];
b = [
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00020;
   0.00020;
   0.00020;
   0.00020;
   0.00020;
   0.01710;
   0.01710;
   0.01710;
   0.01710;
   0.01710;
   0.04500;
   0.04500;
   0.04500;
   0.04500;
   0.04500;
   0.05415;
   0.05415;
   0.05415;
   0.05415;
   0.05415;
   0.04825;
   0.04825;
   0.04825;
   0.04825;
   0.04825;
   0.02250;
   0.02250;
   0.02250;
   0.02250;
   0.02250;
   0.00510;
   0.00510;
   0.00510;
   0.00510;
   0.00510;
   0.00035;
   0.00035;
   0.00035;
   0.00035;
   0.00035;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
   0.00000;
];

    return Dict([("population", population),  ("birth_rate", b),
        ("death_rate", d)])
end


for N in [Float64, Rational{Int}, Float32]
    # random ball
    rand(BallInf)

    # 1D BallInf
    b = BallInf(N[0], N(1))
    # Test Dimension
    @test dim(b) == 1
    # Test Support Vector
    d = N[1]
    @test σ(d, b) == N[1]
    d = N[-1]
    @test σ(d, b) == N[-1]

    # 2D BallInf
    b = BallInf(N[0, 0], N(1))
    # Test Dimension
    @test dim(b) == 2
    # Test Support Vector
    d = N[1, 1]
    @test σ(d, b) == N[1, 1]
    d = N[-1, 1]
    @test σ(d, b) == N[-1, 1]
    d = N[-1, -1]
    @test σ(d, b) == N[-1, -1]
    d = N[1, -1]
    @test σ(d, b) == N[1, -1]

    # 2D BallInf not 0-centered
    b = BallInf(N[1, 2], N(1))
    # Test Dimension
    @test dim(b) == 2
    # Test Support Vector
    d = N[1, 1]
    @test σ(d, b) == N[2, 3]
    d = N[-1, 1]
    @test σ(d, b) == N[0, 3]
    d = N[-1, -1]
    @test σ(d, b) == N[0, 1]
    d = N[0, -1]
    @test σ(d, b) == N[2, 1]

    # 2D BallInf radius =/= 1
    b = BallInf(N[0, 0], N(2))
    # Test Dimension
    @test dim(b) == 2
    # Test Support Vector
    d = N[1, 1]
    @test σ(d, b) == N[2, 2]
    d = N[-1, 1]
    @test σ(d, b) == N[-2, 2]
    d = N[-1, -1]
    @test σ(d, b) == N[-2, -2]
    d = N[1, -1]
    @test σ(d, b) == N[2, -2]

    # support vector for single entry vector
    svec = σ(SingleEntryVector(2, 3, N(2)), BallInf(zeros(N, 3), N(2)))
    @test svec[1] ∈ Interval(N[-2, 2]) && svec[2] == N(2) && svec[3] ∈ Interval(N[-2, 2])

    # support function
    B = BallInf(N[1, 2], N(1))
    @test ρ(N[1, 1], B) == N(5)
    @test ρ(N[1, 0], B) == N(2)
    @test ρ(N[0, 1], B) == N(3)
    @test ρ(N[-1, -1], B) == N(-1)
    @test ρ(N[-1, 0], B) == N(0)
    @test ρ(N[-1, 1], B) == N(3)

    # boundedness
    @test isbounded(b)

    # isempty
    @test !isempty(b)

    # isuniversal
    answer, w = isuniversal(b, true)
    @test !isuniversal(b) && !answer && w ∉ b

    # membership
    b = BallInf(N[1, 1], N(1))
    @test N[0.5, -0.5] ∉ b
    @test N[0.5, 1.5] ∈ b

    # an_element function
    b = BallInf(N[1, 2], N(3))
    @test an_element(b) ∈ b

    # check that vertices_list for zero radius doesn't repeat vertices
    b = BallInf(N[1, 2], N(0))
    vl = vertices_list(b)
    @test vl == [center(b)]

    # high and low
    b = BallInf(N[1, 2], N(1))
    @test high(b) == N[2, 3]
    @test low(b) == N[0, 1]

    # isflat
    @test !isflat(BallInf(N[1, 1], N(1))) && isflat(BallInf(N[1, 1], N(0)))

    # split
    b = BallInf(N[3, 3], N(1))
    @test split(b, [1, 1]) == [Hyperrectangle(N[3, 3], N[1, 1])]
    @test ispermutation(split(b, [2, 1]),
        [Hyperrectangle(N[2.5, 3], N[0.5, 1]),
         Hyperrectangle(N[3.5, 3], N[0.5, 1])])

    # translation
    b = BallInf(N[1, 2], N(1))
    @test translate(b, N[1, 2]) == BallInf(N[2, 4], N(1))

    # affine map
    M = N[2 1; 0 2]
    B = BallInf(N[0, 0], N(1))
    v = N[1, -1]
    am = affine_map(M, B, v)
    @test ispermutation(vertices_list(am),
                        [N[4, 1], N[0, 1], N[-2, -3], N[2, -3]])

    # volume
    B = BallInf(N[0, 0], N(1))
    @test volume(B) ≈ N(4)
    if N <: AbstractFloat
        B = BallInf(zeros(N, 100), N(1/2 + 1e-5))
        @test volume(B) ≈ N(1.0020019812942185)
    end
end

# tests that only work with Float64
for N in [Float64]
    # concrete Minkowski sum
    b = BallInf(N[1, 2], N(1))
    p = minkowski_sum(b, N[2 0; 0 1] * b)
    @test p isa HPolytope{N} && ispermutation(constraints_list(p),
        [HalfSpace(N[0, -1], N(-2)), HalfSpace(N[0, 1], N(6)),
         HalfSpace(N[-1, 0], N(0)), HalfSpace(N[1, 0], N(6))])
end


# Finding X-points and O-points in a 2D magnetic reconnection configuration.
# This example assumes X-Z meridional plane.
# Note:
# 1. The input B field domain matters for computing the flux function, but I'm entirely sure
# why there are differences.
# 2. In identifying the X-points and O-points, we currently provide two methods: method 1
# needs to set the relative tolerance, while method 2 does not.
#
# Hongyang Zhou, hyzhou@umich.edu

using Vlasiator, PyPlot

file = "bulk.0001657.vlsv"

meta = load(file)

ndims(meta) != 2 && @error "Flux function only works for 2D simulations!"

if meta.ncells[3] == 1
   @error "equatorial plane, no reconnection!"
end

b = meta["vg_b_vol"]
b = reshape(b, 3, meta.ncells[1], meta.ncells[3])
dx = [meta.dcoord[1], meta.dcoord[3]]

x = LinRange(meta.coordmin[1], meta.coordmax[1], meta.ncells[1]) ./ Vlasiator.RE
z = LinRange(meta.coordmin[3], meta.coordmax[3], meta.ncells[3]) ./ Vlasiator.RE

xmin_ = searchsortedfirst(x, 7.0)
xmax_ = searchsortedlast(x, 9.5)
zmin_ = searchsortedfirst(z, -5.0)
zmax_ = searchsortedlast(z, 5.0)

# meshgrid for plotting
X = [a for a in x[xmin_:xmax_], _ in z[zmin_:zmax_]]
Z = [b for _ in x[xmin_:xmax_], b in z[zmin_:zmax_]]

flux = compute_flux_function(view(b,:,xmin_:xmax_,zmin_:zmax_), dx)

fig, ax = plt.subplots(subplot_kw=Dict("projection"=>"3d"))

ax.plot_surface(X, Z, flux;
   cmap=matplotlib.cm.turbo,
   linewidth=0, antialiased=false)

indices_x, indices_o = find_reconnection_points(flux; retol=1e-4, method=2)

fig, ax = subplots(figsize=(6,10), constrained_layout=true)

pcolormesh(meta, "proton/vg_v", ax;
   comp=:z, extent=[5, 10, -7.5, 7.5],
   cmap=matplotlib.cm.RdBu_r)
s1 = ax.scatter(x[indices_x[1,:].+xmin_.-1], z[indices_x[2,:].+zmin_.-1];
   s=50, marker="x", color="tab:gray")
s2 = ax.scatter(x[indices_o[1,:].+xmin_.-1], z[indices_o[2,:].+zmin_.-1];
   s=12, marker="o", color="tab:brown")

ax.legend([s1, s2], ["X-point", "O-point"])


export hessian_dense

function hessian_dense(nlp, x)
    return Matrix(hess(nlp, x))
    #return Symmetric(Matrix(hess(nlp, x)))
end


using vOptGeneric, CPLEX

if length(ARGS) == 0
    m = parseMOP("example.MOP", solver = CplexSolver())
else
    m = parseMOP(ARGS[1], solver = CplexSolver())
end

solve(m, method=:epsilon)

printX_E(m)

using PyPlot
Y_N = getY_N(m)
f1, f2 = map(x -> x[1], Y_N), map(x -> x[2], Y_N)
xlabel("z1") ; ylabel("z2")
plot(f1,f2,"bx", markersize = "6")
!isinteractive() && show()


using Printf

include("gen_rmesh_exp.jl")
include("lagrange_interp.jl")
include("radial_interp.jl")
include("rsch_integ_rk4.jl")
include("sch_rk4_step.jl")
include("get_min_idx.jl")
include("get_n_nodes.jl")

function main()
    # Build radial function
    r_min = 1e-9
    r_max = 50.0
    a = 1e9
    Nr = 1000
    rmesh = gen_rmesh_exp(r_min, r_max, a, Nr)

    Z = 1
    l = 0

    P = zeros(Float64,Nr)
    Q = zeros(Float64,Nr)
    V = -Z ./ rmesh

    E = -0.125
    imax = rsch_integ_rk4!( E, Z, l, rmesh, V, P, Q )
    println("imax = ", imax)

    println("Some P and Q")
    for i in 1:5
        @printf("%8d %20.10e %20.10e\n", i, P[i], Q[i])
    end
    @printf("...\n")
    for i in imax-5:imax
        @printf("%8d %20.10e %20.10e\n", i, P[i], Q[i])
    end

    minidx = get_min_idx(imax, P)
    println("minidx = ", minidx)
    println("imax   = ", imax)

    Nnodes = get_n_nodes( minidx-1, P )
    println("Nnodes = ", Nnodes)

    println("Pass here")
end

@time main()
@time main()


n_k = [80; 100; 120]
k = length(n_k)
n = sum(n_k)
p = 3

y, X = sampledata(n_k, p)
M = MOD.classmeans(Val{:row}, X, y)
H = MOD.centerclasses!(Val{:row}, copy(X), M, y)
Σ = H'H/(n-1)
Σ_k =
priors = ones(k)./k

info("Testing ", MOD.classwhiteners!)
for order in (:row, :col)
    for T in FloatingPointTypes
        isrowmajor = order == :row
        H_tst = isrowmajor ? copy(convert(Array{T}, H)) : transpose(convert(Array{T}, H))
        Σ_tst = isrowmajor ? copy(convert(Array{T}, Σ)) : transpose(convert(Array{T}, Σ))
        Σ_k_tst = isrowmajor ? [MOD.covmatrix(Val{:row}, H_tst[y .== i,:]) for i = 1:y.k] :
                               [MOD.covmatrix(Val{:col}, H_tst[:,y .== i]) for i = 1:y.k]

        for γ in (Nullable{T}(), Nullable(convert(T,0.4)), Nullable(one(T)))
            for λ in (Nullable{T}(), Nullable(convert(T,0.4)), Nullable(one(T)))
                W_k = isnull(λ) ? MOD.classwhiteners!(Val{order}, H_tst, y, γ) :
                                  MOD.classwhiteners!(Val{order}, H_tst, y, γ, get(λ))
                for i in eachindex(W_k)
                    S = isnull(λ) ? Σ_k_tst[i] : (1-get(λ))*Σ_k_tst[i] + get(λ)*Σ_tst
                    S = isnull(γ) ? S      : (1-get(γ))*S + get(γ)*trace(S)/p*I
                    if order == :row
                        @test_approx_eq eye(T,p) W_k[i]'*S*W_k[i]
                    else
                        @test_approx_eq eye(T,p) W_k[i]*S*(W_k[i]')
                    end
                end
            end
        end
    end
end

info("Testing ", MOD.qda!)
for order in (:row, :col)
    for T in FloatingPointTypes
        isrowmajor = order == :row
        X_tst = isrowmajor ? copy(convert(Matrix{T}, X)) : transpose(convert(Matrix{T}, X))
        M_tst = isrowmajor ? copy(convert(Matrix{T}, M)) : transpose(convert(Matrix{T}, M))
        H_tst = isrowmajor ? copy(convert(Matrix{T}, H)) : transpose(convert(Matrix{T}, H))

        for γ in (Nullable{T}(), Nullable(convert(T,0.5)))
            for λ in (Nullable{T}(), Nullable(convert(T,0.5)))
                W_k_tst = MOD.qda!(Val{order}, copy(X_tst), copy(M_tst), y, γ, λ)
                W_k_tmp = isnull(λ) ? MOD.classwhiteners!(Val{order}, H_tst, y, γ) :
                                      MOD.classwhiteners!(Val{order}, H_tst, y, γ, get(λ))
                for i in eachindex(W_k_tmp)
                    @test_approx_eq W_k_tst[i] W_k_tmp[i]
                end
            end
        end
    end
end

info("Testing ", MOD.qda)
for order in (:row, :col)
    for T in FloatingPointTypes
        isrowmajor = order == :row
        X_tst = isrowmajor ? copy(convert(Matrix{T}, X)) : transpose(convert(Matrix{T}, X))
        M_tst = isrowmajor ? copy(convert(Matrix{T}, M)) : transpose(convert(Matrix{T}, M))

        for γ in (Nullable{T}(), zero(T), Nullable(convert(T, 0.5)), convert(T, 0.5), one(T))
            for λ in (Nullable{T}(), zero(T), Nullable(convert(T, 0.5)), convert(T, 0.5), one(T))
                model = MOD.qda(X_tst, y, order=Val{order}, gamma=γ, lambda=λ)
                γ_tst = isa(γ, Nullable) ? γ : Nullable(γ)
                λ_tst = isa(λ, Nullable) ? λ : Nullable(λ)
                W_k_tst = MOD.qda!(Val{order}, copy(X_tst), copy(M_tst), y, γ_tst, λ_tst)
                for i in eachindex(model.W_k)
                    @test_approx_eq model.W_k[i] W_k_tst[i]
                end
                @test_approx_eq model.M M_tst
            end
        end
    end
end

info("Testing ", MOD.discriminants_qda)
for T in FloatingPointTypes
    X_tmp = copy(convert(Matrix{T}, X))
    priors_tmp = convert(Vector{T}, priors)

    model1 = qda(X_tmp,  y, order=Val{:row})
    model2 = qda(X_tmp', y, order=Val{:col})
    D = hcat(-[MOD.dotvectors(Val{:row}, (X_tmp .- M[i,:])*model1.W_k[i])/2
                for i in eachindex(priors_tmp)]...) .+ log(priors_tmp)'

    @test_approx_eq D  MOD.discriminants(model1, X_tmp)
    @test_approx_eq D' MOD.discriminants(model2, X_tmp')
end

info("Testing ", MOD.classify)
for order in (:row, :col)
    for T in FloatingPointTypes
        isrowmajor = order == :row
        X_tst = isrowmajor ? copy(convert(Matrix{T}, X)) : transpose(convert(Matrix{T}, X))
        priors_tst = convert(Vector{T}, priors)

        model = qda(X_tst, y, order=Val{order})
        D1 = vec(mapslices(indmax, MOD.discriminants(model, X_tst), isrowmajor ? 2 : 1))
        @test D1 == MOD.classify(model, X_tst)
    end
end

info("Testing ", MOD.ModelQDA)
show(DevNull, qda(X, y))


function mkelemblocks(ne,ns::Int=512)

ns = min(ns,ne);
nb = Int(floor(ne/ns));  # number of blocks
na = Int(round(ne/nb)); # number of elements per block
nk = collect(1:na:ne);
nk = reshape(nk,1,length(nk));
if length(nk)==1
    nm = [1; ne];
else
    tm = 0*nk;
    tm[1:end-1] = nk[2:end].-1;
    tm[end] = ne;
    nm = [nk; tm];
    if (nm[2,end]-nm[1,end]) < (na/2)
        nm[2,end-1] = nm[2,end];
        nm = nm[:,1:(end-1)];
    end
end

nb = size(nm,2);
while minimum(nm[2,:]-nm[1,:])<(ns/2-1) || maximum(nm[2,:]-nm[1,:])>ns
    nb = nb+1;
    na = round(ne/nb); # number of elements per block
    nk = collect(1:na:ne);
    nk = reshape(nk,1,length(nk));
    tm = 0*nk;
    tm[1:end-1] = nk[2:end].-1;
    tm[end] = ne;
    nm = [nk; tm];
    #nm = [nk[1:end]; [nk[2:end]-1 ne]];
    if (nm[2,end]-nm[1,end]) < (na/2)
        nm[2,end-1] = nm[2,end];
        nm = nm[:,1:(end-1)];
    end
end
nm = Int.(nm);
nb = size(nm,2);

if nm[2,end]!=ne
    error("something wrong");
end
if maximum(nm[2,:]-nm[1,:])>ns
    error("something wrong");
end
if minimum(nm[2,:]-nm[1,:])<(ns/2-1)
    error("something wrong");
end

return nm,nb

end


# matrix_eqn.jl

"""
Solves the discrete lyapunov equation.

The problem is given by

    AXA' - X + B = 0

`X` is computed by using a doubling algorithm. In particular, we iterate to
convergence on `X_j` with the following recursions for j = 1, 2,...
starting from X_0 = B, a_0 = A:

    a_j = a_{j-1} a_{j-1}
    X_j = X_{j-1} + a_{j-1} X_{j-1} a_{j-1}'

##### Arguments

- `A::Matrix{Float64}` : An n x n matrix as described above.  We assume in order
for  convergence that the eigenvalues of `A` have moduli bounded by unity
- `B::Matrix{Float64}` :  An n x n matrix as described above.  We assume in order
for convergence that the eigenvalues of `B` have moduli bounded by unity
- `max_it::Int(50)` :  Maximum number of iterations

##### Returns

- `gamma1::Matrix{Float64}` Represents the value X

"""
function solve_discrete_lyapunov(A::ScalarOrArray,
                                 B::ScalarOrArray,
                                 max_it::Int=50)
    # TODO: Implement Bartels-Stewardt
    n = size(A, 2)
    alpha0 = reshape([A;], n, n)
    gamma0 = reshape([B;], n, n)

    alpha1 = zeros(alpha0)
    gamma1 = zeros(gamma0)

    diff = 5
    n_its = 1

    while diff > 1e-15

        alpha1 = alpha0*alpha0
        gamma1 = gamma0 + alpha0*gamma0*alpha0'

        diff = maximum(abs(gamma1 - gamma0))
        alpha0 = alpha1
        gamma0 = gamma1

        n_its += 1

        if n_its > max_it
            error("Exceeded maximum iterations, check input matrices")
        end
    end

    return gamma1
end

"""
Solves the discrete-time algebraic Riccati equation

The prolem is defined as

    X = A'XA - (N + B'XA)'(B'XB + R)^{-1}(N + B'XA) + Q

via a modified structured doubling algorithm.  An explanation of the algorithm
can be found in the reference below.

##### Arguments

- `A` : k x k array.
- `B` : k x n array
- `R` : n x n, should be symmetric and positive definite
- `Q` : k x k, should be symmetric and non-negative definite
- `N::Matrix{Float64}(zeros(size(R, 1), size(Q, 1)))` : n x k array
- `tolerance::Float64(1e-10)` Tolerance level for convergence
- `max_iter::Int(50)` : The maximum number of iterations allowed

Note that `A, B, R, Q` can either be real (i.e. k, n = 1) or matrices.

##### Returns
- `X::Matrix{Float64}` The fixed point of the Riccati equation; a  k x k array
representing the approximate solution

##### References

Chiang, Chun-Yueh, Hung-Yuan Fan, and Wen-Wei Lin. "STRUCTURED DOUBLING
ALGORITHM FOR DISCRETE-TIME ALGEBRAIC RICCATI EQUATIONS WITH SINGULAR CONTROL
WEIGHTING MATRICES." Taiwanese Journal of Mathematics 14, no. 3A (2010): pp-935.

"""
function solve_discrete_riccati(A::ScalarOrArray, B::ScalarOrArray,
                                Q::ScalarOrArray,
                                R::ScalarOrArray,
                                N::ScalarOrArray=zeros(size(R, 1), size(Q, 1));
                                tolerance::Float64=1e-10,
                                max_it::Int=50)
    # Set up
    dist = tolerance + 1
    best_gamma = 0.0

    n = size(R, 1)
    k = size(Q, 1)
    I = eye(k)

    current_min = Inf
    candidates = [0.0, 0.01, 0.1, 0.25, 0.5, 1.0, 2.0, 10.0, 100.0, 10e5]
    BB = B' * B
    BTA = B' * A

    for gamma in candidates
        Z = getZ(R, gamma, BB)
        cn = cond(Z)
        if isfinite(cn)
            Q_tilde = -Q + N' * (Z \ (N + gamma .* BTA)) + gamma .* I
            G0 = B * (Z \ B')
            A0 = (I - gamma .* G0) * A - B * (Z \ N)
            H0 = gamma .* (A' * A0) - Q_tilde
            f1 = cond(Z, Inf)
            f2 = gamma .* f1
            f3 = cond(I + G0 * H0)
            f_gamma = max(f1, f2, f3)

            if f_gamma < current_min
                best_gamma = gamma
                current_min = f_gamma
            end
        end
    end

    if isinf(current_min)
        msg = "Unable to initialize routine due to ill conditioned args"
        error(msg)
    end

    gamma = best_gamma
    R_hat = R + gamma .* BB

    # Initial conditions
    Q_tilde = - Q + N' * (R_hat\(N + gamma .* BTA)) + gamma .* I
    G0 = B * (R_hat\B')
    A0 = (I - gamma .* G0) * A - B * (R_hat\N)
    H0 = gamma .* A'*A0 - Q_tilde
    i = 1

    # Main loop
    while dist > tolerance

        if i > max_it
            msg = "Maximum Iterations reached $i"
            error(msg)
        end

        A1 = A0 * ((I + G0 * H0)\A0)
        G1 = G0 + A0 * G0 * ((I + H0 * G0)\A0')
        H1 = H0 + A0' * ((I + H0*G0)\(H0*A0))

        dist = Base.maxabs(H1 - H0)
        A0 = A1
        G0 = G1
        H0 = H1
        i += 1
    end

    return H0 + gamma .* I  # Return X
end

"""
Simple method to return an element Z in the Riccati equation solver whose type is Float64 (to be accepted by the cond() function)

##### Arguments

- `BB::Float64` : result of B' * B
- `gamma::Float64` : parameter in the Riccati equation solver
- `R::Float64`

##### Returns
- `::Float64` : element Z in the Riccati equation solver

"""
getZ(R::Float64, gamma::Float64, BB::Float64) = R + gamma * BB

"""
Simple method to return an element Z in the Riccati equation solver whose type is Float64 (to be accepted by the cond() function)

##### Arguments

- `BB::Union{Vector, Matrix}` : result of B' * B
- `gamma::Float64` : parameter in the Riccati equation solver
- `R::Float64`

##### Returns
- `::Float64` : element Z in the Riccati equation solver

"""
getZ(R::Float64, gamma::Float64, BB::Union{Vector, Matrix}) = R + gamma * BB[1]

"""
Simple method to return an element Z in the Riccati equation solver whose type is Matrix (to be accepted by the cond() function)

##### Arguments

- `BB::Matrix` : result of B' * B
- `gamma::Float64` : parameter in the Riccati equation solver
- `R::Matrix`

##### Returns
- `::Matrix` : element Z in the Riccati equation solver

"""
getZ(R::Matrix, gamma::Float64, BB::Matrix) = R + gamma .* BB


using InvertedIndices, BenchmarkTools, LinearAlgebra, ProfileView,
    GraphRecipes

## rank 1 update
function sym_inv_update(M_inv, ind_not::Int64, ind)
    h = M_inv[ind_not, ind_not]
    a = M_inv[ind, [ind_not]] / sqrt(h)
    out = M_inv[ind, ind] - LinearAlgebra.BLAS.gemm('N', 'T', 1.0, a, a)
    return out
end

# rank-n-update
function sym_inv_update(M_inv, ind_not, ind)
    A = M_inv[ind_not, ind]
    H = cholesky(M_inv[ind_not, ind_not])
    D = H \ A
    out = M_inv[ind, ind] - LinearAlgebra.BLAS.gemm('T', 'N', 1.0, A, D)
    return out
end

S = rand(30,30)
S = S*S'
S = Symmetric(S)
S_inv = Symmetric(inv(S))

ind = filter(x ->!(x ∈ [5]), 1:size(S, 1))

sym_inv_update(S_inv, 5, ind) ≈ inv(S[Not(5), Not(5)])
sym_inv_update(S_inv, [5], ind) ≈ inv(S[Not(5), Not(5)])

@benchmark inv($S[ind, ind])
@benchmark sym_inv_update($S_inv, 5, ind)
@benchmark sym_inv_update($S_inv, [5], ind)

del = [2, 3, 5, 7, 9, 10, 12, 18, 21]
ind = filter(x ->!(x ∈ del), 1:size(S, 1))

sym_inv_update(S_inv, del, ind) ≈ inv(S[ind, ind])
inv(cholesky(S)) ≈ cholesky(S) \ I

@benchmark cholesky($S[ind, ind]) \ I
@benchmark sym_inv_update($S_inv, del, $ind)

function testf(S_inv, ind_not, ind)
    for i = 1:10000
        sym_inv_update(S_inv, ind_not, ind)
    end
end

ProfileView.@profview testf(S_inv, [5, 7], ind)

a = rand(40, 40)
b = rand(40, 40)

@benchmark a-b

#### random graph generation
using Random, LinearAlgebra, SparseArrays, GraphRecipes, Plots, BenchmarkTools

Random.seed!(78435472956284237434)

nvar = 500
#p_con = 0.1
e_n = 4

p_con = (e_n/(nvar-1))

S = zeros(nvar, nvar)
S[diagind(S)] .= 1.0

# DAG
G = rand(nvar, nvar)
G = G .< p_con
G = LowerTriangular(G)
G[diagind(G)] .= false
G = sparse(G)

C = inv(I-Matrix(G))
C = C .!= 0
(C .& C') == I
@benchmark C = inv(I-Matrix(G))

# Cyclic
G = rand(nvar, nvar)
G = G .< p_con
G = sparse(G)

C = inv(I-Matrix(G))
C = C .!= 0
(C .& C') == I
@benchmark C = inv(I-Matrix(G))

graphplot(G)
print(eigvals(Matrix(G)))

## Fill with effects
ind = findall(G)
effects = rand(size(ind, 1))

A = zeros(nvar, nvar)
A[ind] .= effects
A[ind] .= 0.3

inv(I-A)*S*inv(I-A)'

### generate alternative Graphs

A = [

    ]


using DataFrames, CSV
using PyPlot
using Printf
using SpecialFunctions
include("utilities.jl")

filename = @sprintf("../data/zeta_test_data.csv")
data1 = CSV.read(filename; delim=",", type=String)

data1[!,:s] = parse.(Complex{Float64}, data1[!,:s] )
data1[!,:z] = parse.(Complex{Float64}, data1[!,:z] )
data1[!,:zeta] = parse.(Complex{Float64}, data1[!,:zeta] )

m = size(data1,1)
s = data1[!,:s]
z = complex(data1[!,:z])

S1 = zeros(Complex{Float64}, m)
for i = 1:m
    S1[i] = SpecialFunctions.zeta( 1-s[i], 0.5 + log(z[i])/(2*pi*im) )
end

error1 = S1 - data1[!,:zeta]
rel_error = abs.( error1 ) ./ abs.( data1[!,:zeta] )
maximum(rel_error)


# Symmetric triangular distribution
immutable SymTriangularDist <: ContinuousUnivariateDistribution
    location::Float64
    scale::Float64
    function SymTriangularDist(l::Real, s::Real)
        s > zero(s) || error("scale must be positive")
        new(float64(l), float64(s))
    end
end

SymTriangularDist(location::Real) = SymTriangularDist(location, 1.0)
SymTriangularDist() = SymTriangularDist(0.0, 1.0)

@distr_support SymTriangularDist d.location - d.scale d.location + d.scale

## Properties
mean(d::SymTriangularDist) = d.location
median(d::SymTriangularDist) = d.location
mode(d::SymTriangularDist) = d.location

var(d::SymTriangularDist) = d.scale^2 / 6.0
skewness(d::SymTriangularDist) = 0.0
kurtosis(d::SymTriangularDist) = -0.6

entropy(d::SymTriangularDist) = 0.5 + log(d.scale)

## Functions
function pdf(d::SymTriangularDist, x::Float64)
    if insupport(d, x)
        (1.0 - abs(x-d.location)/d.scale) / d.scale
    else
        0.0
    end
end

function logpdf(d::SymTriangularDist, x::Float64)
    if insupport(d, x)
        log1p(-abs(x-d.location)/d.scale) -log(d.scale)
    else
        -Inf
    end
end

function cdf(d::SymTriangularDist, x::Float64)
    u = (x - d.location) / d.scale
    if u <= -1.0
        0.0
    elseif u <= 0.0
        0.5*(1.0 + u)^2
    elseif u <= 1.0
        1.0 - 0.5*(1.0 - u)^2
    else
        1.0
    end
end
function ccdf(d::SymTriangularDist, x::Float64)
    u = (x - d.location) / d.scale
    if u <= -1.0
        1.0
    elseif u <= 0.0
        1.0 - 0.5*(1.0 + u)^2
    elseif u <= 1.0
        0.5*(1.0 - u)^2
    else
        0.0
    end
end

function logcdf(d::SymTriangularDist, x::Float64)
    u = (x - d.location) / d.scale
    if u <= -1.0
        -Inf
    elseif u <= 0.0
        loghalf + 2.0*log1p(u)
    elseif u <= 1.0
        log1p(-0.5*(1.0 - u)^2)
    else
        0.0
    end
end
function logccdf(d::SymTriangularDist, x::Float64)
    u = (x - d.location) / d.scale
    if u <= -1.0
        0.0
    elseif u <= 0.0
        log1p(- 0.5*(1.0 + u)^2)
    elseif u <= 1.0
        loghalf + 2.0*log1p(-u)
    else
        -Inf
    end
end

function quantile(d::SymTriangularDist, p::Float64)
    @checkquantile p begin
        if p < 0.5
            d.location - d.scale*(1.0 - sqrt(2.0*p))
        else
            d.location + d.scale*(1.0 - sqrt(2.0*(1.0-p)))
        end
    end
end

function cquantile(d::SymTriangularDist, p::Float64)
    @checkquantile p begin
        if p > 0.5
            d.location - d.scale*(1.0 - sqrt(2.0*(1.0-p)))
        else
            d.location + d.scale*(1.0 - sqrt(2.0*p))
        end
    end
end

function invlogcdf(d::SymTriangularDist, lp::Float64)
    @checkinvlogcdf lp begin
        if lp < loghalf
            d.location + d.scale*expm1(0.5*(lp-loghalf))
        else
            d.location + d.scale*(1.0 - sqrt(-2.0*expm1(lp)))
        end
    end
end

function invlogccdf(d::SymTriangularDist, lp::Float64)
    @checkinvlogcdf lp begin
        if lp > loghalf
            d.location - d.scale*(1.0 - sqrt(-2.0*expm1(lp)))
        else
            d.location - d.scale*expm1(0.5*(lp-loghalf))
        end
    end
end

function mgf(d::SymTriangularDist, t::Real)
    a = d.scale*t
    a == zero(a) && return one(a)
    4.0*exp(d.location*t)*(sinh(0.5*a)/a)^2
end
function cf(d::SymTriangularDist, t::Real)
    a = d.scale*t
    a == zero(a) && return complex(one(a))
    4.0*exp(im*d.location*t)*(sin(0.5*a)/a)^2
end

## Sampling
function rand(d::SymTriangularDist)
    ξ1, ξ2 = rand(), rand()
    return d.location + (ξ1 - ξ2) * d.scale
end


using QuasiArrays, IntervalSets, Test
import QuasiArrays: ApplyQuasiArray

@testset "Continuous" begin
    @testset "Inclusion" begin
        # note behaviour changes with ContinuumArrays
        x = Inclusion(0..1)
        @test eltype(x) == Int
        @test axes(x) ≡ (x,)
        @test_throws InexactError x[0.1]
        @test x[1] ≡ 1
        @test_throws BoundsError x[2]
        @test Base.unsafe_getindex(x,2) ≡ 2

        for x in (Inclusion{Float64}(0..1),Inclusion(0.0..1))
            @test Inclusion(x) ≡ Inclusion{Float64}(x) ≡ convert(Inclusion,x) ≡
                    convert(Inclusion{Float64},x) ≡ convert(AbstractQuasiArray,x) ≡
                    convert(AbstractQuasiArray{Float64},x) ≡ convert(AbstractQuasiVector{Float64},x) ≡
                    convert(AbstractQuasiVector,x) ≡ x
            @test axes(x) ≡ (x,)
            @test x[0.1] ≡ 0.1
            @test x[1] ≡ 1.0
            @test_throws BoundsError x[2]
            @test Base.unsafe_getindex(x,2) ≡ 2.0
            @test x[BigFloat(π)/4] == π/4
        end

        @test Inclusion{Float64}(0..1) == Inclusion(0.0..1)
        @test Inclusion{Float64}(0..1) !== Inclusion(0.0..1)

        @test findfirst(isequal(0.1), Inclusion(0.0..1)) ≡ findlast(isequal(0.1), Inclusion(0..1)) ≡ 0.1
        @test findfirst(isequal(2.3), Inclusion(0.0..1)) ≡ findlast(isequal(2.3), Inclusion(0..1)) ≡ nothing
        @test findall(isequal(0.1), Inclusion(0.0..1)) == [0.1]
        @test findall(isequal(2.3), Inclusion(0.0..1)) == Float64[]
    end

    @testset "QuasiDiagonal" begin
        x = Inclusion(0.0..1)
        D = QuasiDiagonal(x)
        @test D[0.1,0.2] ≡ 0.0
        @test D[0.1,0.1] ≡ 0.1
        @test D[1,1] ≡ 1.0
        @test_throws BoundsError D[1.1,1.2]
    end
end


using DReal
import DReal: opensmt_context, opensmt_set_verbosity, opensmt_assert, opensmt_check, opensmt_print_expr, opensmt_mk_context,
              opensmt_mk_bool_var, opensmt_mk_int_var, opensmt_mk_minus, opensmt_mk_leq, opensmt_mk_num_from_string,
              opensmt_push, opensmt_pop, opensmt_reset, opensmt_del_context, opensmt_mk_uminus, opensmt_mk_times_2, opensmt_mk_times_3,
              opensmt_set_precision, opensmt_mk_unbounded_real_var, opensmt_mk_num, opensmt_mk_div, opensmt_mk_sqrt, opensmt_mk_pow,
              opensmt_mk_exp, opensmt_define_ode, opensmt_mk_integral, opensmt_mk_eq, opensmt_mk_geq, opensmt_mk_real_var
using Base.Test
using Compat

ctx = opensmt_mk_context(@compat UInt32(1))
opensmt_set_precision(ctx, 1.0)

# Creating Real variables
x = opensmt_mk_real_var(ctx, "x", -100.0, 100.0)
P = opensmt_mk_real_var(ctx, "P", -100.0, 100.0)
x_0 = opensmt_mk_real_var(ctx, "x_0", -100.0, 100.0)
P_0 = opensmt_mk_real_var(ctx, "P_0", -100.0, 100.0)
x_t = opensmt_mk_real_var(ctx, "x_t", -100.0, 100.0)
P_t = opensmt_mk_real_var(ctx, "P_t", -100.0, 100.0)
time_0 = opensmt_mk_real_var(ctx, "time_0", 0.0, 100.0)

# =======================================================================
# Define ODE (flow)
# d/dt[x] = 1.0
# d/dt[P] = 1.0 / ((2.0 * 3.14159265359) ^ 0.5) * exp (- (x ^ 2) / 2.0)
zero = opensmt_mk_num(ctx, 0.0)
one = opensmt_mk_num(ctx, 1.0)
ten = opensmt_mk_num(ctx, 10.0)
neg_ten = opensmt_mk_uminus(ctx, ten)
two = opensmt_mk_num(ctx, 2.0)
pi  = opensmt_mk_num(ctx, 3.14159265359)
two_pi = opensmt_mk_times_2(ctx, two, pi)
# p_rhs = opensmt_mk_uminus(ctx, x)
p_rhs = opensmt_mk_times_2(
  ctx,
  # 1.0 / sqrt(2 * pi)
  # opensmt_mk_div(ctx, one, opensmt_mk_sqrt(ctx, two_pi)),
  opensmt_mk_div(ctx, one, opensmt_mk_pow(ctx, two_pi, opensmt_mk_num(ctx, 0.5))),
  # exp (- x^2 / 2.0)
  opensmt_mk_exp(
    ctx,
    opensmt_mk_uminus(ctx, opensmt_mk_div(ctx, opensmt_mk_pow(ctx, x, two), two))))
#
vars = [x, P]
rhses = [one, p_rhs]
opensmt_define_ode(ctx, "flow_1", vars, rhses, Cuint(2))

# =======================================================================
## Define initial and final state
opensmt_assert(ctx, opensmt_mk_eq(ctx, x_0, neg_ten))  # x_0 = -10.0
opensmt_assert(ctx, opensmt_mk_eq(ctx, P_0, zero))     # P_0 = 0.0
opensmt_assert(ctx, opensmt_mk_eq(ctx, x_t, ten))      # x_t = 10.0
opensmt_assert(ctx, opensmt_mk_geq(ctx, P_t, zero))    # P_t >= 0.0

# =======================================================================
# Add integral constraint
# [x_t P_t] = integral_0^{time_0} [x_0 P_0] | flow_1

vars_t = [x_t, P_t]
vars_0 = [x_0, P_0]
integral = opensmt_mk_integral(ctx, vars_t, zero, time_0, vars_0, Cuint(2), "flow_1")
opensmt_assert(ctx, integral)
res = opensmt_check(ctx)

println( "res = $(res == -1 ? "unsat" : "sat")")

print( "Deleting context\n" )
opensmt_del_context( ctx )


sumRand(n::Int = 10) = sum(rand(n))


immutable Laplace <: ContinuousUnivariateDistribution
    location::Float64
    scale::Float64
    function Laplace(l::Real, s::Real)
        if s > 0.0
            new(float64(l), float64(s))
        else
            error("scale must be positive")
        end
    end
end

Laplace(location::Real) = Laplace(location, 1.0)
Laplace() = Laplace(0.0, 1.0)

const Biexponential = Laplace

function cdf(d::Laplace, q::Real)
    return 0.5 * (1.0 + sign(q - d.location) *
           (1.0 - exp(-abs(q - d.location) / d.scale)))
end

entropy(d::Laplace) = log(2.0 * e * d.scale)

insupport(d::Laplace, x::Number) = isreal(x) && isfinite(x)

kurtosis(d::Laplace) = 3.0

mean(d::Laplace) = d.location

median(d::Laplace) = d.location

function mgf(d::Laplace, t::Real)
    m, b = d.location, d.scale
    return exp(t * m) / (1.0 - b^2 * t^2)
end

function cf(d::Laplace, t::Real)
    m, b = d.location, d.scale
    return exp(im * t * m) / (1.0 + b^2 * t^2)
end

modes(d::Laplace) = [d.location]

function pdf(d::Laplace, x::Real)
    (1.0 / (2.0 * d.scale)) * exp(-abs(x - d.location) / d.scale)
end

function logpdf(d::Laplace, x::Real)
    -log(2.0 * d.scale) - abs(x - d.location) / d.scale
end

function quantile(d::Laplace, p::Real)
    d.location - d.scale * sign(p - 0.5) * log(1.0 - 2.0 * abs(p - 0.5))
end

# Need to see whether other RNG strategies are more efficient:
# (1) Difference of two Exponential(1/b) variables
# (2) Ratio of logarithm of two Uniform(0.0, 1.0) variables
function rand(d::Laplace)
    u = rand() - 0.5
    return d.location - d.scale * sign(u) * log(1.0 - 2.0 * abs(u))
end

skewness(d::Laplace) = 0.0

std(d::Laplace) = sqrt(2.0) * d.scale

var(d::Laplace) = 2.0 * d.scale^2


function show(io::IO, x::ULogarithmic)
  print(io, "exp(")
  show(io, x.log)
  print(io, ")")
end

function show(io::IO, x::Logarithmic)
  x.signbit && print(io, "-")
  show(io, x.abs)
end

function show(io::IO, x::CLogarithmic)
	print(io, "exp(")
	show(io, log(x))
	print(io, ")")
end

write(io::IO, x::ULogarithmic) = write(io, x.log)
write(io::IO, x::Logarithmic) = write(io, x.abs, x.signbit)
write(io::IO, x::CLogarithmic) = write(io, x.abs, x.angle)

read(io::IO, ::Type{ULogarithmic{T}}) where {T} = uexp(T, read(io, T))
read(io::IO, ::Type{Logarithmic{T}}) where {T} = Logarithmic{T}(read(io, ULogarithmic{T}), read(io, Bool))
read(io::IO, ::Type{CLogarithmic{T}}) where {T} = CLogarithmic{T}(read(io, ULogarithmic{T}), read(io, T))


using Gnuplot
x = y = -15:0.33:15
fz(x,y) = sin.(sqrt.(x.^2 + y.^2))./sqrt.(x.^2+y.^2)
fxy = [fz(x,y) for x in x, y in y]

@gsp(x, y, fxy, "w pm3d", "set ylabel 'y' textcolor rgb 'red'",
     "set xlabel 'x' textcolor rgb 'blue'", "set key off")

# Draw contours using the same palette as the surface plot
z = -2.5  # z coordinate at which contours are drawn
@gsp :- "set cbrange [$(minimum(fxy)):$(maximum(fxy))]"
@gsp :- "set xyplane 0"
levels = contourlines(x, y, fxy)
for i in 1:length(levels)
    @gsp :- levels[i].data "u 1:2:($z):($(levels[i].z)) w l notit lc pal" :-
end
@gsp


using FastTransforms, Compat
using Compat.Test

import FastTransforms: δ

@testset "Gaunt coefficients" begin
    println("Testing Table 2 of Y.-l. Xu, JCAM 85:53–65, 1997.")
    for (m,n) in ((0,2),(1,2),(1,8),(6,8),(3,18),
                  (10,18),(5,25),(-23,25),(2,40),(-35,40),
                  (28,62),(-42,62),(1,99),(90,99),(10,120),
                  (80,120),(23,150),(88,150))
        @test norm(gaunt(m,n,-m,n)[end]./(big(-1.0)^m/(2n+1))-1, Inf) < 400eps()
    end
    println("Testing Table 3 of Y.-l. Xu, JCAM 85:53–65, 1997.")
    for (m,n,μ,ν) in ((0,1,0,5),(0,5,0,10),(0,9,0,10),(0,10,0,12),
                      (0,11,0,15),(0,12,0,20),(0,20,0,45),(0,40,0,80),
                      (0,45,0,100),(3,5,-3,6),(4,9,-4,15),(-8,18,8,23),
                      (-10,20,10,30),(5,25,-5,45),(15,50,-15,60),(-28,68,28,75),
                      (32,78,-32,88),(45,82,-45,100))
        @test norm(sum(gaunt(m,n,μ,ν))-δ(m,0), Inf) < 15000eps()
    end
end


"""
    NegativeBinomial

* Parametrization
parametrized in r, p

* Score

* Fisher Information

* `time_varying_params` map.

* Default link
"""
NegativeBinomial

function score!(score_til::Matrix{T}, y::Int, ::Type{NegativeBinomial}, param::Matrix{T}, t::Int) where T
    score_til[t, 1] = digamma(y + param[t, 1]) - digamma(param[t, 1]) + log(param[t, 2])
    score_til[t, 2] = param[t, 1] / param[t, 2] - y / (one(T) - param[t, 2])
    return
end

function fisher_information!(aux::AuxiliaryLinAlg{T}, ::Type{NegativeBinomial}, param::Matrix{T}, t::Int) where T
    return error("Fisher information not implemented for NegativeBinomial distribution.")
end

function log_likelihood(::Type{NegativeBinomial}, y::Vector{Int}, param::Matrix{T}, n::Int) where T
    loglik = zero(T)
    for t in 1:n
        loglik += loggamma(y[t] + param[t, 1]) - logfactorial(y[t]) - loggamma(param[t, 1]) +
                param[t, 1] * log(param[t, 2]) + y[t] * log(1 - param[t, 2])
    end
    return -loglik
end

# Links
function link!(param_tilde::Matrix{T}, ::Type{NegativeBinomial}, param::Matrix{T}, t::Int) where T
    param_tilde[t, 1] = link(LogLink, param[t, 1], zero(T))
    param_tilde[t, 2] = link(LogitLink, param[t, 2], zero(T), one(T))
    return
end
function unlink!(param::Matrix{T}, ::Type{NegativeBinomial}, param_tilde::Matrix{T}, t::Int) where T
    param[t, 1] = unlink(LogLink, param_tilde[t, 1], zero(T))
    param[t, 2] = unlink(LogitLink, param_tilde[t, 2], zero(T), one(T))
    return
end
function jacobian_link!(aux::AuxiliaryLinAlg{T}, ::Type{NegativeBinomial}, param::Matrix{T}, t::Int) where T
    aux.jac[1] = jacobian_link(LogLink, param[t, 1], zero(T))
    aux.jac[2] = jacobian_link(LogitLink, param[t, 2], zero(T), one(T))
    return
end

# utils
function update_dist(::Type{NegativeBinomial}, param::Matrix{T}, t::Int) where T
    return NegativeBinomial(param[t, 1], param[t, 2])
end

function params_sdm(d::NegativeBinomial)
    return Distributions.params(d)
end

function num_params(::Type{NegativeBinomial})
    return 2
end


gaussian = html_div() do
    html_div(id="fields") do
        html_div(className="container-slider") do
            html_label("Samples: "),
            html_div(
                dcc_slider(
                        id = "n_samples",
                        min = 1,
                        max = 10000,
                        marks = Dict(i => "$i" for i = 0:1000:10000),
                        value = 100,
                ),
            style = (width = "70%", display = "inline-block"),
            className = "slider")
        end,
        html_div(className="container-slider") do
            html_label("Cov: "),
            html_div(
                dcc_slider(
                        id = "cov",
                        min = 0.01,
                        max = 1.0,
                        step = 0.01,
                        marks = Dict(i => "$i" for i = 0.0:0.1:1.0),
                        value = 0.1,
                ),
            style = (width = "70%", display = "inline-block"),
            className = "slider")
        end,
        html_div(className="container-slider") do
            html_label("Classes: "),
            html_div(
                dcc_slider(
                        id = "n_classes",
                        min = 1,
                        max = 30,
                        marks = Dict(i => "$i" for i = 0:1:30),
                        value = 2,
                ),
            style = (width = "70%", display = "inline-block"),
            className = "slider")
        end
    end,
    html_div(className="container-graph") do
        dcc_graph(
            id = "plot_gaussian",
            figure = gen_gaussian(1000, 0.25,2)
        )
    end
end

callback!(app, Output("plot_gaussian", "figure"),[Input("n_samples", "value"), Input("cov", "value"),Input("n_classes", "value")]) do n_samples, cov, n_classes
    gen_gaussian(n_samples, cov,n_classes)
end