"""
    autocor(x,lag)

Returns autocorrelation function for a distance of lag, i.e. between values of
```x[t]``` and ```x[t+lag]```. The autocorrelation function is normalized such
that ```autocor(x,0) = 1```. This is equvivalent to the quantity
``\\Gamma_X(t)`` of equation (4.61) in [Gattringer/Lang](@cite).
"""
function autocor(x, lag)
    # (wasteful in terms of allocations but clear)
    z = x .- mean(x)
    a = sum(z[1+lag:end].*z[1:end-lag])/sum(z.*z)
    return a
end
"""
    autocor(x)

Returns autocorrelation function for a set of lags from ``0`` up to the closest
integer to ``10 \\text{log}_{10}(l)``. Here ``l`` is the number of elements in
```x```.
"""
function autocor(x)
    lx   = length(x)
    lags = collect(0:min(lx-1, round(Int,10*log10(lx))))
    a = zeros(eltype(x),length(lags))
    for i in 1:length(lags)
        a[i] = autocor(x, lags[i])
    end
    return a
end
"""
    autotimeexp(x)

Returns ``\\max(1,\\tau)`` where ``\\tau`` is the exponential autocorrelation
time of a series of measurements ``x``. This is obtained by fitting the
autocorrelation function to an exponential function of the form
``A \\exp(\\frac{t}{\\tau})``.
"""
function autotimeexp(O)
    a = autocor(O)
    @. modelτ(x,p) = abs.(p[2])*exp(-x/p[1])
    x = collect(1:length(a))
    c = curve_fit(modelτ, x, a, ones(2))
    τ = c.param[1]
    return max(one(τ),τ)
end


#=

  https://edu.swi-prolog.org/mod/assign/view.php?id=249
  """
  Medical diagnosis

  Develop an expert system for medical diagnosis.

  Consider three diseases: flu, gastroenteritis and bronchitis.

  A priori, flu has probability 0.3, gastroenteritis 0 0.2, and bronchitis 0.25.

  If you have the flu, you may have the following symptons, associated with their
  probabilities (symptoms are not mutually exclusive):

    fever, 0.8
    cough 0.6
    sore throat 0.5
    headaches 0.4
    aches 0.7

  If you have gastroenteritis, you may have the following symptons, associated with their
  probabilities (symptoms are not mutually exclusive):

    diarrhea 0.8
    aches 0.7
    nausea 0.4
    fatigue 0.3

  If you have bronchitis, you may have the following symptons, associated with their probabilities
  (symptoms are not mutually exclusive):

    cough 0.8
    fatigue 0.7
    fever 0.3

  Compute the probability of each disease given that the patient has the symptoms fever and aches.

  Do the same supposing the patient adds that he also experiences fatigue.
  """

  Cf ~/cplint/course_medical_diagnosis.pl
     ~/blog/medical_diagnosis.blog
     ~/psi/medical_diagnosis.psi
     ~/webppl/medical_diagnosis.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function medical_diagnosis()
    #=
      Consider three diseases: flu, gastroenteritis and bronchitis.

      A priori, flu has probability 0.3, gastroenteritis 0.2, and bronchitis 0.25.
    =#

    flu ~ flip(0.3)
    gastroenteritis ~ flip(0.2)
    bronchitis ~ flip(0.25)

    #=
      If you have the flu, you may have the following symptons, associated with their
      probabilities (symptoms are not mutually exclusive):

      fever, 0.8
      cough 0.6
      sore throat 0.5
      headaches 0.4
    aches 0.7
    =#

    fever ~ flu ? flip(0.8) : ( bronchitis ? flip(0.3) : flip(0.0))
    cough ~ flu ? flip(0.6) : ( bronchitis ? flip(0.8) : flip(0.0))
    sore_throat ~ flu ? flip(0.5) : flip(0.0)
    headaches ~ flu ? flip(0.4) : flip(0.0)
    aches ~ flu ? flip(0.7) : (gastroenteritis ? flip(0.7) : flip(0.0))

    #=
      If you have gastroenteritis, you may have the following symptons, associated with their
      probabilities (symptoms are not mutually exclusive):

      diarrhea 0.8
      aches 0.7
      nausea 0.4
      fatigue 0.3
    =#

    diarrhea ~ gastroenteritis ? flip(0.8) : flip(0.0)
    nausea  ~ gastroenteritis ? flip(0.4) : flip(0.0)
    fatigue ~ gastroenteritis ? flip(0.3) : ( bronchitis ? flip(0.7) : flip(0.0))

    #=
      If you have bronchitis, you may have the following symptons, associated with their probabilities
      (symptoms are not mutually exclusive):

      cough 0.8
      fatigue 0.7
      fever 0.3
    =#
    cough ~ bronchitis ? flip(0.8) : flip(0.0)
    fatigue ~ bronchitis ? flip(0.7) : flip(0.0)
    fever ~ bronchitis ? flip(0.3) : flip(0.0)

    # The observations
    true ~ Dirac(fever)
    true ~ Dirac(aches)

end

model = medical_diagnosis()
num_chns = 4

# HH has problem with this!
# chns = sample(model, MH(), MCMCThreads(), 10_000, num_chns)
# chns = sample(model, MH(), 40_000)

# chns = sample(model, PG(20), MCMCThreads(), 1000, num_chns)
chns = sample(model, SMC(1000), MCMCThreads(), 10_000, num_chns)

# Note: IS don't generate chns the same way as MH, PG, and SMC!
# chns = sample(model, IS(), MCMCThreads(), 1000, num_chns)

display(chns)

show_var_dist_pct(chns,:flu)
show_var_dist_pct(chns,:gastroenteritis)
show_var_dist_pct(chns,:bronchitis)


######################################################################
# test counter and mcmc with a standard normal
######################################################################

"Normal(0,1) density. Keeps track of simulations. For unit tests."
mutable struct StandardNormalModel
    counter::Int
end

logdensity(::StandardNormalModel, θ) = -0.5*θ[1]^2

simulate!(p::StandardNormalModel) = p.counter += 1

@testset "standard normal" begin
    p = StandardNormalModel(0)
    chain, a = mcmc(RWMH(diagm([0.5])), p, [0.0], 1000)
    @test p.counter == 1000
    xs = first.(chain)
    @test mean(xs) ≈ 0 atol = 0.1
    @test std(xs) ≈ 1 atol = 0.2
end

######################################################################
# test mcmc with a transformed Γ
######################################################################

@testset "log transformed Γ" begin
    dist = Gamma(5.0, 1.0)
    # sample log(x), where x ∼ distribution, apply the Jacobian correction
    p = DensityWrapper(θ -> logpdf(dist, exp(θ[1])) + θ[1])

    chain, a = mcmc(RWMH(diagm([2.0])), p, [log(5.0)], 10000)
    xs = exp.(first.(chain))

    @test mean(xs) ≈ mean(dist) atol = 0.2
    test_cdf(dist, xs; ps = 0.1:0.1:0.9, atol = 0.2)
end


export
point_contained,
point_contained!

"""
point_contained(simplex::ST, point::Vector{VT}) where {T <: AbstractSimplex, VT}

Check if `point` is contained in the interior of `simplex`.

The `point` is contained within `simplex` if it can be expressed as a
convex combination of the vertices of `simplex`. To check this condition,
we iteratively check the signed volume of a modified simplex. Specifically,
at the i-th iteration step, substitute the i-th vertex of `simplex` with
`point` and take the signed volume when when one vertex is substituted.
In the next step, replace the substitated vertex by the original, then
replace the (i + 1)th vertex by `point` and repeat the process.
If any two consecutive signs are different, the point is not contained.

When the point is very close to the edge of a simplex, numerical accuracies
will play a role, so the function sometimes (very rarely) puts `point` inside
`simplex`.
"""
function point_contained(simplex::T, point::Vector{VT}) where {T <: AbstractSimplex, VT}
    dim = dimension(simplex)

    # The signs of the linear combination coefficients
    coefficient_signs = zeros(Float64, dim + 1)

    # A temporary simplex into which we substitute `point`
    temp_simplex = MutableSimplex([zeros(Float64, dim) for i = 1:(dim + 1)])

    # Using the temporary simplex, replace the the first vertex by `point`
    # and fill the remaining vertices with the original `simplex` vertices.
    temp_simplex[1] = point
    [temp_simplex[i] = simplex[i] for i = 2:(dim + 1)]

    # The signed volume of the simplex
    coefficient_signs[1] = sign(orientation(temp_simplex))

    # Iteratively replace the other vertices and see if the signed
    # volume changes. It it does, the point can not be expressed as a
    # convex combination of the vertices of `simplex`.
    for κ = 2:dim

        # Replace the i-th vertex of the temporary simplex with `point`
        temp_simplex[κ] = point

        # Replace the remaining vertices of the temporary simplex
        # with the original vertices
        remaining_vertex_idxs = setdiff(1:(dim + 1), κ)
        [temp_simplex[i] = simplex[i] for i in remaining_vertex_idxs]

        # Check the signed volume and stop if it has changed.
        coefficient_signs[κ] = sign(orientation(temp_simplex))

        if !(coefficient_signs[κ-1] == coefficient_signs[κ])
            return false
        end
    end

    # Check when replacing the last point.
    temp_simplex[end] = point
    [temp_simplex[i] = simplex[i] for i = 1:dim]

    coefficient_signs[end] = sign(orientation(temp_simplex))

    if !(coefficient_signs[end-1] == coefficient_signs[end])
        return false
    else
        return true
    end
end

