using NumCME
using Test
using LinearAlgebra: norm

## Bursting gene model
𝕊 = [[-1, 1, 0] [1, -1, 0] [0, 0, 1] [0, 0, -1]]
x₀ = [1, 0, 0]
k₀₁ = 0.05
k₁₀ = 0.1
λ = 5.0
γ = 1.0

θ = [k₀₁, k₁₀, λ, γ]

α₁ = propensity() do x, p
    p[1]*x[1]
end
α₂ = propensity() do x, p
    p[2]*x[2]
end
α₂tv = propensity((t,p) -> max(0.0, 1.0 - sin(π * t / 2))) do x, p
    p[2]*x[2]
end
α₂tvj = propensity() do t, x, p
    max(0.0, 1.0 - sin(π * t / 2))*p[2]*x[2]
end
α₃ = propensity() do x, p
    p[3]*x[2]
end
α₄ = propensity() do x, p
    p[4]*x[3]
end

propensities_ti = [α₁, α₂, α₃, α₄]
propensities_tv = [α₁, α₂tv, α₃, α₄]
propensities_tvj = [α₁, α₂tvj, α₃, α₄]

𝔛 = StateSpaceSparse(𝕊, x₀)
expand!(𝔛, 2)
𝐀 = FspMatrixSparse(𝔛, propensities_ti, parameters=θ)
@test size(𝐀, 1) == get_state_count(𝔛) + get_sink_count(𝔛)
@test size(𝐀, 2) == get_state_count(𝔛) + get_sink_count(𝔛)
𝐯 = ones(Float64, size(𝐀, 1))
𝐰 = 𝐀(1.0) * 𝐯
@test sum(𝐰) ≈ 0.0 atol = 1.0e-14
𝐰 = 𝐀 * 𝐯
@test sum(𝐰) ≈ 0.0 atol = 1.0e-14

# Test mat-vec for time-varying matrix
𝔛 = StateSpaceSparse(𝕊, x₀)
expand!(𝔛, 2)
A1 = FspMatrixSparse(𝔛, propensities_tv, parameters=θ)
@test size(A1, 1) == get_state_count(𝔛) + get_sink_count(𝔛)
@test size(A1, 2) == get_state_count(𝔛) + get_sink_count(𝔛)
𝐯 = ones(Float64, size(A1, 1))
w1 = A1(1.0) * 𝐯
@test sum(w1) ≈ 0.0 atol = 1.0e-14
w1 = A1 * 𝐯
@test sum(w1) ≈ 0.0 atol = 1.0e-14

A2 = FspMatrixSparse(𝔛, propensities_tvj, parameters=θ)
w2 = A2(1.0) * 𝐯
@test sum(w2) ≈ 0.0 atol = 1.0e-14
w2 = A2 * 𝐯

@test sum(w2) ≈ 0.0 atol = 1.0e-14
@test norm(w1 -w2) ≈ 0


include("euler/euler.jl")
using .BigIntegers: BigInteger

compute(n::Int, list_numbers::String)::SubString = sum(map(BigInteger, split(list_numbers, "\n")))[1:n]


include("euler/euler.jl")
using .BigIntegers: BigInteger
using .Math: fast_max
using .Numbers: digits_sum
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

function compute(n::Int)::Int
    result = 0
    for a ∈ n - 5:n - 1
        number = BigInteger(a)
        for b ∈ 1:n - 1
            number *= BigInteger(a)
            result = fast_max(result, digits_sum(number))
        end
    end
    return result
end

compute(100)

@benchmark compute(100)


using SparseArrays
using LinearAlgebra
using Logging


module TestStructuredCospans
using Test

using Catlab, Catlab.Theories, Catlab.Graphs, Catlab.CategoricalAlgebra
using Catlab.Graphs.BasicGraphs: AbstractGraph, TheoryGraph

# Structured cospans of C-sets
##############################

const OpenGraphOb, OpenGraph = OpenCSetTypes(Graph, :V)

# Directed fork as open graph with one input and two outputs.
g0 = Graph(4)
add_edges!(g0, [1,2,2], [2,3,4])
g = OpenGraph(g0, FinFunction([1],4), FinFunction([3,4],4))
@test apex(g) == g0
@test dom.(legs(g)) == [Graph(1), Graph(2)]
@test feet(g) == [FinSet(1), FinSet(2)]
@test dom(g) == OpenGraphOb(FinSet(1))
@test codom(g) == OpenGraphOb(FinSet(2))

# Opposite of previous graph.
h0 = Graph(4)
add_edges!(h0, [1,2,3], [3,3,4])
h = OpenGraph(h0, FinFunction([1,2],4), FinFunction([4],4))

# Structured multicospans.
g′ = OpenGraph(g0, FinFunction([1],4), FinFunction([3],4), FinFunction([4],4))
@test feet(g′) == fill(FinSet(1), 3)
@test bundle_legs(g′, [1, (2,3)]) == g

# Category
#---------

# Composition.
k = compose(g, h)
@test dom(k) == dom(g)
@test codom(k) == codom(h)
k0 = apex(k)
@test (nv(k0), ne(k0)) == (6, 6)
@test (src(k0), tgt(k0)) == ([1,2,2,3,4,5], [2,3,4,5,5,6])

# Identities.
a = OpenGraphOb(FinSet(3))
@test dom(id(a)) == a
@test codom(id(a)) == a
@test compose(g, id(codom(g))) == g
@test compose(id(dom(g)), g) == g

# Symmetric monoidal category
#----------------------------

a, b = OpenGraphOb(FinSet(3)), OpenGraphOb(FinSet(2))
@test otimes(a, b) == OpenGraphOb(FinSet(5))
@test munit(OpenGraphOb) == OpenGraphOb(FinSet(0))
@test dom(braid(a, b)) == a⊗b
@test codom(braid(b, a)) == b⊗a
@test force(braid(a,b)⋅braid(b,a)) == force(id(a⊗b))

k = otimes(g, h)
@test dom(k) == dom(g)⊗dom(h)
@test codom(k) == codom(g)⊗codom(h)
@test (nv(apex(k)), ne(apex(k))) == (8, 6)

# Hypergraph category
#--------------------

@test compose(create(a), mcopy(a)) == dunit(a)
@test compose(mmerge(a), delete(a)) == dcounit(a)

@test dom(dagger(g)) == codom(g)
@test codom(dagger(g)) == dom(g)

# Structured cospan of attributed C-sets
########################################

# Non-attributed boundary
#------------------------

const OpenWeightedGraphOb, OpenWeightedGraph = OpenACSetTypes(WeightedGraph, :V)

g0 = WeightedGraph{Float64}(2)
add_edge!(g0, 1, 2, weight=1.5)
g = OpenWeightedGraph{Float64}(g0, FinFunction([1],2), FinFunction([2],2))

h0 = WeightedGraph{Float64}(3)
add_edges!(h0, [1,1], [2,3], weight=[1.0,2.0])
h = OpenWeightedGraph{Float64}(h0, FinFunction([1],3), FinFunction([2,3],3))
@test dom.(legs(h)) == [WeightedGraph{Float64}(1), WeightedGraph{Float64}(2)]
@test feet(h) == [FinSet(1), FinSet(2)]
@test dom(h) == OpenWeightedGraphOb{Float64}(FinSet(1))
@test codom(h) == OpenWeightedGraphOb{Float64}(FinSet(2))

k = compose(g, h)
k0 = apex(k)
@test src(k0) == [1,2,2]
@test tgt(k0) == [2,3,4]
@test subpart(k0, :weight) == [1.5, 1.0, 2.0]

# Attributed boundary
#--------------------

@present TheoryVELabeledGraph <: TheoryGraph begin
  Label::AttrType
  vlabel::Attr(V,Label)
  elabel::Attr(E,Label)
end

@acset_type VELabeledGraph(TheoryVELabeledGraph, index=[:src,:tgt]) <: AbstractGraph
const VELGraph = VELabeledGraph

const OpenVELGraphOb, OpenVELGraph = OpenACSetTypes(VELGraph, :V)

g0 = VELGraph{Symbol}()
add_vertices!(g0, 2, vlabel=[:x,:y])
add_edge!(g0, 1, 2, elabel=:f)
g = OpenVELGraph{Symbol}(g0, FinFunction([1],2), FinFunction([2],2))

h0 = VELGraph{Symbol}()
add_vertices!(h0, 3, vlabel=[:y,:w,:z])
add_edges!(h0, [1,1], [2,3], elabel=[:g,:h])
h = OpenVELGraph{Symbol}(h0, FinFunction([1],3), FinFunction([2,3],3))
lfoot, rfoot = feet(h)
@test nparts(rfoot, :V) == 2
@test subpart(rfoot, :vlabel) == [:w,:z]
@test dom(h) == OpenVELGraphOb{Symbol}(FinSet(1), vlabel=:y)
@test codom(h) == OpenVELGraphOb{Symbol}(FinSet(2), vlabel=[:w,:z])

# Category
#---------

k = compose(g, h)
k0 = apex(k)
@test src(k0) == [1,2,2]
@test tgt(k0) == [2,3,4]
@test subpart(k0, :vlabel) == [:x, :y, :w, :z]
@test subpart(k0, :elabel) == [:f, :g, :h]

# Incompatible attributes.
set_subpart!(h0, 1, :vlabel, :y′)
h = OpenVELGraph{Symbol}(h0, FinFunction([1],3), FinFunction([2,3],3))
@test_throws ErrorException compose(g, h)

# Symmetric monoidal category
#----------------------------

k = otimes(g, h)
@test dom(k) == dom(g)⊗dom(h)
@test codom(k) == codom(g)⊗codom(h)
@test (nv(apex(k)), ne(apex(k))) == (5, 3)

a = OpenVELGraphOb{Symbol}(FinSet(2), vlabel=[:u,:v])
b = OpenVELGraphOb{Symbol}(FinSet(3), vlabel=[:x,:y,:z])
@test dom(braid(a, b)) == a⊗b
@test codom(braid(a, b)) == b⊗a

# Structured cospan of multiply attributed C-sets
#################################################

# Attributed boundary
#--------------------

@present TheoryMultiplyAttributedGraph <: TheoryGraph begin
  Length::AttrType
  Size::AttrType
  Label::AttrType
  vlabel::Attr(V,Label)
  elabel::Attr(E,Label)
  vsize::Attr(V, Size)
  elength::Attr(E, Length)
end

@acset_type MAGraph(TheoryMultiplyAttributedGraph, index=[:src,:tgt]) <: AbstractGraph
const OpenMAGraphOb, OpenMAGraph = OpenACSetTypes(MAGraph, :V)

g0 = @acset MAGraph{Int, Float64, Symbol} begin
  V = 5
  E = 3
  vlabel = [:a, :b, :c, :d, :e]
  elabel = [:a, :b, :c]
  vsize = [0.0, 1.0, 3.0, 3.0, 4.0]
  elength = [1,2,3]
  src = [1, 2, 3]
  tgt = [3, 4, 5]
end
g = OpenMAGraph{Int, Float64, Symbol}(g0, FinFunction([1],5), FinFunction([3,4],5))

g′ = OpenMAGraph{Int, Float64, Symbol}(g0, FinFunction([1],5), FinFunction([3],5), FinFunction([4],5))
@test bundle_legs(g′, [1, (2,3)]) == g

end


# Second-Order upwind central scheme
# Based on:
# Kurganov A., Noelle S., Petrova G., Semidiscrete Central-Upwind schemes
# for hyperbolic Conservation Laws and Hamilton-Jacobi Equations. SIAM. Sci Comput,
# Vol 23, No 3m pp 707-740. 2001

struct FVCUScheme{ltype <: AbstractSlopeLimiter} <: AbstractFVAlgorithm
  slopeLimiter :: ltype
end

function FVCUScheme(;slopeLimiter=GeneralizedMinmodLimiter())
  FVCUScheme(slopeLimiter)
end

function update_flux_value(uold,∇u,j,mesh,Flux, alg::FVCUScheme,nonscalar::Bool)
    uminus=cellval_at_left(j,uold,mesh)+0.5*cellval_at_left(j,∇u,mesh)
    uplus=cellval_at_right(j,uold,mesh)-0.5*cellval_at_right(j,∇u,mesh)

    if nonscalar
      λm = eigvals(Flux.Dflux(uminus))
      λp = eigvals(Flux.Dflux(uplus))
      aa_plus=maximum((maximum(λm), maximum(λp),0))
      aa_minus=minimum((minimum(λm), minimum(λp),0))
    else
      λm = Flux.Dflux(uminus)
      λp = Flux.Dflux(uplus)
      aa_plus=maximum((λm, λp,0))
      aa_minus=minimum((λm, λp,0))
    end
    # Update numerical fluxes
    if abs(aa_plus-aa_minus) < 1e-8
      return 0.5*(Flux(uminus)+Flux(uplus))
    else
      return (aa_plus*Flux(uminus)-aa_minus*Flux(uplus))/(aa_plus-aa_minus) +
      (aa_plus*aa_minus)/(aa_plus-aa_minus)*(uplus - uminus)
    end
end

# Dissipation Reduced Central upwind Scheme: Second-Order
# Based on:
# Kurganov A., Lin C., On the reduction of Numerical Dissipation in Central-Upwind
# Schemes, Commun. Comput. Phys. Vol 2. No. 1, pp 141-163, Feb 2007.

struct FVDRCUScheme{ltype <: AbstractSlopeLimiter} <: AbstractFVAlgorithm
  slopeLimiter :: ltype
end

function FVDRCUScheme(;slopeLimiter=GeneralizedMinmodLimiter())
  FVDRCUScheme(slopeLimiter)
end

