# Convert vector values to quantile scores
function quantiles(vec::Array; ignorenothing::Bool=true)
    if ignorenothing == true
        qfinder = ecdf(filter(!isnothing, vec))
        qvec = replace(x -> !isnothing(x) ? qfinder(x) : x, vec)
    else
        qfinder = ecdf(vec)
        qvec = qfinder(vec)
    end
    return qvec
end

# Convert & mutate vector values to quantile scores
function quantiles!(vec::Array; ignorenothing::Bool=true)
    if ignorenothing == true
        qfinder = ecdf(filter(!isnothing, vec))
        replace!(x -> !isnothing(x) ? qfinder(x) : x, vec)
    else
        qfinder = ecdf(vec)
        replace!(x -> qfinder(x), vec)
    end
    return vec
end

# Convert SimpleSDMLayer values to quantile scores
function quantiles(layer::SimpleSDMLayer)
    qfinder = ecdf(Float32.(filter(!isnothing, layer.grid)))
    qgrid = replace(x -> isnothing(x) ? x : qfinder(x), layer.grid)
    qlayer = SimpleSDMResponse(qgrid, layer)
    return qlayer
end

# Convert & mutate SimpleSDMLayer values to quantile scores
function quantiles!(layer::SimpleSDMLayer)
    qfinder = ecdf(Float32.(filter(!isnothing, layer.grid)))
    replace!(x -> isnothing(x) ? x : qfinder(x), layer.grid)
    return layer
end


@testset "Theoretical" begin
  Random.seed!(123)
  h = range(0, stop=10, length=50)
  x, y = rand(Point3), rand(Point3)

  # stationary variogram models
  γs = [NuggetEffect(), GaussianVariogram(), ExponentialVariogram(),
        MaternVariogram(), SphericalVariogram(),
        SphericalVariogram(range=2.), CubicVariogram(),
        PentasphericalVariogram(), SineHoleVariogram()]

  # non-stationary variogram models
  γn = [PowerVariogram(), PowerVariogram(exponent=.4)]

  # non-decreasing variogram models
  γnd = [GaussianVariogram(), ExponentialVariogram(),
         MaternVariogram(), SphericalVariogram(),
         SphericalVariogram(range=2.), CubicVariogram(),
         PentasphericalVariogram(), PowerVariogram()]

  # anisotropic variogram models
  γa = [GaussianVariogram(MetricBall((2.,1.))),
        MaternVariogram(MetricBall((3.,2.,1.)))]

  # check stationarity
  @test all(isstationary, γs)
  @test all(!isstationary, γn)

  # check isotropy
  @test all(isisotropic, γs)
  @test all(isisotropic, γn)
  @test all(isisotropic, γnd)
  @test isisotropic(sum(γs) + sum(γn) + sum(γnd))
  @test all(!isisotropic, γa)

  # variograms are symmetric under Euclidean distance
  for γ in (γs ∪ γn ∪ γnd ∪ [sum(γs) + sum(γn) + sum(γnd)])
    @test γ(x, y) ≈ γ(y, x)
  end

  # some variograms are non-decreasing
  for γ in (γnd ∪ [sum(γnd)])
    @test all(γ.(h) .≤ γ.(h.+1))
  end

  # variograms are valid at the origin
  for γ in (γs ∪ γn ∪ γnd)
    @test !isnan(γ(0.)) && !isinf(γ(0.))
  end

  # nugget effect
  γ = NuggetEffect(nugget=0.2)
  @test nugget(γ) == 0.2
  @test sill(γ) == 0.2
  @test range(γ) == 0.0

  # ill-conditioned models and nugget regularization
  # see https://github.com/JuliaEarth/GeoStats.jl/issues/29
  pset = PointSet([93.0 90.0 89.0 94.0 93.0 97.0 95.0 88.0 96.0 98.0
                   40.0 33.0 34.0 36.0 30.0 39.0 39.0 28.0 25.0 35.0])

  # ill-conditioned covariance
  γ = GaussianVariogram(range=20.)
  C = sill(γ) .- Variography.pairwise(γ, pset)
  @test cond(C) > 1000.

  # nugget regularization
  γ = GaussianVariogram(range=20., nugget=0.1)
  C = sill(γ) .- Variography.pairwise(γ, pset)
  @test γ(0) == 0
  @test γ(1e-6) > 0
  @test cond(C) < 100.

  if visualtests
    plt1 = plot()
    for γ ∈ γs
      plot!(plt1, γ)
    end
    plt2 = plot()
    for γ ∈ γn
      plot!(plt2, γ)
    end
    plt = plot(plt1, plt2, size=(600,800), layout=(2,1))
    @test_reference "data/theoretical.png" plt

    @test_reference "data/nugget.png" plot(NuggetEffect(0.1),ylim=(0,1))
  end
end


using IterativeSolvers
using Test
using LinearMaps
using LinearAlgebra
using Random
using SparseArrays

#GMRES
@testset "GMRES" begin

srand(1234321)
n = 10

@testset "Matrix{$T}" for T in (Float32, Float64, ComplexF32, ComplexF64)
    A = rand(T, n, n)
    b = rand(T, n)
    F = lu(A)
    tol = √eps(real(T))

    # Test optimality condition: residual should be non-increasing
    x, history = gmres(A, b, log = true, restart = 3, maxiter = 10, tol = tol);
    @test isa(history, ConvergenceHistory)
    @test all(diff(history[:resnorm]) .<= 0.0)

    # Left exact preconditioner
    x, history = gmres(A, b, Pl=F, maxiter=1, restart=1, tol=tol, log=true)
    @test history.isconverged
    @test norm(F \ (A * x - b)) / norm(b) ≤ tol

    # Right exact preconditioner
    x, history = gmres(A, b, Pl=Identity(), Pr=F, maxiter=1, restart=1, tol=tol, log=true)
    @test history.isconverged
    @test norm(A * x - b) / norm(b) ≤ tol
end

@testset "SparseMatrixCSC{$T}" for T in (Float64, ComplexF64)
    A = sprand(T, n, n, 0.5) + I
    b = rand(T, n)
    F = lu(A)
    tol = √eps(real(T))

    # Test optimality condition: residual should be non-increasing
    x, history = gmres(A, b, log = true, restart = 3, maxiter = 10);
    @test all(diff(history[:resnorm]) .<= 0.0)

    # Left exact preconditioner
    x, history = gmres(A, b, Pl=F, maxiter=1, restart=1, log=true)
    @test history.isconverged
    @test norm(F \ (A * x - b)) / norm(b) ≤ tol

    # Right exact preconditioner
    x, history = gmres(A, b, Pl = Identity(), Pr=F, maxiter=1, restart=1, log=true)
    @test history.isconverged
    @test norm(A * x - b) / norm(b) ≤ tol
end

@testset "Linear operator defined as a function" begin
    A = LinearMap(cumsum!, 100; ismutating=true)
    b = rand(100)
    tol = 1e-5

    x = gmres(A, b; tol=tol, maxiter=2000)
    @test norm(A * x - b) / norm(b) ≤ tol
end
end


using Poirot, Poirot.LAX, Test

double(x) = xla(() -> x+x)

@test double(21) == 42
@test double(4.5) == 9.0

let x = 2, y = 2.0
  @test_broken xla(() -> x+y) == 4.0
end

add(a, b) = a+b
@test @code_xla(add(2, 3)) isa Poirot.Abstract.IR

@test xla(() -> 2+2) == 4

let x = 5
  @test xla(() -> 3x^(1+1) + (2x + 1)) == 86
end

relu(x) = xla(() -> x > 0 ? x : 0)

@test relu(5) == 5
@test relu(-5) == 0
@test_broken relu(5.0) == 5.0

let x = rand(3), y = rand(3)
  @test collect(xla(() -> x+y)) == x+y
end

function updatezero!(env)
  if env[:x] < 0
    env[:x] = 0
  end
end

function relu(x)
  env = Dict()
  env[:x] = x
  updatezero!(env)
  return env[:x]
end

@test xla(() -> relu(5)) == 5
@test xla(() -> relu(-5)) == 0


using FractionalSystems

G = fotf([1], [0], [1, 10, 35, 50, 24], [3.5, 2.8, 2.1, 1.4, 0.7])
rlocus(G)


################################################################################
#                            Planar and Radial Flows                           #
#             Ref: Variational Inference with Normalizing Flows,               #
#               D. Rezende, S. Mohamed(2015) arXiv:1505.05770                  #
################################################################################

###############
# PlanarLayer #
###############

# TODO: add docstring

struct PlanarLayer{T1<:AbstractVector{<:Real}, T2<:Union{Real, AbstractVector{<:Real}}} <: Bijector{1}
    w::T1
    u::T1
    b::T2
end
function Base.:(==)(b1::PlanarLayer, b2::PlanarLayer)
    return b1.w == b2.w && b1.u == b2.u && b1.b == b2.b
end

function PlanarLayer(dims::Int, wrapper=identity)
    w = wrapper(randn(dims))
    u = wrapper(randn(dims))
    b = wrapper(randn(1))
    return PlanarLayer(w, u, b)
end

# all fields are numerical parameters
Functors.@functor PlanarLayer

"""
    get_u_hat(u::AbstractVector{<:Real}, w::AbstractVector{<:Real})

Return a tuple of vector ``û`` that guarantees invertibility of the planar layer, and
scalar ``wᵀ û``.

# Mathematical background

According to appendix A.1, vector ``û`` defined by
```math
û(w, u) = u + (\\log(1 + \\exp{(wᵀu)}) - 1 - wᵀu) \\frac{w}{\\|w\\|²}
```
guarantees that the planar layer ``f(z) = z + û tanh(wᵀz + b)`` is invertible for all ``w, u ∈ ℝᵈ`` and ``b ∈ ℝ``.
We can rewrite ``û`` as
```math
û = u + (\\log(1 + \\exp{(-wᵀu)}) - 1) \\frac{w}{\\|w\\|²}.
```

Additionally, we obtain
```math
wᵀû = wᵀu + \\log(1 + \\exp{(-wᵀu)}) - 1 = \\log(1 + \\exp{(wᵀu)}) - 1.
```

# References

D. Rezende, S. Mohamed (2015): Variational Inference with Normalizing Flows.
arXiv:1505.05770
"""
function get_u_hat(u::AbstractVector{<:Real}, w::AbstractVector{<:Real})
    wT_u = dot(w, u)
    û = u .+ ((LogExpFunctions.log1pexp(-wT_u) - 1) / sum(abs2, w)) .* w
    wT_û = LogExpFunctions.log1pexp(wT_u) - 1
    return û, wT_û
end

# An internal version of the transform in eq. (10) that returns intermediate variables
function _transform(flow::PlanarLayer, z::AbstractVecOrMat{<:Real})
    w = flow.w
    b = first(flow.b)
    û, wT_û = get_u_hat(flow.u, w)
    wT_z = aT_b(w, z)
    transformed = z .+ û .* tanh.(wT_z .+ b)
    return (transformed = transformed, wT_û = wT_û, wT_z = wT_z)
end

(b::PlanarLayer)(z) = _transform(b, z).transformed

#=
Log-determinant of the Jacobian of the planar layer

The log-determinant of the Jacobian of the planar layer ``f(z) = z + û tanh(wᵀz + b)``
is given by
```math
\\log |det ∂f(z)/∂z| = \\log |1 + ûᵀsech²(wᵀz + b)w| = \\log |1 + sech²(wᵀz + b) wᵀû|.
```

Since ``0 < sech²(x) ≤ 1`` and
```math
wᵀû = wᵀu + \\log(1 + \\exp{(-wᵀu)}) - 1 = \\log(1 + \\exp{(wᵀu)}) - 1 > -1,
```
we get
```math
\\log |det ∂f(z)/∂z| = \\log(1 + sech²(wᵀz + b) wᵀû).
```
=#
function forward(flow::PlanarLayer, z::AbstractVecOrMat{<:Real})
    transformed, wT_û, wT_z = _transform(flow, z)

    # Compute ``\\log |det ∂f(z)/∂z|`` (see above).
    b = first(flow.b)
    log_det_jacobian = log1p.(wT_û .* abs2.(sech.(_vec(wT_z) .+ b)))

    return (rv = transformed, logabsdetjac = log_det_jacobian)
end

function (ib::Inverse{<:PlanarLayer})(y::AbstractVecOrMat{<:Real})
    flow = ib.orig
    w = flow.w
    b = first(flow.b)
    û, wT_û = get_u_hat(flow.u, w)

    # Find the scalar ``α`` by solving ``wᵀy = α + wᵀû tanh(α + b)``
    # (eq. (23) from appendix A.1).
    wT_y = aT_b(w, y)
    α = find_alpha.(wT_y, wT_û, b)

    # Compute ``z = y - û tanh(α + b)``.
    z = y .- û .* tanh.(α .+ b)

    return z
end

"""
    find_alpha(wt_y, wt_u_hat, b)

Compute an (approximate) real-valued solution ``α̂`` to the equation
```math
wt_y = α + wt_u_hat tanh(α + b)
```

The uniqueness of the solution is guaranteed since ``wt_u_hat ≥ -1``.
For details see appendix A.1 of the reference.

# Initial bracket

For all ``α``, we have
```math
α - |wt_u_hat| - wt_y \\leq α + wt_u_hat tanh(α + b) - wt_y \\leq α + |wt_u_hat| - wt_y.
```
Thus
```math
α̂ - |wt_u_hat| - wt_y \\leq 0 \\leq α̂ + |wt_u_hat| - wt_y,
```
which implies ``α̂ ∈ [wt_y - |wt_u_hat|, wt_y + |wt_u_hat|]``.
To avoid floating point issues if ``α̂  = wt_y ± |wt_u_hat|``, we use the more conservative
interval ``[wt_y - 2 |wt_u_hat|, wt_y + 2|wt_u_hat|]`` as initial bracket, at the cost of
one additional iteration step.

# References

D. Rezende, S. Mohamed (2015): Variational Inference with Normalizing Flows.
arXiv:1505.05770
"""
function find_alpha(wt_y::Real, wt_u_hat::Real, b::Real)
    # avoid promotions in root-finding algorithm and simplify AD dispatches
    return find_alpha(promote(wt_y, wt_u_hat, b)...)
end
function find_alpha(wt_y::T, wt_u_hat::T, b::T) where {T<:Real}
    # Compute the initial bracket (see above).
    Δ = 2 * abs(wt_u_hat)
    lower = float(wt_y - Δ)
    upper = float(wt_y + Δ)

    # Handle empty brackets (https://github.com/TuringLang/Bijectors.jl/issues/204)
    if lower == upper
        return lower
    end

    # Solve the root-finding problem
    α0 = Roots.find_zero((lower, upper)) do α
        return α + wt_u_hat * tanh(α + b) - wt_y
    end

    return α0
end