"""
    point_contained!(temp_simplex::MutableSimplex, simplex::ST, point::Vector{VT}) where {T <: AbstractSimplex, VT}

Check if `point` is contained in the interior of `simplex`.

The `point` is contained within `simplex` if it can be expressed as a
convex combination of the vertices of `simplex`. To check this condition,
we iteratively check the signed volume of a modified simplex. Specifically,
at the i-th iteration step, substitute the i-th vertex of `simplex` with
`point` and take the signed volume when when one vertex is substituted.
In the next step, replace the substitated vertex by the original, then
replace the (i + 1)th vertex by `point` and repeat the process.
If any two consecutive signs are different, the point is not contained.

When the point is very close to the edge of a simplex, numerical accuracies
will play a role, so the function sometimes (very rarely) puts `point` inside
`simplex`.
"""
function point_contained!(temp_simplex::MutableSimplex, simplex::T, point::Vector{VT}) where {T <: AbstractSimplex, VT}
    dim = dimension(simplex)

    # The signs of the linear combination coefficients
    coefficient_signs = zeros(Float64, dim + 1)

    # Using the temporary simplex, replace the the first vertex by `point`
    # and fill the remaining vertices with the original `simplex` vertices.
    temp_simplex[1] = point
    [temp_simplex[i] = simplex[i] for i = 2:(dim + 1)]

    # The signed volume of the simplex
    coefficient_signs[1] = sign(orientation(temp_simplex))

    # Iteratively replace the other vertices and see if the signed
    # volume changes. It it does, the point can not be expressed as a
    # convex combination of the vertices of `simplex`.
    for κ = 2:dim

        # Replace the i-th vertex of the temporary simplex with `point`
        temp_simplex[κ] = point

        # Replace the remaining vertices of the temporary simplex
        # with the original vertices
        remaining_vertex_idxs = setdiff(1:(dim + 1), κ)
        [temp_simplex[i] = simplex[i] for i in remaining_vertex_idxs]

        # Check the signed volume and stop if it has changed.
        coefficient_signs[κ] = sign(orientation(temp_simplex))

        if !(coefficient_signs[κ-1] == coefficient_signs[κ])
            return false
        end
    end

    # Check when replacing the last point.
    temp_simplex[end] = point
    [temp_simplex[i] = simplex[i] for i = 1:dim]

    coefficient_signs[end] = sign(orientation(temp_simplex))

    if !(coefficient_signs[end-1] == coefficient_signs[end])
        return false
    else
        return true
    end
end


using SparseArrays, DiffEqOperators, LinearAlgebra, Random,
      Test, BandedMatrices, FillArrays

Random.seed!(1234)

# These functions create analytic solutions to upwind operators. In the function name,
# the first number indicates the derivative order, the second number indicates
# the approximation order,and the last three characters "Pos" or "Neg" indicating
# the winding direction. For instance analyticOneTwoPos corresponds to the positive
# upwind second order approximation of the first derivative.

function analyticOneOnePos()
      A = zeros(5,7)
      for i in 1:5
            A[i,i+1:i+2] = [-1 1]
      end
      return A
end

function analyticOneOneNeg()
      A = zeros(5,7)
      for i in 1:5
            A[i,i:i+1] = [-1 1]
      end
      return A
end

function analyticOneTwoPos()
      A = zeros(5,7)
      for i in 1:4
            A[i,i+1:i+3] = [-3/2 2 -1/2]
      end
      A[5,5:7] = [-1/2 0 1/2]
      return A
end

function analyticOneTwoNeg()
      A = zeros(5,7)
      A[1,1:3] = [-1/2 0 1/2]
      for i in 2:5
            A[i,i-1:i+1] = [1/2 -2 3/2]
      end
      return A
end

function analyticTwoTwoPos()
      A = zeros(5,7)
      for i in 1:3
            A[i,i+1:i+4] = [2 -5 4 -1]
      end
      A[4,4:7] = [1 -2 1 0]
      A[5, 4:7] = [0 1 -2 1]
      return A
end

function analyticTwoTwoNeg()
      A = zeros(5,7)
      A[1,1:4] = [1 -2 1 0]
      A[2,1:4] = [0 1 -2 1]
      for i in 3:5
            A[i,i-2:i+1] = [-1 4 -5 2]
      end
      return A
end

function analyticTwoThreePos()
      A = zeros(7,9)
      for i in 1:4
            A[i,i+1:i+5] = [35/12 -104/12 114/12 -56/12 11/12]
      end
      A[5,5:9] = [11/12 -20/12 6/12 4/12 -1/12]
      A[6,5:9] = [-1/12 16/12 -30/12 16/12 -1/12]
      A[7,5:9] = [-1/12 4/12 6/12 -20/12 11/12]
      return A
end

function analyticTwoThreeNeg()
      A = zeros(7,9)
      A[1,1:5] = [11/12 -20/12 6/12 4/12 -1/12]
      A[2,1:5] = [-1/12 16/12 -30/12 16/12 -1/12]
      A[3,1:5] = [-1/12 4/12 6/12 -20/12 11/12]
      for i in 4:7
            A[i,i-3:i+1] = [11/12 -56/12 114/12 -104/12 35/12]
      end
      return A
end

# analytically derived upwind operators for the irregular grid:
# [0.0, 0.08, 0.1, 0.15, 0.19, 0.26, 0.29], and spacing:
# dx = [0.08, 0.02, 0.05, 0.04, 0.07, 0.03]

function analyticOneOnePosIrr()
      A = zeros(5,7)
      A[1,2:3] = [-50, 50]
      A[2,3:4] = [-20, 20]
      A[3,4:5] = [-25, 25]
      A[4,5:6] = [-100/7, 100/7]
      A[5,6:7] = [-100/3, 100/3]
      return A
end

function analyticOneOneNegIrr()
      A = zeros(5,7)
      A[1,1:2] = [-25/2, 25/2]
      A[2,2:3] = [-50, 50]
      A[3,3:4] = [-20, 20]
      A[4,4:5] = [-25, 25]
      A[5,5:6] = [-100/7, 100/7]
      return A
end

function analyticOneTwoPosIrr()
      A = zeros(5,7)
      A[1,2:4] = [-450/7, 490/7, -40/7]
      A[2,3:5] = [-280/9, 405/9, -125/9]
      A[3,4:6] = [-2625/77, 3025/77, -400/77]
      A[4,5:7] = [-510/21, 1000/21, -490/21]
      A[5,5:7] = [-90/21, -400/21, 490/21]
      return A
end

function analyticOneTwoNegIrr()
      A = zeros(5,7)
      A[1,1:3] = [-5/2, -75/2, 80/2]
      A[2,1:3] = [5/2, -125/2, 120/2]
      A[3, 2:4] = [250/7, -490/7, 240/7]
      A[4, 3:5] = [80/9, -405/9, 325/9]
      A[5, 4:6] = [1225/77, -3025/77, 1800/77]
      return A
end

function analyticTwoTwoPosIrr()
      A = zeros(5,7)
      A[1,2:5] = [200000/77, -308000/77, 143000/77, -35000/77]
      A[2,3:6] = [27500/33, -75000/33, 55000/33, -7500/33]
      A[3,4:7] = [72500/77, -137500/77, 120000/77, -55000/77]
      A[4,4:7] = [42500/77, -71500/77, 40000/77, -11000/77]
      A[5,4:7] = [-10000/77, 44000/77, -100000/77, 66000/77]
      return A
end

function analyticTwoTwoNegIrr()
      A = zeros(5,7)
      A[1,1:4] = [1050/7, -1250/7, -1400/7, 1600/7]
      A[2,1:4] = [350/7, 6250/7, -9800/7, 3200/7]
      A[3,1:4] = [-1400/7, 25000/7, -30800/7, 7200/7]
      A[4,2:5] = [-390000/231, 770000/231, -660000/231, 280000/231]
      A[5,3:6] = [-38500/77, 161000/77, -165000/77, 42500/77]
      return A
end

@testset "Test: Derivative Order = 1, Approx Order = 1, Winding = Positive" begin
      N = 5
      L = UpwindDifference(1,1, 1.0, N, 1.0)
      analyticL = analyticOneOnePos()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) == analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Derivative Order = 1, Approx Order = 1, Winding = Negative" begin
      N = 5
      L = UpwindDifference(1,1, 1.0, N, -1.0)
      analyticL = -1*analyticOneOneNeg()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) == analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Derivative Order = 1, Approx Order = 2, Winding = Positive" begin
      N = 5
      L = UpwindDifference(1,2, 1.0, N, 1.0)
      analyticL = analyticOneTwoPos()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) == analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Derivative Order = 1, Approx Order = 2, Winding = Negative" begin
      N = 5
      L = UpwindDifference(1,2, 1.0, N, -1.0)
      analyticL = -1*analyticOneTwoNeg()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) == analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Derivative Order = 2, Approx Order = 2, Winding = Positive" begin
      N = 5
      L = UpwindDifference(2,2, 1.0, N, 1.0)
      analyticL = analyticTwoTwoPos()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) == analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Derivative Order = 2, Approx Order = 2, Winding = Negative" begin
      N = 5
      L = UpwindDifference(2,2, 1.0, N, -1.0)
      analyticL = -1*analyticTwoTwoNeg()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) == analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

# Here the operators are too big for five grid points, so weird corner cases must be accounted for
# We should be able to assume that users will not have cases like this.
@testset "Test: Derivative Order = 2, Approx Order = 3, Winding = Positive" begin
      N = 7
      L = UpwindDifference(2,3, 1.0, N, 1.0)
      analyticL = analyticTwoThreePos()
      x = rand(9)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Derivative Order = 2, Approx Order = 3, Winding = Negative" begin
      N = 7
      L = UpwindDifference(2,3, 1.0, N, -1.0)
      analyticL = -1*analyticTwoThreeNeg()
      x = rand(9)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Derivative Order = 1, Approx Order = 1, Winding = Positive, Grid = Irregular" begin
      N = 5
      # constructor throws an error at the moment
      L = UpwindDifference(1,1, [0.08, 0.02, 0.05, 0.04, 0.07, 0.03], N, 1.0)
      analyticL = analyticOneOnePosIrr()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv
end

@testset "Test: Derivative Order = 1, Approx Order = 1, Winding = Negative, Grid = Irregular" begin
      N = 5
      # constructor throws an error at the moment
      L = UpwindDifference(1,1, [0.08, 0.02, 0.05, 0.04, 0.07, 0.03], N, -1.0)
      analyticL = -1*analyticOneOneNegIrr()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv
end