function update_flux_value(uold,∇u,j,mesh,Flux, alg::FVDRCUScheme,nonscalar::Bool)
  uminus=cellval_at_left(j,uold,mesh)+0.5*cellval_at_left(j,∇u,mesh)
  uplus=cellval_at_right(j,uold,mesh)-0.5*cellval_at_right(j,∇u,mesh)

  if nonscalar
    λm = eigvals(Flux.Dflux(uminus))
    λp = eigvals(Flux.Dflux(uplus))
    aa_plus=maximum((λm[end], λp[end],0))
    aa_minus=minimum((λm[1], λp[1],0))
  else
    λm = Flux.Dflux(uminus)
    λp = Flux.Dflux(uplus)
    aa_plus=maximum((λm, λp,0))
    aa_minus=minimum((λm, λp,0))
  end

  # Update numerical fluxes
  if abs(aa_plus-aa_minus) < 1e-8
    return 0.5*(Flux(uminus)+Flux(uplus))
  else
    flm = Flux(uminus); flp = Flux(uplus)
    wint = 1/(aa_plus-aa_minus)*(aa_plus*uplus-aa_minus*uminus-(flp-flm))
    qj = minmod.((uplus-wint)/(aa_plus-aa_minus),(wint-uminus)/(aa_plus-aa_minus))
    return (aa_plus*flm-aa_minus*flp)/(aa_plus-aa_minus) +
            (aa_plus*aa_minus)*((uplus - uminus)/(aa_plus-aa_minus) - qj)
  end
end

##################3
FVCUSchemes = Union{FVCUScheme,FVDRCUScheme}
"""
compute_fluxes!(hh, Flux, uold, mesh, dt, M, alg::FVCUScheme, ::Type{Val{true}})
Numerical flux of Second-Order upwind central scheme in 1D
"""
function compute_fluxes!(fluxes, Flux, uold, mesh, dt, alg::T, nonscalar::Bool, ::Type{Val{true}}) where {T <: FVCUSchemes}
  slopeLimiter = alg.slopeLimiter
  #update vector
  # 1. slopes
  ∇u = compute_slopes(uold, mesh, slopeLimiter, nonscalar, Val{true})

  Threads.@threads for j in node_indices(mesh)
      if nonscalar
          fluxes[:,j] .= update_flux_value(uold,∇u,j,mesh,Flux,alg,nonscalar)
      else
          fluxes[j] = update_flux_value(uold,∇u,j,mesh,Flux,alg,nonscalar)
      end
  end
end

function compute_fluxes!(fluxes, Flux, uold, mesh, dt, alg::T, nonscalar::Bool, ::Type{Val{false}}) where {T <: FVCUSchemes}
  slopeLimiter = alg.slopeLimiter
  #update vector
  # 1. slopes
  ∇u = compute_slopes(uold, mesh, slopeLimiter, nonscalar, Val{false})

  for j in node_indices(mesh)
      if nonscalar
          fluxes[:,j] .= update_flux_value(uold,∇u,j,mesh,Flux,alg,nonscalar)
      else
          fluxes[j] = update_flux_value(uold,∇u,j,mesh,Flux,alg,nonscalar)
      end
  end
end

# Dissipation Reduced Central upwind Scheme: Fifth-Order
# Based on:
# Kurganov A., Lin C., On the reduction of Numerical Dissipation in Central-Upwind
# Schemes, Commun. Comput. Phys. Vol 2. No. 1, pp 141-163, Feb 2007.
# Kurganov, Liu, New adaptive artificial viscosity method for hyperbolic systems
# of conservation laws

struct FVDRCU5Scheme{ltype <: AbstractSlopeLimiter} <: AbstractFVAlgorithm
  slopeLimiter :: ltype
end

function FVDRCU5Scheme(;slopeLimiter=GeneralizedMinmodLimiter())
  FVDRCU5Scheme(slopeLimiter)
end

function compute_fluxes!(fluxes, Flux, uold, mesh, dt, alg::FVDRCU5Scheme, nonscalar::Bool, ::Type{Val{true}})
  #update vector
  Threads.@threads for j in node_indices(mesh)
      if nonscalar
          fluxes[:,j] .= update_flux_value(uold,j,mesh,Flux,alg,nonscalar)
      else
          fluxes[j] = update_flux_value(uold,j,mesh,Flux,alg,nonscalar)
      end
  end
end

function compute_fluxes!(fluxes, Flux, uold, mesh, dt, alg::FVDRCU5Scheme, nonscalar::Bool, ::Type{Val{false}})
  #update vector
  for j in node_indices(mesh)
      if nonscalar
          fluxes[:,j] .= update_flux_value(uold,j,mesh,Flux,alg,nonscalar)
      else
          fluxes[j] = update_flux_value(uold,j,mesh,Flux,alg,nonscalar)
      end
  end
end

function update_flux_value(uold,j,mesh,Flux,alg::FVDRCU5Scheme, nonscalar)
    slopeLimiter = alg.slopeLimiter
    # A fifth-order piecewise polynomial reconstruction
    @inbounds ulll=cellval_at_left(j-2,uold,mesh)
    @inbounds ull=cellval_at_left(j-1,uold,mesh)
    @inbounds ul=cellval_at_left(j,uold,mesh)
    @inbounds ur=cellval_at_right(j,uold,mesh)
    @inbounds urr=cellval_at_right(j+1,uold,mesh)
    @inbounds urrr=cellval_at_right(j+2,uold,mesh)

    uminus = 1/60*(2*ulll-13*ull+47*ul+27*ur-3*urr)
    uplus = 1/60*(-3*ull+27*ul+47*ur-13*urr+2*urrr)

    #Remark: some extrange bug wth eigvals force me to use Lapack
    if nonscalar
      λm = sort(LAPACK.geev!('N','N',Array(Flux.Dflux(uminus)))[1])
      λp = sort(LAPACK.geev!('N','N',Array(Flux.Dflux(uplus)))[1])
      aa_plus=maximum((λm[end], λp[end],0))
      aa_minus=minimum((λm[1], λp[1],0))
    else
      λm = Flux.Dflux(uminus)
      λp = Flux.Dflux(uplus)
      aa_plus=maximum((λm, λp,0))
      aa_minus=minimum((λm, λp,0))
    end
    # Update numerical fluxes
    if abs(aa_plus-aa_minus) < 1e-8
      return 0.5*(Flux(uminus)+Flux(uplus))
    else
      flm = Flux(uminus); flp = Flux(uplus)
      wint = 1/(aa_plus-aa_minus)*(aa_plus*uplus-aa_minus*uminus-
      (flp-flm))
      qj = minmod.((uplus-wint)/(aa_plus-aa_minus),(wint-uminus)/(aa_plus-aa_minus))
      return (aa_plus*flm-aa_minus*flp)/(aa_plus-aa_minus) +
      (aa_plus*aa_minus)*((uplus - uminus)/(aa_plus-aa_minus) - qj)
    end
end


struct NNStatistics
    units::String
    arealUnits::String
    area::Polygon
    observedMeanDistance::Real
    expectedMeanDistance::Real
    nearestNeighbourIndex::Real
    pointsCount::Real
    zScore::Real
end

"""
    analysis(data::F, area::Union{P, Nothing}=nothing, units::String="kilometers") where {F <: AbstractFeatureCollection, P <: AbstractPolygon}

Nearest Neighbor Analysis calculates an index based the average distances
between points in the dataset, thereby providing inference as to whether the
data is clustered, dispersed, or randomly distributed within the study area.

It returns a Polygon of the study area, with the results of
the analysis attached as part of of the `nearestNeighborAnalysis` property
of the study area's `properties`. The attached
[_z_-score](https://en.wikipedia.org/wiki/Standard_score) indicates how many
standard deviations above or below the expected mean distance the data's
observed mean distance is. The more negative, the more clustered. The more
positive, the more evenly dispersed. A _z_-score between -2 and 2 indicates
a seemingly random distribution. That is, within _p_ of less than 0.05, the
distribution appears statistically significantly neither clustered nor
dispersed.

**Remarks**

- Though the analysis will work on any FeatureCollection type, it
works best with Point collections.

- This analysis is **very** sensitive to the study area provided.
If no Polygon is passed as the study area, the function draws a box
around the data, which may distort the findings. This analysis works best
with a bounded area of interest within with the data is either clustered,
dispersed, or randomly distributed. For example, a city's subway stops may
look extremely clustered if the study area is an entire state. On the other
hand, they may look rather evenly dispersed if the study area is limited to
the city's downtown.
"""
function analysis(data::F, area::Union{P, Nothing}=nothing, units::String="kilometers") where {F <: AbstractFeatureCollection, P <: AbstractPolygon}
    area == nothing && (area = bbox_polygon(bbox(data)[1]))

    feats = []

    for ft in data.features
        push!(feats, centroid(ft.geometry))
    end

    n = length(feats)

    obsDist = []
    for (index, p) in enumerate(feats)
        other = filter((f, i) -> i !== index, feats)
        push!(obsDist, distance(p, nearestpoint(p, other).coordinates, units))
    end

    reduce((a, b) -> a + b, obsDist)

    popDensity = n / convert_area(area(area), "meters", units)
    expDist = 1 / (2 * sqrt(popDensity))
    var = 0.26136 / (sqrt(n * popDensity))

    return NNStatistics(units, units * "²", area, obsDist, expDist, obsDist / expDist, n, (obsDist - expDist) / var)

end


using Test
import ShuffleProofs: decode, encode, Tree, unmarshal, marshal_publickey, unmarshal_publickey, unmarshal_privatekey
import CryptoGroups: ElGamal, PGroup, Dec

CIPHERTEXT_FILE = "$(@__DIR__)/../../ref/demo/ciphertexts"
CIPHERTEXTOUT_FILE = "$(@__DIR__)/../../ref/demo/ciphertextsout"
PUBLIC_KEY = "$(@__DIR__)/../../ref/demo/publicKey"
PRIVATE_KEY = "$(@__DIR__)/../../ref/privateKey"

y, g = let
    bytes = read(PUBLIC_KEY)
    tree = decode(bytes)
    unmarshal_publickey(tree)
end

sk, g′ = let
    bytes = read(PRIVATE_KEY)
    tree = decode(bytes)
    unmarshal_privatekey(tree)
end

@test g′ == g
@test y == g^sk # This way we can be sure that correct inputs have been choosen.

# It would be cool to also to input my own ciphertexts!

G = typeof(g)

#𝓖 = group(g)

𝐞 = let
    bytes = read(CIPHERTEXT_FILE)
    tree = decode(bytes)
    𝐚, 𝐛 = convert(Tuple{Vector{BigInt}, Vector{BigInt}}, tree)
    ElGamal{G}(𝐚, 𝐛)
end

𝐞′ = let
    bytes = read(CIPHERTEXTOUT_FILE)
    tree = decode(bytes)
    convert(ElGamal{G}, tree)
end

dec = Dec(sk)

𝐦 = g .^ (2:11)
@test dec(𝐞) == 𝐦

𝐦′ = dec(𝐞′)
@test sort(𝐦) == sort(𝐦′)


struct Lwplsr
    X::Array{Float64}
    Y::Array{Float64}
    fm
    metric::String
    h::Real
    k::Int
    nlv::Int
    tol::Real
    verbose::Bool
end

"""
    lwplsr(X, Y; nlvdis, metric, h, k, nlv, tol = 1e-4, verbose = false)
k-Nearest-Neighbours locally weighted partial least squares regression (KNN-LWPLSR).
* `X` : X-data.
* `Y` : Y-data.
* `nlvdis` : Number of latent variables (LVs) to consider in the global PLS
    used for the dimension reduction before computing the dissimilarities.
    If `nlvdis = 0`, there is no dimension reduction.
* `metric` : Type of dissimilarity used to select the neighbors.
    Possible values are "eucl" (default; Euclidean distance)
    and "mahal" (Mahalanobis distance).
* `h` : A scalar defining the shape of the weight function. Lower is h,
    sharper is the function. See function `wdist`.
* `k` : The number of nearest neighbors to select for each observation to predict.
* `nlv` : Nb. latent variables (LVs).
* `tol` : For stabilization when very close neighbors.
* `verbose` : If true, fitting information are printed.

Function `lwplsr` fits kNN-LWPLSR models (Lesnoff et al., 2020).
The function uses functions `getknn`, `locw` and a PLSR function;
see the code for details. Many other variants of kNN-LWPLSR pipelines can be built.

The general principles of the method are as follows.

LWPLSR is a particular case of weighted PLSR (WPLSR) (e.g. Schaal et al. 2002).
In WPLSR, a priori weights, different from the usual 1/n (standard PLSR),
are given to the n training observations. These weights are used for calculating
(i) the scores and loadings of the WPLS and (ii) the regression model that fits
(by weighted least squares) the Y-response(s) to the WPLS scores.
The specificity of LWPLSR (compared to WPLSR) is that the weights are computed
from dissimilarities (e.g. distances) between the new observation to predict
and the training observations ("L" in LWPLSR comes from "localized").
Note that in LWPLSR the weights and therefore the fitted WPLSR model
change for each new observation to predict.

In the original LWPLSR, all the n training observations are used for each
observation to predict (e.g. Sicard & Sabatier 2006, Kim et al 2011).
This can be very time consuming, in particular for large n.
A faster and often more efficient strategy is to preliminary select,
in the training set, a number of `k` nearest neighbors to the observation to predict
(= "weighting 1") and then to apply LWPLSR only to this pre-selected
neighborhood (= "weighting 2"). This strategy corresponds to a kNN-LWPLSR
and is the one implemented in function `lwplsr`.

In `lwplsr`, the dissimilarities used for weightings 1 and 2 are
computed from the raw X-data or after a dimension reduction,
depending on argument `nlvdis`. In the last case, global PLS2 scores (LVs) are
computed from {`X`, `Y`} and the dissimilarities are computed over these scores.

In general, for high dimensional X-data, using the Mahalanobis distance requires
preliminary dimensionality reduction of the data.

## References

Kim, S., Kano, M., Nakagawa, H., Hasebe, S., 2011. Estimation of active
pharmaceutical ingredients content using locally weighted partial least squares
and statistical wavelength selection. Int. J. Pharm., 421, 269-274.

Lesnoff, M., Metz, M., Roger, J.-M., 2020. Comparison of locally weighted PLS
strategies for regression and discrimination on agronomic NIR data.
Journal of Chemometrics, e3209. https://doi.org/10.1002/cem.3209

Schaal, S., Atkeson, C., Vijayamakumar, S. 2002. Scalable techniques from nonparametric
statistics for the real time robot learning. Applied Intell., 17, 49-60.

Sicard, E. Sabatier, R., 2006. Theoretical framework for local PLS1 regression
and application to a rainfall data set. Comput. Stat. Data Anal., 51, 1393-1410.
"""
function lwplsr(X, Y; nlvdis, metric, h, k, nlv, tol = 1e-4, verbose = false)
    X = ensure_mat(X)
    Y = ensure_mat(Y)
    if nlvdis == 0
        fm = nothing
    else
        fm = plskern(X, Y; nlv = nlvdis)
    end
    return Lwplsr(X, Y, fm, metric, h, k, nlv, tol, verbose)