logabsdetjac(flow::PlanarLayer, x) = forward(flow, x).logabsdetjac
isclosedform(b::Inverse{<:PlanarLayer}) = false


# Pizza Slices
# Let's try this out on a problem. Let's say you want to order pizzas for 10 people (people = 10) and each person wants 2.5 slices on average (avg = 2.5). A pizza has 8 slices per pizza (slices = 8). How many pizzas should you order (pizzas = ?)? So we have the following

# Meaning	Variable
# Number of people	people
# Average number of slices each person eats	avg
# Number of slices on a piece of pizza	slices

people = 10;
avg = 2.5;
slices = 8;

pizzas = people * avg / slices;

println(pizzas);

pizzasRounded = ceil(pizzas);

println(pizzasRounded);


include("low_volumefraction.jl")

using JLD
using Plots
plotly()

strfs = ["f_slow1","f_slow2","f_fast1","f_fast2"];
f = strfs[4]
species = load("data/$f.jld")["species"]
ωs = load("data/$f.jld")["ωs"]
medium = load("data/$f.jld")["medium"]

kTs_arr = [
    [ sqrt(wavenumber_low_volumefraction(ω,medium, sps)) for ω in ωs]
for sps in [[species[1]],[species[2]],species]];

speed_arr = [ ωs./real(kTs) for kTs in kTs_arr]
atten_arr = imag(kTs_arr)

labs = ["void" "stone" "mix"];
p1 = plot(ωs./real(medium.c),speed_arr, xlabel="k", ylabel="wave speed", labels=labs,);
p2 = plot(ωs./real(medium.c),atten_arr, xlabel="k", ylabel="attenuation", labels=labs,);
plot(p1,p2)


function jt04()
#
#  test polynomial suggested by Jenkins and Traub
#
   p = poly([0.1;0.1;0.1;0.5;0.6;0.7]);
   z = [0.5  1; 0.6  1; 0.7  1; 0.1  3];
   p, PolyZeros(z)
end


function demoBurglarSimple()
    # This demo is the same as demoBurglar.jl, except that here we do not make use of the available potential indexing methods and also a potential dictionary (as in demoBurglarDictPot).

    burglar,earthquake,alarm,radio=1,2,3,4 # variables must start from 1 and be contiguous
    yes, no= 1,2 # define states, starting from 1. These do need to be contiguous, starting from 1, since currently the probability tables are defined by arrays of contiguous states.

    # This defines the meaning of the variables and their states.
    # It is not necessary for the computation, but is useful for
    # printing the results in a meaningful form:
    DictVariable=Dict{Integer,DiscreteVariable}()
    DictVariable[burglar]=DiscreteVariable("burglar",["yes","no"])
    DictVariable[earthquake]=DiscreteVariable("earthquake",["yes","no"])
    DictVariable[alarm]=DiscreteVariable("alarm",["yes","no"])
    DictVariable[radio]=DiscreteVariable("radio",["yes","no"])

    # Define the discrete probability tables:
    PotBurglar=PotArray(burglar,[0.01 0.99])
    PotEarthquake=PotArray(earthquake,[0.000001 1-0.000001])

    table=zeros(2,2)
    table[yes,yes]=1
    table[no,yes]=0
    table[yes,no]=0
    table[no,no]=1
    PotRadio=PotArray([radio earthquake],table)

    table=zeros(2,2,2)
    table[yes,yes,yes]=0.9999
    table[yes,yes,no]=0.99
    table[yes,no,yes]=0.99
    table[yes,no,no]=0.0001
    table[no,:,:]=1-table[yes,:,:]
    PotAlarm=PotArray([alarm burglar earthquake],table)

    jointpot=prod([PotBurglar PotRadio PotAlarm PotEarthquake])

    println("p(burglar|alarm=yes):")
    show(condpot(setpot(jointpot,alarm,yes),burglar),DictVariable)

    println("p(burglar|alarm=yes,radio=yes):")
    show(condpot(setpot(jointpot,[alarm radio],[yes yes]),burglar),DictVariable)

end


"""
    CoamoebaFiber2D(f::AbstractPolynomial, θ=(0., 0.))

Construct the fiber `CoamoebaFiber2D` of the bivariate polynomial `f` at `θ`.
"""
function CoamoebaFiber2D(p::MP.AbstractPolynomialLike{S}, θ=(0.0, 0.0)) where S
    vars = MP.variables(p)
    @assert (length(vars) == 2) "Expected a bivariate polynomial, but got $(p)."

    T = realified_coefficient_type(S)
    f_re, f_im = SP.Polynomial.(one(T) .* realify(p))

    sincosθ = _sincosθ(θ)
    U = zeros(2, 4)
    x = SVector(0.0, 0.0, 0.0, 0.0)
    CoamoebaFiber2D(f_re, f_im, sincosθ, U, x)
end

_sincosθ(θ) = (sin(θ[1]), sin(θ[2]), cos(θ[1]), cos(θ[2]))

function update_fiber!(F::CoamoebaFiber2D, θ)
    F.sincosθ = _sincosθ(θ)
    F
end

function phi!(fiber::CoamoebaFiber2D, w)
    s1, s2, c1, c2 = fiber.sincosθ
    v1, v2 = exp(w[1]), exp(w[2])
    x = SVector(v1 * c1, v2 * c2, v1 * s1, v2 * s2)
    fiber.x = x
    x
end

function evaluate(fiber::CoamoebaFiber2D, w)
    x = phi!(fiber, w)
    SVector(SP.evaluate(fiber.f_re, x), SP.evaluate(fiber.f_im, x))
end

function jacobian(fiber::CoamoebaFiber2D, w, precomputed=false)
    x = precomputed ? fiber.x : phi!(fiber, w)
    U = fiber.U
    @inbounds begin
        ∇f_re = SP.gradient(fiber.f_re, x)
        for i=1:4
            U[1, i] = ∇f_re[i]
        end
        U[2, 3] = U[1, 1]
        U[2, 4] = U[1, 2]
        U[2, 1] = -U[1, 3]
        U[2, 2] = -U[1, 4]
    end

    @inbounds out = begin
        ∂re∂w1 = muladd(U[1, 1], x[1], U[1, 3] * x[3])
        ∂re∂w2 = muladd(U[1, 2], x[2], U[1, 4] * x[4])
        ∂im∂w1 = muladd(U[2, 1], x[1], U[2, 3] * x[3])
        ∂im∂w2 = muladd(U[2, 2], x[2], U[2, 4] * x[4])

        SMatrix{2,2}(∂re∂w1, ∂im∂w1, ∂re∂w2, ∂im∂w2)
    end
    out
end


using Dolfin

import Dolfin:dot, Function, det
import PyCall: PyDict

parameters["form_compiler"]["cpp_optimize"] = true

ffc_options = PyDict()
ffc_options["optimize"]= true
ffc_options["eliminate_zeros"]= true
ffc_options["precompute_basis_const"]= true
ffc_options["precompute_ip_const"]=true

# Create mesh and define function space
mesh = UnitCubeMesh(24, 16, 16)
V = VectorFunctionSpace(mesh, "Lagrange", 1)

# Mark boundary subdomians
left =  CompiledSubDomain("near(x[0], side) && on_boundary", side = 0.0)
right = CompiledSubDomain("near(x[0], side) && on_boundary", side = 1.0)

# Define Dirichlet boundary (x = 0 or x = 1)
c = Expression(("0.0", "0.0", "0.0"))
r = Expression(("scale*0.0",
                "scale*(y0 + (x[1] - y0)*cos(theta) - (x[2] - z0)*sin(theta) - x[1])",
                "scale*(z0 + (x[1] - y0)*sin(theta) + (x[2] - z0)*cos(theta) - x[2])"),
                scale = 0.5, y0 = 0.5, z0 = 0.5, theta = pi/3)

bcl = DirichletBC(V, c, left)
bcr = DirichletBC(V, r, right)
bcs = [bcl, bcr]

# Define functions
du = TrialFunction(V)            # Incremental displacement
v  = TestFunction(V)             # Test function
u  = Function(V)                 # Displacement from previous iteration
B  = Constant((0.0, -0.5, 0.0))  # Body force per unit volume
T  = Constant((0.1,  0.0, 0.0))  # Traction force on the boundary

# Kinematics
d = u["geometric_dimension"]()
I = Identity(d)             # Identity tensor
F = I + grad(u)             # Deformation gradient
C = F["T"]*F                   # Right Cauchy-Green tensor

# Invariants of deformation tensors
Ic = tr(C)
J  = det(F)

# Elasticity parameters
E, nu = 10.0, 0.3
mu, lmbda = Constant(E/(2*(1 + nu))), Constant(E*nu/((1 + nu)*(1 - 2*nu)))

# Stored strain energy density (compressible neo-Hookean model)
# change ** to ^
psi = (mu/Constant(2))*(Ic - Constant(3)) - mu*ln(J) + (lmbda/Constant(2))*(ln(J)) * (ln(J))

# Total potential energy
Pi = psi*dx - dot(B, u)*dx - dot(T, u)*ds

# Compute first variation of Pi (directional derivative about u in the direction of v)
F = derivative(Pi, u, v)

# Compute Jacobian of F
J = derivative(F, u, du)

# Solve variational problem
solve(F == Constant(0), u, bcs, J=J, form_compiler_parameters=ffc_options)

# Save solution in VTK format
# file = File("displacement.pvd");
# file << u;

# Plot and hold solution
plot(u, mode = "displacement", interactive = true)


@snn_kw struct NoisyIFParameter{FT=Float32} <: AbstractIFParameter
    σ::FT = 0
end

@snn_kw mutable struct NoisyIF{VFT=Vector{Float32}} <: AbstractIF
    param::NoisyIFParameter = NoisyIFParameter()
    N::Int32 = 100
    randncache::VFT = randn(N)
end

"""
Noisy Integrate-And-Fire Neuron
"""
NoisyIF

function integrate!(p::NoisyIF, param::NoisyIFParameter, dt::Float32)
    @unpack N, randncache = p
    @unpack σ = param
    randn!(randncache)
    @inbounds for i = 1:N
        v[i] += dt * (ge[i] + gi[i] - (v[i] - El) + I[i] + σ / √dt * randncache[i]) / τm
        ge[i] += dt * -ge[i] / τe
        gi[i] += dt * -gi[i] / τi
    end
    @inbounds for i = 1:N
        fire[i] = v[i] > Vt
        v[i] = ifelse(fire[i], Vr, v[i])
    end
end


using BCBIStats
using StatsBase
using FreqTables
using ARules
using SparseArrays

mutable struct Stats
    pmid_count::Int

    mesh_names::Vector{String}

    topn_mesh_labels::Vector{String}
    topn_mesh_counts::Vector{Integer}

    pmi_sp::SparseMatrixCSC{Float64,Integer}
    top_coo_sp::SparseMatrixCSC{Int64, Int64}
    corrcoef::Array{Float64}

    mh_rules_pmid_count::Int
    mh_rules_df::DataFrame

end

Stats() = Stats(0, [], [], [], Array{Any,2}(undef, 0, 0), Array{Any,2}(undef, 0, 0), [], 0, DataFrame(lhs = [], rhs = [], supp = [], conf = [], lift = []))

"""
    get_plotting_inputs(mesh_df, topn)

Given a dataframe with PMIDs and Mesh Terms, returns various statistics for plotting
"""
function mesh_stats(mesh_df::DataFrame, topn::Int=50)

    # convert dataframe to frequency table
    mesh_frequencies = freqtable(string.(mesh_df[:pmid]), mesh_df[:descriptor])

    # total number of articles
    pmid_count = length(unique(mesh_df[:pmid]))

    # find the counts by mesh term, find the top mesh headings, and names
    mesh_counts = vec(sum(mesh_frequencies, 1))
    count_perm = sortperm(mesh_counts, rev=true)
    mesh_names = collect(keys(mesh_frequencies.dicts[2]))

    # check if there are topn frequencies, if not use length
    topn = min(topn, length(mesh_names))

    #RETURN LABELS FOR ALL CHARTS
    top_mesh_labels = mesh_names[count_perm[1:topn]]

    #RETURN FOR BAR CHART
    topn_mesh = mesh_counts[count_perm[1:topn]]

    #co-occurrance matrix - only for top MeSH
    top_occ = mesh_frequencies.array[:, count_perm[1:topn]]
    top_occ_sp = sparse(top_occ)
    top_coo_sp = top_occ_sp' * top_occ_sp

    # RETURN FOR PLOT MATRICES
    #Point Mutual Information
    pmi_sp = BCBIStats.COOccur.pmi_mat(top_coo_sp, pmid_count)
    #chi2
    # top_chi2= BCBIStats.COOccur.chi2_mat(top_occ, min_freq=0) NOT USED RIGHT NOW
    #correlation
    corrcoef = BCBIStats.COOccur.corrcoef(top_occ);

    # convert to BitArray
    mh_occ = convert(BitArray{2}, mesh_frequencies.array)

    # ARules

    top_mesh_df = filter(row -> row[:descriptor] in top_mesh_labels, mesh_df)

    fsets = Array{Array{String,1},1}()
    idx = Dict{Int32,Int}()

    for i = 1:size(top_mesh_df)[1]
      if haskey(idx, top_mesh_df[i,:pmid])
        push!(fsets[idx[top_mesh_df[i,:pmid]]], top_mesh_df[i,:descriptor])
      else
        push!(fsets, [top_mesh_df[i,:descriptor]])
        idx[top_mesh_df[i,:pmid]] = length(fsets)
      end
    end

    mh_rules_pmid_count = idx.count

    # find ARules - not actually used in plots?
    mh_rules = apriori(fsets, supp = 0.001, conf = 0, maxlen = 9)

    return PubMedMiner.Stats(
        pmid_count,
        mesh_names,
        top_mesh_labels,
        topn_mesh,
        pmi_sp,
        top_coo_sp,
        corrcoef,
        # mh_rules,
        # rules_df,
        # root,
        mh_rules_pmid_count,
        mh_rules
        # s,
        # t,
        # v
    )

end


"""
`OrthogonalPolynomials` is the abstract supertype of all univariate orthogonal
polynomials.
"""
abstract type OrthogonalPolynomials{T} <: PolynomialBasis{T}
end

const OPS{T} = OrthogonalPolynomials{T}

abstract type OrthogonalPolynomial{T} <: Polynomial{T} end

approx_length(b::OPS, n::Int) = n

# Taking the derivative of a polynomial lowers its degree. This results in a
# rectangular differentiation matrix. With the reducedegree option set to false,
# the degree is maintained (resulting in a square differentiation matrix).
function derivative_dict(dict::OPS, order::Int; reducedegree = true, options...)
	@assert order >= 0
	reducedegree ? resize(dict, length(dict)-order) : dict
end