@testset "Test: Derivative Order = 1, Approx Order = 2, Winding = Positive, Grid = Irregular" begin
      N = 5
      # constructor throws an error at the moment
      L = UpwindDifference(1,2, [0.08, 0.02, 0.05, 0.04, 0.07, 0.03], N, 1.0)
      analyticL = analyticOneTwoPosIrr()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Derivative Order = 1, Approx Order = 2, Winding = Negative, Grid = Irregular" begin
      N = 5
      # constructor throws an error at the moment
      L = UpwindDifference(1,2, [0.08, 0.02, 0.05, 0.04, 0.07, 0.03], N, -1.0)
      analyticL = -1*analyticOneTwoNegIrr()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Derivative Order = 2, Approx Order = 2, Winding = Positive, Grid = Irregular" begin
      N = 5
      # constructor throws an error at the moment
      L = UpwindDifference(2,2, [0.08, 0.02, 0.05, 0.04, 0.07, 0.03], N, 1.0)
      analyticL = analyticTwoTwoPosIrr()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Derivative Order = 2, Approx Order = 2, Winding = Negative, Grid = Irregular" begin
      N = 5
      # constructor throws an error at the moment
      L = UpwindDifference(2,2, [0.08, 0.02, 0.05, 0.04, 0.07, 0.03], N, -1.0)
      analyticL = -1*analyticTwoTwoNegIrr()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Scaling by dx and Derivative Order in Uniform Case" begin
      N = 5
      # constructor throws an error at the moment
      L = UpwindDifference(1,2, 0.1, N, 1.0)
      analyticL = 10.0*analyticOneTwoPos()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

      L = UpwindDifference(1,2, 0.1, N, -1.0)
      analyticL = -10.0*analyticOneTwoNeg()

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

      L = UpwindDifference(2,2, 0.1, N, 1.0)
      analyticL = 100.0*analyticTwoTwoPos()

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

      L = UpwindDifference(2,2, 0.1, N, -1.0)
      analyticL = -100.0*analyticTwoTwoNeg()

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Non-Trivial Coefficient Handling in Uniform Grid Case" begin
      N = 5
      # constructor throws an error at the moment
      L = UpwindDifference(2,2, 1.0, N, 4.56)
      analyticL = 4.56*analyticTwoTwoPos()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

      L = UpwindDifference(2,2, 1.0, N, -4.56)
      analyticL = -4.56*analyticTwoTwoNeg()

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: dx and Derivative Order Scaling and Non-Trivial Coefficient Handling in Uniform Grid Case" begin
      N = 5
      # constructor throws an error at the moment
      L = UpwindDifference(2,2, 0.1, N, 4.56)
      analyticL = 4.56*100.0*analyticTwoTwoPos()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

      L = UpwindDifference(2,2, 0.1, N, -4.56)
      analyticL = -4.56*100.0*analyticTwoTwoNeg()

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end

@testset "Test: Coefficient Handling in Non-Uniform Grid Case" begin
      N = 5
      # constructor throws an error at the moment
      L = UpwindDifference(2,2, [0.08, 0.02, 0.05, 0.04, 0.07, 0.03], N, 4.56)
      analyticL = 4.56*analyticTwoTwoPosIrr()
      x = rand(7)

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

      L = UpwindDifference(2,2, [0.08, 0.02, 0.05, 0.04, 0.07, 0.03], N, -4.56)
      analyticL = -4.56*analyticTwoTwoNegIrr()

      # Test that multiplication agrees with analytic multiplication
      @test L*x ≈ analyticL*x

      # Test that concretized multiplication agrees with analytic multiplication
      @test Array(L)*x ≈ analyticL*x

      # Test that matrix-free multiplication agrees with concretized multiplication
      @test L*x ≈ Array(L)*x

      # Test that concretized matrix agrees with analytic matrix
      @test Array(L) ≈ analyticL

      # Test Banded and Sparse concretizations
      @test Array(L) == BandedMatrix(L)
      @test Array(L) == SparseMatrixCSC(L)
      @test typeof(BandedMatrix(L)) <: BandedMatrix
      @test typeof(SparseMatrixCSC(L)) <: SparseMatrixCSC

      # Test convolutions against bpv
      bpv = PeriodicBC(Float64)*x[1:N]
      @test L*bpv ≈ analyticL*bpv

end


using Revise
using AdFem
using PyPlot

n = 20
mmesh = Mesh(n, n, 1/n, degree=2)
bdedge = bcedge((x1,y1,x2,y2)->((y1<1e-5) && (y2<1e-5)), mmesh)
bdnode = bcnode((x,y)->!(y<1e-5 && x>1e-5 && x<1-1e-5), mmesh)

F1 = eval_f_on_gauss_pts((x,y)->3.0, mmesh)
F2 = eval_f_on_gauss_pts((x,y)->-1.0, mmesh)
F = compute_fem_source_term(F1, F2, mmesh)

t1 = eval_f_on_boundary_edge((x,y)->-x-y, bdedge, mmesh)
t2 = eval_f_on_boundary_edge((x,y)->2y, bdedge, mmesh)
T = compute_fem_traction_term(t1, t2, bdedge, mmesh)

D = constant(diagm(0=>[1,1,0.5]))
K = compute_fem_stiffness_matrix(D, mmesh)
rhs = T - F
bdval = [eval_f_on_boundary_node((x,y)->x^2+y^2, bdnode, mmesh);
        eval_f_on_boundary_node((x,y)->x^2-y^2, bdnode, mmesh)]
DOF = [bdnode;bdnode .+ mmesh.ndof]
K, rhs = impose_Dirichlet_boundary_conditions(K, rhs, DOF, bdval)
u = K\rhs
sess = Session(); init(sess)
S = run(sess, u)

figure(figsize=[10,4])
subplot(121)
visualize_scalar_on_fem_points(S[1:mmesh.nnode], mmesh)
subplot(122)
visualize_scalar_on_fem_points(S[mmesh.ndof+1:mmesh.ndof + mmesh.nnode], mmesh)
savefig("numerical.png")

figure(figsize=[10,4])
x = LinRange(0, 1, 50)
X = zeros(50, 50)
Y = zeros(50, 50)
for i = 1:50
    for j = 1:50
        X[i,j] = x[i]
        Y[i,j] = x[j]
    end
end
subplot(121)
X = mmesh.nodes[:,1]
Y = mmesh.nodes[:,2]
U = (@. X^2+Y^2)
V = (@. X^2-Y^2)
subplot(121)
visualize_scalar_on_fem_points(U, mmesh)
subplot(122)
visualize_scalar_on_fem_points(V, mmesh)
savefig("exact.png")

figure(figsize=[10,4])
subplot(121)
visualize_scalar_on_fem_points(U-S[1:mmesh.nnode], mmesh)
subplot(122)
visualize_scalar_on_fem_points(V-S[mmesh.ndof+1:mmesh.ndof + mmesh.nnode], mmesh)
savefig("difference.png")


# 泡利算符

const σ0 = [1 0; 0 1]
const σx = [0 1; 1 0]
const σy = [0 -im; im 0]
const σz = [1 0; 0 -1]
const σv = σx + σy + σz

"""
    ⊗(x::AbstractVector, y::AbstractVector)
张量积
"""
⊗(x::AbstractVector, y::AbstractVector) = begin
    T = promote_type(eltype(x), eltype(y))
    lx = length(x)
    ly = length(y)
    z = zeros(T, lx * ly)
    for i = 1:lx
        for j = 1:ly
            z[(i-1)*ly+j] = x[i] * y[j]
        end
    end
    return z
end

"""
    ⊗(x::AbstractMatrix, y::AbstractMatrix)
张量积
"""
⊗(x::AbstractMatrix, y::AbstractMatrix) = begin
    T = promote_type(eltype(x), eltype(y))
    xrows, xcols = size(x)
    yrows, ycols = size(y)
    z = zeros(T, xrows * yrows, xcols * ycols)
    for xj = 1:xcols
        for yj = 1:ycols
            for xi = 1:xrows
                for yi = 1:yrows
                    z[(xi-1)*yrows+yi, (xj-1)*ycols+yj] = x[xi, xj] * y[yi, yj]
                end
            end
        end
    end
    return z
end


"""
    circular_layout(g) -> x, y

Return two vectors representing the positions of the nodes of graph `g`
when placed on the unit circonference of radius one centered at the origin.
"""
function circular_layout(g::AGraphOrDiGraph)
    # Discard the extra angle since it matches 0 radians.
    θ = range(0, stop=2pi, length=nv(g)+1)[1:end-1]
    return cos.(θ), sin.(θ)
end

"""
    spring_layout(g;
                  k = 1 / √nv(g),
                  maxiter = 100,
                  inittemp = 2.0,
                  [x0, y0]) -> x, y

Return the positions of the nodes in graph `g` according to
Fruchterman and Reingold's spring/repulsion model.
The forces as function of the distance `d` beetween two nodes
are given by

    f_a(d) =  d / k # attractive force
    f_r(d) = -k^2 / d^2 # repulsive force:

`maxiter` is the number of updates of the positions.
`inittemp` controls displacement per iteration.

Initial positions can passed though the argument `x0` and `y0`.

The positions are rescaled to fit the [-1, +1]^2 box.

This function is adapted from [GraphLayout.jl](https://github.com/IainNZ/GraphLayout.jl).
"""
function spring_layout(g::AGraphOrDiGraph;
                       x0 = 2*rand(nv(g)) .- 1.0,
                       y0 = 2*rand(nv(g)) .- 1.0,
                       k = 1/√nv(g),
                       maxiters = 100,
                       inittemp = 2.0)

    x = copy(x0)
    y = copy(y0)
    nvg = nv(g)
    adj_matrix = adjacency_matrix(g)
    k² = k * k

    # Store forces and apply at end of iteration all at once
    force_x = zeros(nvg)
    force_y = zeros(nvg)

    # Iterate MAXITER times
    @inbounds for iter = 1:maxiters
        # Calculate forces
        for i = 1:nvg
            force_vec_x = 0.0
            force_vec_y = 0.0
            for j = 1:nvg
                i == j && continue
                d_x = x[j] - x[i]
                d_y = y[j] - y[i]
                dist²  = (d_x * d_x) + (d_y * d_y)
                dist = sqrt(dist²)

                if !(iszero(adj_matrix[i,j]) && iszero(adj_matrix[j,i]) )
                    # Attractive + repulsive force
                    # F_d = dist² / k - k² / dist # original FR algorithm
                    F_d = dist / k - k² / dist²
                else
                    # Just repulsive
                    # F_d = -k² / dist  # original FR algorithm
                    F_d = -k² / dist²
                end
                force_vec_x += F_d * d_x
                force_vec_y += F_d * d_y
            end
            force_x[i] = force_vec_x
            force_y[i] = force_vec_y
        end
        # Cool down
        temp = inittemp / iter
        # Now apply them, but limit to temperature
        for i = 1:nvg
            fx = force_x[i]
            fy = force_y[i]
            force_mag  = sqrt((fx * fx) + (fy * fy))
            scale      = min(force_mag, temp) / force_mag
            x[i] += force_x[i] * scale
            y[i] += force_y[i] * scale
        end
    end

    # Scale to unit square
    min_x, max_x = minimum(x), maximum(x)
    min_y, max_y = minimum(y), maximum(y)
    function scaler(z, a, b)
        2.0*((z - a) / (b - a)) - 1.0
    end
    map!(z -> scaler(z, min_x, max_x), x, x)
    map!(z -> scaler(z, min_y, max_y), y, y)

    return x, y
end