end

"""
    predict(object::Lwplsr, X)
Compute the Y-predictions from the fitted model.
* `object` : The fitted model.
* `X` : X-data for which predictions are computed.
"""
function predict(object::Lwplsr, X; nlv = nothing)
    X = ensure_mat(X)
    m = size(X, 1)
    a = object.nlv
    isnothing(nlv) ? nlv = a : nlv = (max(minimum(nlv), 0):min(maximum(nlv), a))
    # Getknn
    if isnothing(object.fm)
        res = getknn(object.X, X; k = object.k, metric = object.metric)
    else
        res = getknn(object.fm.T, transform(object.fm, X); k = object.k, metric = object.metric)
    end
    listw = copy(res.d)
    for i = 1:m
        w = wdist(res.d[i]; h = object.h)
        w[w .< object.tol] .= object.tol
        listw[i] = w
    end
    # End
    pred = locwlv(object.X, object.Y, X;
        listnn = res.ind, listw = listw, fun = plskern, nlv = nlv,
        verbose = object.verbose).pred
    (pred = pred, listnn = res.ind, listd = res.d, listw = listw)
end


NumericalTricks.volume(eM::eMesh{T1,Nothing}) where {T1} = error("meshes without tets have no volume")
function NumericalTricks.volume(eM::eMesh{T1,Tet}) where {T1}
    vol = 0.0
    for k = 1:n_tet(eM)
        vol += volume(eM.point[eM.tet[k]])
    end
    return vol
end

NumericalTricks.area(eM::eMesh{Nothing,T2}) where {T2} = error("meshes without triangles have no area")
function NumericalTricks.area(eM::eMesh{Tri,T2}) where {T2}
    area_cum = 0.0
    for k = 1:n_tri(eM)
        area_cum += area(vertex_pos_for_tri_ind(eM, k))
    end
    return area_cum
end

NumericalTricks.area(a::BB)   where {BB <: BoundingBox} = 8 * dot(a.e, SVector{3,Float64}(a.e[2], a.e[3], a.e[1]))
NumericalTricks.volume(a::BB) where {BB <: BoundingBox} = 8 * prod(a.e)


# These tests have been adapted from Criterim clojure library

macro fact_max_err(expected, actual, max_err)
    :(@fact abs(($expected) - ($actual)) < $max_err => true)
end

facts("bootstrap estimate") do
    C = Criterion
    testdata = rand(10_000)
    est_mean  = C.bootstrap_bca(testdata, mean, 0.05, 1000)
    @fact (est_mean.point  - 0.5) < 1e-2 => true
    @fact_max_err 0.5 est_mean.lbound  0.02
    @fact_max_err 0.5 est_mean.ubound  0.02
    @fact C.scale_bootstrap_estimate((4.5, 5.5), 0.1) => (0.45, 0.55)
end

facts("bootstrap bca") do
    C = Criterion
    context("mean") do
        test_data = ones(20)
        est_mean  = C.bootstrap_bca(test_data, mean, 0.05, 100)
        @fact est_mean.point  => roughly(1.0)
        @fact est_mean.lbound => roughly(1.0)
        @fact est_mean.ubound => roughly(1.0)
    end

    context("variance") do
        test_data = ones(20)
        est_var   = C.bootstrap_bca(test_data, std, 0.05, 100)
        @fact est_var.point  => roughly(0.0)
        @fact est_var.lbound => roughly(0.0)
        @fact est_var.ubound => roughly(0.0)
    end
end

facts("normal quantile") do
    nq = Criterion.normal_quantile

    @fact nq(0.5001) > 0 => true
    @fact nq(0.4999) < 0 => true
    @fact 2e-8 < (nq(0.999) - nq(0.001)) => true
    let max_error = 1e-7
        @fact_max_err  1.2815515655446   nq(0.9)  max_error
        @fact_max_err  0.674489750196082 nq(0.75) max_error
        @fact_max_err -1.03643338949379  nq(0.15) max_error
        @fact_max_err -2.32634787404084  nq(0.01) max_error
    end
end

facts("normal cdf test") do
    normal_cdf = Criterion.normal_cdf
    max_error = 1.5e-7
    @fact_max_err 0.99865010196837  normal_cdf(3.0) max_error
    @fact_max_err 0.977249868051821 normal_cdf(2.0) max_error
    @fact_max_err 0.841344746068543 normal_cdf(1.0) max_error
    @fact_max_err 0.691462461274013 normal_cdf(0.5) max_error
    @fact_max_err 0.5 normal_cdf(0.0) max_error
    @fact_max_err 0.158655253931457 normal_cdf(-1.0) max_error
    @fact_max_err 0.00134989803163009 normal_cdf(-3.0) max_error
end


@testset "Example 1" begin
    Mod.@polyvar x y
    p = 2x + 3.0x*y^2 + y
    @test differentiate(p, x) == 2 + 3y^2
    @test differentiate.(p, (x, y)) == (2 + 3y^2, 6x*y + 1)
    @test p((x, y)=>(y, x)) == 2y + 3y*x^2 + x
    @test subs(p, y=>x^2) == 2x + 3x^5 + x^2
    @test p(x=>1, y=>2) == 2 + 3*4 + 2
end


# Min Trust LAMB (https://github.com/Smerity/pytorch-lamb/blob/master/pytorch_lamb/lamb.py)
#
# It has been proposed in Large Batch Optimization for Deep Learning: Training BERT in 76 minutes.
# https://arxiv.org/abs/1904.00962
#
# Adapted from BertAdam https://github.com/OsmanMutlu/BERT.jl/blob/master/src/optimizer.jl

import Knet: update!
using AutoGrad: full
using LinearAlgebra
using CUDA

warmup_cosine(x, warmup=0.002) = x < warmup ? x/warmup : 0.5 * (1.0 + cos(π * x))
warmup_constant(x, warmup=0.002) = x < warmup ? x/warmup : 1.0
warmup_linear(x, warmup=0.002) = x < warmup ? x/warmup : 1.0 - x

mutable struct MinTrustLAMB
    lr::AbstractFloat
    beta1::AbstractFloat
    beta2::AbstractFloat
    eps::AbstractFloat
    w_decay::AbstractFloat
    min_trust::AbstractFloat
    gclip::AbstractFloat
    t::Int
    schedule
    warmup
    t_total
    fstm
    scndm
end

MinTrustLAMB(; min_trust=0.25, lr=1e-3, gclip=0.25, beta1=0.9, beta2=0.999, eps=1e-6, w_decay=0.0, schedule="warmup_constant", warmup=-1, t_total=-1) = MinTrustLAMB(lr, beta1, beta2, eps, w_decay, min_trust, gclip, 0, schedule, warmup, t_total, nothing, nothing)

function initlamb!(model, t_total; min_trust=0.25, lr=0.001, warmup=0.1, schedule="warmup_constant")
    for par in params(model)
        if length(size(value(par))) === 1
            par.opt = MinTrustLAMB(;min_trust=min_trust, lr=lr, warmup=warmup, schedule=schedule, t_total=t_total, w_decay=1.2e-6)
        else
            par.opt = MinTrustLAMB(;min_trust=min_trust, lr=lr, warmup=warmup, schedule=schedule, t_total=t_total)
        end
    end
end

for T in (Array{Float32},Array{Float64},KnetArray{Float32},KnetArray{Float64},CuArray{Float32},CuArray{Float64})
    @eval begin
        function update!(w::$T, g, p::MinTrustLAMB)
            Knet.Train20.gclip!(g, p.gclip)
            g = full(g)
            if p.fstm===nothing; p.fstm=zero(w); p.scndm=zero(w); end
            p.t += 1

            # Decay the first and second moment running average coefficient
            # scndm : v_t, fstm : m_t
            lmul!(p.beta1, p.fstm)
            axpy!(1-p.beta1, g, p.fstm)
            lmul!(p.beta2, p.scndm)
            axpy!(1-p.beta2, g .* g, p.scndm)

            # Get Learning rate from scheduler
            if p.t_total !== -1
                schedule_func = eval(Meta.parse(p.schedule))
                lr_scheduled = p.lr * schedule_func(p.t/p.t_total, p.warmup)
            else
                lr_scheduled = p.lr
            end

            # Calculate Opt. Step
            if p.w_decay > 0.0
                adam_step = (p.fstm ./ (sqrt.(p.scndm) .+ p.eps)) .+ (p.w_decay * w)
            else
                adam_step = (p.fstm ./ (sqrt.(p.scndm) .+ p.eps))
            end

            # Calculate Trust Ratio
            weight_norm = clamp(norm(w), 0, 10)
            adam_norm = norm(adam_step)
            if weight_norm == 0 || adam_norm == 0
                trust_ratio = 1
            else
                trust_ratio = weight_norm / adam_norm
            end
            trust_ratio = max(trust_ratio, p.min_trust)

            # Update weight matrix
            axpy!(-lr_scheduled * trust_ratio, adam_step, w)
        end
    end
end;

nothing