function antiderivative_dict(s::OPS, order::Int; options...)
	@assert order >= 0
	resize(s, length(s)+order)
end

size(o::OrthogonalPolynomials) = (o.n,)

p0(::OPS{T}) where {T} = one(T)

hasextension(b::OPS) = true

extension(::Type{T}, src::O, dest::O; options...) where {T,O <: OPS} = IndexExtension{T}(src, dest, 1:length(src))
restriction(::Type{T}, src::O, dest::O; options...) where {T,O <: OPS} = IndexRestriction{T}(src, dest, 1:length(dest))

include("recurrence.jl")

# Default evaluation of an orthogonal polynomial: invoke the recurrence relation
unsafe_eval_element(b::OPS, idx::PolynomialDegree, x) =
    recurrence_eval(b, idx, x)
function unsafe_eval_element_derivative(b::OPS, idx::PolynomialDegree, x, order)
	@assert order == 1
    recurrence_eval_derivative(b, idx, x)
end

hasmeasure(dict::OPS) = true

measure(ϕ::OrthogonalPolynomial) = measure(dictionary(ϕ))

weightfun(b::OPS, x) = weightfun(measure(b), x)

function gram1(T, dict::OPS, m; options...)
    isorthonormal(dict, m) && return IdentityOperator{T}(dict)
	if isorthogonal(dict, m)
		diagonal_gram(T, dict, m; options...)
	else
		default_gram(T, dict, m; options...)
	end
end

diagonal_gram(::Type{T}, dict::OPS, measure; options...) where {T} =
    default_diagonal_gram(T, dict::OPS, measure; options...)

function symmetric_jacobi_matrix(b::OPS)
    T = codomaintype(b)
    n = length(b)
    α, β = monic_recurrence_coefficients(b)
    J = zeros(T, n, n)
    for k in 1:n
        J[k,k] = α[k]
        if k > 1
            J[k,k-1] = sqrt(β[k])
        end
        if k < n
            J[k,k+1] = sqrt(β[k+1])
        end
    end
    J
end

function roots(b::OPS{T}) where {T<:Number}
    J = symmetric_jacobi_matrix(b)
    eigen(J).values
end

function roots(b::OPS{T}) where {T<:Union{BigFloat}}
    J = symmetric_jacobi_matrix(b)
    # assuming the user has imported GenericLinearAlgebra.jl
    sort(real(eigvals!(J)))
end

interpolation_grid(b::OPS) = roots(b)
hasinterpolationgrid(dict::OPS) = true
opsorthogonal(dict, measure) = length(dict) -issymmetric(dict) <= length(points(measure))

"Return the first moment, i.e., the integral of the weight function."
function first_moment(b::OPS)
    @warn "implement first_moment for $b"
end

"""
Compute the Gaussian quadrature rule using the roots of the orthogonal polynomial.
"""
function gauss_rule(b::OPS{T}) where {T <: Real}
    J = symmetric_jacobi_matrix(b)
    x,v = eigen(J)
    b0 = first_moment(b)
    # In the real-valued case it is sufficient to use the first element of the
    # eigenvector. See e.g. Gautschi's book, "Orthogonal Polynomials and Computation".
    w = b0 * v[1,:].^2
    x,w
end

# In the complex-valued case, we have to compute the weights by summing explicitly
# over the full eigenvector.
function gauss_rule(b::OPS{T}) where {T <: Complex}
    J = symmetric_jacobi_matrix(b)
    x,v = eigen(J)
    b0 = first_moment(b)
    w = similar(x)
    for i in 1:length(w)
        w[i] = b0 * v[1,i]^2 / sum(v[:,i].^2)
    end
    x,w
end

import FastGaussQuadrature: gaussjacobi
using GaussQuadrature: jacobi
gaussjacobi(n::Integer,α::T,β::T) where T<:AbstractFloat =
    jacobi(n, α, β)

"""
Compute the weights of Gaussian quadrature rule from the given roots of the
orthogonal polynomial and using the formula:
```
w_j = 1 / \\sum_{k=0}^{n-1} p_k(x_j)^2
```
This formula only holds for the orthonormal polynomials.
"""
function gauss_weights_from_points(b::OPS, roots)
    pk = zeros(eltype(roots), length(b))
    w = zeros(eltype(roots), length(b))
    for j in 1:length(w)
        recurrence_eval_orthonormal!(pk, b, roots[j])
        w[j] = 1/sum(pk.^2)
    end
    w
end

function leading_order_coefficient(b::OPS{T}, idx) where {T}
    @assert 1 <= idx <= length(b)
    γ = one(T)
    for k in 0:idx-2
        γ *= rec_An(b, k)
    end
    γ
end

"""
Stieltjes Algorithm

Given nodes (t_k) and weights (w_k) of a quadrature rule, the Stieltjes Algorithm
calculates the N first recurrence coefficients of the monic polynomials (p_k) satisfying

```
p_{k+1}(t) = (t-a_k)p_k(t)-b_kp_{k-1}(t)
p_0(t) = 1, p_{-1}(t) = 0
```

and orthogonal with respect to the inner product induced by the quadrature

```
\\sum w_k p_m(t_k)q_n(t_k) = \\delta_{m,n}
```

See e.g. Gautschi's book, "Orthogonal Polynomials and Computation" and its
accompanying code `https://www.cs.purdue.edu/archives/2002/wxg/codes/OPQ.html`

Note: The vectors used have indices starting at `1` (such that `a_n` above is
given by `a[n+1]`). The last value is `a_{n-1} = a[n]`.
"""
function stieltjes(N, x::Array{RT},w::Array{T}) where {T,RT}
    @assert real(T) == RT
    # ####
    # Remove data with zero weights done in, e.g, Gautschi's book, "Orthogonal Polynomials and Computation".
    # ####
    # indices = sortperm(w)
    # x = x[indices]
    # w = w[indices]
    # index = minimum(find(!(w==0)))
    # w = w[index:length(w)]
    # x = x[index:length(w)]
    #
    # indices = sortperm(x)
    # x = x[indices]
    # w = w[indices]

    Ncap=length(x)
    # ####
    if N < 1 || N > Ncap
        error("Too little quadrature points")
    end

    a = zeros(T,N)
    b = zeros(T,N)

    p1 = zeros(T,Ncap)
    p2 = ones(T,Ncap)
    # The following scaling has to be adopted in case the
    # orthogonal polynomials underflow or overflow. In the
    # former case, c has to be chosen sufficiently large
    # positive, in the latter case sufficiently small
    # positive. (The example below relates to Table2_9.m.)
    #
    # if N==320
    #     c=1e50;
    #     p2=c*p2;
    #     s0=c^2*s0;
    # end
    p0 = similar(p1)
    scratch = similar(p1)
    stieltjes!(a,b,N,x,w,p0,p1,p2,scratch)
    a,b
end

"""
In-place Stieltjes Algorithm

see `stieltjes`

`p0`, `p1`, `p2`, and `scratch` are vectors of size `N`
"""
function stieltjes!(a::Array{T},b::Array{T},N,x::Array{RT},w::Array{T},p0::Array{T},p1::Array{T},p2::Array{T},scratch::Array{T}) where {T, RT}
    @assert real(T) == RT
    tiny = 10*typemin(real(T))
    huge = .1*typemax(real(T))

    s0 = sum(w)
    a[1] = dot(x,w)/s0
    b[1] = s0
    for k=1:N-1
        copyto!(p0,p1)
        copyto!(p1,p2)

        p2 .= (x .- a[k]) .* p1 .- b[k] .* p0
        scratch .= (p2 .^ 2) .* w
        s1 = sum(scratch)
        scratch .= w .* (p2 .^ 2) .* x
        s2 = sum(scratch)
        scratch .= abs.(p2)
        if (maximum(abs.(scratch))>huge || abs(s2)>huge)
            error("impending overflow in stieltjes for k=$(k)")
        end
        if abs(s1)<tiny
            error("impending underflow in stieltjes for k=$(k)")
        end
        a[k+1] = s2/s1
        b[k+1] = s1/s0
        s0 = s1
    end
end

"""
Modified Chebyshev Algorithm

Implements the map of the modified moments

```
\\mu_k = \\int p_k(x) dλ(x), k=0,...,2N-1
```
where (p_k) are some monic orthogonal polynomials satisfying the recurrence relation
```
p_{k+1}(t) = a_k(t_k)p_k(t) - b_kp{k-1}(t)
p_0(t) = 1, p_{-1}(t) = 0
```
to the first recurrence coefficients (α_k), (β_k), with k=0,...,N-1, such that
the monic polynomials (q_k) satisfying
```
q_{k+1}(t) = α_k(t_k)q_k(t) - β_kq{k-1}(t)
q_0(t) = 1, q_{-1}(t) = 0
```
are orthogonal with respect to the measure λ.

By default a, and b are zero vectors and therefore the moments are assumed to be
```
\\mu_k = \\int x^k dλ(x), k=0,...,2N-1
```
This procedure is however not stable.

See e.g. Gautschi's book, "Orthogonal Polynomials and Computation". Algorithm 2.1, p77

Note: The vectors used have indices starting at `1` (such that `α_n` above is
given by `α[n+1]`). The last value is `α_{n-1} = α[n]`.
"""
function modified_chebyshev(m::Array{T}, a=zeros(T,length(m)), b=zeros(T,length(m))) where {T}
    L = length(m)
    n = L>>1
    α = Array{T}(undef, n)
    β = Array{T}(undef, n)
    σmone = Array{T}(undef, L)
    σzero = Array{T}(undef, L)
    σ = Array{T}(undef, L)

    modified_chebyshev!(α,β,m,a,b,σ,σzero,σmone,n,1)
    α, β
end

"""
Modified Chebyshev Algorithm

see `chebyshev`

`σ`, `σzero`, and `σmone` are vectors of size or larger then `2n`
`n` is the number of coefficients required and `os` is the offset, i.e.,
the initial index that can be used in `σ`, `σzero`, and `σmone`.
"""
function modified_chebyshev!(α,β,m,a,b,σ,σzero,σmone,n=length(α),os=1)
    # Initialisation
    α[1] = a[1]+m[2]/m[1]
    β[1] = m[1]
    fill!(σmone,0)
    copyto!(σzero,os,m,os,2n)
    # continue
    for k=1:n-1
        for l in k+os:2n-k-1+os
            σ[l] = σzero[l+1]-(α[k]-a[l]).*σzero[l]-β[k].*σmone[l]+b[l].*σzero[l-1]
        end
        α[k+1] = a[k+1]
        α[k+1] += σ[k+2]/σ[k+1]
        α[k+1] = a[k+1]+σ[k+2]/σ[k+1]-σzero[k+1]/σzero[k]
        β[k+1] = σ[k+1]/σzero[k]
        copyto!(σmone,os,σzero,os,2n)
        copyto!(σzero,os,σ,os,2n)
    end
    nothing
end

"""
Calculates the first `n` recurrence coefficients (`α_k`) and (`β_k`) of the monic orthogonal polynomials
up to the tolerance `tol`
given a quadrature rule `my_quadrature_rule` with a relative degree of exactness
```
d(M)/M = δ + O(1/M) as M goes to infinity
```
with d(M) the degree of exactness. Thus δ is 2 for Gaussian quadrature.

See equation page 101 from Gautschi's book, "Orthogonal Polynomials and Computation"
"""
function adaptive_stieltjes(n,my_quadrature_rule::Function; tol = 1e-12, δ = 1, maxits = 20, quadrature_size=false)
    M = 1 + floor(Int, (2n-1)/δ)
    nodes, weights = my_quadrature_rule(M)
    ELT = eltype(nodes)

    α0,β0 = stieltjes(n,nodes,weights)

    M = M+1
    nodes, weights = my_quadrature_rule(M)
    α1,β1 = stieltjes(n,nodes,weights)

    no_its = 2
    while reduce(&, abs.(β1-β0) .> tol.*abs.(β1))
        if no_its > maxits
            @warn("accuracy of Stieltjes is not obtained, degree of Quadrature is $(M) and error is $(maximum(abs.(β1-β0)./abs.(β1)))")
            break
        end
        M = M+2^floor(Int,no_its/5)
        no_its = no_its + 1
        copyto!(α0,α1)
        copyto!(β0,β1)

        nodes, weights = my_quadrature_rule(M)
        α1,β1 = stieltjes(n,nodes,weights)
    end
    quadrature_size ? (α1, β1, M) : (α1, β1)
end

"""
Transforms the `N` recurrence coefficients (α_k) and (β_k) of monic
orthogonal polynomials (p_k), which satisfy the three-term recurrence relation
```
p_{k+1}(t) = (t-α_k)p_k(t)-β_kp_{k-1}(t)
```
to the `N-1` first recurrence coefficients (a_k), (b_k), and (c_k) the associated
monic polynomials (q_k), such that

```
q_{k+1}(t) = (a_kt+b_k)q_k(t)-c_kq_{k-1}(t)
```
"""
function monic_to_orthonormal_recurrence_coefficients(α::Array{T}, β::Array{T}) where {T}
    n = length(α)
    a = Array{T}(undef, n-1)
    b = Array{T}(undef, n-1)
    c = Array{T}(undef, n-1)
    monic_to_orthonormal_recurrence_coefficients!(a,b,c,α,β)
    a,b,c
end

"""
See `monic_to_orthonormal_recurrence_coefficients`
"""
function monic_to_orthonormal_recurrence_coefficients!(a::Array{T},b::Array{T},c::Array{T},α::Array{T},β::Array{T}) where {T}
    a .= 1 ./ sqrt.(view(β,2:length(β)))
    b .= -1.0.*view(α,1:length(α)-1)./sqrt.(view(β,2:length(β)))
    c .= sqrt.(view(β,1:length(β)-1)./view(β,2:length(β)))
    a,b,c
end


type SoftMaxLayer <: Layer
    inputActivation::Vol
    outputActivation::Vol

    function SoftMaxLayer(numOutputs::Int)
        dimensions = (1, 1, numOutputs)
        new(Vol(dimensions, 0.0), Vol(dimensions, 0.0))
    end
end

function forward(self::SoftMaxLayer, inputActivation::Vol)
    @assert inputSize(self) == outputSize(self)
    self.inputActivation = inputActivation
    exponentials = exp(self.inputActivation.w .- maximum(self.inputActivation.w))
    exponentials /= sum(exponentials)

    self.outputActivation.w = exponentials
    @assert abs(sum(self.outputActivation.w) - 1.0) < 1E-5
    return self.outputActivation
end

function backward(self::SoftMaxLayer, class::Int)
    @assert inputSize(self) == outputSize(self)
    fill!(self.inputActivation.dw, 0.0)
    for i in 1:length(self.outputActivation)
        indicator = ((i == class) ? 1.0 : 0.0)
        mul = -(indicator - self.outputActivation.w[1, 1, i])
        self.inputActivation.dw[1, 1, i] = mul
    end
    @assert self.inputActivation.dw |> sum |> abs < 1E-5
    return -log(self.outputActivation.w[1, 1, class])