"""
    shell_layout(g, nlist) -> x, y

Position the nodes of `g` in concentric circles.

`nlist` is a vector of vectors containing the nodes
for each shell.
"""
function shell_layout(g::AGraphOrDiGraph, nlist::Vector{Vector{T}}) where T<:Integer
    if nv(g) == 1
        return [0.0], [0.0]
    end
    radius = length(nlist[1]) > 1 ? 1.0 : 0.0
    x = Float64[]
    y = Float64[]
    for nodes in nlist
        θ = range(0, stop=2pi, length=length(nodes)+1)[1:end-1]
        append!(x, radius * cos.(θ))
        append!(y, radius * sin.(θ))
        radius += 1.0
    end
    return x, y
end


export UpdateColley

"""
    UpdateColley

 Colley update rule, e.g., see "Whos's #1", Langville and Meyer, p.21

## Parameters (none)
```
"""
struct UpdateColley <: UpdateRule
# this is a batch calculation that takes no account of past ratings, and
# has no parameters
end

function update_info( rule::UpdateColley )
    info = Dict(
                :name => "Colley",
                :reference => "\"Whos's #1\", Langville and Meyer, p.21",
                :computation => "simultaneous",
                :state_model => "none",
                :input => "outcome",
                :output => "deterministic",
                :model => "single",
                :ties => false,
                :factors => false,
                :parameters => [],
                :record => false
                )
    return info
end
update_info( ::Type{UpdateColley} ) =  update_info( UpdateColley() )

function update_ratings( rule::UpdateColley,
                         input_ratings::RatingsList,
                         input_competitions::DataFrame)
    n = size(input_competitions,1)
    m = length( input_ratings.players )
    I = player_indexes( input_ratings.players )

    # construct Colley matrices and vectors
    C = diagm(0 => 2*ones(Int,m))
    wins = zeros(Int, m)
    losses = zeros(Int, m)
    d = input_competitions
    for i=1:n
        if d[i,Outcome] == 1
            wins[I[d[i,PlayerA]]] += 1
            losses[I[d[i,PlayerB]]] += 1
        elseif d[i,Outcome] == -1
            wins[I[d[i,PlayerB]]] += 1
            losses[I[d[i,PlayerA]]] += 1
        end
        C[ I[d[i,PlayerA]], I[d[i,PlayerB]] ] -= 1
        C[ I[d[i,PlayerB]], I[d[i,PlayerA]] ] -= 1
        C[ I[d[i,PlayerA]], I[d[i,PlayerA]] ] += 1
        C[ I[d[i,PlayerB]], I[d[i,PlayerB]] ] += 1
    end

    # solve Colley's equation
    b = 1 .+ 0.5*(wins .- losses)
    r = C \ b
    ratings = Dict{String, Float64}()
    for player in input_ratings.players
        ratings[player] = r[ I[player] ]
    end

    # output ratings list
    output_ratings = RatingsList(input_ratings.players, ratings )
    return output_ratings
end

# Colley's predictions are deterministic, unless you apply some model after the fact
function predict_outcome(rule::UpdateColley,
                         ratingA::Real, ratingB::Real,
                         factorA::Union{Missing,Real}, factorB::Union{Missing,Real})
    rating_diff = ratingA - ratingB
    s = sign(rating_diff)
    if s==1
        return (1.0, 0.0, 0.0)
    elseif s==-1
        return (0.0, 1.0, 0.0)
    else
        return (0.0, 0.0, 1.0)
    end
end


using Plots, LaTeXStrings, Statistics

function Deposition(;len, tot_time, time_steps)
    Time = ceil.(Int, exp.(2:(tot_time-2)/(time_steps):tot_time))
    surf = [0 for i=1:len]
    VarList = [0.0 for i=1:time_steps]
    for n in 2:time_steps+1
        randsurf = rand(1:len,(Time[n]-Time[n-1]))
        for i in randsurf
            MAX = FindMax(surf, i, len)
            surf[i] = MAX
        end
        VarList[n-1] = std(surf)
    end
    return VarList
end

function sides(n, L)
    if n == L
        return n-1 , 1
    elseif n == 1
        return L , n+1
    else
        return n-1, n+1
    end
end

function FindMax(surface, index_,L_surf)
    i1 , i2 = sides(index_, L_surf)
    maxlen = max(surface[i1],surface[index_] + 1,surface[i2])
    return maxlen
end

function Linear_fit(;Time, VarList, time_steps)
    A = [hcat(log.(Time[1:time_steps])) reshape(ones(time_steps), time_steps, 1)]
    b = reshape(log.(VarList[1:time_steps]), time_steps, 1)
    line = (A \ b)
    return line
end

# function FindLine(tot_steps, Time, VarList)
#     retline = [0.0 0.0]
#     for i in 5:tot_steps
#         Paraline = Dict(
#                         :Time => Time,
#                         :VarList => VarList,
#                         :time_steps => i
#                             )
#         Line = Linear_fit(;Paraline...)
#         absl = abs(Line[1]*Time[i] + Line[2] - VarList[i])
#         standev = std(Line[1].*Time .+ Line[2] .- VarList)
#         if absl > 0.2 && standev > 1 && retline == [0.0 0.0]
#             retline = Line
#         end
#         if absl > 1 && standev > 3
#             return retline, i
#         end
#     end
# end

theme(:dark)
gr()

iternum = 1000
Parameters = Dict(:len => 200,
                    :tot_time => 13,
                        :time_steps => 50)
allVar = [ [0.0 for i in 1:Parameters[:time_steps]] for j = 1:iternum]
meanVar = [0.0 for i in 1:Parameters[:time_steps]]
vars = [0.0 for i in 1:Parameters[:time_steps]]
for i in 1:iternum
    VarList = Deposition(;Parameters...)
    allVar[i] = VarList
    meanVar += VarList
    print("\r$i")
end
meanVar /= iternum
for i in 1:Parameters[:time_steps]
    vars[i] = std(log.(hcat(allVar...))[i,:])
end

Time = exp.(0:(Parameters[:tot_time])/(Parameters[:time_steps]-1):Parameters[:tot_time])
scatter(log.(Time[3:end]), log.(meanVar[3:end]),
    # xlims = (1, Parameters[:tot_time]),
    xlabel= L"Log\ Time",
    ylabel= L"Log\ W_{(t)}",
    title= L"Log-Log\ Plot\ ~W_{(t)}-Time~\ (L = %$(Parameters[:len]))",
    label = L"Data\ point",
    yerror = vars,
    legend = 150)

Paraline = Dict(
                :Time => Time[3:end],
                :VarList => meanVar[3:end],
                :time_steps => 13)
Line = Linear_fit(;Paraline...)
X = log.(Time[3]:Time[30])
Y = X .* Line[1] .+ Line[2]

plot!(X,Y,label = L"y = %$(round(Line[1],digits= 2))x + %$(round(Line[2],digits= 2))")

savefig("C:\\Users\\Yaghoub\\Documents\\GitHub\\Ballistic-Deposition\\Deposition\\Ballistic-Deposition-Kardar-Parisi-Zhang\\Fig\\W-t(L=$(Parameters[:len])).png")


### A Pluto.jl notebook ###
# v0.14.7

using Markdown
using InteractiveUtils

# ╔═╡ 5057c590-c5a4-11eb-2621-4b67f12ca95d
md"""
# K Nearest Neighbors
"""

# ╔═╡ 74f4afa1-dd5a-4931-ab48-a1d4fc77bbe5
md"""
[Reference 1](https://scikit-learn.org/stable/modules/neighbors.html)\
[Reference 2](https://en.wikipedia.org/wiki/Nearest_neighbor_search)\
[Reference 3](https://booking.ai/k-nearest-neighbours-from-slow-to-fast-thanks-to-maths-bec682357ccd)

### Naive Approach

No transforming on original dataset (No training)\
For prediction, iter through the original dataset, and find the nearest K data, by a given metrics

Can use euclidean distance:

$$\text{dist}(X_1, X_2) = \|X_1 - X_2\|$$

A better approach is cosine similarity:

$$\text{sim}(X_1, X_2) = \frac{X_1 \cdot X_2}{\|X_1\| \|X_2\|}$$

which computes the cos value between two vectors\
1 for 0 degree, and less than 1 for $(0, \pi]$
"""

# ╔═╡ cb1dd9ac-ca78-4169-b073-943f64160dbc
module tools include("../tools.jl") end

# ╔═╡ f1fa76da-f656-4f12-97d7-a8b847ca0733
JuTools = tools.JuTools

# ╔═╡ c3ad6021-8080-4870-8350-1337d548f23a
import Statistics

# ╔═╡ 83864e6e-4b62-45d5-962e-39e36bbcc43d
import Random

# ╔═╡ 0d7f31a9-7ae2-4be5-81b2-313ac58af62d
import LinearAlgebra

# ╔═╡ 73b4e7e0-7294-46dc-9545-9728277df652
function cosine_sim(X1::Array{T} where T<:Number, X2::Array{T} where T<:Number)::AbstractFloat
    @assert size(X1) == size(X2)
    @assert ndims(X1) == ndims(X2) == 1
    product = LinearAlgebra.dot(X1, X2)
    X1_norm = LinearAlgebra.norm(X1, 2)
    X2_norm = LinearAlgebra.norm(X2, 2)
    return product / (X1_norm * X2_norm)
end

# ╔═╡ 19b1c000-6fa3-42b7-9fc7-983c7caf44b9
begin
X_data, Y_data = JuTools.data_generate_linear_2d()
size(X_data), size(Y_data)
end

# ╔═╡ 063389c7-07b6-4050-9241-ea953c4f4338
X_data[1:2, :]

# ╔═╡ cf291226-c0a1-4206-a8cb-53bedb67b139
cosine_sim(X_data[1, :], X_data[2, :])

# ╔═╡ fbf7a012-dcf8-498f-ad29-ec8c4854406c
cosine_sim(X_data[1, :], X_data[3, :])

# ╔═╡ 2cf4e990-f788-44e3-94cc-9e210e47c762
# define majority vote function
function majority_vote(y::Array{T} where T<:Number)::Number
    @assert ndims(y) == 1
    unique_votes = Dict{Number, Integer}()
    for y_val in y
        if !haskey(unique_votes, y_val)
            push!(unique_votes, y_val => 1)
        else
            unique_votes[y_val] += 1
        end
    end
    result = sort(collect(unique_votes), by=m->m[2])
    return result[end][1]
end

# ╔═╡ c9d27318-5e3d-4c9c-9ba5-556c649baaf9
majority_vote([1,1,0])

# ╔═╡ ba24b3c5-f350-4774-a018-9acf1f4597ad
md"""
Output ordering is affected by input ordering
"""

# ╔═╡ 5e564f3b-1bda-41b1-8b22-bdf6d91314d2
majority_vote([1,1,0,0])

# ╔═╡ ed575b52-7509-4140-b5e7-295920f73e3d
majority_vote([1,1,0,0,0])