"""
    calc_deviance_resids(;model, df, y)

Return deviance residuals (unstandardised) from a fitted linear model.

...
# Arguments
- `model`: Variable containing a fitted linear model (e.g. general linear model [GLM]).
- `df::DataFrame=df`: DataFrame array containing response and predictor variable(s)/covariates.
- `y`: Column name from `df` containing the response variable.
...

"""
function calc_deviance_resids(;model, df, y)

    ###################################################
    # Section #1: Extract response vector and model residuals
    ###################################################

    # Extract fitted/predicted values from model object
    pred = GLM.predict(model)

    # Extract vector containing response variable
    response_y = df[!, y]

    ###################################################
    # Section #2: Calculate deviance residuals
    #             - requires conditional statements to
    #               loop through different data distributions
    ###################################################

    # Extract statistical distribution specified in model
    distType = typeof(model).parameters[1].parameters[1].parameters[2]
    distType = string.(distType)

    # Use if-elseif-else statements to calculate deviance residuals depending
    # on which errors were specified during model specification

        # If model was specified with Poisson errors
        if distType == "Poisson{Float64}"
            sign.(response_y .- pred) .* sqrt.(GLM.devresid.(Distributions.Poisson(), response_y, pred))

        # If model was specified with negative binomial errors
        elseif distType == "NegativeBinomial{Float64}"
            sign.(response_y .- pred) .* sqrt.(GLM.devresid.(Distributions.NegativeBinomial(), response_y, pred))

        elseif distType == "Normal{Float64}"
            sign.(response_y .- pred) .* sqrt.(GLM.devresid.(Distributions.Normal(), response_y, pred))

        # If model was specified with Gamma errors
        elseif distType == "Gamma{Float64}"
            sign.(response_y .- pred) .* sqrt.(GLM.devresid.(Distributions.Gamma(), response_y, pred))

        # If model was specified with Bernoulli errors
        elseif distType == "Bernoulli{Float64}"
            sign.(response_y .- pred) .* sqrt.(GLM.devresid.(Distributions.Bernoulli(), response_y, pred))

        # If model was specified with Binomial errors
        elseif distType == "Binomial{Float64}"
            sign.(response_y .- pred) .* sqrt.(GLM.devresid.(Distributions.Binomial(), response_y, pred))

        # If model was specified with unsupported error distribution (e.g. tweedie)
        else
            println("ERROR: Unsupported error distribution specified.
                    Please refer to GLM.jl documentation for supported error distributions.")

        end
end


@testset "Sphere" begin
    # Make some points between (0,0,0) and (1,1,1) and put in a point cloud
    N = 1000
    pos = rand(SVector{3,Float32}, N)
    sphere = Sphere(0.5, 0.5, 0.5, 0.1)

    sphere_search_correct = true
    for i = 1:length(pos)
        v = pos[i]
        dv = v - SVector(0.5,0.5,0.5)
        distance = sum(abs2, dv)
        expected = distance < 0.01
        actual = v ∈ sphere
        sphere_search_correct &= actual == expected
    end
    @test sphere_search_correct

    @test boundingbox(sphere) == BoundingBox(0.4, 0.4, 0.4, 0.6, 0.6, 0.6)

    @test volume(sphere) ≈ 0.04188790204786391

    @test might_intersect(
        BoundingBox(0.0, 0.0, 0.0, 0.1, 0.1, 0.1),
        Sphere(0.5, 0.5, 0.5, 0.5)
    ) == false

    @test might_intersect(
        BoundingBox(0.0, 0.0, 0.0, 1.0, 1.0, 1.0),
        Sphere(0.5, 0.5, 0.5, 0.2)
    ) == true

    @test might_intersect(
        BoundingBox(0.0, 0.0, 0.0, 0.4, 0.4, 0.4),
        Sphere(0.5, 0.5, 0.5, 0.2)
    ) == true
end


using WignerFunctions
using Test
using Random

d222(β::AbstractFloat) = (1 + cos(β))^2 / 4
d221(β::AbstractFloat) = -sin(β) * (1 + cos(β)) / 2
d220(β::AbstractFloat) = sin(β) * sin(β) * √(3 / 8)
d221̄(β::AbstractFloat) = -sin(β) * (1 - cos(β)) / 2
d222̄(β::AbstractFloat) = (1 - cos(β))^2 / 4
d211(β::AbstractFloat) = (2cos(β) * cos(β) + cos(β) - 1) / 2
d210(β::AbstractFloat) = -sin(2β) * √(3 / 8)
d211̄(β::AbstractFloat) = (-2cos(β) * cos(β) + cos(β) + 1) / 2
d200(β::AbstractFloat) = (3cos(β) * cos(β) - 1) / 2
WignerFunctions.Index(::typeof(d222)) = WignerFunctions.Index((2, 2, 2))
WignerFunctions.Index(::typeof(d221)) = WignerFunctions.Index((2, 2, 1))
WignerFunctions.Index(::typeof(d220)) = WignerFunctions.Index((2, 2, 0))
WignerFunctions.Index(::typeof(d221̄)) = WignerFunctions.Index((2, 2, -1))
WignerFunctions.Index(::typeof(d222̄)) = WignerFunctions.Index((2, 2, -2))
WignerFunctions.Index(::typeof(d211)) = WignerFunctions.Index((2, 1, 1))
WignerFunctions.Index(::typeof(d210)) = WignerFunctions.Index((2, 1, 0))
WignerFunctions.Index(::typeof(d211̄)) = WignerFunctions.Index((2, 1, -1))
WignerFunctions.Index(::typeof(d200)) = WignerFunctions.Index((2, 0, 0))

@testset "Naive" begin
    naive = WignerFunctions.naive
    @testset "diagonal: β=-0, l=$l, m₁=$m₁, m₂=$m₂" for l ∈ 0:4, m₁ ∈ -l:l, m₂ ∈ -l:l
        @test naive(BigFloat(0), l, m₁, m₂) ≈ (m₁ == m₂ ? 1 : 0)
    end

    functions = [d222, d221, d220, d221̄, d222̄, d211, d210, d211̄, d211̄, d200]
    @testset "function $func" for func ∈ functions
        β = 2π*rand(10)
        l, m, m′ = WignerFunctions.Index(func)
        @test all(func.(β) .≈ naive.(BigFloat.(β), l, m, m′))
        @test all(func.(β) .≈ ((m - m′) % 2 == 0 ? 1 : -1) .* naive.(BigFloat.(β),  l, m′, m))
        @test all(naive.(BigFloat.(β), l, m′, m) .≈ naive.(BigFloat.(β),  l, -m, -m′))
    end
end

@testset "$method" for method ∈ (WignerFunctions.Risbo.nosym,)
    dmatrix = method(0)
    @testset "diagonal: β=-0, l=$l, m₁=$m₁, m₂=$m₂" for l ∈ 0:4, m₁ ∈ -l:l, m₂ ∈ -l:l
        @test dmatrix(l, m₁, m₂) ≈ (m₁ == m₂ ? 1 : 0)
    end

    l = rand(2:20, 50)
    m = rand.(range.(-l, l, step=1))
    m′ = rand.(range.(-l, l, step=1))
    β = 2π * rand(length(l))
    @testset "Symmetries" begin
        @testset "β=$(β[i]), l=$(l[i]), m=$(m[i]), m′=$(m′[i]) " for i ∈ 1:length(l)
            k, m₁, m₂ = l[i], m[i], m′[i]
            dmatrix = method(β[i])
            @test dmatrix(k, m₁, m₂) ≈ ((m₁ - m₂) % 2 == 0 ? 1 : -1) * dmatrix(k, m₂, m₁)
            @test dmatrix(k, m₁, m₂) ≈ ((m₁ - m₂) % 2 == 0 ? 1 : -1) * dmatrix(k, -m₁, -m₂)
            @test dmatrix(k, m₁, m₂) ≈ dmatrix(k, -m₂, -m₁)
        end
    end

    naive = WignerFunctions.naive
    oldprec = precision(BigFloat)
    setprecision(512)
    @testset "Against naive β=$(β[i]), l=$(l[i]), m=$(m[i]), m′=$(m′[i]) " for i ∈ 1:length(l)
        @test method(β[i])(l[i], m[i], m′[i]) ≈ float(naive(BigFloat(β[i]), l[i], m[i], m′[i]))
    end
    setprecision(oldprec)
end
nothing

@testset "Delta recurrence" begin
    recurrence = WignerFunctions.Trapani.delta_recurrence
    @test recurrence(0, 0, 0) == 1
    @test recurrence(3, 4, 0) == 0
    @test recurrence(3, 0, -4) == 0

    ls = rand(2:20, 50)
    ms = rand.(range.(-ls, ls, step=1))
    ms′ = rand.(range.(-ls, ls, step=1))
    oldprec = precision(BigFloat)
    setprecision(512)
    naive = (l, m₁, m₂) -> float(WignerFunctions.naive(BigFloat(pi)/2, l, m₁, m₂))
    @testset "Against naive l=$(ls[i]), m=$(ms[i]), m′=$(ms′[i]) " for i ∈ 1:length(ls)
        l, m₁, m₂ = ls[i], ms[i], ms′[i]
        @test recurrence(l, m₁, m₂) ≈ naive(l, m₁, m₂)
    end
    setprecision(oldprec)
end


using Test
using JOLI
using Zygote
using JOLI4Flux

x = rand(Float32, 3)

W = randn(Float32, 3, 3)
W_JOLI = joMatrix(W)

g = gradient(x -> sum(W*x), x)[1]
g_JOLI = gradient(x -> sum(W_JOLI*x), x)[1]

@test isapprox(sum((g - g_JOLI).^2)/sum(g_JOLI.^2), 0f0; atol=1f-8)


using Flux, Flux.Data.MNIST, Statistics
using Flux: onehotbatch, onecold, crossentropy, throttle
using Base.Iterators: repeated, partition

# Classify MNIST digits with a convolutional network

imgs = MNIST.images()

labels = onehotbatch(MNIST.labels(), 0:9)

# Partition into batches of size 1,000
train = [(cat(float.(imgs[i])..., dims = 4), labels[:,i])
         for i in partition(1:60_000, 1000)]

train = gpu.(train)

# Prepare test set (first 1,000 images)
tX = cat(float.(MNIST.images(:test)[1:1000])..., dims = 4) |> gpu
tY = onehotbatch(MNIST.labels(:test)[1:1000], 0:9) |> gpu

m = Chain(
  Conv((2,2), 1=>16, relu),
  x -> maxpool(x, (2,2)),
  Conv((2,2), 16=>8, relu),
  x -> maxpool(x, (2,2)),
  x -> reshape(x, :, size(x, 4)),
  Dense(288, 10), softmax) |> gpu

m(train[1][1])

loss(x, y) = crossentropy(m(x), y)

accuracy(x, y) = mean(onecold(m(x)) .== onecold(y))

evalcb = throttle(() -> @show(accuracy(tX, tY)), 10)
opt = ADAM(params(m))

Flux.train!(loss, train, opt, cb = evalcb)


module LineSearch

using Printf
using Stopping
using NLPModels
using NLPModelsJuMP
using Printf
using LinearAlgebra
using OptimizationProblems

include("LSFunctionMetamod.jl")
include("LineModel2mod.jl")

include("phi_dphi.jl")
include("pred_ared.jl")
include("update_H.jl")

include("armijo_backtracking.jl")
include("one-step-size.jl")
include("shamanskii-line-search.jl")

include("TR/TR_generic_ls.jl")
include("TR/TR-ls-step-computation.jl")
include("TR/TR-Nwt-LS.jl")
include("TR/TR-Sec-LS.jl")
include("TR/TR-SecA-LS.jl")

include("ARC/ARC_generic_ls.jl")
include("ARC/ARC-Nwt-LS.jl")
include("ARC/ARC-Sec-LS.jl")
include("ARC/ARC-SecA-LS.jl")
include("ARC/ARC_direction_computation.jl")

include("Bisection/find_interval.jl")
include("Bisection/Biss-Nwt.jl")
include("Bisection/Biss-Sec.jl")
include("Bisection/Biss-SecA.jl")
include("Bisection/Biss-Cub.jl")

include("other-ls/struct-ls.jl")
include("other-ls/hagerzhang2.jl")
include("other-ls/hagerzhang.jl")

end # module


using CalibrationErrors, Distances

const x1, y1 = 1, 3
const x2, y2 = [1, -1], [2, 1]

@testset "Squared exponential kernel" begin
    @test_throws ArgumentError SquaredExponentialKernel(-1)

    @test SquaredExponentialKernel() === SquaredExponentialKernel(1)
    @test SquaredExponentialKernel() === SquaredExponentialKernel(Euclidean())
    @test SquaredExponentialKernel() === SquaredExponentialKernel(1, Euclidean())

    for γ in (1, 2)
        kernel = SquaredExponentialKernel(γ, TotalVariation())
        @test kernel(x1, y1) ≈ exp(- γ)
        @test kernel(x2, y2) ≈ exp(- 2.25 * γ)

        kernel2 = SquaredExponentialKernel(γ, Cityblock())
        @test kernel2(x1, y1) ≈ exp(- 4 * γ)
        @test kernel2(x2, y2) ≈ exp(- 9 * γ)

        kernel3 = SquaredExponentialKernel(γ, Euclidean())
        @test kernel3(x1, y1) ≈ exp(- 4 * γ)
        @test kernel3(x2, y2) ≈ exp(- 5 * γ)
    end
end

@testset "Exponential kernel" begin
    @test_throws ArgumentError ExponentialKernel(-1)

    @test ExponentialKernel() === ExponentialKernel(1)
    @test ExponentialKernel() === ExponentialKernel(Euclidean())
    @test ExponentialKernel() === ExponentialKernel(1, Euclidean())

    for γ in (1, 2)
        kernel = ExponentialKernel(γ, TotalVariation())
        @test kernel(x1, y1) ≈ exp(- γ)
        @test kernel(x2, y2) ≈ exp(- 1.5 * γ)

        kernel2 = ExponentialKernel(γ, Cityblock())
        @test kernel2(x1, y1) ≈ exp(- 2 * γ)
        @test kernel2(x2, y2) ≈ exp(- 3 * γ)

        kernel3 = ExponentialKernel(γ, Euclidean())
        @test kernel3(x1, y1) ≈ exp(- 2 * γ)
        @test kernel3(x2, y2) ≈ exp(- sqrt(5) * γ)
    end
end


using Random, LinearAlgebra

# Standardised tests for the eigen decomposition of a square kronecker product
function eigen_tests(rng, C::AbstractKroneckerProduct)

    # Check approximate correctness of decomposition
    λ, Γ = eigen(C)
    @test Γ * Diagonal(λ) * Γ' ≈ C

    # Check approximate correctness of shifted decomposition
    σ² = abs2(randn(rng))
    λ′, Γ′ = (eigen(C) + σ² * I)
    @test Γ′ * Diagonal(λ′) * Γ′' ≈ Matrix(C + σ² * I)
    @test σ² * I + C == C + σ² * I

    # Test with decomposition
    v = randn(rng, length(λ))
    @test eigen(C) \ v ≈ Float64.(Matrix(C)) \ v

    # Test various linear algebra operations with decomposition
    C_dense = Float64.(Matrix(C))
    @test det(eigen(C)) ≈ det(C_dense)
    @test logdet(eigen(C)) ≈ logdet(C_dense)
    @test Matrix(inv(eigen(C))) ≈ inv(C_dense)
end

@testset "eigen" begin
    rng, P, Q, R = MersenneTwister(123456), 3, 5, 7

    # Generate some positive definite matrices so that logdet can be tested.
    A_, B_, C_ = randn(rng, P, P), randn(rng, Q, Q), randn(rng, R, R)
    A, B, C = Symmetric(A_ * A_' + I), Symmetric(B_ * B_' + I), Symmetric(C_ * C_' + I)

    D = kronecker(A, B)
    eigen_tests(rng, D)
    eigen_tests(rng, kronecker(D, C))
    eigen_tests(rng, kronecker(C, D))
    eigen_tests(rng, kronecker(kronecker(A, B), kronecker(B, A)))
    eigen_tests(rng, kronecker(A, 4))

    function test_non_square_extension()
        local n, m, A, B, K, M
        n, m = 3, 5
        # 1. K is square, while A, B aren't
        A = randn(n, m)
        B = randn(m, n)
        K = A ⊗ B
        M = Matrix(K)
        E = eigen(K)
        F = eigen(M)
        @test E.values ≈ F.values
        @test det(E) ≈ det(F)
        @test Matrix(E) ≈ K

        # 2. K is not square
        A = randn(n, m)
        B = randn(m, m)
        K = A ⊗ B
        @test_throws DimensionMismatch eigen(K)
    end
    test_non_square_extension()
end


using FunctionalData, QuickShiftClustering
data = @p map unstack(1:10) (x->10*randn(2,1).+randn(2,1000)) | flatten
@time a = quickshift(data)
Profile.clear_malloc_data()
@time a = quickshift(data)
exit()


using Statistics

@time begin
    data = [mean(rand(5*10^2)) for _ in 1:10^6]
    println("98% of the means lie in the estimated range: ",
                    (quantile(data,0.01),quantile(data,0.99)) )
end


function local_metric_jacobian(M::MetricManifold, x)
    n = size(x, 1)
    ∂g = reshape(ForwardDiff.jacobian(x -> local_metric(M, x), x), n, n, n)
    return ∂g
end

function christoffel_symbols_second_jacobian(M::MetricManifold, x)
    n = size(x, 1)
    ∂Γ = reshape(
        ForwardDiff.jacobian(x -> christoffel_symbols_second(M, x), x),
        n, n, n, n
    )
    return ∂Γ
end


n_samples = 100_000

for d in [MultivariateNormal([0.0, 0.0], [1.0 0.9; 0.9 1.0]),
          Dirichlet([2.0, 1.0]),
          Multinomial(1, [0.5, 0.5])]
    # Test rand
    s = rand(d)
    s = rand(d, 10)
    A = Array(eltype(s), 2, n_samples)
    rand!(d, A)

    # Test mean under CLT
    mu_hat = mean(A, 2)
    @assert norm(mu_hat - mean(d), Inf) < 0.1

    # Test cov under CLT
    sigma_hat = cov(A')
    @assert norm(sigma_hat - var(d), Inf) < 0.1

    # Test entropy and logpdf using AEP
    en_hat = -mean(logpdf(d, A))
    en = entropy(d)
    @assert abs(en_hat - en) < 0.1

    # Test entropy and pdf using AEP
    en_hat = -mean(log(pdf(d, A)))
    en = entropy(d)
    @assert abs(en_hat - en) < 0.1
end


using Test
using OpenQuantumSystems
using Random, SparseArrays, LinearAlgebra, StableRNGs

@testset "trace" begin

    Random.seed!(StableRNG(0), 1)

    D(op1::Array, op2::Array) = abs(norm(op1 - op2))
    D(x1::StateVector, x2::StateVector) = norm(x2 - x1)
    D(op1::AbstractOperator, op2::AbstractOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))
    D(op1::AbstractSuperOperator, op2::AbstractSuperOperator) =
        abs(tracedistance_nh(dense(op1), dense(op2)))

    # TODO: change to macro
    mode1 = Mode(0.2, 1.0)
    Energy = [0.0, 200.0]
    mol1 = Molecule([mode1], 3, Energy)
    mol2 = Molecule([mode1], 3, Energy)
    agg = Aggregate([mol1, mol2])
    aggInds = getIndices(agg)
    vibindices = getVibIndices(agg, aggInds)
    aggIndsLen = length(aggInds)
    basis = GenericBasis([aggIndsLen])
    FCFact = getFranckCondonFactors(agg, aggInds)

    Ham = getAggHamiltonian(agg, aggInds, FCFact)
    Ham_I = getAggHamInteraction(agg, aggInds, FCFact)
    Ham_0 = getAggHamSystemBath(agg, aggInds, FCFact)

    data = Matrix(Hermitian(rand(StableRNG(0), ComplexF64, aggIndsLen, aggIndsLen)))
    rho0 = DenseOperator(basis, basis, data)
    normalize!(rho0)

    t = 0.0
    U_op = evolutionOperator(Ham, t)
    rho = U_op * rho0 * U_op'
    rho_traced_ref = ComplexF64[0.30545089951602383 + 0.0im 0.14375197756961597 + 0.21658981002618588im 0.23626964940425038 + 0.1914303091850988im; 0.14375197756961597 - 0.21658981002618588im 0.30379617012935617 + 1.5178830414797062e-18im 0.15887255193342634 + 0.11195793936579877im; 0.23626964940425035 - 0.19143030918509882im 0.15887255193342636 - 0.11195793936579874im 0.4026345547478164 - 1.734723475976807e-18im]
    rho_traced = trace_bath_slow(rho, agg, FCFact, aggInds, vibindices)
    @test 1e-7 > D(rho_traced.data, rho_traced_ref)
    rho_traced =
        trace_bath_slow(rho.data, agg, FCFact, aggInds, vibindices)
    @test 1e-7 > D(rho_traced, rho_traced_ref)

    FCProd = getFCProd(agg, FCFact, aggInds, vibindices)
    rho_traced = trace_bath(rho, agg, FCProd, aggInds, vibindices)
    @test 1e-7 > D(rho_traced.data, rho_traced_ref)
    rho_traced = trace_bath(rho.data, agg, FCProd, aggInds, vibindices)
    @test 1e-7 > D(rho_traced, rho_traced_ref)

    for a = 1:3, b = 1:3
        rho_traced_ab = trace_bath(rho, a, b, agg, FCProd, aggInds, vibindices)
        @test 1e-7 > abs(rho_traced_ref[a, b] - rho_traced_ab)
    end

    rho_bath = get_rho_bath(rho0, agg, FCProd, aggInds, vibindices)
    rho_traced = trace_bath(rho_bath, agg, FCProd, aggInds, vibindices)
    rho_traced_ref = [1.0 1.0 1.0; 1.0 1.0 1.0; 1.0 1.0 1.0]
    @test 1e-10 > D(rho_traced_ref, rho_traced.data)

    rho_traced = trace_bath(rho0, agg, FCProd, aggInds, vibindices)
    rho0_ad = ad(rho_traced, rho_bath, agg, FCProd, aggInds, vibindices)
    @test 1e-7 > D(rho0, rho0_ad)

    t = 1.0
    U_op = evolutionOperator(Ham, t)
    rho = U_op * rho0 * U_op'
    rho_traced_ref = ComplexF64[0.30545089951602605 + 3.0357660829594124e-18im 0.19173860810181675 - 0.0844643501820273im 0.19018921460047641 - 0.20310674141599672im; 0.19173860810181684 + 0.0844643501820273im 0.3499470437099394 + 8.50622960880644e-18im 0.13057918573136992 + 0.08614578235028769im; 0.19018921460047636 + 0.20310674141599672im 0.13057918573136984 - 0.08614578235028776im 0.411399271707198 + 6.1416735243015536e-18im]
    rho_traced = trace_bath_slow(rho, agg, FCFact, aggInds, vibindices)
    @test 1e-7 > D(rho_traced.data, rho_traced_ref)
    rho_traced =
        trace_bath_slow(rho.data, agg, FCFact, aggInds, vibindices)
    @test 1e-7 > D(rho_traced, rho_traced_ref)

    FCProd = getFCProd(agg, FCFact, aggInds, vibindices)
    rho_traced = trace_bath(rho, agg, FCProd, aggInds, vibindices)
    @test 1e-7 > D(rho_traced.data, rho_traced_ref)
    rho_traced = trace_bath(rho.data, agg, FCProd, aggInds, vibindices)
    @test 1e-7 > D(rho_traced, rho_traced_ref)

    for a = 1:3, b = 1:3
        rho_traced_ab = trace_bath(rho, a, b, agg, FCProd, aggInds, vibindices)
        @test 1e-7 > abs(rho_traced_ref[a, b] - rho_traced_ab)

        rho_ab = take_el_part(rho.data, a, b, vibindices)
        rho_traced_ab = trace_bath_part(
            rho_ab,
            a,
            b,
            agg,
            FCProd,
            aggInds,
            vibindices
        )
        @test 1e-7 > abs(rho_traced_ref[a, b] - rho_traced_ab)
    end

    rho_bath = get_rho_bath(rho0, agg, FCProd, aggInds, vibindices)
    rho_traced = trace_bath(rho_bath, agg, FCProd, aggInds, vibindices)
    rho_traced_ref = [1.0 1.0 1.0; 1.0 1.0 1.0; 1.0 1.0 1.0]
    @test 1e-10 > D(rho_traced_ref, rho_traced.data)

    rho_traced = trace_bath(rho0, agg, FCProd, aggInds, vibindices)
    rho0_ad = ad(rho_traced, rho_bath, agg, FCProd, aggInds, vibindices)
    @test 1e-7 > D(rho0, rho0_ad)

    rho_bath = get_rho_bath(rho0, agg, FCProd, aggInds, vibindices)

    t = 0.0
    Ham_II_t = getInteractionHamIPicture(Ham_0, Ham_I, t)
    prod = Ham_II_t * Ham_I * rho_bath
    corr_ref = trace_bath(prod, agg, FCProd, aggInds, vibindices)
    corr = correlation_function(
        t,
        rho_bath,
        Ham_0,
        Ham_I,
        agg,
        FCProd,
        aggInds,
        vibindices
    )
    @test 1e-7 > D(corr, corr_ref)

    t = 1.0
    Ham_II_t = getInteractionHamIPicture(Ham_0, Ham_I, t)
    prod = Ham_II_t * Ham_I * rho_bath
    corr_ref = trace_bath(prod, agg, FCProd, aggInds, vibindices)
    corr = correlation_function(
        t,
        rho_bath,
        Ham_0,
        Ham_I,
        agg,
        FCProd,
        aggInds,
        vibindices
    )
    @test 1e-7 > D(corr, corr_ref)

end


function combine_spouses_data(spouse_value_matrix)
        n_ins = length(spouse_value_matrix[:,1])

        combined_data = Array(Any,n_ins)
        for i = 1 : n_ins
                combined_data[i] = tuple(spouse_value_matrix[i,:]...)
        end

        return combined_data
end

function parent_class_combine(data_matrix,parent_set)
        n_parent = length(parent_set)
        N = length(data_matrix[:,1])

        parent_data_set = Array(Int64,N,n_parent)

        for i = 1 : n_parent
                parent_data_set[:,i] = data_matrix[:,parent_set[i]]
        end

        parent_data_set = combine_spouses_data(parent_data_set)

        class_combine = length(unique(parent_data_set))
        return class_combine
end

function log_prob_single_edge_last_term(class)
        n = length(class)
        uniq_cl = unique(class)
        val_cl = length(uniq_cl)

        distr_table = Array(Array,n)

        # Survey distribution for each data point

        for index_data = 1 : n
                class_cl = findfirst(uniq_cl,class[index_data])
                z = zeros(Int64,val_cl)
                z_modify = zeros(Int64,val_cl)
                z_modify[class_cl] = 1
                distr_table[index_data] = z_modify
        end

        # Survey class distribution between point i and point j

        distr_intval_table = Array(Array,n,n)

        for ind_init = 1 : n
                for ind_end = ind_init : n
                        if ind_init == ind_end
                                current_distr = distr_table[ind_init]
                        else
                                current_distr = distr_intval_table[ind_init,ind_end-1] +
                                                distr_table[ind_end]
                        end
                        distr_intval_table[ind_init,ind_end] = current_distr
                end
        end
        # println(distr_intval_table[1,n])
        # Get 1/P(D|M) for single edge case and count only last term in objective func

        inv_p = Array(Float64,n,n)

        for ind_init = 1 : n
                for ind_end = 1 : n
                        if ind_end < ind_init
                                inv_p[ind_init,ind_end] = Inf
                        else
                                distr_in_intval = distr_intval_table[ind_init,ind_end]
                                current = lfact(ind_end-ind_init+1)
                                for ind_cl = 1 : val_cl
                                        current -= lfact(distr_in_intval[ind_cl])
                                end

                                inv_p[ind_init,ind_end] = current
                        end
                end
        end

        return inv_p
end

function log_prob_spouse_child_data(child,spouse)
        n = length(child)
        uniq_sp = unique(spouse)
        uniq_ch = unique(child)
        val_sp = length(uniq_sp)
        val_ch = length(uniq_ch)

        distr_table = Array(Array,n,val_sp)

        # Survey distribution for each data point

        for index_data = 1 : n
                class_ch = findfirst(uniq_ch,child[index_data])
                class_sp = findfirst(uniq_sp,spouse[index_data])
                z = zeros(Int64,val_ch)
                z_modify = zeros(Int64,val_ch)
                z_modify[class_ch] = 1

                for class_sp_index = 1 : val_sp
                        if class_sp_index == class_sp
                                distr_table[index_data,class_sp_index] = z_modify
                        else
                                distr_table[index_data,class_sp_index] = z
                        end
                end
        end

        # Survey class distribution in between point i and point j

        distr_intval_table = Array(Array,n,n,val_sp)

        for ind_val_sp = 1 : val_sp
                for ind_init = 1 : n
                for ind_end = ind_init : n
                        if ind_init == ind_end
                                current_distr = distr_table[ind_end,ind_val_sp]
                        else
                                current_distr = distr_intval_table[ind_init,ind_end-1,ind_val_sp] +
                                                distr_table[ind_end,ind_val_sp]
                        end
                        #println(current_distr)
                        distr_intval_table[ind_init,ind_end,ind_val_sp] = current_distr
                end
                end
        end

        # Get the 1/P(D|M) for this child-spouse set

        inv_p = Array(Float64,n,n)

        for ind_init = 1 : n
        for ind_end = 1 : n
                if ind_end < ind_init
                        inv_p[ind_init,ind_end] = Inf
                else
                        current_val = 0.0
                        for ind_val_sp = 1 : val_sp
                                distr_in_intval = distr_intval_table[ind_init,ind_end,ind_val_sp]
                                total_num = sum(distr_in_intval)

                                # The forth term in objective function
                                current = lfact(total_num)
                                for ind_val_ch = 1 : val_ch
                                        current -= lfact(distr_in_intval[ind_val_ch])
                                end

                                # The third term in objective function
                                current += lfact(total_num + val_ch -1) - lfact(val_ch-1) - lfact(total_num)

                                # The 3rd and 4th terms for this value of spouse
                                current_val += current
                        end
                        inv_p[ind_init,ind_end] = current_val
                end
        end
        end

        return inv_p
end

function prior_of_intval(continuous,lambda)
        N = length(continuous)
        d_1_N = continuous[N] - continuous[1]
        prior = Array(Float64,N)
        for i = 1 : N-1
                d_i = continuous[i+1] - continuous[i]
                prior[i] = 1 - exp(-lambda * d_i / d_1_N)
        end
        prior[N] = 1

        return prior
end

function largest_class_value(data_matrix)
        n = length(data_matrix[1,:])
        largest = 0
        for i = 1 : n
                class_value = length(unique(data_matrix[:,i]))
                if class_value > largest
                        largest = class_value
                end
        end

        return largest
end

function BN_discretizer_p_data_model(data_matrix,parent_set,child_spouse_set,approx = true)

        N = length(data_matrix[:,1])
        n = length(data_matrix[1,:])
        n_p = length(parent_set)
        n_c = length(child_spouse_set)
        n_r = n_p + n_c

        parent_class_number = 1
        if length(parent_set) > 0
                for i = 1 : n_p
                        multi = length(unique(data_matrix[:,parent_set[i]]))
                        parent_class_number = parent_class_number * multi
                end

                parent_class_number_2 = parent_class_combine(data_matrix,parent_set)
        end

        # -log(P(D|M)) part:

        log_P_data_model = zeros(Float64,N,N)
        nearest_var_set = [parent_set;child_spouse_set]
        if approx == true
        for ind_parent = 1 : n_p
                table = 0
                single_variable_data = data_matrix[:,parent_set[ind_parent]]
                table = log_prob_single_edge_last_term(single_variable_data)
                log_P_data_model += table
        end

        else
                if n_p > 0
                table = 0
                parent_data = Array(Int64,N,n_p)
                for p = 1 : n_p
                        parent_data[:,p] = data_matrix[:,parent_set[p]]
                end
                combined_parent = combine_spouses_data(parent_data)
                table = log_prob_single_edge_last_term(combined_parent)
                log_P_data_model += table
                end
        end

        for ind_child = 1 : n_c
                table = 0
                if length(child_spouse_set[ind_child]) == 1
                        child_data = data_matrix[:,child_spouse_set[ind_child]]
                        spouse_data = zeros(Int64,N)
                        table = log_prob_spouse_child_data(child_data,spouse_data)

                elseif length(child_spouse_set[ind_child]) == 2
                        (child,spouse) = child_spouse_set[ind_child]
                        child_data = data_matrix[:,child]
                        spouse_data = data_matrix[:,spouse]
                        table = log_prob_spouse_child_data(child_data,spouse_data)
                else
                        # A child has more than 2 parents
                        child = child_spouse_set[ind_child][1]
                        child_data = data_matrix[:,child]
                        spouse_set = child_spouse_set[ind_child][2:end]
                        spouse_matrix = Array(Int64,N,length(spouse_set))
                        for spouse_index = 1 : length(spouse_set)
                                spouse_matrix[:,spouse_index] = data_matrix[:,spouse_set[spouse_index]]
                        end
                        spouse_data = combine_spouses_data(spouse_matrix)
                        table = log_prob_spouse_child_data(child_data,spouse_data)
                end

                log_P_data_model += table
        end

        # Add -log(p_class_distribution_for_parent)

        for i = 1 : N
                for j = i : N
                        log_P_data_model[i,j] += lfact(j - i + parent_class_number) -
                                                 lfact(j-i+1) - lfact(parent_class_number-1)
                end
        end

        return log_P_data_model

end

function BN_discretizer_free_number_rep(continuous,data_matrix,parent_set,child_spouse_set,approx = true)
        p_data_model = BN_discretizer_p_data_model(data_matrix,parent_set,child_spouse_set,approx)
        #lambda = div(largest_class_value(data_matrix),2)
        lambda = largest_class_value(data_matrix)
        split_on_intval = prior_of_intval(continuous,lambda)
        N = length(continuous)
        not_split_on_intval = ones(Float64,N) - split_on_intval

        conti_norep = unique(continuous)
        conti_head = Array(Int64,length(conti_norep))
        conti_tail = Array(Int64,length(conti_norep))
        N_norep = length(conti_norep)

        conti_head[1] = 1
        index_head = 1
        for i = 2 : N
                if continuous[i] != continuous[i-1]
                        index_head += 1
                        conti_head[index_head] = i
                end
        end

        conti_tail[end] = N
        index_tail = length(conti_norep)
        for i = N-1 : -1 : 1
                if continuous[i] != continuous[i+1]
                        index_tail -= 1
                        conti_tail[index_tail] = i
                end
        end

        #println(conti_tail==conti_head)
        smallest_value = Array(Float64,N_norep)
        optimal_disc = Array(Array,N_norep)
        length_conti_data = continuous[N] - continuous[1]

        for a = 1 : N_norep
                if a == 1
                        smallest_value[1] = -log(split_on_intval[conti_tail[1]]) +
                                                        p_data_model[conti_head[1],conti_tail[1]]
                        optimal_disc[1] = [conti_tail[1]]
                else
                        current_value = Inf
                        current_disc = [0]
                        for b = 1 : a
                                if b == a
                                        value =
                                        ( ( (continuous[conti_tail[a]] - continuous[1])/
                                                length_conti_data)*lambda) -
                                        log(split_on_intval[conti_tail[a]]) +
                                        p_data_model[1,conti_tail[a]]
                                else

                                        value = smallest_value[b] +
                                        ( ( (continuous[conti_tail[a]] - continuous[conti_head[b+1]])/
                                        length_conti_data)*lambda) -
                                        log(split_on_intval[conti_tail[a]]) +
                                        p_data_model[conti_head[b+1],conti_tail[a]]
                                end

                                if value < current_value
                                        current_value = value
                                        if b == a
                                                current_disc = [conti_tail[a]]
                                        else
                                                current_disc = [optimal_disc[b];conti_tail[a]]
                                        end
                                end
                        end
                        smallest_value[a] = current_value
                        optimal_disc[a] = current_disc
                end
        end

        #println(optimal_disc[end])
        optimal_disc_value = Array(Float64,length(optimal_disc[end])+1)

        for i = 0 : length(optimal_disc[end])
                if i == 0
                        optimal_disc_value[i+1] = continuous[1]
                elseif i == length(optimal_disc[end])
                        optimal_disc_value[i+1] = continuous[end]
                else
                        optimal_disc_value[i+1] = 0.5 * (continuous[optimal_disc[end][i]]
                                                      +continuous[optimal_disc[end][i]+1])
                end
        end
        #println(smallest_value[end])
        return optimal_disc_value

end

# Equal width discretization
function equal_width_disc(continuous,m)

        N = length(continuous)
        min_value = minimum(continuous)
        max_value = maximum(continuous)
        span = (max_value - min_value)/m
        class_list = Array(Int64,N)

        for i = 1 : N
                class = div((continuous[i] - min_value),span)
                if class >= m
                        class -= 1
                end
                class_list[i] = class + 1
        end
        return class_list
end

function equal_width_edge(continuous,m)
        min_value = continuous[1]
        max_value = continuous[end]
        span = max_value - min_value
        edge = [continuous[1]]
        for i = 1 : m
                edge = [edge, min_value + (span/m)*i]
        end
        return edge
end

#######################################################################################
#######################################################################################

function cartesian_product(product_set)
        N = length(product_set)
        M = 1
        for i = 1 : N
                M *= length(product_set[i])
        end

        product_returned = [product_set[1]...]

        for class_index = 2 : N
                class_number = length(product_set[class_index])
                current_length = length(product_returned[:,1])

                enlarged_matrix = Array(Int64,current_length*class_number,class_index)

                if class_number == 1
                        enlarged_matrix[:,1:class_index-1] = product_returned
                        for i = 1 : current_length
                                enlarged_matrix[i,class_index] = product_set[class_index][1]
                        end
                else

                        for enlarge_times = 1 : class_number
                                enlarged_matrix[(enlarge_times-1)*current_length+1:enlarge_times*current_length,1:class_index-1] =
                                product_returned
                        end

                        for i = 1 : class_number * current_length
                                item_index = div(i-1,current_length) + 1
                                enlarged_matrix[i,class_index] = product_set[class_index][item_index]
                        end
                end
                product_returned = enlarged_matrix

        end

        return product_returned
end

function find_intval(x,disc)
    if x < disc[2]
        return 1
    end
    if x >= disc[end-1]
        return length(disc)-1
    end

    for i = 2 : length(disc)-2
           if (x >= disc[i])&(x < disc[i+1])
                   return i
           end
    end
end

function continuous_to_discrete(data,bin_edge_1)
        #bin_edge_extend
        bin_edge = copy(bin_edge_1)
        bin_edge[1] = -Inf
        bin_edge[end] = Inf
        data_discrete = Array(Int64,length(data))
        for i = 1 : length(data)
                index = 0
                for j = 2 : length(bin_edge)
                        if (data[i] > bin_edge[j-1])&(data[i] <= bin_edge[j])
                                index = j-1
                        end
                end
                data_discrete[i] = index
        end
        return data_discrete
end

function sort_disc_by_vorder(continuous_order,disc_edge)
      reorder_disc_edge = Array(Any,length(continuous_order))
      for i = 1 : length(continuous_order)
            num_less = 1
            for j = 1 : length(continuous_order)
                  if continuous_order[j] < continuous_order[i]
                        num_less += 1
                  end
            end
            reorder_disc_edge[num_less] = disc_edge[i]
    end
    return reorder_disc_edge
end

function rand_seq(N)

        seq = Array(Int64,N)

        i = 1
        seq[1] = 1 + round(Int64, div(N * rand(),1) )

        while i < N
                number = 1 + round(Int64,div( N*rand(),1))
                if ~(number in seq[1:i])
                        i += 1
                        seq[i] = number
                end
        end
        return seq
end

#######################################################################################
#######################################################################################

function disc_intval_seq(table)
        N = length(table)
        seq = Array(Int64,N)
        for i = 1 : N
                seq[i] = length(table[i])
        end
        return seq
end

function graph_to_reverse_order(graph)
        order = Array(Int64,length(graph))
        for i = 1 : length(graph)
                order[length(graph) - i + 1] = graph[i][end]
        end
        return order
end

function graph_to_reverse_conti_order(graph,continuous_index)
        order = graph_to_reverse_order(graph)
        Order = Array(Int64,length(continuous_index))
        ind = 0
        for i = 1 : length(order)
                if order[i] in continuous_index
                        ind += 1
                        Order[ind] = order[i]
                end
        end
        return Order
end

function graph_to_markov(graph,target)
        parent_set = []; child_spouse_set = [];
        for i = 1 : length(graph)

                condi_graph = graph[i]
                if length(condi_graph) == 1
                        if target == condi_graph
                                parent_set = []
                        end

                elseif target in condi_graph

                        index = findfirst(condi_graph,target)
                        if index == length(condi_graph)
                                parent_set = [condi_graph[1:end-1]...]
                        else

                                child = condi_graph[end]
                                spouse = []

                                for j = 1 : length(condi_graph)-1
                                        if j != index
                                                spouse = [spouse,condi_graph[j]]
                                        end
                                end
                                child_spouse_set = [child_spouse_set,tuple([child,spouse]...)] ###### change to ; ######

                        end
                else

                end
        end

        for i = 1 : length(child_spouse_set)
                if length(child_spouse_set[i]) == 1
                        child_spouse_set = [child_spouse_set[1:i-1],child_spouse_set[i][1],child_spouse_set[i+1:end]]
                end
        end

        return (parent_set,child_spouse_set)
end

function one_iteration(data,data_integer,graph,discrete_index,continuous_index,lcard,approx = true)

        # Save discretization edge
        disc_edge_collect = Array(Any,length(continuous_index))

        for i = 1 : length(continuous_index)
                target = continuous_index[i]

                increase_order = sortperm(data[:,target])
                conti = data[:,target][increase_order]
                data_integer_sort = Array(Int64,size(data))

                # sort data_integer properly
                for j = 1 : length(data_integer[1,:])
                        data_integer_sort[:,j] = data_integer[:,j][increase_order]
                end
                sets = graph_to_markov(graph,target)
                parent_set = sets[1]; child_spouse_set = sets[2];

                if (length(sets[1]) + length(sets[2])) == 0
                        disc_edge = equal_width_edge(conti,lcard)
                        disc_edge_collect[i] = disc_edge
                else
                        disc_edge = BN_discretizer_free_number_rep(conti,data_integer_sort,parent_set,child_spouse_set,approx)
                        disc_edge_collect[i] = disc_edge
                end

                # Update by current discretization
                data_integer[:,target] = continuous_to_discrete(data[:,target],disc_edge)
        end

        return (data_integer,disc_edge_collect)
end

function BN_discretizer_iteration_converge(data,graph,discrete_index,continuous_index,cut_time,approx = true)
        # intital the first data_integer
        l_card = 0
        for i = 1 : length(discrete_index)
                card = length(unique(data[:,discrete_index[i]]))
                if card > l_card
                        l_card = card
                end
        end

        # pre-equal-width-discretize on continuous variables
        data_integer = zeros(Int64,size(data))
        for i = 1 : length(discrete_index)
                index = discrete_index[i]
                data_integer[:,index] = data[:,index]
        end

        for i = 1 : length(continuous_index)
                index = continuous_index[i]
                data_integer[:,index] = equal_width_disc(data[:,index],l_card)
        end

        disc_edge_collect = Array(Any,length(continuous_index))

        disc_edge_previous = Array(Any,length(continuous_index))

        for i = 1 : length(continuous_index)
                disc_edge_previous[i] = []
                disc_edge_collect[i] = [0]
        end

        times = 0
        while (disc_edge_previous != disc_edge_collect)&(times<cut_time)

            times += 1
            #println(("iteration times = ",times))
            disc_edge_previous = disc_edge_collect

            X = one_iteration(data,data_integer,graph,discrete_index,continuous_index,l_card,approx)
            data_integer = X[1]
            disc_edge_collect = X[2]

        end

        return (data_integer,disc_edge_collect)
end

function K2_f(x,parent)
        N = length(x)
        n = length(parent[1,:])

        # build carti product
        carti_set = [tuple(unique(parent[:,1])...)]

        for i = 2 : n
                carti_set = [carti_set;tuple(unique(parent[:,i])...)]
        end

        uni_x = unique(x); r = length(uni_x)
        carti_set = [carti_set;tuple(uni_x...)]

        CP = cartesian_product(carti_set)

        r_p = length(CP[:,1]) / r
        r_p = round(Int64,r_p)

        # Survey distribution
        count = zeros(Int64,length(CP[:,1]))

        for i = 1 : N
                reform_data = [parent[i,:] x[i]]
                for j = 1 : length(CP[:,1])
                        if reform_data == CP[j,:]
                                count[j] += 1
                        end
                end
        end

        # Sum over same parent value
        count_p = zeros(Int64,r_p)
        for i = 1 : r_p
                count_pp = 0
                for j = 1 : r
                        index = i + r_p * (j-1)
                        count_pp += count[index]
                end
                count_p[i] = count_pp
        end

        # Evaluate
        f = 0
        for i = 1 : r_p
                f += lfact(count_p[i]+r-1) - lfact(r-1)
        end
        for i = 1 : length(CP[:,1])
                f -= lfact(count[i])
        end

        return f
end

function parent_table_to_graph(parent_table)
        m = length(parent_table)
        # Parent_table -> graph
        graph = []
        for i = 1 : m
                child = i
                if length(parent_table[i]) == 0
                        graph = [graph;child]
                else
                parent = []
                for j = 1 : length(parent_table[i])
                        parent = [parent;parent_table[i][j]]
                end

                add = [parent,child]
                graph = [graph;tuple(add...)]
                end
        end
        return graph
end

function K2_one_iteration_discretization(order,u,data_matrix,continuous_index,cut_time,approx = true)
        N = length(data_matrix[:,1])
        m = length(data_matrix[1,:]) # number of variables
        mc = length(continuous_index)

        # continuous/discrete index in the new order
        conti_index = Array(Int64,mc)
        disc_index = Array(Int64,m-mc)
        for i = 1 : mc
                conti_index[i] = findfirst(order,continuous_index[i])
        end
        conti_index = sort(conti_index, rev=true)

        index_disc = 0
        for i = 1 : m
                if ~(i in continuous_index)
                        index_disc += 1
                        disc_index[index_disc] = findfirst(order,i)
                end
        end

        # data that is sorted by order
        data = Array(Any,size(data_matrix))
        for i = 1 : m
                data[:,i] = data_matrix[:,order[i]]
                #println(i)
        end

        # Initialize graph structure
        initial_lcard= 0
        data_discretized = Array(Int64,size(data_matrix))
        if (m - length(conti_index)) == 0
            initial_lcard = round(Int64, log(m)+1)
        else
            for i = 1 : m
                    if ~(i in conti_index)
                            data_discretized[:,i] = data[:,i]
                            current_initial_lcard = length(unique(data[:,i]))

                            if current_initial_lcard > initial_lcard
                                    initial_lcard = current_initial_lcard
                            end
                    end
            end
        end
        for i = 1 : m
                if i in conti_index
                        data_discretized[:,i] = equal_width_disc(data[:,i],initial_lcard)
                end
        end

        # Initialize graph structure
        parent_table = Array(Any,m)
        for i = 1 : m
                parent_table[i] = [] # No parent in the begining
        end
        # store discretization edges
        disc_edge = 0

        # Iterate for all nodes
        for i = 1 : m
                p_old = K2_f(data_discretized[:,i],zeros(Int64,N))
                OKToProceed = true
                current_parent = []

                while OKToProceed & (length(parent_table[i]) < u)
                        # find node z for iteration
                        iteration_list = []
                        for j = 1 : i-1
                                if ~(j in current_parent)
                                        iteration_list = [iteration_list;j]
                                end
                        end

                        # iteration to find most probable
                        current_value = Inf
                        current_best_parent = []

                        for j = 1 : length(iteration_list)

                                iteration_parent = [current_parent;iteration_list[j]]
                                iteration_parent_data = data_discretized[:,iteration_parent[1]]
                                #if i ==4 ; println(iteration_parent); end;

                                for k = 2 : length(iteration_parent)
                                        iteration_parent_data = [iteration_parent_data data_discretized[:,iteration_parent[k]]]
                                end
                                this_iteration_value = K2_f(data_discretized[:,i],iteration_parent_data)

                                if this_iteration_value < current_value
                                        current_value = this_iteration_value
                                        current_best_parent = iteration_parent
                                end
                        end

                        if current_value < p_old
                                p_old = current_value
                                current_parent = current_best_parent
                                parent_table[i] = current_parent
                                ###### Rediscretize ######
                                #println((i,"th varaible"))
                                current_graph = parent_table_to_graph(parent_table)
                                disc_result = BN_discretizer_iteration_converge(data,current_graph,
                                                                               disc_index,conti_index,cut_time,approx)

                                ##### Replace data_discretized #####

                                data_discretized = disc_result[1]
                                disc_edge = disc_result[2]

                        else

                                OKToProceed = false

                        end

                end
                #println(("parent of",i,"is",parent_table[i]))
        end

        score = 0
        # Calculate score
        for i = 1 : m
                if length(parent_table[i]) == 0
                        score += K2_f(data_discretized[:,i],zeros(Int64,N))
                else
                        parent_data_final = Array(Int64,N,length(parent_table[i]))
                        for j = 1 : length(parent_table[i])
                                parent_data_final[:,j] = data_discretized[:,parent_table[i][j]]
                        end
                        score += K2_f(data_discretized[:,i],parent_data_final)
                end
        end

        # Parent_table -> graph
        graph = []
        for i = 1 : m
                child = order[i]
                if length(parent_table[i]) == 0
                        graph = [graph,child]
                else
                parent = []
                for j = 1 : length(parent_table[i])
                        parent = [parent,order[parent_table[i][j]]]
                end

                add = [parent,child]
                graph = [graph,tuple(add...)]
                end
        end

        # Disc_edge is modified to correct order
        conti_index_correspond = Array(Any,mc,3)
        for i = 1 : mc
              conti_index_correspond[i,2] = continuous_index[i]
              conti_index_correspond[i,1] = findfirst(order,continuous_index[i])
        end

        zz = sortperm(conti_index_correspond[:,1], rev = true)
        conti_index_correspond[:,1] = conti_index_correspond[:,1][zz]
        conti_index_correspond[:,2] = conti_index_correspond[:,2][zz]

        for i = 1 : mc
              conti_index_correspond[i,3] = disc_edge[i]
        end

        zzz = sortperm(conti_index_correspond[:,2])
        disc_edge_result = conti_index_correspond[:,3][zzz]

        return (score,graph,disc_edge_result)
end

function K2_w_discretization(data_matrix,u,continuous_index,times,cut_time,approx = true)

        score = Inf
        graph = 0
        disc_edge = 0

        for time = 1 : times
                # Produce random sequence of indexes
                println(("Times of K2 along with discretization method:",time))
                order = rand_seq(length(data_matrix[1,:]))
                iteration_result = K2_one_iteration_discretization(order,u,data_matrix,
                                                          continuous_index,cut_time,approx)

                if iteration_result[1] < score
                        score = iteration_result[1]
                        graph = iteration_result[2]
                        disc_edge = iteration_result[3]
                end
        end

        return (score,graph,disc_edge)
end

function Discretize_All(data_matrix,graph,continuous_index,cut_time)
    discrete_index = Array(Int64,length(data_matrix[1,:]) - length(continuous_index))
    num = 0
    for i = 1 : length(data_matrix[1,:])
        if ~(i in continuous_index)
            num += 1
            discrete_index[num] = i
        end
    end
    sort_continuous = graph_to_reverse_conti_order(graph,continuous_index)
    X = BN_discretizer_iteration_converge(data_matrix,graph,discrete_index,sort_continuous,cut_time,false)
    edge = X[2]
    reorder_edge = sort_disc_by_vorder(sort_continuous,edge)

    return reorder_edge
end

function Learn_DVBN(data_matrix,continuous_index,order,u,cut_time)
    X = K2_one_iteration_discretization(order,u,data_matrix,continuous_index,cut_time,false)
    return (X[2],X[3])
end

function Learn_Discrete_Bayesian_Net(data_matrix,continuous_index,u,cut_time,times)
    X = K2_w_discretization(data_matrix,u,continuous_index,times,cut_time,false)
    return (X[2],X[3],X[1])
end


# Copyright (c) 2021, Eric Sabo
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

include("linearcode.jl")

abstract type AbstractReedMullerCode <: AbstractLinearCode end

mutable struct ReedMullerCode <: AbstractReedMullerCode
    F::Union{FqNmodFiniteField, Nemo.GaloisField, AbstractAlgebra.GFField{Int64}}
    n::Integer
    k::Integer
    d::Union{Integer, Missing}
    r::Integer
    m::Integer
    G::Union{gfp_mat, fq_nmod_mat}
    Gorig::Union{gfp_mat, fq_nmod_mat, Missing}
    H::Union{gfp_mat, fq_nmod_mat}
    Horig::Union{gfp_mat, fq_nmod_mat, Missing}
    Gstand::Union{gfp_mat, fq_nmod_mat}
    Hstand::Union{gfp_mat, fq_nmod_mat}
    weightenum::Union{WeightEnumerator, Missing}
end

r(C::ReedMullerCode) = C.r
m(C::ReedMullerCode) = C.m

function show(io::IO, C::AbstractReedMullerCode)
    if get(io, :compact, false)
        println(io, "[$(length(C)), $(dimension(C)), $(minimumdistance(C))]_$(order(field(C))) Reed Muller code RM($(r(C)), $(m(C))).")
    else
        println(io, "[$(length(C)), $(dimension(C)), $(minimumdistance(C))]_$(order(field(C))) Reed Muller code RM($(r(C)), $(m(C))).")
        println(io, "Generator matrix: $(dimension(C)) × $(length(C))")
        for i in 1:dimension(C)
            print(io, "\t")
            for j in 1:length(C)
                if j != length(C)
                    print(io, "$(C.G[i, j]) ")
                elseif j == length(C) && i != dimension(C)
                    println(io, "$(C.G[i, j])")
                else
                    print(io, "$(C.G[i, j])")
                end
            end
        end
    end
end

function ReedMullergeneratormatrix(q::Integer, r::Integer, m::Integer)
    (0 ≤ r ≤ m) || error("Reed Muller codes require 0 ≤ r ≤ m, received r = $r and m = $m.")

    if q == 2
        F, _ = FiniteField(2, 1, "α")
        if r == m
            M = MatrixSpace(F, 2^m, 2^m)
            return M(1)
        elseif r == 0
            M = MatrixSpace(F, 1, 2^m)
            return M([1 for i in 1:2^m])
        else
            Grm1 = ReedMullergeneratormatrix(q, r, m - 1)
            Gr1m1 = ReedMullergeneratormatrix(q, r - 1, m - 1)
            M = MatrixSpace(F, size(Gr1m1, 1), size(Gr1m1, 2))
            return vcat(hcat(Grm1, Grm1), hcat(M(0), Gr1m1))
        end
    else
        error("Nonbinary Reed Muller codes have not yet been implemented.")
    end
end

function ReedMullerCode(q::Integer, r::Integer, m::Integer, verify::Bool=true)
    (1 ≤ r < m) || error("Reed Muller codes require 1 ≤ r < m, received r = $r and m = $m.")
    m < 64 || error("This Reed Muller code requires the implmentation of BigInts. Change if necessary.")
    q == 2 || error("Nonbinary Reed Muller codes have not yet been implemented.")

    G = ReedMullergeneratormatrix(q, r, m)
    H = ReedMullergeneratormatrix(q, m - r - 1, m)
    Gstand, Hstand = _standardform(G)

    if verify
        size(G, 2) == 2^m || error("Generator matrix computed in ReedMuller has the wrong number of columns; received: $(size(G, 2)), expected: $(BigInt(2)^m).")
        k = sum([binomial(m, i) for i in 0:r])
        size(G, 1) == k || error("Generator matrix computed in ReedMuller has the wrong number of rows; received: $(size(G, 1)), expected: $k.")
        if size(H) == (2^m - k, k)
            H = deepcopy(H')
        end
        for r in 1:size(Gstand, 1)
            iszero(Gstand[r, :] * H') || error("Column swap appeared in _standardform.")
        end
    end

    return ReedMullerCode(base_ring(G), size(G, 2), size(G, 1), 2^(m - r), r, m, G, missing,
        H, missing, Gstand, Hstand, missing)
end

function dual(C::ReedMullerCode)
    # return ReedMullerCode(C.q, C.m - C.r - 1, C.m)
    return ReedMullerCode(field(C), length(C), length(C) - dimension(C), 2^(r(C) + 1),
        m(C) - r(C) - 1, m(C), paritycheckmatrix(C), originalparitycheckmatrix(C), generatormatrix(C),
        originalgeneratormatrix(C), paritycheckmatrix(C, true), generatormatrix(C, true), missing)
end


# Compare with tabulated values,  classical hypergeometric formulas

@testset "Bessel" begin

    T = Float64

    for α in  (3/2, 1/2, 1, 2)
        P = Bessel{α,T}
        β =  2

        p0,p1,p2,p3,p4 = basis.(P, 0:4)

        # Krall and Frink
        for x in range(0, 1, length=5)
            @test p0(x) ≈ 1
            @test p1(x) ≈ 1 + α*(x/β)
            @test p2(x) ≈ 1 + 2(α+1)*(x/β) + (α+1)*(α+2)*(x/β)^2
            @test p3(x) ≈ 1 + 3(α+2)*(x/β) + 3(α+2)*(α+3)*(x/β)^2 +  (α+2)*(α+3)*(α+4)*(x/β)^3
            @test p4(x) ≈ 1 + 4(α+3)*(x/β) + 6(α+3)*(α+4)*(x/β)^2 +  4(α+3)*(α+4)*(α+5)*(x/β)^3 + (α+3)*(α+4)*(α+5)*(α+6)*(x/β)^4
        end
    end

    x = variable(Polynomial)
    for  α  ∈ (1/4,1/2,3/4, 1, 2)
        P =  Bessel{α}
        for n in 0:5
            p  = basis(P,n)
            @test  p(x) ≈ SP.classical_hypergeometric(P, n,  x)
        end
    end

end

@testset "ChebyshevT" begin

    p0 = Chebyshev([1])  # 1
    p1 = Chebyshev([0,1]) # x
    p2 = Chebyshev([0,0,1]) #
    p3 = Chebyshev([0,0,0,1]) #
    p4 = Chebyshev([0,0,0,0,1]) #

    for x in range(-1, 1, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ x
        @test p2(x) ≈ 2x^2 - 1
        @test p3(x) ≈ 4x^3  - 3x
        @test p4(x) ≈ 8x^4 - 8x^2 + 1
    end

end

@testset "ChebyshevU" begin

    p0 = ChebyshevU([1])  # 1
    p1 = ChebyshevU([0,1]) # x
    p2 = ChebyshevU([0,0,1]) #
    p3 = ChebyshevU([0,0,0,1]) #
    p4 = ChebyshevU([0,0,0,0,1]) #

    for x in range(-1, 1, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ 2x
        @test p2(x) ≈ 4x^2 - 1
        @test p3(x) ≈ 8x^3 - 4x
        @test p4(x) ≈ 16x^4 - 12x^2 +  1
    end

end

@testset "Hermite" begin

    p0 = Hermite([1])  # 1
    p1 = Hermite([0,1]) # x
    p2 = Hermite([0,0,1]) #
    p3 = Hermite([0,0,0,1]) #
    p4 = Hermite([0,0,0,0,1]) #

    for x in range(-1, 1, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ 2x
        @test p2(x) ≈ 4x^2 - 2
        @test p3(x) ≈ 8x^3 - 12x
        @test p4(x) ≈ 16x^4 - 48x^2 + 12
    end

    x = variable(LaurentPolynomial)
    P =  Hermite
    for n in 2:5
        p  = basis(P,n)
        @test  p(x) ≈ SP.classical_hypergeometric(P, n,  x)
    end

end

@testset "ChebyshevHermite" begin

    p0 = ChebyshevHermite([1])  # 1
    p1 = ChebyshevHermite([0,1]) # x
    p2 = ChebyshevHermite([0,0,1]) #
    p3 = ChebyshevHermite([0,0,0,1]) #
    p4 = ChebyshevHermite([0,0,0,0,1]) #

    for x in range(-1, 1, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ x
        @test p2(x) ≈ x^2 - 1
        @test p3(x) ≈ x^3 - 3x
        @test p4(x) ≈ x^4 - 6x^2 + 3
    end

    x = variable(LaurentPolynomial)
    P =  ChebyshevHermite
    for n in 2:5
        p  = basis(P,n)
        @test  p(x) ≈ SP.classical_hypergeometric(P, n,  x)
    end

end

@testset "Legendre" begin

    p0 = Legendre([1])  # 1
    p1 = Legendre([0,1]) # x
    p2 = Legendre([0,0,1]) #
    p3 = Legendre([0,0,0,1]) #
    p4 = Legendre([0,0,0,0,1]) #

    for x in range(-1, stop=1, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ x
        @test p2(x) ≈ 1/2*(3x^2-1)
        @test p3(x) ≈ 1/2*(5x^3-3x)
        @test p4(x) ≈ 1/8*(35x^4 - 30x^2 + 3)
    end

end

@testset "Shifted Legendre" begin

    T = Float64
    P = ShiftedLegendre{T}

    p0 = P([1])  # 1
    p1 = P([0,1]) # x
    p2 = P([0,0,1]) #
    p3 = P([0,0,0,1]) #
    p4 = P([0,0,0,0,1]) #

    for x in range(0, stop=1, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ 2x - 1
        @test p2(x) ≈ 6x^2 - 6x + 1
        @test p3(x) ≈ 20x^3 - 30x^2+ 12x - 1
        @test p4(x) ≈ 70x^4 - 140x^3 + 90x^2 - 20x + 1
    end

end

@testset "Laguerre" begin

    P = Laguerre{0}
    p0 = P([1])  # 1
    p1 = P([0,1]) # x
    p2 = P([0,0,1]) #
    p3 = P([0,0,0,1]) #
    p4 = P([0,0,0,0,1]) #

    for x in range(0, stop=5, length=5)
        @test p0(x) ≈ 1
        @test p1(x) ≈ -x + 1
        @test p2(x) ≈ 1/2*(x^2 - 4x + 2)
        @test p3(x) ≈ 1/6*(-x^3 + 9x^2 -18x + 6)
        @test p4(x) ≈ 1/24 * (x^4 - 16x^3 + 72x^2 - 96x + 24)
    end

    x = variable(Polynomial)
    for  α  ∈ (0, 1/4,1/2,3/4)
        P =  Laguerre{α}
        for n in 2:5
            p  = basis(P,n)
            @test  p(x) ≈ SP.classical_hypergeometric(P, n,  x)
        end
    end

end

@testset "Gegenbauer" begin

    for α in (1/4, 1/3, 1/2, 1, 2)
        P = Gegenbauer{α, Float64}
        p0 = P([1])  # 1
        p1 = P([0,1]) # x
        p2 = P([0,0,1]) #
        p3 = P([0,0,0,1]) #
        p4 = P([0,0,0,0,1]) #

        for x in range(-1, stop=1, length=5)
            @test p0(x) ≈ 1
            @test p1(x) ≈ 2α*x
            @test p2(x) ≈ -α + 2α*(1+α)*x^2
            @test p3(x) ≈ -2α * (1 + α) *x + 4/3*α*(1+α) * (2 + α)*x^3
        end
    end

    x = variable(Polynomial)
    for  α  ∈ (1/4,1/2,3/4, 1)
        P =  Gegenbauer{α}
        for n in 2:5
            p  = basis(P,n)
            @test  p(x) ≈ SP.classical_hypergeometric(P, n,  x)
        end
    end

end

@testset "Jacobi" begin

    T = Float64
    x = variable()

    # ChebyshevT = J(-1/2, -1/2)
    for (alpha_beta, Q) in  (((-1/2, -1/2), Chebyshev{T}),
                             ((1/2,1/2), ChebyshevU{T}),
                             ((0,0), Legendre{T})
                             )

        P = Jacobi{alpha_beta..., T}
        for i in  2:6
            p = basis(P,i)
            q = basis(Q,i)
            @test SP.monic(p)(x) ≈ SP.monic(q)(x)
        end

    end

    x = variable(Polynomial)
    for  (α,β)  ∈ ((1/2,1/2),  (-1/2,1/2), (1/2,-1/2), (-1/2, -1/2))
        P =  Jacobi{α,β}
        for n in 2:5
            p  = basis(P,n)
            @test  p(x) ≈ SP.classical_hypergeometric(P, n,  x)
        end
    end

    ## We have `jacobi_eval` for α and β outside [-1, ∞)
    α,β = 1/2, -1/2
    for n in 2:5
        x0 = 1/2
        @test SP.jacobi_eval(α, β, n, x0) ≈ basis(Jacobi{α, β},n)(x0)
    end

end

## CDOP

@testset "Charlier" begin

    x = variable()
    for μ ∈ (1/4, 1/2, 1,2)
        P = Charlier{μ}
        @test basis(P,0)(x) ≈ one(x)
        @test basis(P,1)(x) ≈ 1 - x/μ
        @test basis(P,2)(x) ≈ (x^2  +  μ^2 - x *(1 + 2μ))/μ^2

    end

    for μ ∈ (1/4, 1/2, 1, 2)
        P = Charlier{μ}
        for i in 0:5
            @test basis(P,i)(x) ≈  SP.classical_hypergeometric(P,i,x)
        end
    end

end

@testset  "Hahn"  begin

    x  = variable()
    for  α ∈  (1/4, 1/2)
        for β  ∈ (1/4, 1/2)
            𝐍  = 5

            P = Hahn{α,β,𝐍}
            for i in 0:𝐍
                @test basis(P,i)(x) ≈  SP.classical_hypergeometric(P,i,x)
            end
        end
    end

end

@testset "Krawchouk" begin

    x = variable()
    for p ∈ (1/4, 1/2, 1, 2)
        for N ∈ (3,6)
            P = Krawchouk{p,N}
            @test basis(P,0)(x) ≈ one(x)
            @test basis(P,1)(x) ≈ -N*p + x
            @test basis(P,2)(x) ≈ 1/2*(N^2*p^2 + x*(2p + x - 1) - N*p*(p+2x))
        end

    end

    for p ∈ (1/4, 1/2, 1, 2)
        for N ∈ (3,6)
            P = Krawchouk{p,N}
            for i in 0:N
                @test basis(P,i)(x) ≈  SP.classical_hypergeometric(P,i,x)
            end
        end
    end

end

@testset  "Meixner" begin

    x = variable()
    for γ ∈ (1/4, 1/2, 3/4)
        for μ ∈ (1/4, 1/2, 3/4)
            P = Meixner{γ,μ}
            for i in 0:4
                @test basis(P,i)(x) ≈  SP.classical_hypergeometric(P,i,x)
            end
        end
    end

end


# ------------------------------------------------------------------
# Copyright (c) 2015, Júlio Hoffimann Mendes <juliohm@stanford.edu>
# Licensed under the ISC License. See LICENCE in the project root.
# ------------------------------------------------------------------

function dijkstra_cut(A1::AbstractArray, A2::AbstractArray, dim::Integer)
  # permute dimensions so that the algorithm is
  # the same for cuts in x, y and z directions
  B = abs.(A1 - A2)
  if dim == 1
    B = permutedims(B, [1,2,3])
  elseif dim == 2
    B = permutedims(B, [2,1,3])
  elseif dim == 3
    B = permutedims(B, [3,2,1])
  end

  mx, my, mz = size(B)

  # accumulation cube and overlap mask
  E = zeros(size(B)); M = falses(size(B))

  # pad accumulation cube with +inf
  Epad(i,j,k) = all(0 .< [i,j,k] .≤ [mx,my,mz]) ? E[i,j,k] : Inf

  # forward accumulation along 3D cube
  E[:,:,1] = B[:,:,1]
  for k=2:mz, i=1:mx, j=1:my
    square = [Epad(i-1 , j-1 , k-1) , Epad(i-1 , j , k-1) , Epad(i-1 , j+1 , k-1),
              Epad(i   , j-1 , k-1) , Epad(i   , j , k-1) , Epad(i   , j+1 , k-1),
              Epad(i+1 , j-1 , k-1) , Epad(i+1 , j , k-1) , Epad(i+1 , j+1 , k-1)]
    E[i,j,k] = B[i,j,k] + minimum(square)
  end

  # forward accumulation along 2D square (last slice of the cube)
  zslice = view(E,:,:,mz)
  for j=2:my
    zslice[1,j] += minimum(zslice[1:2,j-1])
    for i=2:mx-1
      zslice[i,j] += minimum(zslice[i-1:i+1,j-1])
    end
    zslice[mx,j] += minimum(zslice[mx-1:mx,j-1])
  end

  # backward search along last slice
  idx = argmin(zslice[:,my])
  mslice = view(M,:,:,mz)
  mslice[1:idx,my] = trues(idx)
  idxvec = zeros(Int, my); idxvec[my] = idx # keep track of indexes
  width = isodd(mx) ? mx+1 : mx # avoid zig-zag artifact
  for j=my-1:-1:1
    for i=1:width
      if idx < mx && minimum(zslice[max(idx-1,1):idx+1,j]) == zslice[idx+1,j]
        idx += 1
      elseif idx > 1 && zslice[idx-1,j] ≤ zslice[idx,j]
        idx -= 1
      end
    end
    mslice[1:idx,j] .= true
    idxvec[j] = idx
  end

  # backward search along cube
  for j=1:my
    yslice = view(E,:,j,:)
    idx = idxvec[j]
    for k=mz-1:-1:1
      for i=1:width
        if idx < mx && minimum(yslice[max(idx-1,1):idx+1,k]) == yslice[idx+1,k]
          idx += 1
        elseif idx > 1 && yslice[idx-1,k] ≤ yslice[idx,k]
          idx -= 1
        end
      end
      M[1:idx,j,k] .= true
    end
  end

  # permute back to original shape
  if dim == 1
    M = permutedims(M, [1,2,3])
  elseif dim == 2
    M = permutedims(M, [2,1,3])
  elseif dim == 3
    M = permutedims(M, [3,2,1])
  end

  M
end