end

gradientStates(self::SoftMaxLayer) = []


module BTGModel

using Reexport
using SpecialFunctions
using FastGaussQuadrature
using Distributions
using LinearAlgebra

include("./BTGFuncs.jl")

@reexport using .BTGFuncs

export
    Model,
    btgdensity,
    btgdistribution,
    computeweights

"""
    BTGModel

Settings for the Bayesian Transformed Gaussian model.
"""
struct Model{K<:Correlation,G<:Transform,V<:Covariate}
    k::K
    g::G
    f::V
    X::Array{Float64, 2}
    Y::Array{Float64, 1}
    knumgauss::Int
    gnumgauss::Int
end

function gaussweights(d::Uniform, n)
    x, w = gausslegendre(n)
    return ((d.b - d.a) .* (x .+ 1) ./ 2 .+ d.a), w ./ 2
end

function gaussweights(d::Normal, n)
    x, w = gausshermite(n)
    return (sqrt(2) * d.σ) .* x .+ d.μ, w ./ sqrt(π)
end

function gaussweights(d::Gamma, n)
    x, w = gausslaguerre(n, d.α - 1)
    return d.θ .* x, w ./ gamma(d.α)
end

function gaussweights(d::Exponential, n)
    x, w = gausslaguerre(n)
    return d.θ .* x, w
end

function weighttensor(f::Parametric, n)
    tups = (gaussweights(getfield(f, s), n) for s in fieldnames(typeof(f)))
    x = collect(Iterators.product((t[1] for t in tups)...))
    w = map(x -> *(x...), Iterators.product((t[2] for t in tups)...))
    return reshape(x, length(x)), reshape(w, length(w))
end

function computeweights(mdl::Model)
    kgauss = weighttensor(mdl.k, mdl.knumgauss)
    ggauss = weighttensor(mdl.g, mdl.gnumgauss)
    return (kgauss, ggauss)
end

struct HSSSolver end # TODO Yuanxi Solver

# function Base.size(s::HSSSolver) end # TODO

# function Base.:\(s::HSSSolver, y) end # TODO

# function LinearAlgebra.det(s::HSSSolver) end # TODO

function makesolver!(A)
    return cholesky!(A) # TODO Placeholder
end