# ╔═╡ 8406faec-66f4-46aa-9582-964ebdbecfba
begin
X_train, X_test, Y_train, Y_test = JuTools.split_data(X_data, Y_data, shuffle=true, ratio=0.3)
size(X_train), size(X_test), size(Y_train), size(Y_test)
end

# ╔═╡ 3001da4a-3288-4168-b514-c0b12f552fdd
# define predict function, naive approach
function predict_naive(X_predict::Array{T} where T<:Number, K::Integer, X_data::Array{T} where T<:Number, Y_data::Array{T} where T<:Number)::Array
    @assert ndims(X_data) == 2
    @assert ndims(Y_data) == 1
    @assert size(X_data)[1] == size(Y_data)[1]
    @assert 0 < ndims(X_predict) <= 2
    @assert 0 < K < size(X_data)[1]
    if ndims(X_predict) < 2
        X_predict = reshape(X_predict, (1, size(X_predict)[1]))
    end
    @assert size(X_predict)[2] == size(X_data)[2]
    result = Array{Number}(undef, size(X_predict)[1])
    sim = Array{Tuple{Integer, AbstractFloat}}(undef, size(X_data)[1])
    for i in 1:size(X_predict)[1]
        vec_predict = X_predict[i, :]
        for j in 1:size(X_data)[1]
            vec_data = X_data[j, :]
            vec_similarity = cosine_sim(vec_predict, vec_data)
            sim[j] = (j, vec_similarity)
        end
        sort!(sim, by=m->m[2], rev=true)
        K_nearest_votes = Y_data[[m[1] for m in sim[1:K]]]
        result[i] = majority_vote(K_nearest_votes)
    end
    return result
end

# ╔═╡ 30320c6e-7ada-41b6-a938-f1fd19cc120a
Y_predict = predict_naive(X_test, 5, X_train, Y_train)

# ╔═╡ 27b26afb-a91c-47c4-94ef-1d64761f291d
JuTools.compute_accuracy(Y_predict, Y_test)

# ╔═╡ 0e1b5c16-f444-41dd-9e44-999969ba647c
# what about dist similarity?
function dist_sim(X1::Array{T} where T<:Number, X2::Array{T} where T<:Number)::AbstractFloat
    @assert size(X1) == size(X2)
    @assert ndims(X1) == ndims(X2) == 1
    return sqrt(sum((X1 .- X2).^2))
end

# ╔═╡ c42f2b8a-a88d-4e8c-a1c0-1a371ae6a669
dist_sim(X_data[1, :], X_data[2, :])

# ╔═╡ 158f77bf-2eb8-490a-a062-6364a6ec70e9
dist_sim(X_data[2, :], X_data[3, :])

# ╔═╡ 323e10a0-c68f-433a-afc5-0ac2514a90d1
md"""
It's greatly affected by the scale of data!
"""

# ╔═╡ 52010e3f-2e4f-4778-b53e-4fad04dd75bf
function predict_naive_fun(X_predict::Array{T} where T<:Number, K::Integer, X_data::Array{T} where T<:Number, Y_data::Array{T} where T<:Number)::Array
    @assert ndims(X_data) == 2
    @assert ndims(Y_data) == 1
    @assert size(X_data)[1] == size(Y_data)[1]
    @assert 0 < ndims(X_predict) <= 2
    @assert 0 < K < size(X_data)[1]
    if ndims(X_predict) < 2
        X_predict = reshape(X_predict, (1, size(X_predict)[1]))
    end
    @assert size(X_predict)[2] == size(X_data)[2]
    result = Array{Number}(undef, size(X_predict)[1])
    sim = Array{Tuple{Integer, AbstractFloat}}(undef, size(X_data)[1])
    for i in 1:size(X_predict)[1]
        vec_predict = X_predict[i, :]
        for j in 1:size(X_data)[1]
            vec_data = X_data[j, :]
            vec_similarity = dist_sim(vec_predict, vec_data)
            sim[j] = (j, vec_similarity)
        end
        sort!(sim, by=m->m[2])
        K_nearest_votes = Y_data[[m[1] for m in sim[1:K]]]
        result[i] = majority_vote(K_nearest_votes)
    end
    return result
end

# ╔═╡ 36ceae2a-4c04-4f67-9a6d-19a44cb2fa5b
JuTools.compute_accuracy(predict_naive_fun(X_test, 5, X_train, Y_train), Y_test)

# ╔═╡ 2c7f7263-5ba0-4786-8b68-8ca8fb3ffc26
md"""
Although it (`predict_naive`) may be slow on large dataset, it is easy to implement and it works as expected
"""

# ╔═╡ b2aac2a9-cf86-44f0-ae97-b2a52327b0f9
md"""
### K-Dimensional Tree (K-d tree) Approach

A space partitioning technique\
Treat each data row as a point in `k`-dimensional space\
[Wikipedia](https://en.wikipedia.org/wiki/K-d_tree)
"""

# ╔═╡ b127297e-cb0c-46f5-8551-3ba2c4dae4c6
mutable struct KdTree
    X_data::Array{T} where T<:Number # 1d vector
    Y_data::Number                   # number
    child_l::Union{KdTree,Nothing}
    child_r::Union{KdTree,Nothing}
end

# ╔═╡ c3c4fe2c-6e92-4c98-bce3-e77f394c405d
# K-d tree generator function
function create_kdtree(X_data::Array{T} where T<:Number, Y_data::Array{T} where T<:Number)::KdTree
    @assert ndims(X_data) == 2
    @assert ndims(Y_data) == 1
    @assert size(X_data)[1] == size(Y_data)[1]
    function kdtree_recursive_generate(X_data::Array, Y_data::Array, depth::Integer, n_axes::Integer)::KdTree
        curr_axis = mod(depth, n_axes) + 1 # array starts from 1
        data_combined = hcat(X_data, Y_data)
        data_combined = sortslices(data_combined, by=m->m[curr_axis], dims=1)
        X_data = data_combined[:, 1:end-1]
        Y_data = data_combined[:, end]
        i_mid = div(size(X_data)[1], 2) + 1
        node_X_data = X_data[i_mid, :]
        node_Y_data = Y_data[i_mid]
        node = KdTree(node_X_data, node_Y_data, nothing, nothing)
        if i_mid > 1
            node.child_l = kdtree_recursive_generate(X_data[1:i_mid-1,:], Y_data[1:i_mid-1], depth+1, n_axes)
        end
        if i_mid < size(X_data)[1]
            node.child_r = kdtree_recursive_generate(X_data[i_mid+1:end,:], Y_data[i_mid+1:end], depth+1, n_axes)
        end
        return node
    end
    return kdtree_recursive_generate(X_data, Y_data, 0, size(X_data)[2])
end

# ╔═╡ a76dcf16-060e-4ece-9b8d-76c45deb38c4
kdtree = create_kdtree(X_data, Y_data)

# ╔═╡ 4c3444dd-f2cb-4954-8dad-d2e9466252e8
kdtree_test = create_kdtree(reshape([30,5,10,70,50,35], (6, 1)), [1,1,1,1,1,1])

# ╔═╡ 94c43601-cff7-4c7f-a116-8a82aefed5b0
# inspired from https://stackoverflow.com/questions/1627305/nearest-neighbor-k-d-tree-wikipedia-proof/37107030#37107030
# note that for kdtree search, we use euclidean distance
function predict_kdtree(X_predict::Array{T} where T<:Number, kdtree::KdTree; K::Integer=5)::Array
    @assert K > 0
    @assert 0 < ndims(X_predict) <= 2
    if ndims(X_predict) == 1
        X_predict = reshape(X_predict, (1, size(X_predict)[1]))
    end
    @assert size(X_predict)[2] == size(kdtree.X_data)[1]

    function kdtree_closest_max(kdtree_closest::Array{Union{KdTree, Nothing}},
                kdtree_closest_val::Array{AbstractFloat})::Tuple{Integer, AbstractFloat}
        default = (0, 0.0)
        for i in 1:size(kdtree_closest)[1]
            if kdtree_closest[i] === nothing
                break
            elseif default[1] == 0 || (kdtree_closest_val[i] > default[2])
                default = (i, kdtree_closest_val[i])
            end
        end
        return default
    end

    function kdtree_update_nearest!(X_vec::Array, kdtree::KdTree, kdtree_closest::Array{Union{KdTree, Nothing}},
                kdtree_closest_val::Array{AbstractFloat})
        @assert size(kdtree_closest) == size(kdtree_closest_val)
        distance = dist_sim(kdtree.X_data, X_vec)
        if nothing in kdtree_closest
            for i in 1:size(kdtree_closest)[1]
                if kdtree_closest[i] === nothing
                    kdtree_closest[i] = KdTree(kdtree.X_data, kdtree.Y_data, nothing, nothing)
                    kdtree_closest_val[i] = distance
                    break
                end
            end
        else
            curr_max = kdtree_closest_max(kdtree_closest, kdtree_closest_val)
            if distance < curr_max[2]
                kdtree_closest[curr_max[1]] = KdTree(kdtree.X_data, kdtree.Y_data, nothing, nothing)
                kdtree_closest_val[curr_max[1]] = distance
            end
        end
    end

    function kdtree_recursive_search!(X_vec::Array, kdtree::KdTree, depth::Integer, n_axes::Integer,
                kdtree_closest::Array{Union{KdTree, Nothing}}, kdtree_closest_val::Array{AbstractFloat})
        @assert size(kdtree_closest) == size(kdtree_closest_val)
        # check current node
        kdtree_update_nearest!(X_vec, kdtree, kdtree_closest, kdtree_closest_val)
        # run on children
        curr_axis = mod(depth, n_axes) + 1 # array starts from 1
        if X_vec[curr_axis] < kdtree.X_data[curr_axis]
            if kdtree.child_l !== nothing
                kdtree_recursive_search!(X_vec, kdtree.child_l, depth+1, n_axes, kdtree_closest, kdtree_closest_val)
            end
            if (X_vec[curr_axis] + kdtree_closest_max(kdtree_closest, kdtree_closest_val)[2] >= kdtree.X_data[curr_axis]) && kdtree.child_r !== nothing
                kdtree_recursive_search!(X_vec, kdtree.child_r, depth+1, n_axes, kdtree_closest, kdtree_closest_val)
            end
        else
            if kdtree.child_r !== nothing
                kdtree_recursive_search!(X_vec, kdtree.child_r, depth+1, n_axes, kdtree_closest, kdtree_closest_val)
            end
            if (X_vec[curr_axis] - kdtree_closest_max(kdtree_closest, kdtree_closest_val)[2] <= kdtree.X_data[curr_axis]) && kdtree.child_l !== nothing
                kdtree_recursive_search!(X_vec, kdtree.child_l, depth+1, n_axes, kdtree_closest, kdtree_closest_val)
            end
        end
    end

    result = Array{Number}(undef, size(X_predict)[1])
    for i in 1:size(X_predict)[1]
        kdtree_closest = Array{Union{KdTree, Nothing}}(nothing, K)
        kdtree_closest_val = Array{AbstractFloat}(undef, K)
        kdtree_recursive_search!(X_predict[i, :], kdtree, 0, size(X_predict)[2], kdtree_closest, kdtree_closest_val)
        K_nearest_votes = Number[]
        for i in 1:K
            if kdtree_closest[i] === nothing
                break
            else
                push!(K_nearest_votes, kdtree_closest[i].Y_data)
            end
        end
        result[i] = majority_vote(K_nearest_votes)
    end
    return result
end

# ╔═╡ 569da607-dea8-4994-9ffe-05959071a855
kdtree_train = create_kdtree(X_train, Y_train)

# ╔═╡ 0ea4cb86-6be5-49f1-b9b8-56f6a8cc8ed0
JuTools.compute_accuracy(predict_kdtree(X_test, kdtree_train, K=10), Y_test)

# ╔═╡ 98be952f-0591-45c6-be83-3c0f6cc5d26b
md"""
### Ball Tree Approach

A better space partition approach\
More efficient than K-d Tree when searching\
[Wikipedia](https://en.wikipedia.org/wiki/Ball_tree)
"""

# ╔═╡ dc218b5e-73f5-4dc8-b523-62bed53a71e7
mutable struct BallTree
    X_data::Array{T} where T<:Number
    Y_data::Number
    pivot::Union{Array{T},Nothing} where T<:Number # defines pivot point of hypersphere
    radius::AbstractFloat                            # defines radius of hypersphere
    child_l::Union{BallTree,Nothing}
    child_r::Union{BallTree,Nothing}
end

# ╔═╡ d7cc4b1c-f953-4f64-95bb-09e1d51a38f5
# inspired from https://gist.github.com/jakevdp/5216193
# ball tree generator function
function create_balltree(X_data::Array{T} where T<:Number, Y_data::Array{T} where T<:Number)::BallTree
    @assert ndims(X_data) == 2
    @assert ndims(Y_data) == 1
    @assert size(X_data)[1] == size(Y_data)[1]
    balltree = nothing
    if size(X_data)[1] == 1
        balltree = BallTree(X_data[1, :], Y_data[1], nothing, 0.0, nothing, nothing)
    else
        # find pivot
        pivot = vec(sum(X_data, dims=1)) ./ size(X_data)[1]
        # find radius
        radius = 0.0
        for i in 1:size(X_data)[1]
            X_vec = X_data[i, :]
            dist = dist_sim(pivot, X_vec)
            if dist > radius
                radius = dist
            end
        end
        # find greatest spread dimension
        d_greatest_spread = 1
        n_spread = 0.0
        for i in 1:size(X_data)[2]
            X_vec = X_data[:, i]
            current_spread = abs(maximum(X_vec) - minimum(X_vec))
            if current_spread > n_spread
                d_greatest_spread = i
                n_spread = current_spread
            end
        end
        data_combined = hcat(X_data, Y_data)
        data_combined = sortslices(data_combined, by=m->m[d_greatest_spread], dims=1)
        X_data = data_combined[:, 1:end-1]
        Y_data = data_combined[:, end]
        i_mid = div(size(X_data)[1], 2) + 1
        node_X_data = X_data[i_mid, :]
        node_Y_data = Y_data[i_mid]
        balltree = BallTree(node_X_data, node_Y_data, pivot, radius, nothing, nothing)
        if i_mid > 1
            balltree.child_l = create_balltree(X_data[1:i_mid-1,:], Y_data[1:i_mid-1])
        end
        if i_mid < size(X_data)[1]
            balltree.child_r = create_balltree(X_data[i_mid+1:end,:], Y_data[i_mid+1:end])
        end
    end
    return balltree
end

# ╔═╡ ffb15055-c887-47a9-88d7-71a33a3c2dba
balltree = create_balltree(X_data, Y_data)

# ╔═╡ 77ed169b-d426-4f3f-8362-809cb0c8dd22
balltree.X_data

# ╔═╡ 79ddb157-8ee4-4c44-b8bf-b4f36f88acce
balltree.Y_data

# ╔═╡ 6c589940-e00a-40c9-b2d9-c10384e4ac62
balltree.pivot

# ╔═╡ 5b5c5f78-75c2-4a85-b918-e7485f492d9e
balltree.radius

# ╔═╡ 75368c67-8774-4f67-9163-d31475f4e289
balltree_test = create_balltree(reshape([30,5,10,70,50,35], (6, 1)), [1,1,1,1,1,1])

# ╔═╡ a14bd64d-4188-4682-8a47-c3999afb0d09
# search function for ball tree
# similar to kdtree search
function predict_balltree(X_predict::Array{T} where T<:Number, balltree::BallTree; K::Integer=5)::Array
    @assert K > 0
    @assert 0 < ndims(X_predict) <= 2
    if ndims(X_predict) == 1
        X_predict = reshape(X_predict, (1, size(X_predict)[1]))
    end
    @assert size(X_predict)[2] == size(balltree.X_data)[1]

    function balltree_closest_max(balltree_closest::Array{Union{BallTree, Nothing}},
                balltree_closest_val::Array{AbstractFloat})::Tuple{Integer, AbstractFloat}
        # find the maximum value in balltree_closest_val
        # return its index and value
        default = (0, 0.0)
        for i in 1:size(balltree_closest)[1]
            if balltree_closest[i] === nothing
                break
            elseif default[1] == 0 || (balltree_closest_val[i] > default[2])
                default = (i, balltree_closest_val[i])
            end
        end
        return default
    end

    function balltree_update_nearest!(X_vec::Array, balltree::BallTree, balltree_closest::Array{Union{BallTree, Nothing}},
                balltree_closest_val::Array{AbstractFloat})
        # update current node by distance
        @assert size(balltree_closest) == size(balltree_closest_val)
        distance = dist_sim(balltree.X_data, X_vec)
        if nothing in balltree_closest
            for i in 1:size(balltree_closest)[1]
                if balltree_closest[i] === nothing
                    balltree_closest[i] = balltree
                    balltree_closest_val[i] = distance
                    break
                end
            end
        else
            curr_max = balltree_closest_max(balltree_closest, balltree_closest_val)
            if distance < curr_max[2]
                balltree_closest[curr_max[1]] = balltree
                balltree_closest_val[curr_max[1]] = distance
            end
        end
    end

    function balltree_recursive_search!(X_vec::Array, balltree::BallTree, balltree_closest::Array{Union{BallTree, Nothing}},
            balltree_closest_val::Array{AbstractFloat})
        # recursively search a balltree for K nearest neighbors
        @assert size(balltree_closest) == size(balltree_closest_val)
        if (!(nothing in balltree_closest) && (balltree.pivot !== nothing)
                && (dist_sim(X_vec, balltree.pivot) - balltree.radius >=
                    balltree_closest_max(balltree_closest, balltree_closest_val)[2]))
            return nothing
        end
        # check current node
        balltree_update_nearest!(X_vec, balltree, balltree_closest, balltree_closest_val)
        # run on children
        if balltree.child_l === nothing || balltree.child_r === nothing
            if balltree.child_l !== nothing
                balltree_recursive_search!(X_vec, balltree.child_l, balltree_closest, balltree_closest_val)
            end
            if balltree.child_r !== nothing
                balltree_recursive_search!(X_vec, balltree.child_r, balltree_closest, balltree_closest_val)
            end
        else
            dist_left = dist_sim(X_vec, balltree.child_l.X_data)
            dist_right = dist_sim(X_vec, balltree.child_r.X_data)
            if dist_left < dist_right
                balltree_recursive_search!(X_vec, balltree.child_l, balltree_closest, balltree_closest_val)
                balltree_recursive_search!(X_vec, balltree.child_r, balltree_closest, balltree_closest_val)
            else
                balltree_recursive_search!(X_vec, balltree.child_r, balltree_closest, balltree_closest_val)
                balltree_recursive_search!(X_vec, balltree.child_l, balltree_closest, balltree_closest_val)
            end
        end
    end

    result = Array{Number}(undef, size(X_predict)[1])
    for i in 1:size(X_predict)[1]
        balltree_closest = Array{Union{BallTree, Nothing}}(nothing, K)
        balltree_closest_val = Array{AbstractFloat}(undef, K)
        balltree_recursive_search!(X_predict[i, :], balltree, balltree_closest, balltree_closest_val)
        K_nearest_votes = Number[]
        for i in 1:K
            if balltree_closest[i] === nothing
                break
            else
                push!(K_nearest_votes, balltree_closest[i].Y_data)
            end
        end
        result[i] = majority_vote(K_nearest_votes)
    end
    return result
end

# ╔═╡ 06e6d38d-b08e-4583-863a-e6e06726143b
balltree_train = create_balltree(X_train, Y_train)

# ╔═╡ 41e118e2-2172-4910-ba52-12e199737c6a
JuTools.compute_accuracy(predict_balltree(X_test, balltree_train, K=20), Y_test)