# TODO Clean this function up, make it faster
function computedist(mdl::Model, x::Array{Float64, 1}, y::Float64, ((kx, kw), (gx, gw)), f)
    n, p = size(mdl.X)

    c = 0
    acc = 0

    for i = 1:length(kx)
        θ = kx[i]
        α = kw[i]

        cholΣ = makesolver!(Symmetric(kernelmatrix(mdl.k, θ, mdl.X, mdl.X)))
        cholXΣX = makesolver!(Symmetric(mdl.X' * (cholΣ \ mdl.X)))
        B = kernelmatrix(mdl.k, θ, reshape(x, 1, length(x)), mdl.X)
        D = 1.0 - (B * (cholΣ \ B'))[1]
        H = x .- B * (cholΣ \ mdl.X)
        C = D + (H * (cholXΣX \ H'))[1]

        for j = 1:length(gx)
            λ = gx[j]
            α′ = gw[j]

            Z = mdl.g.(λ, mdl.Y)
            z = mdl.g(λ, y)

            β = cholXΣX \ (mdl.X' * (cholΣ \ Z))
            q = ((Z - mdl.X * β)' * (cholΣ \ (Z - mdl.X * β)))[1]
            J = abs(prod(prime.(Ref(mdl.g), λ, mdl.Y)))
            m = (B * (cholΣ \ Z))[1] + (H * β)[1]

            h = q^(-(n - p) / 2) * J^(1 - p / n) / sqrt(det(cholΣ) * det(cholXΣX))
            T = LocationScale(m, sqrt(q * C), TDist(n - p))
            P = f(T, θ, λ, z)

            c += α * α′ * h
            acc += α * α′ * h * P
        end
    end

    return acc / c
end

"""
    btgdensity(mdl, x, y)

The probability density of the value `y` at the location `x` given a BTG model.

```math
\\mathcal{P}(y_x\\lvert m)
```
"""
function btgdensity(mdl::Model, x::Array{Float64, 1}, y::Float64, weights)
    return computedist(mdl, x, y, weights, function (T, θ, λ, z)
                       return pdf(T, z) * abs(prime(mdl.g, λ, y))
                       end)
end

"""
    btgdistribution(mdl, x, y)

The cumulative probability of the value `y` at the location `x` given a BTG model.

```math
\\Phi(y_x\\lvert m)
```
"""
function btgdistribution(mdl::Model, x::Array{Float64, 1}, y::Float64, weights)
    return computedist(mdl, x, y, weights, (T, _, _, z) -> cdf(T, z))
end

end # module


struct _bi01IP
    C::Matrix{Int} # objective functions, k=1..2, j=1..n
    A::Matrix{Int} # matrix of constraints, i=1..m, j=1..n
    b::Vector{Int} # right-hand side, i=1..m
end

mutable struct solution
    val_objectif::Vector{Int64}
    sol::Vector{Int64}
    cout::Vector{Int64}
end


using Qsosed, Test

@testset "Test corona properties" begin
    # comparison with Xspec model

    bh = BlackHole(1e8 * M_SUN, 0.5, 0.0)
    corona =
        Corona(bh, hard_xray_fraction = 0.02, electron_energy = 100, reprocessing = true)
    Ledd = compute_eddington_luminosity(bh)

    computed_radius = compute_corona_radius(corona)
    @test computed_radius ≈ 8.9008961989970832 rtol = 1e-3

    temp_corona = compute_disk_temperature(bh, computed_radius)
    @test temp_corona ≈ 107102.98423666063 rtol = 1e-1

    dissip_lumin = compute_corona_dissipated_luminosity(corona) / Ledd
    @test dissip_lumin ≈ 2.0183853632627917e-2 rtol = 1e-2

    corona_luminosity = compute_corona_luminosity(corona) / Ledd
    @test corona_luminosity ≈ 3.6390444796974425e-2 rtol = 1e-1

    corona_photon_index = compute_corona_photon_index(corona)
    @test corona_photon_index ≈ 2.2826826043284854 rtol = 1e-2

    # Second black hole

    bh = BlackHole(1e6 * M_SUN, 0.025, 0.0)
    @test gravity_radius(bh) ≈ 1161.0642639605840 rtol = 1e-2
    corona =
        Corona(bh, hard_xray_fraction = 0.02, electron_energy = 100, reprocessing = true)
    Ledd = compute_eddington_luminosity(bh)

    computed_radius = compute_corona_radius(corona)
    @test computed_radius ≈ 89.731710842601927  rtol = 1e-2

    temp_corona = compute_disk_temperature(bh, computed_radius)
    @test temp_corona ≈ 47643.787634112188 rtol = 1e-1

    dissip_lumin = compute_corona_dissipated_luminosity(corona) / Ledd
    @test dissip_lumin ≈ 2.0001856311471888e-2 rtol = 1e-2

    corona_luminosity = compute_corona_luminosity(corona) / Ledd
    @test corona_luminosity ≈ 2.0621184497821025e-2 rtol = 1e-1

    corona_photon_index = compute_corona_photon_index(corona)
    @test corona_photon_index ≈ 1.6483970267857697  rtol = 1e-1

    warm = Warm(corona, electron_energy = 0.2, photon_index = 2.5)
    energy_range = 10 .^ range(-5, 3, length = 500);
    corona_sed = compute_corona_spectral_luminosity(corona, warm, energy_range)
    integ_lumin = sum(corona_sed[2:end] .* diff(energy_range))
    target_lumin = compute_corona_luminosity(corona)
    @test integ_lumin ≈ target_lumin rtol = 1e-5

end


"""
    `CalcVSCoefs(boundary, etas; U_slip, sigma)`

Calculates the vortex sheet RBF coefficients.

# ARGUMENTS
* `boundary::Boundary`                       : Boundary object
* `etas::Array{T}`                           : RBF coefficients for rho
* `U_slip::Union{Nothing, Function}=nothing` : Slip velocity field function: U_slip(X)
* `sigma = 0.2`                              : Gaussian width
where {T<:Real}

# RETURNS
* `alphas::Array{Float64}`                   : Array of RBF coefficients
"""
function CalcVSCoefs(boundary::Boundary,
                     etas::Array{T};
                     U_slip::Union{Nothing, Function}=nothing,
                     sigma=0.2) where {T<:Real}

    # === Extract geometry from boundary ===
    X_body = [[point[1], point[2], 0.0] for point in boundary.bodyPTS]
    n_hats = [[n_hat[1], n_hat[2], 0.0] for n_hat in boundary.nHats]
    t_hats = [[t_hat[1], t_hat[2], 0.0] for t_hat in boundary.tHats]
    NPTS = boundary.NPTS_BODY

    # === Function definitions ===

    # i -> RBF index
    # j -> Integration index
    # k -> Evaluation index

    X_i(i) = X_body[i]                                  # Set_i of points on S
    X_j(j) = X_body[j]                                  # Set_j of points on S
    X_jp1(j) = (j != NPTS) ? X_body[j+1] : X_body[1]    # Point j + 1
    X_k(k) = X_body[k]                                  # Set_k of points on S

    R_kj(k,j) = X_k(k) - X_j(j)                         # X_k - X_j
    R_ji(j,i) = X_j(j) - X_i(i)                         # X_j - X_i
    R_ki(k,i) = X_k(k) - X_i(i)                         # X_k - X_i
    r_kj(k,j) = norm(R_kj(k,j))                         # | X_k - X_j |
    r_ji(j,i) = norm(R_ji(j,i))                         # | X_j - X_i |
    r_ki(k,i) = norm(R_ki(k,i))                         # | X_k - X_i |

    del_S(j) = norm(X_j(j) - X_jp1(j))                  # ΔS_j
    L = sum( [del_S(j) for j in 1:NPTS] )               # Length of surface

    #etas = CalcRhoCoefs(boundary)                       # Coefficients for rho RBF
    rho_1(k) = CalcRho(boundary, etas, X_k(k))          # ρ_1 of Eigendecomposition

    # === Coefficent functions ===
    function Theta_ki(k,i)
        res = 0
        for j in 1:NPTS
            if j != k
                res = res + (1/(2*pi)) * (dot(R_kj(k,j), n_hats[j]) * RBF_gauss(r_ji(j,i), sigma=sigma) * del_S(j))/(r_kj(k,j)^2)
            end
        end
        return res
    end
    #Theta_ki(k,i) = sum([((k != j) ? (1/pi) * (dot(R_kj(k,j), n_hats[j]) * RBF_gauss(r_ji(j,i)) * del_S(j))/(r_kj(k,j)^2) : 0) for j in 1:NPTS] )
    Lambda_ki(k,i) = sum( [RBF_gauss(r_ji(j,i), sigma=sigma) * (rho_1(k))/(L) for j in 1:NPTS] )
    phi_ki(k,i) = RBF_gauss(r_ki(k,i), sigma=sigma)

    # === Calculate coefficient matrix ===
    A = zeros(NPTS,NPTS)
    for k in 1:NPTS
        for i in 1:NPTS
            A[k,i] = phi_ki(k,i) - Theta_ki(k,i) + Lambda_ki(k,i)
        end
    end
    #A = [ [phi_ki(k,i) - Theta_ki(k,i) + Lambda_ki(k,i) for k in 1:NPTS] for i in 1:NPTS]
    #println("A: ",A)

    # === Calculate RHS vector ===
    b = zeros(NPTS)
    for i in 1:NPTS
        b[i] = dot(U_slip(X_i(i)), t_hats[i])
    end
    #b = [dot(U_slip(X_i(i)), t_hats[i]) for i in 1:NPTS]
    #println("b: ",b)

    # === Calculate alpha coefficients ===
    alphas = A\b
    return alphas

end

"""
    `CalcVS(boundary, alpha, X_eval)`

Evaluates the RBF interpolation function: gamma(X) = sum(alpha_i * phi(X - X_i)).

# ARGUMENTS
* `boundary::Boundary` : Boundary object
* `alpha::Array{T}`    : Array containing RBF coefficients
* `X_eval::Array{T}`   : Point to evaluate function at
* `sigma = 0.2`        : Gaussian width
where {T<:Real}

# RETURNS
* `gamma::Float64`     : Evaluated function value
"""
function CalcVS(boundary::Boundary,
                alpha::Array{T},
                X_eval::Array{T};
                sigma=0.2) where {T<:Real}

    # === Unpack geometry ===
    XP = [[point[1], point[2], 0.0] for point in boundary.bodyPTS]
    NPTS = boundary.NPTS_BODY

    # === Summation over all points for gamma ===
    gamma = -sum( [alpha[i] * RBF_gauss(norm(X_eval - XP[i]), sigma=sigma) for i in 1:NPTS ] )
    return gamma

end

# """
#
# """
# function CalcVSVelocityCirlce(boundary::Boundary,
#                               alpha::Array{T},
#                               X_eval::Array{T};
#                               radius=1.0) where {T<:Real}
#
#     # === Unpack geometry ===
#     X_body = [[point[1], point[2], 0.0] for point in boundary.bodyPTS]
#     NPTS = boundary.NPTS_BODY
#
#     # === Function definitions ===
#     X_j(j) = X_body[j]                                       # Set_j of points on S
#     X_jp1(j) = (j != NPTS) ? X_body[j+1] : X_body[1]         # Point j + 1
#     R_xj(j) = X_eval - X_j(j)                                # X_eval - X_j
#     r_xj(j) = norm(R_xj(j))                                  # | X_eval - X_j |
#
#     gamma(j) = CalcVS(boundary, alpha, X_j(j)) .* [0, 0, 1]  # Vortex sheet strength
#     del_S(j) = norm(X_j(j) - X_jp1(j))                       # ΔS_j
#
#     num(j) = cross(X_eval - X_j(j), gamma(j)) * del_S(j)     # Numerator
#     den(j) = 2*pi * r_xj(j)^2                                # Denominator
#
#     # === Calculate velocity ===
#     U = sum( [((r_xj(j) < 1e-12) ? zeros(3) : num(j) ./ den(j)) for j in 1:NPTS] )
#     return U
#
# end


using SemiseparableMatrices, LazyArrays, MatrixFactorizations, Test

@testset "Inverse of banded" begin
    @testset "Bidiagonal" begin
        @testset "SemiseparableMatrix" begin
            n = 5
            L = Bidiagonal(randn(n), randn(n-1), :L)
            U = Bidiagonal(randn(n), randn(n-1), :U)
            Li = SemiseparableMatrix(ApplyArray(inv, L))
            Ui = SemiseparableMatrix(ApplyArray(inv, U))
            @test Li ≈ inv(L)
            @test Ui ≈ inv(U)
        end

        @testset "Derivation" begin
            n = 5
            L = Bidiagonal(ones(n), randn(n-1), :L)
            Li = ApplyArray(inv, L)
            ℓ = [1; cumprod(-L.ev)]
            @test tril(ℓ * inv.(ℓ)') ≈ Li

            L = Bidiagonal(randn(n), randn(n-1), :L)
            L̃ = Bidiagonal(ones(n), L.ev ./ L.dv[1:end-1], :L)
            # inv(L) == inv(L̃ * D) == inv(D) * inv(L̃)
            ℓ = [1; cumprod(-L.ev./L.dv[1:end-1])]
            @test tril(ℓ * inv.(ℓ)') ≈ inv(L̃)
            @test tril((inv.(L.dv) .* ℓ) * inv.(ℓ)') ≈ inv(L)

            n = 5
            U = Bidiagonal(ones(n), randn(n-1), :U)
            u = [1; cumprod(-U.ev)]
            @test triu(inv.(u) * u') ≈ inv(U)

            U = Bidiagonal(randn(n), randn(n-1), :U)
            u = [1; cumprod(-U.ev./U.dv[1:end-1])]
            @test triu(inv.(u) * (inv.(U.dv) .* u)') ≈ inv(U)
        end

        @testset "Tridiagonal inverse" begin
            n = 10
            A = Tridiagonal(randn(n-1), randn(n), randn(n-1)) + 10I
            U,L = ul(A, Val(false))
            Li = SemiseparableMatrix(ApplyArray(inv, L))
            Ui = SemiseparableMatrix(ApplyArray(inv, U))
            Ai = inv(A)
            @test Li * Ui ≈ Ai
            @test rank(triu(Ai)[1:5,5:end]) == 1
            a,b = Li.L.args; b = b'
            c,d = Ui.U.args; d = d'
            @test Ai[1,1] ≈ a[1]b[1]c[1]d[1]
            @test Ai[1,2] ≈ a[1]b[1]c[1]d[2]
            @test Ai[1,3] ≈ a[1]b[1]c[1]d[3]
            @test Ai[1,5] ≈ a[1]b[1]c[1]d[5]
            @test Ai[2,1] ≈ a[2]b[1]c[1]d[1]
            @test Ai[2,2] ≈ a[2]b[1]c[1]d[2]+a[2]b[2]c[2]d[2]
            @test Ai[2,3] ≈ a[2]b[1]c[1]d[3]+a[2]b[2]c[2]d[3]
            @test Ai[3,1] ≈ a[3]b[1]c[1]d[1]
            @test Ai[3,2] ≈ a[3]b[1]c[1]d[2]+a[3]b[2]c[2]d[2]
            @test Ai[3,3] ≈ a[3]b[1]c[1]d[3]+a[3]b[2]c[2]d[3]+a[3]b[3]c[3]d[3]
            @test Ai[3,4] ≈ a[3]b[1]c[1]d[4]+a[3]b[2]c[2]d[4]+a[3]b[3]c[3]d[4]

            @test Ai[1:3,3:5] ≈ [a[1]b[1]c[1]; a[2]*(b[1]c[1]+b[2]c[2]); a[3]*(b[1]c[1]+b[2]c[2]+b[3]c[3])]*d[3:5]'
            @test Ai[3:5,1:3] ≈ a[3:5] * [b[1]c[1]d[1] (b[1]c[1]+b[2]c[2])*d[2] (b[1]c[1]+b[2]c[2]+b[3]c[3])*d[3]]

            b̃ = cumsum(b .* c) .* d
            c̃ = a .* cumsum(b .* c)
            @test tril(a*b̃') ≈ tril(Ai)
            @test triu(c̃*d') ≈ triu(Ai)
        end
    end
end


function u_pcm(x, c, η; m)
    d_squared = FuzzySets.euclidean_distance(x, c; squared=true)
    u = 1 / (1 + (d_squared / η)^(1 / (m - 1)))
    u
end


using Test

flip(bit) = bit == '0' ? '1' : '0'
most_common(bits) = 2count(==('0'), bits) > length(bits) ? '0' : '1'
least_common(bits) = flip.(most_common(bits))
lines_to_bits(lines) = hcat(Vector{Char}.(lines)...)
parse_bits(chars::Vector{Char}) = parse(Int, join(chars, ""); base = 2)

codes = readlines("day03.in")
test_codes = ["00100",
              "11110",
              "10110",
              "10111",
              "10101",
              "01111",
              "00111",
              "11100",
              "10000",
              "11001",
              "00010",
              "01010"]

function q1(codes)
    codebits = lines_to_bits(codes)
    γ = most_common.(eachrow(codebits))
    ε = flip.(γ)
    return parse_bits(γ) * parse_bits(ε)
end

function rating(criteria, codebits)
    matching = reduce(eachrow(codebits), init = 1:size(codebits, 2)) do m, row
        return length(m) == 1 ? m : m[row[m] .== criteria(row[m])]
    end
    return vec(codebits[:, matching])
end

function q2(codes)
    codebits = lines_to_bits(codes)
    oxygen_rating = rating(most_common, codebits)
    co2_rating = rating(least_common, codebits)
    return parse_bits(oxygen_rating) * parse_bits(co2_rating)
end

@test q1(test_codes) == 198
@test q2(test_codes) == 230

println("Q1: ", q1(codes))
println("Q2: ", q2(codes))


function bonferroni(pValues)
  @check_pValues pValues
  min(1, pValues * length(pValues))
end


# ---
# title: 810. Chalkboard XOR Game
# id: problem810
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Math
# link: <https://leetcode.com/problems/chalkboard-xor-game/description/>
# hidden: true
# ---
#
# We are given non-negative integers nums[i] which are written on a chalkboard.
# Alice and Bob take turns erasing exactly one number from the chalkboard, with
# Alice starting first.  If erasing a number causes the bitwise XOR of all the
# elements of the chalkboard to become 0, then that player loses.  (Also, we'll
# say the bitwise XOR of one element is that element itself, and the bitwise XOR
# of no elements is 0.)
#
# Also, if any player starts their turn with the bitwise XOR of all the elements
# of the chalkboard equal to 0, then that player wins.
#
# Return True if and only if Alice wins the game, assuming both players play
# optimally.
#
#
#
#     **Example:**
#     Input: nums = [1, 1, 2]
#     Output: false
#     Explanation:
#     Alice has two choices: erase 1 or erase 2.
#     If she erases 1, the nums array becomes [1, 2]. The bitwise XOR of all the elements of the chalkboard is 1 XOR 2 = 3. Now Bob can remove any element he wants, because Alice will be the one to erase the last element and she will lose.
#     If Alice erases 2 first, now nums becomes [1, 1]. The bitwise XOR of all the elements of the chalkboard is 1 XOR 1 = 0. Alice will lose.
#
#
#
# **Notes:**
#
#   * `1 <= N <= 1000`.
#   * `0 <= nums[i] <= 2^16`.
#
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


@testset "k_cliques tests" begin

    @testset "triangle tests" begin
        @test length(kcliques(MatrixNetwork(sparse([0 1; 1 0])),3)[2]) == 0
        @test length(kcliques(MatrixNetwork(sparse(fill(0,10,10))),3)[2]) == 0
        @test length(kcliques(MatrixNetwork(sparse(fill(0,10,10))),3)[2]) == 0
        @test length(kcliques(MatrixNetwork(sparse(fill(0,10,10))),3)[2]) == 0
        @test length(kcliques(MatrixNetwork(sparse([0 1 1; 1 0 1; 1 1 0])), 3)[2]) == 1
        @test length(kcliques(load_data("facebook_data"), 3)[2]) == 1612010
        @test length(kcliques(MatrixNetwork(load_matrix_network("clique-10")), 3)[2]) == 120
    end

end


## 3D Panels
#==========================================================================================#

abstract type AbstractPanel3D <: AbstractPanel end

"""
    Panel3D(p1, p2, p3, p4)

A composite type consisting of 4 Cartesian coordinates `p1, p2, p3, p4` representing corners of a panel in 3 dimensions. The following commutative diagram depicts the order:

```
z → y
↓
x
        p1 —→— p4
        |       |
        ↓       ↓
        |       |
        p2 —→— p3
```
"""
struct Panel3D{T <: Real} <: AbstractPanel3D
    p1 :: SVector{3,T}
    p2 :: SVector{3,T}
    p3 :: SVector{3,T}
    p4 :: SVector{3,T}
end

Panel3D(p1 :: FieldVector{3,T}, p2 :: FieldVector{3,T}, p3 :: FieldVector{3,T}, p4 :: FieldVector{3,T}) where T <: Real = Panel3D{T}(p1, p2, p3, p4)

average_chord(panel :: Panel3D) = (p2(panel) - p1(panel) + p3(panel) - p4(panel)) / 2
average_width(panel :: Panel3D) = (p4(panel) - p1(panel) + p3(panel) - p2(panel)) / 2

"""
    panel_area(panel :: Panel3D)

Compute the area of a Panel3D.
"""
panel_area(panel :: Panel3D) = (norm ∘ cross)(average_chord(panel), average_width(panel))

wetted_area(panels :: Matrix{<: Panel3D}) = sum(panel -> panel_area(panel), panels)

"""
    panel_coords(panel :: Panel3D)

Compute the coordinates of a Panel3D.
"""
panel_coords(panel :: Panel3D) = structtolist(panel)

"""
    transform(panel :: Panel3D, rotation, translation)

Perform an affine transformation on the coordinates of a Panel3D given a rotation matrix and translation vector.
"""
function transform(panel :: Panel3D, rotation, translation)
    p1, p2, p3, p4 = (Translation(translation) ∘ LinearMap(rotation)).(panel_coords(panel))
    Panel3D(p1, p2, p3, p4)
end

"""
    midpoint(panel :: Panel3D)

Compute the midpoint of a Panel3D.
"""
midpoint(panel :: Panel3D) = (p1(panel) + p2(panel) + p3(panel) + p4(panel)) / 4

"""
    panel_normal(panel :: Panel3D)

Compute the normal vector of a Panel3D.
"""
panel_normal(panel :: Panel3D) = let p31 = p3(panel) - p1(panel), p42 = p4(panel) - p2(panel); p31 × p42 end

"""
    transform_normal(panel :: Panel3D, h_l, g_l)

Compute the normal ``n_l``, the normal ``n_0`` of a `Panel3D` perturbed by the control gain ``\\delta_l`` about the hinge axis ``h_l``.
"""
transform_normal(panel :: Panel3D, h_l, g_l) = g_l * cross(h_l, panel_normal(panel))


struct RouletteWheelSelection end
# selection(pop::Vector{<:GACreature}, n::Integer, rng)
function selection(::RouletteWheelSelection,
                   pop::Vector{<:GACreature}, n::Integer, rng=Base.GLOBAL_RNG)
    wmin,wmax = extrema(fitness(c) for c in pop)
    weight = wmax - wmin
    function stochasticpick()
        while true
            i = rand(rng,1:length(pop))
            if wmin + weight * rand(rng,typeof(weight)) <= fitness(pop[i])
                return i
            end
        end
    end
    parents = Vector{Tuple{Int,Int}}(undef, n)
    for k = 1:n
        i = stochasticpick()
        j = i
        while i==j
            j = stochasticpick()
        end
        parents[k] = (i,j)
    end
    parents
end

struct TournamentSelection
    k::Int # if k=2 then binary tournament selection
    TournamentSelection(k=2) = new(k)
end
function selection(sel::TournamentSelection,
                   pop::Vector{<:GACreature}, n::Integer, rng=Base.GLOBAL_RNG)
    function stochasticpick()
        si = rand(rng,1:length(pop))
        weighti = fitness(pop[si])
        for _ = 1:sel.k-1
            sj = rand(rng,1:length(pop))
            weightj = fitness(pop[sj])
            if weightj > weighti
                si = sj
                weighti = weightj
            end
        end
        si
    end
    parents = Vector{Tuple{Int,Int}}(undef, n)
    for k = 1:n
        i = stochasticpick()
        j = i
        while i==j
            j = stochasticpick()
        end
        parents[k] = (i,j)
    end
    parents
end


@with_kw Jet{Tf, Tdf, Td2f, Td3f, Tj, Td}
	δ::Td = 1e-8
	F::Tf
	dF::Tdf = (x,p,dx) -> (F(x .+ δ .* dx, p) - F(x, p)) / δ
	d2F::Td2f = (x,p,dx1,dx2) -> (dF(x .+ δ .* dx2, p, dx1) - dF(x, p, dx1)) / δ
	d3F::Td3f = (x,p,dx1,dx2) -> (d2F(x .+ δ .* dx3, p, dx1, dx2) - d2F(x, p, dx1, dx2)) / δ
	J::Tj = nothing
end

Jet(f) = Jet(F = f)
(j::Jet)(x, p) = j.F(x, p)
(j::Jet)(::Val{Jacobian}, x, p) = j.J(x, p)
(j::Jet)(x, p, dx) = j.dF(x, p, dx)
(j::Jet)(x, p, dx1, dx2) = j.d2F(x, p, dx1, dx2)
(j::Jet)(x, p, dx1, dx2, dx3) = j.d2F(x, p, dx1, dx2, dx3)


using DiffEqJump, DiffEqBase
using Test
using LinearAlgebra
using LightGraphs

function get_mean_sol(jump_prob, Nsims, saveat)
    sol = solve(jump_prob, SSAStepper(), saveat = saveat).u
    for i in 1:Nsims-1
        sol += solve(jump_prob, SSAStepper(), saveat = saveat).u
    end
    sol/Nsims
end

# assume sites are labeled from 1 to num_sites(spatial_system)
function discrete_laplacian_from_spatial_system(spatial_system, hopping_rate)
    sites = 1:num_sites(spatial_system)
    laplacian = zeros(length(sites), length(sites))
    for site in sites
        laplacian[site,site] = -outdegree(spatial_system, site)
        for nb in neighbors(spatial_system, site)
            laplacian[site, nb] = 1
        end
    end
    laplacian .*= hopping_rate
    laplacian
end

# problem setup
majumps = JumpSet(nothing)
tf = 0.5
u0 = [100]
num_species = 1

domain_size = 1.0 #μ-meter
linear_size = 5
diffusivity = 0.1
dim = 2
dims = Tuple([linear_size for i in 1:dim])
num_nodes = prod(dims)

# Starting state setup
starting_state = zeros(Int, length(u0), num_nodes)
center_node = trunc(Int,(num_nodes+1)/2)
starting_state[:,center_node] = copy(u0)
tspan = (0.0, tf)
prob = DiscreteProblem(starting_state,tspan, [])

hopping_rate = diffusivity * (linear_size/domain_size)^2
hopping_constants = [hopping_rate for i in starting_state]

# analytic solution
lap = discrete_laplacian_from_spatial_system(LightGraphs.grid(dims), hopping_rate)
evals, B = eigen(lap) # lap == B*diagm(evals)*B'
Bt = B'
analytic_solution(t) = B*diagm(ℯ.^(t*evals))*Bt * reshape(prob.u0, num_nodes, 1)

num_time_points = 10
Nsims = 10000
rel_tol = 0.01
times = 0.0:tf/num_time_points:tf

algs = [NSM(), DirectCRDirect()]
grids = [CartesianGridRej(dims), CartesianGridIter(dims), LightGraphs.grid(dims)]
jump_problems = JumpProblem[JumpProblem(prob, algs[2], majumps, hopping_constants=hopping_constants, spatial_system = grid, save_positions=(false,false)) for grid in grids]
sizehint!(jump_problems, length(algs) * 3 + length(jump_problems))

# flattenned jump prob
push!(jump_problems, JumpProblem(prob, NRM(), majumps, hopping_constants=hopping_constants, spatial_system = grids[1], save_positions=(false,false)))

# hop rates of form L_{s,i,j}
hop_constants = Matrix{Vector{Float64}}(undef, size(hopping_constants))
for ci in CartesianIndices(hop_constants)
    (species, site) = Tuple(ci)
    hop_constants[ci] = repeat([hopping_constants[species, site]], outdegree(grids[1], site))
end
for alg in algs
    push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants=hop_constants, spatial_system=grids[1], save_positions=(false,false)))
    push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants=hop_constants, spatial_system=grids[3], save_positions=(false,false)))
end

# hop rates of form D_s * L_{i,j}
species_hop_constants = [hopping_rate]
site_hop_constants = Vector{Vector{Float64}}(undef, num_nodes)
for site in 1:num_nodes
    site_hop_constants[site] = repeat([1.0], DiffEqJump.outdegree(grids[1], site))
end
for alg in algs
    push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants=Pair(species_hop_constants, site_hop_constants), spatial_system=grids[1], save_positions=(false,false)))
    push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants=Pair(species_hop_constants, site_hop_constants), spatial_system=grids[3], save_positions=(false,false)))
end

# hop rates of form D_{s,i} * L_{i,j}
species_hop_constants = hopping_rate * ones(1, num_nodes)
site_hop_constants = Vector{Vector{Float64}}(undef, num_nodes)
for site in 1:num_nodes
    site_hop_constants[site] = repeat([1.0], DiffEqJump.outdegree(grids[1], site))
end
for alg in algs
    push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants=Pair(species_hop_constants, site_hop_constants), spatial_system=grids[1], save_positions=(false,false)))
    push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants=Pair(species_hop_constants, site_hop_constants), spatial_system=grids[3], save_positions=(false,false)))
end

# testing
for (j,spatial_jump_prob) in enumerate(jump_problems)
    mean_sol = get_mean_sol(spatial_jump_prob, Nsims, tf/num_time_points)
    for (i,t) in enumerate(times)
        local diff = analytic_solution(t) - reshape(mean_sol[i], num_nodes, 1)
        @test abs(sum(diff[1:center_node])/sum(analytic_solution(t)[1:center_node])) < rel_tol
    end
end

# testing non-uniform hopping rates
dims = (2,2)
num_nodes = prod(dims)
grid = LightGraphs.grid(dims)
hopping_constants = Matrix{Vector{Float64}}(undef, 1, num_nodes)
for ci in CartesianIndices(hopping_constants)
    (species, site) = Tuple(ci)
    hopping_constants[species, site] = zeros(outdegree(grid, site))
    for (n, nb) in enumerate(neighbors(grid, site))
        if nb < site
            hopping_constants[species, site][n] = 1.0
        end
    end
end
starting_state = 25*ones(Int, length(u0), num_nodes)
tspan = (0.0, 10.0)
prob = DiscreteProblem(starting_state,tspan, [])

jp=JumpProblem(prob, NSM(), majumps, hopping_constants=hopping_constants, spatial_system = grid, save_positions=(false,false))
sol = solve(jp, SSAStepper())

@test sol.u[end][1,1] == sum(sol.u[end])


__precompile__()

module MultivariateTests

using Compat
using Compat.LinearAlgebra
using StatsBase
using Distributions
using HypothesisTests

using HypothesisTests: HypothesisTest

export
    genvar,
    totalvar,
    PartialCorTest,
    partialcor,
    OneSampleHotellingT2,
    EqualCovHotellingT2,
    UnequalCovHotellingT2,
    BartlettsTest

## Common utility functions

function checkdims(X::AbstractMatrix, Y::AbstractMatrix)
    nx, p = size(X)
    ny, q = size(Y)
    p == q || throw(DimensionMismatch("Inconsistent number of variables: $p, $q"))
    (nx > 0 && ny > 0) || throw(ArgumentError("Inputs must be non-empty"))
    return (p, nx, ny)
end

# Pool the given covariance matrices, trashing them both
function _poolcov!(Sx::AbstractMatrix, nxm1::Int, Sy::AbstractMatrix, nym1::Int)
    Sx .*= nxm1
    Sy .*= nym1
    Sx .+= Sy
    Sx ./= nxm1 + nym1
    return Sx
end

# Benchmarks suggest that this is faster than using pivoted Cholesky when the input
# is not positive definite
function trychol(X::AbstractMatrix)
    S = Symmetric(X)
    return isposdef(S) ? cholfact(S) : factorize(S)
end

# Helper function for computing A'B⁻¹A when B is a covariance matrix, taking advantage
# of the fact that covariance matrices are symmetric, positive (semi-)definite. We'll
# check for positive definiteness and use Cholesky, otherwise let Julia pick Bunch-
# Kaufman or LU in the case of failure.
At_Binv_A(A::AbstractArray, B::AbstractArray) = A'*(trychol(B) \ A)

include("dispersion.jl")
include("partialcor.jl")
include("hotelling.jl")
include("covariance.jl")

end # module


LazyBox(T2::Type{T}) where {T<:Real} = LazyBox{T2}(Dict())

"""
    getindex(o::LazyBox{D}, key::Int) where D

Gets nth dim of box, creates it as unit-interval if it doesn't exist
"""
function getindex(o::LazyBox{D}, key::Int) where D
  if haskey(o.intervals,key)
    o.intervals[key]
  else
    o.intervals[key] = unit(Interval{D})
  end
end

setindex!(b::LazyBox{T}, val::Interval{T}, key::Int) where T = b.intervals[key] = val
ndims(b::LazyBox) = length(keys(b.intervals))
dims(b::LazyBox) = keys(b.intervals)
hasdim(b::LazyBox, i::Int) = haskey(b.intervals,i)

## Conversion
## ==========
convert(::Type{Vector{Interval{T}}}, b::LazyBox{T}) where T = collect(values(b.intervals))
convert(::Type{Vector{Interval}}, b::LazyBox{T}) where T = collect(values(b.intervals))
convert(::Type{Vector{Interval{T}}}, b::LazyBox{T}, dims::Vector) where T = Interval{T}[b[d] for d in dims]
convert(::Type{Vector{Interval}}, b::LazyBox{T}, dims::Vector) where T = Interval{T}[b[d] for d in dims]
convert(::Type{HyperBox}, l::LazyBox) = HyperBox(convert(Vector{Interval},l))

# ## Splitting
# ## =========
mid(b::LazyBox{T}) where T =
  Dict(dim => mid(interval) for (dim,interval) in b.intervals)

function split_box(b::LazyBox{T}, split_point::Dict{Int,T}) where T
  ks = collect(keys(b.intervals))
  intervals_set = product_boxes(b,split_point)
  [LazyBox(Dict{Int,Interval{T}}(zip(ks,intervals))) for intervals in intervals_set]
end

partial_split_box(b::LazyBox{T}, split_point::Dict{Int,T}) where T =
  split_box(b,split_point)

## Rand
"""
    struct LazyRandomVector{T<:Real}

A Vector of whose values are sampled uniformly from [0,1], but are not created until accessed (hence Lazy).
"""
struct LazyRandomVector{T<:Real}
  samples::Dict{Int64,T}
end
LazyRandomVector(T1::Type{T}) where {T<:Real} = LazyRandomVector(Dict{Int64,T1}())

function getindex(o::LazyRandomVector{T}, key::Int) where T
  if haskey(o.samples,key)
    o.samples[key]
  else
    i = rand(T)
    o.samples[key] = i
    i
  end
end

function setindex!(o::LazyRandomVector{T}, val::T, key::Int) where T
  o.samples[key] = val
end

function rand(b::LazyBox{T}) where T <: Real
  l = LazyRandomVector(T)
  for (dim,interval) in b.intervals
    l[dim] = rand(interval)
  end
  l
end

## Print
## =====
string(b::LazyBox) = b.intervals
print(io::IO, b::LazyBox) = print(io, string(b))
show(io::IO, b::LazyBox) = print(io, string(b))
showcompact(io::IO, b::LazyBox) = print(io, string(b))


using VehicleCalculate,GLMakie
include("zlsb250.jl")
zs650zhudong = PrimTransmission(44,72,"曲轴")
zs650zhuzhou = DrivingTransmission(15,16,20,21,23,27,"主轴")
zs650fuzhou = DrivenTransmission(39,33,34,31,31,34,"副轴")
zs650flywheel = Flywheel(14,41,"大小飞")

zs650chuan = Transmission(zs650zhudong,zs650zhuzhou,zs650fuzhou,zs650flywheel)

zs650 = vehiclespeed("ZS650",zs650chuan)
zs650.second_ratio
zlsb250.second_ratio

fg = Figure(resolution = (600, 400))
ax1 = Axis(fg[1,2])
ax2 = Axis(fg[1,1])
x_zs650 = 1:6
y_zs650 = zs650.second_ratio
lines!(ax1,x_zs650,y_zs650,label = "zs650")
lines!(ax1,x_sl250,y_sl250,label = "sl250")

scatter!(ax1,x_zs650,y_zs650)
scatter!(ax1,x_sl250,y_sl250)
xb = [1,2,3,4]
yb = [4,4,4,4]
hist!(ax2,yb,direction=:x,bins = 4,offset = 0.1)
fg

h = ["$i" * "档" for i in ["一","二","三","四","五","六"]]
v = ["zs650","SL250"]
c = hcat(zs650.second_ratio,zlsb250.second_ratio)

tableplot(h,v,c,reso = (600,400))


using Base.Test
using QuantumOptics.sparsematrix

# SparseMatrix = quantumoptics.sparsematrix.SparseMatrix
typealias SparseMatrix SparseMatrixCSC{Complex128}

# Set up test matrices
A = rand(Complex128, 5, 5)
A_sp = sparse(A)

B = eye(Complex128, 5)
B_sp = speye(Complex128, 5)

C = rand(Complex128, 3, 3)
C[2,:] = 0
C_sp = sparse(C)

R_sp = A_sp + B_sp
R = A + B

# Test arithmetic
@test_approx_eq 0. norm(full(R_sp) - R)
@test_approx_eq 0. norm(full(Complex128(0.5,0)*A_sp) - 0.5*A)
@test_approx_eq 0. norm(full(A_sp/2) - A/2)
@test_approx_eq 0. norm(full(A_sp*B_sp) - A*B)

# Test kronecker product
@test_approx_eq 0. norm(full(kron(A_sp, C_sp)) - kron(A, C))
@test_approx_eq 0. norm(full(kron(A_sp, B_sp)) - kron(A, B))


# See https://github.com/blegat/SwitchedSystems.jl

# Example 2.8 of
# P. Parrilo and A. Jadbabaie
# Approximation of the joint spectral radius using sum of squares
# Linear Algebra and its Applications, Elsevier, 2008, 428, 2385-2402
# Inspired from the construction of:
# Ando, T. and Shih, M.-h.
# Simultaneous Contractibility.
# SIAM Journal on Matrix Analysis & Applications, 1998, 19, 487
# The JSR is √2

@testset "[PJ08] Example 2.8 with $(factory.constructor)" for factory in sdp_factories
    isscs(factory) && continue
    @polyvar x[1:2]
    A1 = [1 0; 1 0]
    A2 = [0 1; 0 -1]
    expected_ub = [√2, 1]
    function testlyap(d, γ, feasible::Bool)
        m = SOSModel(factory)
        @variable m p Poly(monomials(x, 2d))
        # p strictly positive
        q = sum(x.^(2*d))
        @constraint m p >= q
        c1 = @constraint m p(x => A1 * vec(x)) <= γ^(2*d) * p
        c2 = @constraint m p(x => A2 * vec(x)) <= γ^(2*d) * p

        JuMP.optimize!(m)

        if feasible
            @test JuMP.primal_status(m) == MOI.FeasiblePoint
        else
            @test JuMP.dual_status(m) == MOI.InfeasibilityCertificate
            μ1 = JuMP.dual(c1)
            μ2 = JuMP.dual(c2)

            # The dual constraint should work on any polynomial.
            # Let's test it with q
            lhs = dot(μ1, q(x => A1 * vec(x))) + dot(μ2, q(x => A2 * vec(x)))
            rhs = dot(μ1, q) + dot(μ2, q)
            @test 1e-6 * max(abs(lhs), abs(rhs)) + lhs >= rhs
        end
    end
    testlyap(1, √2 - 1e-1, false)
    testlyap(1, √2 + 1e-1, true)
    testlyap(2, 1 - 1e-1, false)
    testlyap(2, 1 + 1e-1, true)
end


using Gadfly, DataArrays

spy(randn((10, 10)), Coord.cartesian(fixed=true))


function mean_of_zvmean(m)
  nsets = length(m)
  lzv = isnan((m[1].means)[1, 2]) ? false : true
  qzv = isnan((m[1].means)[1, 3]) ? false : true
  cols = (lzv ? (qzv ? (1:3) : (1:2)) : (qzv ? (1:2:3) : (1:1)))
  nmethods = length(cols)
  mm = zeros(Float64, size(m[1].means, 1), nmethods)

  for i in 1:nsets
    mm += (m[i].means)[:, cols]
  end

  mm/nsets
end

function var_of_zvmean(m; means::Union(Matrix{Float64}, Nothing)=nothing)
  nsets = length(m)
  lzv = isnan((m[1].means)[1, 2]) ? false : true
  qzv = isnan((m[1].means)[1, 3]) ? false : true
  cols = (lzv ? (qzv ? (1:3) : (1:2)) : (qzv ? (1:2:3) : (1:1)))
  nmethods = length(cols)
  v = zeros(Float64, size(m[1].means, 1), nmethods)

  mm = means == nothing ? mean_of_zvmean(m) : means
  for i in 1:nsets
    v += ((m[i].means)[:, cols]-mm).^2
  end

  v/(nsets-1)
end

function vrf(v)
  npars, nmethods = size(v)
  vrfs = Array(Float64, npars, nmethods-1)

  for i in 1:nmethods-1
    vrfs[:, i] = v[:, 1]./v[:, i+1]
  end

  vrfs
end

collect_diagnostic(m; diagnostic::ASCIIString="rate") = [(m[i].diagnostics)[diagnostic] for i in 1:length(m)]


struct IsingSampler <: Sampleable{Multivariate,Discrete}
    p::Int64                 # number of vertices
    θ::Vector{Float64}       # parameter matrix
    nbs::Vector{Vector{Int64}}
    state::BitVector
    thin::Int64
end

function IsingSampler(θ::Vector{Float64}; burn::Int64=3000, thin::Int64=1000)
    m = length(θ)
    p = convert(Int, ceil((1. + sqrt(1. + 8. * m)) / 2.))

    # Find neighbors
    nbs = Vector{Vector{Int64}}(undef, p)
    for node=1:p
        nbs[node] = Vector{Int64}(undef, 0)
    end
    for node=1:p
        for nb=1:p
            if node==nb
                continue
            end
            if !iszero(θ[trimap(node, nb)])
                push!(nbs[node], nb)
            end
        end
    end

    # initial state
    state = rand(Bernoulli(), p) .== 1.
    spl = IsingSampler(p, θ, nbs, state, thin)
    for j=1:burn
        _sample_one!(spl)
    end
    spl
end

Base.length(s::IsingSampler) = s.p
Base.eltype(s::IsingSampler) = Bool

Distributions.rand(s::IsingSampler) = Distributions._rand!(s, BitVector(undef, length(s)))
Distributions.rand(s::IsingSampler, n::Int) = Distributions._rand!(s, BitMatrix(undef, length(s), n))

function Distributions._rand!(s::IsingSampler, x::BitMatrix)
    n = size(x, 2)
    for i=1:n
        for j=1:s.thin
            _sample_one!(s)
        end
        _sample_one!(s)
        x[:, i] .= s.state
    end
    x
end

function Distributions._rand!(s::IsingSampler, x::BitVector)
    for j=1:s.thin
        _sample_one!(s)
    end
    _sample_one!(s)
    x .= s.state
end

function _sample_one!(spl::IsingSampler)
    for node=1:length(spl)
        _sample_one_coordinate!(spl, node)
    end
    spl
end

function _sample_one_coordinate!(
    spl::IsingSampler,
    node::Int64)

    θ, nbs, y = spl.θ, spl.nbs, spl.state

    v = 0.
    for nb in spl.nbs[node]
        k = trimap(node, nb)
        v += y[nb] ? θ[k] : -θ[k]
    end
    Q = exp(2. * v)
    y[node] = rand() <= Q / (Q + 1.)
end

### graph construction

# lenC --- length of a chain
# the graph consists of p / lenC chains
# sgn +1, 0, -1 --- whether edges are positive, mixed, or negative
# lb and ub are assumed positive
function chain(p::Integer; lenC::Integer=10, lb::Float64=0.1, ub::Float64=0.3, sgn::Integer=1)
    mod(p, lenC) == 0 || throw(ArgumentError("p=$p should be divisible by lenC=$lenC"))
    m = div(p*(p-1), 2)
    θ = zeros(Float64, m)

    du = Uniform(lb, ub)
    for ic=1:div(p, lenC)
        for j=1:lenC-1
            col = (ic - 1) * lenC + j
            row = col + 1

            v = rand(du)
            if sgn == -1
                v *= -1.
            elseif sgn == 0
                v *= 2. * rand(Bernoulli()) - 1.
            end
            θ[trimap(row, col)] = v
        end
    end
    θ
end

# k --- # of children
# the graph is a k-ary tree on p nodes
# sgn +1, 0, -1 --- whether edges are positive, mixed, or negative
# lb and ub are assumed positive
function kary_tree(p::Integer; k::Integer=3, lb::Float64=0.1, ub::Float64=0.3, sgn::Integer=1)
    k > 1 || throw(ArgumentError("k=$k should be at least 2"))
    m = div(p * (p - 1), 2)
    θ = zeros(Float64, m)

    du = Uniform(lb, ub)

    depth = ceil(Int64, log(k, (k - 1) * p + 1) -1)

    youngest = cumsum(k.^(0:depth))
    youngest[end] = p

    children = (youngest[1] + 1):youngest[2]
    for j in children
        v = rand(du)
        if sgn == -1
            v *= -1.
        elseif sgn == 0
            v *= 2. * rand(Bernoulli()) - 1.
        end
        θ[trimap(1, j)] = v
    end
    for d = 2:depth
        children = (youngest[d] + 1):youngest[d + 1]
        for j in children
            i = fld(j - youngest[d] - 1, k) + youngest[d-1] + 1
            v = rand(du)
            if sgn == -1
                v *= -1.
            elseif sgn == 0
                v *= 2. * rand(Bernoulli()) - 1.
            end
            θ[trimap(i, j)] = v
        end
    end
    θ
end

function removeEdges!(θ, numChanges::Integer=4)
    ind_change = sample(findall(!iszero, θ), numChanges; replace=false)
    for i in ind_change
        θ[i] = 0.
    end
    θ
end

function addEdges!(θ, numChanges::Integer=4, lb::Float64=0.2, ub::Float64=0.4, sgn::Integer=1)
    ind_change = sample(findall(iszero, θ), numChanges; replace=false)
    du = Uniform(lb, ub)
    for i in ind_change
        v = rand(du)
        if sgn == -1
            v *= -1.
        elseif sgn == 0
            v *= 2. * rand(Bernoulli()) - 1.
        end
        θ[i] = v
    end
    θ
end


# row_major_reshape(X::AbstractArray, size...) = permutedims(reshape(X, reverse([size...])...), length(size):-1:1)

"""
    row_major_reshape(Q::AbstractArray, shapes)

Reshaping arrays in a different order.
"""
function row_major_reshape(Q::AbstractArray, shapes)
    shapes = reverse(shapes)
    perm = collect(length(shapes):-1:1)
    return permutedims(reshape(Q, shapes), perm)
end

function chop_op(O::AbstractArray, tol = 1e-8)
    tmp_r = (abs.(real.(O)) .> tol) .* real.(O)
    tmp_i = (abs.(imag.(O)) .> tol) .* imag.(O)
    return tmp_r + 1im .* tmp_i
end

"""
    gaussian(x, mu, sig)

Gaussian function.
"""
function gaussian(x, mu, sig)
    return exp.(- 0.5 * (x .- mu).^2 / sig^2)
end

function gaussian_derivative(x, mu, sig)
    return - (x .- mu) ./ sig^2 .* exp.(- 0.5 * (x .- mu).^2 / sig^2)
end

function trunc_op(op::AbstractArray, states)
    N_trunc = size(states)[1]
    # qstates = [qtp.Qobj(states[i], dims = [[N_s, N_s], [1, 1]]) for i in range(N_trunc)]
    res = spzeros(N_trunc, N_trunc)
    for i in range(1, N_trunc, step = 1)
        for j in range(1, N_trunc, step = 1)
            res += (adjoint(states[i, :]) * op * states[j, :]) * projection(N_trunc, i, j, true)
        end
    end
    return chop_op(res)
end

function eigensystem(A::AbstractArray; k = 6, v0 = nothing, sigma = nothing, krylovdim = 30)
    is_A_hermitian = ishermitian(A)

    if v0 === nothing
        v0 = rand(eltype(A), size(A, 1))
    end
    v0 /= norm(v0)

    if sigma === nothing
        vals, vecs, info = eigsolve(A, v0, k, ishermitian = is_A_hermitian)
    else
        # fac = factorize(A - sigma * I)
        # vals, vecs, info = eigsolve(x -> fac \ x, v0, k, ishermitian = is_A_hermitian)
        # vals = (1 .+ sigma * vals) ./ vals

        A_s = A - sigma * I

        P_cpu = ilu(A_s, τ = 0.01)
        vals, vecs, info = eigsolve(x -> cg(A_s, x; Pl = P_cpu, maxiter = 500), v0, k, ishermitian = true, krylovdim = krylovdim)
        vals = (1 .+ sigma * vals) ./ vals
    end
    if is_A_hermitian
        vals = real.(vals)
    end
    vals = vals[1:k]
    vecs = hcat(vecs...)[:, 1:k]
    return vals, vecs
end

function ptrace(Q::AbstractArray, sel, rd)
    nd = length(rd)
    dkeep = rd[sel]
    qtrace = filter(e->e∉sel,1:nd)
    dtrace = rd[qtrace]

    vmat = row_major_reshape(Q, (prod(rd), 1))
    vmat = row_major_reshape(vmat, rd)
    topermute = []
    append!(topermute, sel)
    append!(topermute, qtrace)
    vmat = permutedims(vmat, topermute)
    vmat = permutedims(vmat, nd:-1:1)
    vmat = row_major_reshape(vmat, (prod(dkeep), prod(dtrace)))
    return vmat * adjoint(vmat)
end

function entropy_vn(rho::AbstractArray, base = 0, tol = 1e-15)
    vals = eigvals(rho)
    indexes = abs.(vals) .> tol
    if 1 ∈ indexes
        nzvals = vals[indexes]
    #     nzvals = vals[vals .!= 0]
        if base != 0
            logvals = log.(base, Complex.(nzvals))
        else
            logvals = log.(Complex.(nzvals))
        end
        return -real(sum(nzvals .* logvals))
    else
        return 0
    end
end

function entanglement(psi, subspaces, size)
    norm = sqrt(sum(abs.(psi).^2))
    ψ = deepcopy(psi)
    if norm > 1e-5
        ψ = ψ ./ norm
    else
        return 0
    end
    rho_rabi = ptrace(ψ, subspaces, size)
    entropy = entropy_vn(rho_rabi)
    return (entropy > 0) * entropy
end


#=
Absolute:
- Julia version: 1.4.2
- Author: akhil
- Date: 2020-05-28
=#
"""
Evaluate Absolute Value of a number

# Examples
```jldoctest
julia> Maths.abs_value(5)
5
```
"""
function abs_value(x)
    return x>0 ? x : -x
end


function partsort!(A, b)
    A1 = []
    A2 = []
    A3 = []
    for x in A
        if x < b
            push!(A1, x)
        elseif x == b
            push!(A2, x)
        else
            push!(A3, x)
        end
    end
    i = length(A1)
    j = length(A1) + length(A2)
    A = append!(A1, A2)
    A = append!(A, A3)
    return A, i, j
end

function partsort2!(A, b)
    A1 = Vector{eltype(A)}(undef, length(A))
    A2 = Vector{eltype(A)}(undef, length(A))
    A3 = Vector{eltype(A)}(undef, length(A))
    i, j, k = 1, 1, 1
    for x in A
        if x < b
            A1[i] = x
            i+=1
        elseif x == b
            A2[j] = x
            j+=1
        else
            A3[k] = x
            k += 1
        end
    end
    i -= 1
    j -= 1
    k -= 1
    for count in 1:i
        A[count] = A1[count]
    end
    for count in i+1:i+j
        A[count] = A2[count - i]
    end
    for count in i+j+1:i+j+k
        A[count] = A3[count - i - j]
    end
    return A, i, i+j
end

function quicksort!(a)
    if length(a) < 2
        return a
    end
    a, i, j = partsort!(a, a[(1+length(a)) ÷ 2])
    a[1:i] = quicksort!(a[1:i])
    a[j+1 : end] = quicksort!(a[j+1:end])
    return a
end

function quicksort_!(a)
    if length(a) < 2
        return a
    end
    a, i, j = partsort2!(a, a[(1+length(a)) ÷ 2])
    a[1:i] = quicksort_!(a[1:i])
    a[j+1 : end] = quicksort_!(a[j+1:end])
    return a
end


using JGraph
using Test
using Statistics

@testset "JGraph.core" begin
    g = Graph();
    @test ne(g) == 0
    @test nv(g) == 0

    # add vertices
    g = Graph();
    add_vertex!(g, 10)
    @test ne(g) == 0
    @test nv(g) == 10

    # add a edge
    g = Graph();
    add_edge!(g, 3, 4);
    @test ne(g) == 1
    @test nv(g) == 4
    @test has_edge(g, 3, 4) == true
    @test has_edge(g, 4, 3) == true

    # test if the edge is already there
    @test add_edge!(g, 3, 4) == false

    # test the handshake lemma
    @test sum(degree(g)) == 2*ne(g)

    # test the neighbors function
    g = Graph();
    @test add_edge!(g, 1, 2) == true
    @test add_edge!(g, 2, 3) == true
    @test add_edge!(g, 3, 1) == true
    @test neighbors(g, 1) == [2, 3]

    # test itertor
    g = Graph();
    add_vertex!(g, 10);
    @test collect(vertices(g)) == collect(1:10)

    # test remove edge
    g = Graph()
    add_edge!(g, 1, 2)
    add_edge!(g, 2, 3)
    add_edge!(g, 3, 1)

    remove_edge!(g, 3, 1)

    @test !has_edge(g, 3, 1)
    @test !has_edge(g, 1, 3)

    e = Edge((1, 2))
    remove_edge!(g, e)

    @test !has_edge(g, 2, 1)
    @test !has_edge(g, 1, 2)

    #  test remove all_edges incident to a node
    g = Graph()
    add_edge!(g, 1, 2)
    add_edge!(g, 2, 3)
    add_edge!(g, 3, 1)

    remove_vertex!(g, 2)
    @test !has_edge(g, 2, 1)
    @test !has_edge(g, 1, 2)
    @test !has_edge(g, 2, 3)
    @test !has_edge(g, 3, 2)
end

@testset "JGraph.edge" begin
    e1 = Edge((1, 2))
    e2 = Edge(Pair(2, 1))
    e3 = Edge((2, 3))
    e4 = Edge((3, 2))
    e5 = Edge((1, 2))
    results = [e4, e3, e2, e1];

    @test src(e1) == 1
    @test dst(e1) == 2

    @test src(e2) == 2
    @test dst(e2) == 1
    @test e1 == e5

    ## test iterator
    g = Graph();
    add_edge!(g, 1, 2);
    add_edge!(g, 2, 3);

    @test length(edges(g)) == 2

    for e in edges(g)
        @test e == pop!(results)
    end
end

@testset "JGraph.generator" begin
    g = erdos_renyi(10, 0.2)
    @test nv(g) == 10

    # test the handshake lemma
    @test sum(degree(g)) == 2*ne(g)

    # test the mean degree
    n = 10000;
    p = 0.0003;
    kmean = (n-1) * p
    g = erdos_renyi(n, p);
    @test mean(degree(g)) ≈ kmean atol=1e-1

    g1 = erdos_renyi(n, p, seed=5);
    g2 = erdos_renyi(n, p, seed=5);
    @test nv(g1) == nv(g2)
    @test ne(g1) == ne(g2)
end

@testset "JGraph.traversal" begin
    g = Graph();
    @test add_edge!(g, 1, 2) == true
    @test add_edge!(g, 2, 3) == true
    @test add_edge!(g, 3, 1) == true
    @test bfs(g, 1) == Set([2, 1, 3])

    g = Graph();
    add_vertex!(g, 6);
    add_edge!(g, 1, 2);
    add_edge!(g, 2, 3);
    add_edge!(g, 1, 3);

    add_edge!(g, 4, 5);
    add_edge!(g, 5, 6);
    add_edge!(g, 4, 6);

    cc = connected_components(g);
    for elem in cc
        @test length(elem) == 3
    end

    @test bfs(g, 1) == Set([2, 3, 1])
    @test bfs(g, 4) == Set([4, 5, 6])
end


abstract type AbstractRegression end

struct LocalLinear <: AbstractRegression

    k::Int64
    ivar::Vector{Int64}
    ivar_neighboor::Vector{Int64}
    res::Array{Float64,2}
    beta::Array{Float64,2}
    Xm::Array{Float64,2}
    Xr::Array{Float64,2}
    Cxx::Array{Float64,2}
    Cxx2::Array{Float64,2}
    Cxy::Array{Float64,2}
    pred::Array{Float64,2}
    X0r::Array{Float64,2}

    function LocalLinear(k, ivar, ivar_neighboor)

        #X = zeros(Float64,(nvn, k))
        #Y = zeros(Float64,(nv, k))
        #w = zeros(Float64, k)

        nv = length(ivar)
        nvn = length(ivar_neighboor)

        res = zeros(Float64, (nv, k))
        beta = zeros(Float64, (nv, nvn + 1))
        Xm = zeros(Float64, (nvn, 1))
        Xr = ones(Float64, (nvn + 1, k))
        Cxx = zeros(Float64, (nvn + 1, nvn + 1))
        Cxx2 = zeros(Float64, (nvn + 1, nvn + 1))
        Cxy = zeros(Float64, (nv, nvn + 1))
        pred = zeros(Float64, (nv, k))
        X0r = ones(Float64, (nvn + 1, 1))

        new(k, ivar, ivar_neighboor, res, beta, Xm, Xr, Cxx, Cxx2, Cxy, pred, X0r)

    end

end

function compute(ll::LocalLinear, x, xf_tmp, xf_mean, ip, X, Y, w)

    ivar = ll.ivar
    ivar_neighboor = ll.ivar_neighboor

    ll.Xm .= sum(X .* w', dims = 2)
    ll.Xr[2:end, :] .= X .- ll.Xm
    mul!(ll.Cxx, (ll.Xr .* w'), ll.Xr')
    mul!(ll.Cxx2, (ll.Xr .* w' .^ 2), ll.Xr')
    mul!(ll.Cxy, (Y .* w'), ll.Xr')

    ll.Cxx .= pinv(ll.Cxx, rtol = 0.01)
    ll.Cxx2 .= ll.Cxx2 * ll.Cxx
    # regression on principal components
    mul!(ll.beta, ll.Cxy, ll.Cxx)
    ll.X0r[2:end, :] .= x[ivar_neighboor, ip] .- ll.Xm
    # weighted mean
    xf_mean[ivar, ip] = ll.beta * ll.X0r
    mul!(ll.pred, ll.beta, ll.Xr)
    Y .-= ll.pred
    xf_tmp[ivar, :] .= xf_mean[ivar, ip] .+ Y
    # weigthed covariance
    cov_xf = (Y * (w .* Y')) ./ (1 .- tr(ll.Cxx2))
    cov_xf .= Symmetric(cov_xf .* (1 .+ tr(ll.Cxx2 * ll.X0r * ll.X0r' * ll.Cxx)))

    return cov_xf

end


"""
    fit_mvnormals(θs, ∇logpθs; history_length=5)

Fit a multivariate-normal distribution to each point on the trajectory `θs`.

Given `θs` with gradients `∇logpθs`, construct LBFGS inverse Hessian approximations with
the provided `history_length`. The inverse Hessians approximate a covariance. The
covariances and corresponding means that define multivariate normal approximations per
point are returned.
"""
function fit_mvnormals(θs, ∇logpθs; history_length, kwargs...)
    if length(θs) < history_length
        error("Insufficient history in trajectory.")
    end
    Σs = lbfgs_inverse_hessians(θs, ∇logpθs; kwargs...)
    l = length(Σs)
    μs = @views muladd.(Σs, ∇logpθs[1:l], θs[1:l])
    return Distributions.MvNormal.(μs, Σs)
end

# faster than computing `logpdf` and `rand` independently
function rand_and_logpdf(rng, dist::Distributions.MvNormal, ndraws)
    μ = dist.μ
    Σ = dist.Σ
    N = length(μ)

    # draw points
    u = Random.randn!(rng, similar(μ, N, ndraws))
    unormsq = vec(sum(abs2, u; dims=1))
    x = PDMats.unwhiten!(u, Σ, u)
    x .+= μ

    # compute log density at each point
    logpx = (muladd(N, log2π, logdet(Σ)) .+ unormsq) ./ -2

    return x, logpx
end


using Wavelets

wf = waveletfilter(WT.db6)
N = 1024*32;
x0 = rand(N);
L = int(log2(N)-2)
tn = 1000

println("iwt (N=",N,"), ", L, " levels")
f(x0, L, wf) = for i = 1:tn; idwt(x0, wf, L); end
f(x0, L, wf);
@time f(x0, L, wf);


norm(A::Array; dims::Int) = sqrt.(sum(x->x^2, A, dims=dims))

function compute_r_nearest_image!(mdsim::MDSimulation)
	P = mdsim.params
	x = mdsim.state.x .% P.L

	mdsim.state.r_nearest_image = reshape(x', (1, P.d, P.N)) .- reshape(x, (P.N, P.d, 1))
	mdsim.state.r_nearest_image[mdsim.state.r_nearest_image .> P.L / 2]  .-= P.L
	mdsim.state.r_nearest_image[mdsim.state.r_nearest_image .< -P.L / 2] .+= P.L

	mdsim.state.r_nearest_image_magnitude = norm(mdsim.state.r_nearest_image, dims=2)
	mdsim.state.r_nearest_image_magnitude[P.diag_indices] .= Inf
end

function compute_row_vector_magnitudes(r)
	if size(r, 2) == 2
		return @. sqrt(r[:, 1]^2 + r[:, 2]^2)
	elseif size(r, 2) == 3
		return @. sqrt(r[:, 1]^2 + r[:, 2]^2 + r[:, 3]^2)
	end
end

function compute_energies!(mdsim::MDSimulation)
	P = mdsim.params
	α = (P.σ_LJ ./ mdsim.state.r_nearest_image_magnitude).^6
	mdsim.state.E_pot = sum((4 * P.ε_LJ) .* (α.^2 .- α)) / 2
	mdsim.state.E_kin = 0.5P.m * sum(mdsim.state.v.^2)
end

function compute_accelerations!(mdsim::MDSimulation)
	P = mdsim.params
	mdsim.state.a_prev = mdsim.state.a
	compute_r_nearest_image!(mdsim)

	α = (P.σ_LJ ./ mdsim.state.r_nearest_image_magnitude).^6
	mdsim.state.a = sum(
		((48 * P.ε_LJ / P.m) * (α.^2 - 0.5α) ./ mdsim.state.r_nearest_image_magnitude.^2)
		.* mdsim.state.r_nearest_image,
		dims=1
	)[1, :, :]'
end


struct LowLevelJointGains{T}
    k_q_p::T # corresponds to kp_position in drcsim API
    k_q_i::T # corresponds to ki_position in drcsim API
    k_qd_p::T # corresponds to kp_velocity in drcsim API
    k_f_p::T
    ff_qd::T # maps to kd_position in drcsim API (there isnt an equivalent gain in the bdi api)
    ff_qd_d::T
    ff_f_d::T
    ff_const::T
end

struct JointState{T}
    position::T
    velocity::T
    effort::T
end

function command_effort(gains::LowLevelJointGains, state::JointState, state_des::JointState, Δt::Number)
    effort = gains.ff_const
    effort += gains.k_q_p * (state_des.position - state.position)
    effort += gains.k_q_i * (state_des.position - state.position) * Δt
    effort += gains.k_qd_p * (state_des.velocity - state.velocity)
    effort += gains.k_f_p * (state_des.effort - state.effort)
    effort += gains.ff_qd * state.velocity
    effort += gains.ff_qd_d * state_des.velocity
    effort += gains.ff_f_d * state_des.effort
    effort
end


function init_sketch(n, R; rseed=0)
    # Random.seed!(rseed)
    return randn(n, R), zeros(n, R)
end

# cache should have length R
function rank_one_update!(S, Ω, v, η; cache::AbstractVector=nothing)
    @assert length(cache) == size(S, 2)
    S .= (1 - η) .* S
    if isnothing(cache)
        tmp = Ω'*v
        BLAS.ger!(η, v, tmp, S)
    else
        mul!(cache, Ω', v)
        BLAS.ger!(η, v, cache, S)
    end
end

# TODO: maybe should make this unallocating for experiments?
# Numerical stable version of eq (5.3)
#  X̂ = S*pinv(Ω'*S)*S' = U*Λ*U'
function reconstruct(Ω, S; correction=false)
    n = size(S, 1)
    σ = sqrt(n)*eps()*norm(S)
    Sσ = S + σ .* Ω
    M = Ω'*S
    M .= 0.5 .* (M .+ M')
    F = nothing

    # Some error handling -- PosDefException on cholesky for first few iters
    # TODO: figure out why this occurs
    try
        F = cholesky(M)
    catch e
        !isa(e, PosDefException) && error("Sketch reconstruction error")
        M[diagind(M)] .+= minimum(eigvals(M)) .+ σ
        F = cholesky(M)
    end
    U, Σ, _ = svd(Sσ / F.U)
    Λ = Diagonal(max.(0, Σ.^2 .- σ))

    # Implements trace correction (Remark 6.1)
    # Note: correction doubles the error at worse; usually works better in practice
    if correction
        R = size(S, 2)
        correction_factor = (1.0 - tr(Λ))/R
        Λ.diag .= Λ.diag .+ correction_factor
    end

    return U, Λ
end

function construct_Xhat(soln::SCGALResults)
    return soln.UT*soln.ΛT*soln.UT'
end

# Avoids overhead of constructing Xhat fully
# cache is a n x R matrix
function compute_objective(C, U, Λ; cache=nothing)
    n = size(C, 1)
    if isnothing(cache)
        cache = C*U
    else
        mul!(cache, C, U)
    end

    obj_val = 0
    @views for i in 1:size(Λ, 1)
        obj_val += Λ.diag[i] * dot(U[:,i], cache[:,i])
    end

    return obj_val
end

# Avoids overhead of constructing Xhat fully
# cache is a n x R matrix
function compute_primal_infeas_mc(U, Λ; cache=nothing)
    n, R = size(U)
    @assert size(cache) == size(U)
    if isnothing(cache)
        cache = U*Λ
    else
        mul!(cache, U, Λ)
    end

    infeas_val = 0
    @views for i in 1:n
        infeas_val += (dot(cache[i,:], U[i,:]) - 1/n)^2
    end
    infeas_val = sqrt(infeas_val)
    return infeas_val / (1 + sqrt(1/n))
end

# MAXCUT objective with GW randomized rounding
# C = -0.25*(Diagonal(G*ones(n)) - G)
function max_cut_val(C, U, Λ)
    n, R = size(U)
    hp = randn(R)
    x = sign.(U*sqrt.(Λ)*hp)
    return -dot(x, C*x)
end


using Distributions, Gadfly, Random
set_default_plot_size(6.6inch, 3.3inch)

Random.seed!(123)
d = rand(MvNormal([2, 2],[1.0 0.7; 0.7 1.0]), 50)'

pa= plot(x=d[:,1], y=d[:,2], Geom.point, layer(Stat.ellipse, Geom.polygon(preserve_order=true)))
pb= plot(x=d[:,1], y=d[:,2], Geom.point, Geom.ellipse(levels=[0.95, 0.99]))
hstack(pa,pb)


"""
Initialize harmonic external potential for electronic system.
"""
function init_pot_harm_3d( LF::LF3dGrid, ω, center )
  Npoints = LF.Nx * LF.Ny * LF.Nz
  Vpot = zeros( Float64,Npoints )
  for ip = 1:Npoints
    Vpot[ip] = 0.5 * ω^2 * norm(LF.lingrid[:,ip] - center[:])^2
  end
  return Vpot
end

"""
Initialize ion-electron potential according to pure Coulombic
potential.
This potential is singular at r = 0.
"""
function init_pot_Hcoul( LF::LF3dGrid, center )
  Npoints = LF.Nx * LF.Ny * LF.Nz
  Vpot = zeros( Float64,Npoints )
  for ip = 1:Npoints
    r = norm( LF.lingrid[:,ip] - center[:] )
    Vpot[ip] = -1.0/r
  end
  return Vpot
end

"""
Initialize ion-electron potential for hydrogen atom using
analytic form and parameters used in Gygi's paper.
(need reference)
"""
function init_pot_Hps( LF::LF3dGrid, center )
  Npoints = LF.Nx * LF.Ny * LF.Nz
  Vpot = zeros( Float64,Npoints )

  const rc1 = 0.25
  const rc2 = 0.284
  const aa  = -1.9287
  const bb  = 0.3374

  for ip=1:Npoints
    r = norm( LF.lingrid[:,ip] - center[:] )
    Vpot[ip] = -1.0/r * erf( r/rc1 ) + (aa + bb*r^2)*exp(-(r/rc2)^2)
  end
  return Vpot
end

"""
Initialize local ion-electron potential for hydrogen atom according
to HGH pseudopotential parameters.

This function is used only for testing purpose.

This function should not be used to initialize ion-electron potential for
periodic system.
"""
function init_pot_Hps_HGH( LF::LF3dGrid, center )
  Npoints = LF.Nx * LF.Ny * LF.Nz
  Vpot = zeros( Float64,Npoints )

  const Zion = 1
  const rloc = 0.2
  const C1 = -4.0663326
  const C2 = 0.6678322

  # TODO Add journal reference
  for ip = 1:Npoints
    r = norm( LF.lingrid[:,ip] - center[:] )
    if r < 1e-7
      println("WARNING: small r:", r);
      println("Using limiting value:")
      Vpot[ip] = -2*Zion/(sqrt(2*pi)*rloc) + C1
    else
      rrloc = r/rloc
      Vpot[ip] = -Zion/r * erf( r/(sqrt(2.0)*rloc) ) +
               (C1 + C2*rrloc^2)*exp(-0.5*(rrloc)^2)
    end
  end
  return Vpot
end