# ╔═╡ Cell order:
# ╟─5057c590-c5a4-11eb-2621-4b67f12ca95d
# ╟─74f4afa1-dd5a-4931-ab48-a1d4fc77bbe5
# ╠═cb1dd9ac-ca78-4169-b073-943f64160dbc
# ╠═f1fa76da-f656-4f12-97d7-a8b847ca0733
# ╠═c3ad6021-8080-4870-8350-1337d548f23a
# ╠═83864e6e-4b62-45d5-962e-39e36bbcc43d
# ╠═0d7f31a9-7ae2-4be5-81b2-313ac58af62d
# ╠═73b4e7e0-7294-46dc-9545-9728277df652
# ╠═19b1c000-6fa3-42b7-9fc7-983c7caf44b9
# ╠═063389c7-07b6-4050-9241-ea953c4f4338
# ╠═cf291226-c0a1-4206-a8cb-53bedb67b139
# ╠═fbf7a012-dcf8-498f-ad29-ec8c4854406c
# ╠═2cf4e990-f788-44e3-94cc-9e210e47c762
# ╠═c9d27318-5e3d-4c9c-9ba5-556c649baaf9
# ╟─ba24b3c5-f350-4774-a018-9acf1f4597ad
# ╠═5e564f3b-1bda-41b1-8b22-bdf6d91314d2
# ╠═ed575b52-7509-4140-b5e7-295920f73e3d
# ╠═8406faec-66f4-46aa-9582-964ebdbecfba
# ╠═3001da4a-3288-4168-b514-c0b12f552fdd
# ╠═30320c6e-7ada-41b6-a938-f1fd19cc120a
# ╠═27b26afb-a91c-47c4-94ef-1d64761f291d
# ╠═0e1b5c16-f444-41dd-9e44-999969ba647c
# ╠═c42f2b8a-a88d-4e8c-a1c0-1a371ae6a669
# ╠═158f77bf-2eb8-490a-a062-6364a6ec70e9
# ╟─323e10a0-c68f-433a-afc5-0ac2514a90d1
# ╠═52010e3f-2e4f-4778-b53e-4fad04dd75bf
# ╠═36ceae2a-4c04-4f67-9a6d-19a44cb2fa5b
# ╟─2c7f7263-5ba0-4786-8b68-8ca8fb3ffc26
# ╟─b2aac2a9-cf86-44f0-ae97-b2a52327b0f9
# ╠═b127297e-cb0c-46f5-8551-3ba2c4dae4c6
# ╠═c3c4fe2c-6e92-4c98-bce3-e77f394c405d
# ╠═a76dcf16-060e-4ece-9b8d-76c45deb38c4
# ╠═4c3444dd-f2cb-4954-8dad-d2e9466252e8
# ╠═94c43601-cff7-4c7f-a116-8a82aefed5b0
# ╠═569da607-dea8-4994-9ffe-05959071a855
# ╠═0ea4cb86-6be5-49f1-b9b8-56f6a8cc8ed0
# ╟─98be952f-0591-45c6-be83-3c0f6cc5d26b
# ╠═dc218b5e-73f5-4dc8-b523-62bed53a71e7
# ╠═d7cc4b1c-f953-4f64-95bb-09e1d51a38f5
# ╠═ffb15055-c887-47a9-88d7-71a33a3c2dba
# ╠═77ed169b-d426-4f3f-8362-809cb0c8dd22
# ╠═79ddb157-8ee4-4c44-b8bf-b4f36f88acce
# ╠═6c589940-e00a-40c9-b2d9-c10384e4ac62
# ╠═5b5c5f78-75c2-4a85-b918-e7485f492d9e
# ╠═75368c67-8774-4f67-9163-d31475f4e289
# ╠═a14bd64d-4188-4682-8a47-c3999afb0d09
# ╠═06e6d38d-b08e-4583-863a-e6e06726143b
# ╠═41e118e2-2172-4910-ba52-12e199737c6a


using DifferentialEquations, ParameterizedFunctions
ball! = @ode_def BallBounce begin
  dy =  v
  dv = -g
end g

function condition(u,t,integrator)
  u[1]
end

function affect!(integrator)
    integrator.u[2] = -integrator.p[2] * integrator.u[2]
end

bounce_cb = ContinuousCallback(condition,affect!)

u0 = [50.0,0.0]
tspan = (0.0,15.0)
p = (9.8,0.9)
prob = ODEProblem(ball!,u0,tspan,p,callback=bounce_cb)

sol = solve(prob,Tsit5())
using Plots; gr()
plot(sol)

function condition_kick(u,t,integrator)
    t == 2
end

function affect_kick!(integrator)
    integrator.u[2] += 50
end

kick_cb = DiscreteCallback(condition_kick,affect_kick!)
u0 = [50.0,0.0]
tspan = (0.0,10.0)
p = (9.8,0.9)
prob = ODEProblem(ball!,u0,tspan,p,callback=kick_cb)

sol = solve(prob,Tsit5(),tstops=[2.0])
plot(sol)

cb = CallbackSet(bounce_cb,kick_cb)

u0 = [50.0,0.0]
tspan = (0.0,15.0)
p = (9.8,0.9)
prob = ODEProblem(ball!,u0,tspan,p,callback=cb)
sol = solve(prob,Tsit5(),tstops=[2.0])
plot(sol)

u0 = [1.,0.]
harmonic! = @ode_def HarmonicOscillator begin
   dv = -x
   dx = v
end
tspan = (0.0,10.0)
prob = ODEProblem(harmonic!,u0,tspan)
sol = solve(prob)
plot(sol)

function terminate_affect!(integrator)
    terminate!(integrator)
end

function terminate_condition(u,t,integrator)
    u[2]
end
terminate_cb = ContinuousCallback(terminate_condition,terminate_affect!)

sol = solve(prob,callback=terminate_cb)
plot(sol)

sol.t[end]

terminate_upcrossing_cb = ContinuousCallback(terminate_condition,terminate_affect!,nothing)

sol = solve(prob,callback=terminate_upcrossing_cb)
plot(sol)

tspan = (0.0,10000.0)
prob = ODEProblem(harmonic!,u0,tspan)
sol = solve(prob)
gr(fmt=:png) # Make it a PNG instead of an SVG since there's a lot of points!
plot(sol,vars=(1,2))

plot(sol,vars=(0,1),denseplot=false)

plot(sol.t,[u[2]^2 + u[1]^2 for u in sol.u]) # Energy ~ x^2 + v^2

function g(resid,u,p,t)
  resid[1] = u[2]^2 + u[1]^2 - 1
  resid[2] = 0
end

cb = ManifoldProjection(g)
sol = solve(prob,callback=cb)
plot(sol,vars=(1,2))

plot(sol,vars=(0,1),denseplot=false)

u1,u2 = sol[500]
u2^2 + u1^2

prob = ODEProblem((du,u,p,t)->du.=u,rand(1000,1000),(0.0,1.0))

saved_values = SavedValues(Float64, Tuple{Float64,Float64})

using LinearAlgebra
cb = SavingCallback((u,t,integrator)->(tr(u),norm(u)), saved_values)

sol = solve(prob, Tsit5(), callback=cb, save_everystep=false, save_start=false, save_end = false) # Turn off normal saving

saved_values.t

saved_values.saveval

saved_values = SavedValues(Float64, Tuple{Float64,Float64}) # New cache
cb = SavingCallback((u,t,integrator)->(tr(u),norm(u)), saved_values, saveat = 0.0:0.1:1.0)
sol = solve(prob, Tsit5(), callback=cb, save_everystep=false, save_start=false, save_end = false) # Turn off normal saving

saved_values.t

saved_values.saveval

using DiffEqTutorials
DiffEqTutorials.tutorial_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


module SimpleML

using Optim
using OptionsMod

export predictLogistic, predictLogisticAll, learnLogistic, learnLogisticAll, predict2nn, learn2nn

sigmoid(x) = (1 + exp (-1 * x)) .^ -1

sigmoidGradient(z) = sigmoid(z) .* (1-sigmoid(z))

#Cost function for a logistic regression, with regularisation
function logisticCost(theta, X, y, grad, lambda)
    m = length(y) #size of training set
    n = length(theta) #number of features

    h = sigmoid(X * theta)

    sum =0
    for i = 1:m
        sum = sum + ( -1 * y[i] * log(h[i]) - (1-y[i]) * log (1-h[i]))
    end
    for j=2:n
        sum = sum + (lambda/2) * (theta[j] ^ 2)
    end

    if !(grad === nothing)
        gptr = pointer(grad)
        for j=1:n
            grad[j]=0
            for i = 1:m
                grad[j] =  grad[j]  + (h[i] - y[i])*X[i,j]
            end

           if j==1
                grad[j] = grad[j]/m
           else
                grad[j] = ( grad[j] + lambda* theta[j]) / m
           end
        end
        @assert gptr == pointer(grad)
    end

    return sum / m
end

#Learn via a logistic regression of two states
#  X : features
#  y : lables (binary)
#  lambda : regularisation parameter
function learnLogistic(X, y, lambda)
    initial_theta = zeros(size(X,2))
    ops = @options itermax=50 tol=1e-5

    results, fval, fcount, converged = cgdescent((g,t)->logisticCost(t, X, y, g, lambda),  initial_theta, ops)
    #println("Converged: $converged @ $(fval[end]) with coefficients $results in $fcount iterations. ")

    return results
end

#Learn via a logistic regression
#  X : features
#  y : lables
#  nlables : number of lables
#  lambda : regularisation parameter
function learnLogisticAll(X, y, lambda, nlables)

    m = size(X, 1); #number of data points
    n = size(X, 2); #number of features

    all_theta = zeros(nlables, n)

    for c=1:nlables
        initial_theta = zeros(n , 1)
        all_theta[c,:]=learnLogistic(X,int((y.==c)),lambda)[:]
    end

     return all_theta

end

function predictLogistic(theta, X )
    m = size(X, 1)
    p = zeros(m, 1)
    h = sigmoid(X * theta)

    for i=1:m
        if h[i] < 0.5
            p[i] = 0
        else
            p[i] = 1
        end
    end
    return p
end

function predictLogisticAll(all_theta, X)
    m = size(X, 1) #Number of training examples
    n = size(X, 2) #Number of features
    nlables = size(all_theta, 1)
    all_h = zeros(nlables, m)
    for c = 1: nlables
        theta = all_theta[c,:]
        h = sigmoid(X * theta')
        all_h[c,:] = h[:]
    end

    p = zeros(m)

    for i = 1:m
        _ , p[i] = findmax(all_h[:,i])
    end

    return p
end

#Learn a 2 layer neural network
#  X : features
#  y : lables
#  input_sz: size of input layer
#  hidden_sz: size of hidden layer
#  nlables : number of lables
#  lambda : regularisation parameter
function learn2nn(X, y, input_sz, hidden_sz, nlables, lambda)

    initial_theta1 = SimpleML.randomInitialWeights(input_sz, hidden_sz)
    initial_theta2 = SimpleML.randomInitialWeights(hidden_sz, nlables)
    initial_params = [initial_theta1[:] ; initial_theta2[:]]

    ops = @options itermax=50 tol=1e-5

    results, fval, fcount, converged = cgdescent((g,t)->nn2Cost(t, X, y,input_sz, hidden_sz, nlables, lambda, g),  initial_params, ops)

    #println("Converged: $converged @ $(fval[end]) with coefficients $results in $fcount iterations. ")

    theta1 = reshape(results[1:hidden_sz * (input_sz + 1)], hidden_sz, (input_sz + 1))
    theta2 = reshape(results[(1 + (hidden_sz * (input_sz + 1))):end], nlables, (hidden_sz + 1))
    return theta1 ,theta2

end

#Predict a 2 layer neural network
function predict2nn(theta1, theta2, X)
    m = size(X, 1)
    nlables = size(theta2, 1)
    p = zeros(size(X, 1), 1)
    X = [ones(m, 1) X]
    a2 = sigmoid(X * theta1');
    a2 = [ones(m, 1) a2]
    a3 = sigmoid (a2 * theta2')

    for i=1:m
        _, p[i] = findmax(a3[i,:])
    end

    return p
end

#Cost function for a 2 layer neural network
function nn2Cost(params, X, y, input_sz,  hidden_sz, nlables,  lambda, grad)

    theta1 = reshape(params[1:hidden_sz * (input_sz + 1)], hidden_sz, (input_sz + 1))
    theta2 = reshape(params[(1 + (hidden_sz * (input_sz + 1))):end], nlables, (hidden_sz + 1))

    m = size(X, 1)

    J = 0
    theta1_grad = zeros(size(theta1))
    theta2_grad = zeros(size(theta2))

    X = [ones(m, 1) X]

    z2 = X * theta1'
    a2 = sigmoid(z2)

    a2 = [ones(m, 1) a2]
    a3 = sigmoid (a2 * theta2')
    s=0
    for i=1:m
        yvector = int([1:nlables] .== y[i])
        s=s+sum(-yvector' .* log(a3[i,:]) - (1-yvector') .* log (1-a3[i,:]));
    end

    temp1 = theta1
    temp2 = theta2
    temp1[:, 1] = zeros(size(temp1, 1), 1)
    temp2[:, 1] = zeros(size(temp2, 1), 1)

    if !(grad === nothing)
        gptr = pointer(grad)
        for t=1:m
            yvector = int([1:nlables] .== y[t])

            delta3 = zeros(1,nlables)
            delta3 = a3[t,:] - yvector'

            delta2 = zeros (1, hidden_sz)
            delta2 = (theta2' * delta3' )[2:end] .* sigmoidGradient(z2[t,:])'

            theta2_grad = theta2_grad + delta3' * a2[t,:]
            theta1_grad = theta1_grad + delta2 * X[t,:]
        end

        theta1_grad = theta1_grad / m
        theta2_grad = theta2_grad / m

        theta1_grad = theta1_grad + (lambda/m) * temp1
        theta2_grad = theta2_grad + (lambda/m) * temp2

        grad[:] = [theta1_grad[:] ; theta2_grad[:]]
        @assert gptr == pointer(grad)

    end
    return s/m  + (lambda/(2*m))*(sum(sum(temp1 .* temp1)) + sum(sum(temp2 .* temp2)))
end

function randomInitialWeights(l_in, l_out)
    ep = 0.12
    rand(l_out, 1 + l_in) * 2 * ep - ep
end

end


const k_pio2 = Double64(1.5707963267948966, 6.123233995736766e-17)

function atan(x::DoubleFloat{T}) where {T<:IEEEFloat}
   signbit(x) && return -atan(abs(x))
   isinf(x) && return DoubleFloat{T}(pi)/2
   iszero(x) && return x
   y = DoubleFloat{T}(atan(x.hi))
   s, c = sin(y), cos(y)
   c2 = square(c)
   t = s/c
   z = y - (t - x)*c2 # z - (tan(x)-x)*(cos(x)^2)
   s, c = sin(z), cos(z)
   t = s/c
   z = z - (t - x)*c2 # z - (tan(x)-x)*(cos(x)^2)
   return z
end

function atan(y::DoubleFloat{T}, x::DoubleFloat{T}) where {T<:IEEEFloat}
   iszero(y) && iszero(x) && return x
   iszero(x) && return copysign(pio2(DoubleFloat{T}), y)
   iszero(y) && return signbit(x) ? pio1(DoubleFloat{T}) : zero(DoubleFloat{T})
   atanyx = atan(y/x)
   if x > 0
      atanyx
   elseif x < 0
      if y >= 0
         atanyx + T(pi)
      else
         atanyx - T(pi)
      end
   elseif y > 0
       T(pi)/2
   elseif y < 0
      -T(pi)/2
   else
      zero(T)
   end
end

function asin(x::DoubleFloat{T}) where {T<:IEEEFloat}
   abs(x) > 1.0 && throw(DomainError("$x"))
   signbit(x) && return -asin(abs(x))
   y = x
   y = y / (1.0 + sqrt(1.0 - square(y)))
   z = atan(y)
   return DoubleFloat{T}(z.hi+z.hi, z.lo+z.lo)
end

function acos(x::DoubleFloat{T}) where {T<:IEEEFloat}
   abs(x) > 1.0 && throw(DomainError("$x"))
   signbit(x) && return DoubleFloat{T}(onepi - acos(abs(x)))
   y = x
   y = sqrt(1.0 - square(y)) / (1.0 + y)
   z = atan(y)
   return DoubleFloat{T}(z.hi+z.hi, z.lo+z.lo)
end

acsc(x::DoubleFloat{T}) where {T<:IEEEFloat} = asin(inv(x))
asec(x::DoubleFloat{T}) where {T<:IEEEFloat} = acos(inv(x))

function acot(x::DoubleFloat{T}) where {T<:IEEEFloat}
   signbit(x) && return DoubleFloat{T}(onepi - acot(abs(x)))
   iszero(x) && return k_pio2
   z = k_pio2 - atan(abs(x))
   if signbit(x.hi)
      z = -z
   end
   return z
end


const tfn = pso.TestFunctions

@testset "TestFunctions" begin

    @testset "sphere function" begin
        @test tfn.sphere([0]) == 0
        @test tfn.sphere([0.0]) == 0
        @test tfn.sphere(fill(0, 10)) == 0
        @test tfn.sphere([3.0, 4.0]) == 25
    end

    @testset "euclidean_distance function" begin
        @test tfn.euclidean_distance([0]) == 0
        @test tfn.euclidean_distance([0.0]) == 0
        @test tfn.euclidean_distance(fill(0, 10)) == 0
        @test tfn.euclidean_distance([3.0, 4.0]) == 5
    end

    @testset "rastrigin function" begin
        @test tfn.rastrigin([0]) == 0
        @test tfn.rastrigin([0.0]) == 0
        @test tfn.rastrigin(fill(0, 10)) == 0
        @test tfn.rastrigin([0.0, 0.0], A=-2.3) == 0.0
    end

    @testset "ackley's function" begin
        @test tfn.ackley([0,0]) == 0
        @test tfn.ackley([0.0, 0.0]) == 0
    end

    @testset "rosenbrock's function" begin
        @test tfn.rosenbrock([1.0]) == 0
        @test tfn.rosenbrock([1, 1]) == 0
        @test tfn.rosenbrock(fill(1, 10)) == 0.0
    end

    @testset "eggholder function" begin
        @test tfn.eggholder(512, 404.2319) ≈ -959.6407 atol=0.0001
        @test tfn.eggholder([512, 404.2319]) ≈ -959.6407 atol=0.0001
    end

    @testset "beale function" begin
        @test tfn.beale(3, 0.5) == 0
        @test tfn.beale([3, 0.5]) == 0
    end

    @testset "square_movingpeak function" begin
        @test tfn.square_movingpeak([0, 0]) == 0
        @test tfn.square_movingpeak([3, 4]) == 25
        @test tfn.square_movingpeak([1, 1], optimum=[1, 1]) == 0
        @test tfn.square_movingpeak([2, 2], optimum=[1, 1]) ≈ 2
    end

    @testset "ellipse function" begin
        @test tfn.ellipse(-0.5, -1.5) == 0
        @test tfn.ellipse([-0.5, -1.5]) == 0
    end

    @testset "stybliski function" begin
        for i in 1:5
            @test tfn.stybliski(fill(-2.903534, i)) ≈ -39.16616*i atol=0.0001
        end
    end

    @testset "mccormick function" begin
        @test tfn.mccormick(0, 0) == 1
        @test tfn.mccormick([0, 0]) == 1
        @test tfn.mccormick(-0.54719, -1.54719) ≈ -1.9133 atol=0.0001
    end

    @testset "himmelblau function" begin
        @test tfn.himmelblau(-0.270845, -0.923039) ≈ 181.617 atol=0.001
        @test tfn.himmelblau(3, 2) ≈ 0
        @test tfn.himmelblau([3, 2]) ≈ 0
        @test tfn.himmelblau(-2.805118,  3.131312) ≈ 0 atol=0.0001
        @test tfn.himmelblau(-3.77931 , -3.283186) ≈ 0 atol=0.0001
        @test tfn.himmelblau( 3.584428, -1.848126) ≈ 0 atol=0.0001
    end

    @testset "griewank function" begin
        @test tfn.griewank([0, 0]) == 0
    end

    @testset "schafferf6 function" begin
        @test tfn.schafferf6(0, 0) == 0
        @test tfn.schafferf6([0, 0]) == 0
    end

end


#    Problem : GROUP A
#    *********
#    The ENGVAL1 problem.
#    This problem is a sum of 2n-2 groups, n-1 of which contain 2 nonlinear
#    elements
#
#    Original SIF Source: problem 31 in
#    Ph.L. Toint,
#    "Test problems for partially separable optimization and results
#    for the routine PSPMIN",
#    Report 83/4, Department of Mathematics, FUNDP (Namur, B), 1983.
#
#    See also Buckley#172 (p. 52)
#    SIF input: Ph. Toint and N. Gould, Dec 1989.
#
#    Implementation translated from Source:
#    http://eprints.tsu.ge/234/14/Tests%20collection-K-F.pdf
#
#    ENGVAL1.SIF classification OUR2-AN-V-0
#
#    N is the number of variables
#
# Daniel Henderson, 08/2021

f = (x) -> begin
    fx = 0.0
    for i in firstindex(x):lastindex(x)-1
        γ = x[i]^2 + x[i+1]^2
        fx += γ^2 + (3.0 - 4.0x[i])
    end
    return fx
end

g! = (g, x) -> begin
    for i in firstindex(x):lastindex(x)-1
        γ = x[i]^2 + x[i+1]^2
        g[i] += 4.0γ*x[i] - 4.0
        g[i+1] += 4.0γ*x[i+1]
    end
    return g
end

fg! = (g, x) -> begin
    fx = 0.0
    for i in firstindex(x):lastindex(x)-1
        γ = x[i]^2 + x[i+1]^2
        fx += γ^2 + (3.0 - 4.0x[i])
        g[i] += 4.0γ*x[i] - 4.0
        g[i+1] += 4.0γ*x[i+1]
    end
    return fx, g
end

init = (n::Int=5000) -> begin
	n < 2 && @warn("ENGVAL1: number of variables must be ≥ 2")
	n = max(n, 2)

    return n, 2.0*ones(n)
end

TestSet["ENGVAL1"] = UncProgram("ENGVAL1",  f, g!, fg!, init)