#=
  Decomposition of global constraint global_cardinality in Julia ConstraintSolver.

  See Global Constraint Catalog
  http://www.emn.fr/x-info/sdemasse/gccat/Cglobal_cardinality.html
  """
  Example:
  <3,3,8,6>,
   (val-3 noccurrence-2,
    val−5 noccurrence-0,
    val-6 noccurrence-1)

  The global_cardinality constraint holds since values 3, 5 and 6
  respectively occur 2, 0 and 1 times within the collection
  <3,3,8,6> and since no constraint was specified for value 8.
  """

  Notes:
     - The version here is limited. See below for details.
     - Picat has already a global_cardinality/2, but the details
       of the parameters are different.

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK
const CS = ConstraintSolver
include("constraints_utils.jl")

#
# Note: global_cardinality_count is defined in constraints_utils.jl
#
function gcc_test(n=5,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>true,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS, # <-
                                                            # "traverse_strategy"=>:DBFS,

                                                            "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            # "branch_split"=>:InHalf, # <-

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>16,

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                        ))

    @variable(model, 1 <= a[1:n] <= n, Int)
    @variable(model, 0 <= gcc[1:n] <= n, Int)

    global_cardinality_count(model, a, gcc)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                ax = convert.(Integer,JuMP.value.(a; result=sol))
                gccx = convert.(Integer,JuMP.value.(gcc; result=sol))
                # println("position:$positionx\n")
                println("a:$ax gcc:$gccx\n")
            end
        end
    else
        println("status:$status")
    end
    return status
end

@time gcc_test(5)


using FractalTools
using Makie
using GeometryBasics

# Construct a mesh
pts = [
    [0., 0.],
    [1., 0.],
    [1., 1.]
]
msh = trimesh(pts...)
trange = 1 : length(faces(msh))

# Plot mesh
fig = Figure()
ax, plt = mesh(fig[1, 1], msh, color=last.(msh.position))
wireframe!(ax, msh)
cb = Colorbar(fig[1, 2], plt, width=20)
button1 = Button(fig[2, 1], tellwidth=false, label="Increase")
button2 = Button(fig[2, 2], tellwidth=false, label="Decrease")
label = Label(fig[3, :], tellwidth=false)

# Plot a triangle
coords = coordinates(msh)
trigs = faces(msh)
idx = Node(1)
x = @lift(getindex.(coords[trigs[$idx]], 1) |> collect)
y = @lift(getindex.(coords[trigs[$idx]], 2) |> collect)
scatter!(ax, x, y)

on(button1.clicks) do n
    idx[] += 1
    label.text = "idx = $(idx[]), x = $(x[]), y = $(y[])"
end

on(button2.clicks) do n
    idx[] -= 1
    label.text = "idx = $(idx[]), x = $(x[]), y = $(y[])"
end

# Bind keybord interacton
scene = fig.scene
on(scene.events.keyboardbuttons) do button
    if ispressed(button, Keyboard.up)
        idx[] += 1
        label.text = "idx = $(idx[]), x = $(x[]), y = $(y[])"
    end
    if ispressed(button, Keyboard.down)
        idx[] -= 1
        label.text = "idx = $(idx[]), x = $(x[]), y = $(y[])"
    end
end

display(fig)


"""
    Interpolation{dim, ref_shape, order}()

Return an `Interpolation` of given dimension `dim`, reference shape
(see see [`AbstractRefShape`](@ref)) `ref_shape` and order `order`.
`order` corresponds to the highest order term in the polynomial.
The interpolation is used to define shape functions to interpolate
a function between nodes.

The following interpolations are implemented:

* `Lagrange{1,RefCube,1}`
* `Lagrange{1,RefCube,2}`
* `Lagrange{2,RefCube,1}`
* `Lagrange{2,RefCube,2}`
* `Lagrange{2,RefTetrahedron,1}`
* `Lagrange{2,RefTetrahedron,2}`
* `Lagrange{3,RefCube,1}`
* `Serendipity{2,RefCube,2}`
* `Serendipity{3,RefCube,2}`
* `Lagrange{3,RefTetrahedron,1}`
* `Lagrange{3,RefTetrahedron,2}`

# Examples
```jldoctest
julia> ip = Lagrange{2,RefTetrahedron,2}()
Ferrite.Lagrange{2,Ferrite.RefTetrahedron,2}()

julia> getnbasefunctions(ip)
6
```
"""
abstract type Interpolation{dim,shape,order} end

"""
Return the dimension of an `Interpolation`
"""
@inline getdim(ip::Interpolation{dim}) where {dim} = dim

"""
Return the reference shape of an `Interpolation`
"""
@inline getrefshape(ip::Interpolation{dim,shape}) where {dim,shape} = shape

"""
Return the polynomial order of the `Interpolation`
"""
@inline getorder(ip::Interpolation{dim,shape,order}) where {dim,shape,order} = order

"""
Compute the value of the shape functions at a point ξ for a given interpolation
"""
function value(ip::Interpolation{dim}, ξ::Vec{dim,T}) where {dim,T}
    [value(ip, i, ξ) for i in 1:getnbasefunctions(ip)]
end

"""
Compute the gradients of the shape functions at a point ξ for a given interpolation
"""
function derivative(ip::Interpolation{dim}, ξ::Vec{dim,T}) where {dim,T}
    [gradient(ξ -> value(ip, i, ξ), ξ) for i in 1:getnbasefunctions(ip)]
end

#####################
# Utility functions #
#####################

"""
Return the number of base functions for an [`Interpolation`](@ref) or `Values` object.
"""
getnbasefunctions

# struct that gathers all the information needed to distribute
# dofs for a given interpolation.
struct InterpolationInfo
    # TODO: Can be smaller than `Int` if that matters...
    nvertexdofs::Int
    nedgedofs::Int
    nfacedofs::Int
    ncelldofs::Int
    dim::Int
    InterpolationInfo(interpolation::Interpolation{dim}) where {dim} =
        new(nvertexdofs(interpolation), nedgedofs(interpolation),
            nfacedofs(interpolation),   ncelldofs(interpolation), dim)
end

# The following functions are used to distribute the dofs. Definitions:
#   vertexdof: dof on a "corner" of the reference shape
#   facedof: dof in the dim-1 dimension (line in 2D, surface in 3D)
#   edgedof: dof on a line between 2 vertices (i.e. "corners") (3D only)
#   celldof: dof that is local to the element

# Fallbacks for the interpolations which are used to distribute the dofs correctly
nvertexdofs(::Interpolation) = 0
nedgedofs(::Interpolation)   = 0
nfacedofs(::Interpolation)   = 0
ncelldofs(::Interpolation)   = 0

# Needed for distrubuting dofs on shells correctly (face in 2d is edge in 3d)
edges(ip::Interpolation{2}) = faces(ip)
nedgedofs(ip::Interpolation{2}) = nfacedofs(ip)

# Fallbacks for vertices
vertices(::Interpolation{2,RefCube}) = (1,2,3,4)
vertices(::Interpolation{3,RefCube}) = (1,2,3,4,5,6,7,8)
vertices(::Interpolation{2,RefTetrahedron}) = (1,2,3)
vertices(::Interpolation{3,RefTetrahedron}) = (1,2,3,4)

############
# Lagrange #
############
struct Lagrange{dim,shape,order} <: Interpolation{dim,shape,order} end

getlowerdim(::Lagrange{dim,shape,order}) where {dim,shape,order} = Lagrange{dim-1,shape,order}()
getlowerorder(::Lagrange{dim,shape,order}) where {dim,shape,order} = Lagrange{dim,shape,order-1}()

##################################
# Lagrange dim 1 RefCube order 1 #
##################################
getnbasefunctions(::Lagrange{1,RefCube,1}) = 2
nvertexdofs(::Lagrange{1,RefCube,1}) = 1

faces(::Lagrange{1,RefCube,1}) = ((1,), (2,))

function reference_coordinates(::Lagrange{1,RefCube,1})
    return [Vec{1, Float64}((-1.0,)),
            Vec{1, Float64}(( 1.0,))]
end

function value(ip::Lagrange{1,RefCube,1}, i::Int, ξ::Vec{1})
    ξ_x = ξ[1]
    i == 1 && return (1 - ξ_x) * 0.5
    i == 2 && return (1 + ξ_x) * 0.5
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end

##################################
# Lagrange dim 1 RefCube order 2 #
##################################
getnbasefunctions(::Lagrange{1,RefCube,2}) = 3
nvertexdofs(::Lagrange{1,RefCube,2}) = 1
ncelldofs(::Lagrange{1,RefCube,2}) = 1

faces(::Lagrange{1,RefCube,2}) = ((1,), (2,))

function reference_coordinates(::Lagrange{1,RefCube,2})
    return [Vec{1, Float64}((-1.0,)),
            Vec{1, Float64}(( 1.0,)),
            Vec{1, Float64}(( 0.0,))]
end

function value(ip::Lagrange{1,RefCube,2}, i::Int, ξ::Vec{1})
    ξ_x = ξ[1]
    i == 1 && return ξ_x * (ξ_x - 1) * 0.5
    i == 2 && return ξ_x * (ξ_x + 1) * 0.5
    i == 3 && return 1 - ξ_x^2
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end

##################################
# Lagrange dim 2 RefCube order 1 #
##################################
getnbasefunctions(::Lagrange{2,RefCube,1}) = 4
nvertexdofs(::Lagrange{2,RefCube,1}) = 1

faces(::Lagrange{2,RefCube,1}) = ((1,2), (2,3), (3,4), (4,1))

function reference_coordinates(::Lagrange{2,RefCube,1})
    return [Vec{2, Float64}((-1.0, -1.0)),
            Vec{2, Float64}(( 1.0, -1.0)),
            Vec{2, Float64}(( 1.0,  1.0,)),
            Vec{2, Float64}((-1.0,  1.0,))]
end

function value(ip::Lagrange{2,RefCube,1}, i::Int, ξ::Vec{2})
    ξ_x = ξ[1]
    ξ_y = ξ[2]
    i == 1 && return (1 - ξ_x) * (1 - ξ_y) * 0.25
    i == 2 && return (1 + ξ_x) * (1 - ξ_y) * 0.25
    i == 3 && return (1 + ξ_x) * (1 + ξ_y) * 0.25
    i == 4 && return (1 - ξ_x) * (1 + ξ_y) * 0.25
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end

##################################
# Lagrange dim 2 RefCube order 2 #
##################################
getnbasefunctions(::Lagrange{2,RefCube,2}) = 9
nvertexdofs(::Lagrange{2,RefCube,2}) = 1
nfacedofs(::Lagrange{2,RefCube,2}) = 1
ncelldofs(::Lagrange{2,RefCube,2}) = 1

faces(::Lagrange{2,RefCube,2}) = ((1,2,5), (2,3,6), (3,4,7), (4,1,8))

function reference_coordinates(::Lagrange{2,RefCube,2})
    return [Vec{2, Float64}((-1.0, -1.0)),
            Vec{2, Float64}(( 1.0, -1.0)),
            Vec{2, Float64}(( 1.0,  1.0)),
            Vec{2, Float64}((-1.0,  1.0)),
            Vec{2, Float64}(( 0.0, -1.0)),
            Vec{2, Float64}(( 1.0,  0.0)),
            Vec{2, Float64}(( 0.0,  1.0)),
            Vec{2, Float64}((-1.0,  0.0)),
            Vec{2, Float64}(( 0.0,  0.0))]
end

function value(ip::Lagrange{2,RefCube,2}, i::Int, ξ::Vec{2})
    ξ_x = ξ[1]
    ξ_y = ξ[2]
    i == 1 && return (ξ_x^2 - ξ_x) * (ξ_y^2 - ξ_y) * 0.25
    i == 2 && return (ξ_x^2 + ξ_x) * (ξ_y^2 - ξ_y) * 0.25
    i == 3 && return (ξ_x^2 + ξ_x) * (ξ_y^2 + ξ_y) * 0.25
    i == 4 && return (ξ_x^2 - ξ_x) * (ξ_y^2 + ξ_y) * 0.25
    i == 5 && return (1 - ξ_x^2) * (ξ_y^2 - ξ_y) * 0.5
    i == 6 && return (ξ_x^2 + ξ_x) * (1 - ξ_y^2) * 0.5
    i == 7 && return (1 - ξ_x^2) * (ξ_y^2 + ξ_y) * 0.5
    i == 8 && return (ξ_x^2 - ξ_x) * (1 - ξ_y^2) * 0.5
    i == 9 && return (1 - ξ_x^2) * (1 - ξ_y^2)
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end

#########################################
# Lagrange dim 2 RefTetrahedron order 1 #
#########################################
getnbasefunctions(::Lagrange{2,RefTetrahedron,1}) = 3
getlowerdim(::Lagrange{2, RefTetrahedron, order}) where {order} = Lagrange{1, RefCube, order}()
nvertexdofs(::Lagrange{2,RefTetrahedron,1}) = 1

vertices(::Lagrange{2,RefTetrahedron,1}) = (1,2,3)
faces(::Lagrange{2,RefTetrahedron,1}) = ((1,2), (2,3), (3,1))

function reference_coordinates(::Lagrange{2,RefTetrahedron,1})
    return [Vec{2, Float64}((1.0, 0.0)),
            Vec{2, Float64}((0.0, 1.0)),
            Vec{2, Float64}((0.0, 0.0))]
end

function value(ip::Lagrange{2,RefTetrahedron,1}, i::Int, ξ::Vec{2})
    ξ_x = ξ[1]
    ξ_y = ξ[2]
    i == 1 && return ξ_x
    i == 2 && return ξ_y
    i == 3 && return 1. - ξ_x - ξ_y
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end

#########################################
# Lagrange dim 2 RefTetrahedron order 2 #
#########################################
getnbasefunctions(::Lagrange{2,RefTetrahedron,2}) = 6
nvertexdofs(::Lagrange{2,RefTetrahedron,2}) = 1
nfacedofs(::Lagrange{2,RefTetrahedron,2}) = 1

vertices(::Lagrange{2,RefTetrahedron,2}) = (1,2,3)
faces(::Lagrange{2,RefTetrahedron,2}) = ((1,2,4), (2,3,5), (3,1,6))

function reference_coordinates(::Lagrange{2,RefTetrahedron,2})
    return [Vec{2, Float64}((1.0, 0.0)),
            Vec{2, Float64}((0.0, 1.0)),
            Vec{2, Float64}((0.0, 0.0)),
            Vec{2, Float64}((0.5, 0.5)),
            Vec{2, Float64}((0.0, 0.5)),
            Vec{2, Float64}((0.5, 0.0))]
end

function value(ip::Lagrange{2,RefTetrahedron,2}, i::Int, ξ::Vec{2})
    ξ_x = ξ[1]
    ξ_y = ξ[2]
    γ = 1. - ξ_x - ξ_y
    i == 1 && return ξ_x * (2ξ_x - 1)
    i == 2 && return ξ_y * (2ξ_y - 1)
    i == 3 && return γ * (2γ - 1)
    i == 4 && return 4ξ_x * ξ_y
    i == 5 && return 4ξ_y * γ
    i == 6 && return 4ξ_x * γ
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end

#########################################
# Lagrange dim 3 RefTetrahedron order 1 #
#########################################
getnbasefunctions(::Lagrange{3,RefTetrahedron,1}) = 4
nvertexdofs(::Lagrange{3,RefTetrahedron,1}) = 1

faces(::Lagrange{3,RefTetrahedron,1}) = ((1,2,3), (1,2,4), (2,3,4), (1,4,3))
edges(::Lagrange{3,RefTetrahedron,1}) = ((1,2), (2,3), (3,1), (1,4), (2,4), (3,4))

function reference_coordinates(::Lagrange{3,RefTetrahedron,1})
    return [Vec{3, Float64}((0.0, 0.0, 0.0)),
            Vec{3, Float64}((1.0, 0.0, 0.0)),
            Vec{3, Float64}((0.0, 1.0, 0.0)),
            Vec{3, Float64}((0.0, 0.0, 1.0))]
end

function value(ip::Lagrange{3,RefTetrahedron,1}, i::Int, ξ::Vec{3})
    ξ_x = ξ[1]
    ξ_y = ξ[2]
    ξ_z = ξ[3]
    i == 1 && return 1.0 - ξ_x - ξ_y - ξ_z
    i == 2 && return ξ_x
    i == 3 && return ξ_y
    i == 4 && return ξ_z
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end

#########################################
# Lagrange dim 3 RefTetrahedron order 2 #
#########################################
getnbasefunctions(::Lagrange{3,RefTetrahedron,2}) = 10
nvertexdofs(::Lagrange{3,RefTetrahedron,2}) = 1
nedgedofs(::Lagrange{3,RefTetrahedron,2}) = 1

faces(::Lagrange{3,RefTetrahedron,2}) = ((1,2,3,5,6,7), (1,2,4,5,9,8), (2,3,4,6,10,9), (1,4,3,8,10,7))
edges(::Lagrange{3,RefTetrahedron,2}) = ((1,5,2), (2,6,3), (3,7,1), (1,8,4), (2,9,4), (3,10,4))

function reference_coordinates(::Lagrange{3,RefTetrahedron,2})
    return [Vec{3, Float64}((0.0, 0.0, 0.0)),
            Vec{3, Float64}((1.0, 0.0, 0.0)),
            Vec{3, Float64}((0.0, 1.0, 0.0)),
            Vec{3, Float64}((0.0, 0.0, 1.0)),
            Vec{3, Float64}((0.5, 0.0, 0.0)),
            Vec{3, Float64}((0.5, 0.5, 0.0)),
            Vec{3, Float64}((0.0, 0.5, 0.0)),
            Vec{3, Float64}((0.0, 0.0, 0.5)),
            Vec{3, Float64}((0.5, 0.0, 0.5)),
            Vec{3, Float64}((0.0, 0.5, 0.5))]
end

# http://www.colorado.edu/engineering/CAS/courses.d/AFEM.d/AFEM.Ch09.d/AFEM.Ch09.pdf
# http://www.colorado.edu/engineering/CAS/courses.d/AFEM.d/AFEM.Ch10.d/AFEM.Ch10.pdf
function value(ip::Lagrange{3,RefTetrahedron,2}, i::Int, ξ::Vec{3})
    ξ_x = ξ[1]
    ξ_y = ξ[2]
    ξ_z = ξ[3]
    i == 1  && return (-2 * ξ_x - 2 * ξ_y - 2 * ξ_z + 1) * (-ξ_x - ξ_y - ξ_z + 1)
    i == 2  && return ξ_x * (2 * ξ_x - 1)
    i == 3  && return ξ_y * (2 * ξ_y - 1)
    i == 4  && return ξ_z * (2 * ξ_z - 1)
    i == 5  && return ξ_x * (-4 * ξ_x - 4 * ξ_y - 4 * ξ_z + 4)
    i == 6  && return 4 * ξ_x * ξ_y
    i == 7  && return 4 * ξ_y * (-ξ_x - ξ_y - ξ_z + 1)
    i == 8  && return ξ_z * (-4 * ξ_x - 4 * ξ_y - 4 * ξ_z + 4)
    i == 9  && return 4 * ξ_x * ξ_z
    i == 10 && return 4 * ξ_y * ξ_z
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end

##################################
# Lagrange dim 3 RefCube order 1 #
##################################
getnbasefunctions(::Lagrange{3,RefCube,1}) = 8
nvertexdofs(::Lagrange{3,RefCube,1}) = 1

faces(::Lagrange{3,RefCube,1}) = ((1,4,3,2), (1,2,6,5), (2,3,7,6), (3,4,8,7), (1,5,8,4), (5,6,7,8))
edges(::Lagrange{3,RefCube,1}) = ((1,2), (2,3), (3,4), (4,1), (1,5), (2,6), (3,7), (4,8), (5,6), (6,7), (7,8), (8,5))

function reference_coordinates(::Lagrange{3,RefCube,1})
    return [Vec{3, Float64}((-1.0, -1.0, -1.0)),
            Vec{3, Float64}(( 1.0, -1.0, -1.0)),
            Vec{3, Float64}(( 1.0,  1.0, -1.0)),
            Vec{3, Float64}((-1.0,  1.0, -1.0)),
            Vec{3, Float64}((-1.0, -1.0,  1.0)),
            Vec{3, Float64}(( 1.0, -1.0,  1.0)),
            Vec{3, Float64}(( 1.0,  1.0,  1.0)),
            Vec{3, Float64}((-1.0,  1.0,  1.0))]
end

function value(ip::Lagrange{3,RefCube,1}, i::Int, ξ::Vec{3})
    ξ_x = ξ[1]
    ξ_y = ξ[2]
    ξ_z = ξ[3]
    i == 1 && return 0.125(1 - ξ_x) * (1 - ξ_y) * (1 - ξ_z)
    i == 2 && return 0.125(1 + ξ_x) * (1 - ξ_y) * (1 - ξ_z)
    i == 3 && return 0.125(1 + ξ_x) * (1 + ξ_y) * (1 - ξ_z)
    i == 4 && return 0.125(1 - ξ_x) * (1 + ξ_y) * (1 - ξ_z)
    i == 5 && return 0.125(1 - ξ_x) * (1 - ξ_y) * (1 + ξ_z)
    i == 6 && return 0.125(1 + ξ_x) * (1 - ξ_y) * (1 + ξ_z)
    i == 7 && return 0.125(1 + ξ_x) * (1 + ξ_y) * (1 + ξ_z)
    i == 8 && return 0.125(1 - ξ_x) * (1 + ξ_y) * (1 + ξ_z)
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end

###############
# Serendipity #
###############
struct Serendipity{dim,shape,order} <: Interpolation{dim,shape,order} end

#####################################
# Serendipity dim 2 RefCube order 2 #
#####################################
getnbasefunctions(::Serendipity{2,RefCube,2}) = 8
getlowerdim(::Serendipity{2,RefCube,2}) = Lagrange{1,RefCube,2}()
getlowerorder(::Serendipity{2,RefCube,2}) = Lagrange{2,RefCube,1}()
nvertexdofs(::Serendipity{2,RefCube,2}) = 1
nfacedofs(::Serendipity{2,RefCube,2}) = 1

faces(::Serendipity{2,RefCube,2}) = ((1,2,5), (2,3,6), (3,4,7), (4,1,8))

function reference_coordinates(::Serendipity{2,RefCube,2})
    return [Vec{2, Float64}((-1.0, -1.0)),
            Vec{2, Float64}(( 1.0, -1.0)),
            Vec{2, Float64}(( 1.0,  1.0)),
            Vec{2, Float64}((-1.0,  1.0)),
            Vec{2, Float64}(( 0.0, -1.0)),
            Vec{2, Float64}(( 1.0,  0.0)),
            Vec{2, Float64}(( 0.0,  1.0)),
            Vec{2, Float64}((-1.0,  0.0))]
end

function value(ip::Serendipity{2,RefCube,2}, i::Int, ξ::Vec{2})
    ξ_x = ξ[1]
    ξ_y = ξ[2]
    i == 1 && return (1 - ξ_x) * (1 - ξ_y) * 0.25(-ξ_x - ξ_y - 1)
    i == 2 && return (1 + ξ_x) * (1 - ξ_y) * 0.25( ξ_x - ξ_y - 1)
    i == 3 && return (1 + ξ_x) * (1 + ξ_y) * 0.25( ξ_x + ξ_y - 1)
    i == 4 && return (1 - ξ_x) * (1 + ξ_y) * 0.25(-ξ_x + ξ_y - 1)
    i == 5 && return 0.5(1 - ξ_x * ξ_x) * (1 - ξ_y)
    i == 6 && return 0.5(1 + ξ_x) * (1 - ξ_y * ξ_y)
    i == 7 && return 0.5(1 - ξ_x * ξ_x) * (1 + ξ_y)
    i == 8 && return 0.5(1 - ξ_x) * (1 - ξ_y * ξ_y)
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end

##################################
# Serendipity dim 3 RefCube order 2 #
##################################
getnbasefunctions(::Serendipity{3,RefCube,2}) = 20
getlowerdim(::Serendipity{3,RefCube,2}) = Serendipity{2,RefCube,2}()
getlowerorder(::Serendipity{3,RefCube,2}) = Lagrange{3,RefCube,1}()
nvertexdofs(::Serendipity{3,RefCube,2}) = 1
nedgedofs(::Serendipity{3,RefCube,2}) = 1

faces(::Serendipity{3,RefCube,2}) = ((1,4,3,2,12,11,10,9), (1,2,6,5,9,18,13,17), (2,3,7,6,10,19,14,18), (3,4,8,7,11,20,15,19), (1,5,8,4,17,16,20,12), (5,6,7,8,13,14,15,16))

function reference_coordinates(::Serendipity{3,RefCube,2})
    return [Vec{3, Float64}((-1.0, -1.0, -1.0)),
            Vec{3, Float64}(( 1.0, -1.0, -1.0)),
            Vec{3, Float64}(( 1.0,  1.0, -1.0)),
            Vec{3, Float64}((-1.0,  1.0, -1.0)),
            Vec{3, Float64}((-1.0, -1.0,  1.0)),
            Vec{3, Float64}(( 1.0, -1.0,  1.0)),
            Vec{3, Float64}(( 1.0,  1.0,  1.0)),
            Vec{3, Float64}((-1.0,  1.0,  1.0)),
            Vec{3, Float64}((0.0, -1.0, -1.0)),
            Vec{3, Float64}((1.0, 0.0, -1.0)),
            Vec{3, Float64}((0.0, 1.0, -1.0)),
            Vec{3, Float64}((-1.0, 0.0, -1.0)),
            Vec{3, Float64}((0.0, -1.0, 1.0)),
            Vec{3, Float64}((1.0, 0.0, 1.0)),
            Vec{3, Float64}((0.0, 1.0, 1.0)),
            Vec{3, Float64}((-1.0, 0.0, 1.0)),
            Vec{3, Float64}((-1.0, -1.0, 0.0)),
            Vec{3, Float64}((1.0, -1.0, 0.0)),
            Vec{3, Float64}((1.0, 1.0, 0.0)),
            Vec{3, Float64}((-1.0, 1.0, 0.0)),]
end

function value(ip::Serendipity{3,RefCube,2}, i::Int, ξ::Vec{3})
    ξ_x = ξ[1]
    ξ_y = ξ[2]
    ξ_z = ξ[3]
    i == 1 && return 0.125(1 - ξ_x) * (1 - ξ_y) * (1 - ξ_z) - 0.5(value(ip,12,ξ) + value(ip,9,ξ) + value(ip,17,ξ))
    i == 2 && return 0.125(1 + ξ_x) * (1 - ξ_y) * (1 - ξ_z) - 0.5(value(ip,9,ξ) + value(ip,10,ξ) + value(ip,18,ξ))
    i == 3 && return 0.125(1 + ξ_x) * (1 + ξ_y) * (1 - ξ_z) - 0.5(value(ip,10,ξ) + value(ip,11,ξ) + value(ip,19,ξ))
    i == 4 && return 0.125(1 - ξ_x) * (1 + ξ_y) * (1 - ξ_z) - 0.5(value(ip,11,ξ) + value(ip,12,ξ) + value(ip,20,ξ))
    i == 5 && return 0.125(1 - ξ_x) * (1 - ξ_y) * (1 + ξ_z) - 0.5(value(ip,16,ξ) + value(ip,13,ξ) + value(ip,17,ξ))
    i == 6 && return 0.125(1 + ξ_x) * (1 - ξ_y) * (1 + ξ_z) - 0.5(value(ip,13,ξ) + value(ip,14,ξ) + value(ip,18,ξ))
    i == 7 && return 0.125(1 + ξ_x) * (1 + ξ_y) * (1 + ξ_z) - 0.5(value(ip,14,ξ) + value(ip,15,ξ) + value(ip,19,ξ))
    i == 8 && return 0.125(1 - ξ_x) * (1 + ξ_y) * (1 + ξ_z) - 0.5(value(ip,15,ξ) + value(ip,16,ξ) + value(ip,20,ξ))
    i == 9 && return 0.25(1 - ξ_x^2) * (1 - ξ_y) * (1 - ξ_z)
    i == 10 && return 0.25(1 + ξ_x) * (1 - ξ_y^2) * (1 - ξ_z)
    i == 11 && return 0.25(1 - ξ_x^2) * (1 + ξ_y) * (1 - ξ_z)
    i == 12 && return 0.25(1 - ξ_x) * (1 - ξ_y^2) * (1 - ξ_z)
    i == 13 && return 0.25(1 - ξ_x^2) * (1 - ξ_y) * (1 + ξ_z)
    i == 14 && return 0.25(1 + ξ_x) * (1 - ξ_y^2) * (1 + ξ_z)
    i == 15 && return 0.25(1 - ξ_x^2) * (1 + ξ_y) * (1 + ξ_z)
    i == 16 && return 0.25(1 - ξ_x) * (1 - ξ_y^2) * (1 + ξ_z)
    i == 17 && return 0.25(1 - ξ_x) * (1 - ξ_y) * (1 - ξ_z^2)
    i == 18 && return 0.25(1 + ξ_x) * (1 - ξ_y) * (1 - ξ_z^2)
    i == 19 && return 0.25(1 + ξ_x) * (1 + ξ_y) * (1 - ξ_z^2)
    i == 20 && return 0.25(1 - ξ_x) * (1 + ξ_y) * (1 - ξ_z^2)
    throw(ArgumentError("no shape function $i for interpolation $ip"))
end


#=
  GLM Link Functions
=#

using Distributions

abstract type AbstractMatrixType end
struct RegularData <: AbstractMatrixType end
struct Block1D <: AbstractMatrixType end
struct Block1DParallel <: AbstractMatrixType end

# Defining the standard normal for further use
const standardNormal = Normal()

# Standard Normal pdf
function pdfSN(eta::T)::T where {T <: AbstractFloat}
  return exp(-(eta^2)/2)/sqrt(2*pi)
end

# Compliment of epsilon 1 - eps(Float)
function ceps(T::Type{<: AbstractFloat})
  return 1 - eps(T)
end

#=
  Types for Link Functions
=#
abstract type AbstractLink end
struct IdentityLink <: AbstractLink end
struct LogLink <: AbstractLink end
struct InverseLink <: AbstractLink end
struct NegativeBinomialLink{T <: AbstractFloat} <: AbstractLink
  Alpha::T
end

# abstract type AbstractProbabilityLink <: AbstractLink end
struct LogitLink <: AbstractLink end
struct ProbitLink <: AbstractLink end
struct CauchitLink <: AbstractLink end
struct OddsPowerLink{T <: AbstractFloat} <: AbstractLink
  Alpha::T
end
struct LogComplimentLink <: AbstractLink end
struct LogLogLink <: AbstractLink end
struct ComplementaryLogLogLink <: AbstractLink end
struct PowerLink{T <: AbstractFloat} <: AbstractLink
  Alpha::T
end

#====================================================================
  Copy constructors for link types
====================================================================#
import Base.copy
copy(::IdentityLink) = IdentityLink()
copy(::LogLink) = LogLink()
copy(::InverseLink) = InverseLink()
function copy(distrib::NegativeBinomialLink{T}) where {T}
  return NegativeBinomialLink(distrib.Alpha)
end
copy(::LogitLink) = LogitLink()
copy(::ProbitLink) = ProbitLink()
copy(::CauchitLink) = CauchitLink()
function copy(distrib::OddsPowerLink{T}) where {T}
  return OddsPowerLink(distrib.Alpha)
end
copy(::LogComplimentLink) = LogComplimentLink()
copy(::LogLogLink) = LogLogLink()
copy(::ComplementaryLogLogLink) = ComplementaryLogLogLink()
function copy(distrib::PowerLink{T}) where {T}
  return PowerLink(distrib.Alpha)
end
#===================================================================#

#=
  Implementation of the groups of link functions
=#

# IdentityLink Functions
@inline function linkfun(::IdentityLink, mu::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return copy(mu)
end
@inline function deta_dmu(::IdentityLink, mu::Array{T, 1}, eta::Array{T})::Array{T, 1} where {T <: AbstractFloat}
  return fill(T(1.0), size(mu))
end
@inline function linkinv(::IdentityLink, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return copy(eta)
end
# LogLink Functions
@inline function linkfun(::LogLink, mu::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return log.(mu)
end
@inline function deta_dmu(::LogLink, mu::Array{T, 1}, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return mu.^(-1)
end
@inline function linkinv(::LogLink, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return exp.(eta)
end
# InverseLink Functions
@inline function linkfun(::InverseLink, mu::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return mu.^(-1)
end
@inline function deta_dmu(::InverseLink, mu::Array{T, 1}, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return -(mu.^-2)
end
@inline function linkinv(::InverseLink, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return eta.^(-1)
end
# LogitLink Functions
@inline function linkfun(::LogitLink, mu::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return log.(mu ./(1 .- mu))
end
@inline function deta_dmu(::LogitLink, mu::Array{T, 1}, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return (mu .* (1 .- mu)).^(-1)
end
@inline function linkinv(::LogitLink, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  tmp = exp.(eta)
  return tmp ./ (1 .+ tmp)
end
# CauchitLink Functions
@inline function linkfun(::CauchitLink, mu::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return tan.(pi .* (mu .- (1/2)))
end
@inline function deta_dmu(::CauchitLink, mu::Array{T, 1}, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return pi .* (cos.(pi .* (mu .- 0.5)).^(-2))
end
@inline function linkinv(::CauchitLink, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return (atan.(eta)./pi) .+ 0.5
end
# ProbitLink Functions
@inline function linkfun(::ProbitLink, mu::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return quantile.(standardNormal, mu)
end
@inline function deta_dmu(::ProbitLink, mu::Array{T, 1}, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return pdf.(standardNormal, eta).^-1
end
@inline function linkinv(::ProbitLink, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return cdf.(standardNormal, eta)
end
# PowerLink{Alpha} Functions
@inline function linkfun(link::PowerLink, mu::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  if link.Alpha == T(0)
    return log.(mu)
  end
  return mu .^link.Alpha
end
@inline function deta_dmu(link::PowerLink, mu::Array{T, 1}, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  if link.Alpha == T(0)
    return mu.^(-1)
  end
  return link.Alpha .* (mu .^(link.Alpha - 1))
end
@inline function linkinv(link::PowerLink, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  if link.Alpha == T(0)
    return exp.(eta)
  end
  return eta .^(link.Alpha^-1)
end
# OddsPowerLink Functions
#=
  Reference from:
  1. Stata Documentation "rglm - Robust variance estimates for generalized linear models"
  2. Multivariate Methods in Epidemiology by Theodore Holford, Appendix 5 , Table A5-1
=#
# OddsPowerLink Functions
@inline function linkfun(link::OddsPowerLink, mu::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  if link.Alpha == 0
    return linkfun(LogitLink(), mu)
  end
  return ((mu ./ (1 .- mu)).^link.Alpha .- 1)./link.Alpha
end
@inline function deta_dmu(link::OddsPowerLink, mu::Array{T, 1}, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  if link.Alpha == 0
    return deta_dmu(LogitLink(), mu, eta)
  end
  return (mu .^ (link.Alpha - 1))./((1 .- mu).^(link.Alpha + 1))
end
@inline function linkinv(link::OddsPowerLink, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  if link.Alpha == 0
    return linkinv(LogitLink(), eta)
  end
  tmp = ((eta .* link.Alpha .+ 1).^(1/link.Alpha))
  return min.(max.(tmp ./ (1 .+ tmp), eps(T)), ceps(T))
end
# LogComplement Link Functions - needs further work
@inline function linkfun(::LogComplimentLink, mu::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return log.(max.(1 .- mu, eps(T)))
end
@inline function deta_dmu(::LogComplimentLink, mu::Array{T, 1}, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return -(max.(1 .- mu, eps(T))).^(-1)
end
@inline function linkinv(::LogComplimentLink, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return min.(max.(-expm1.(eta), eps(T)), ceps(T))
end
# LogLogLink Link Functions - needs further work
@inline function linkfun(::LogLogLink, mu::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return -log.(-log.(mu))
end
@inline function deta_dmu(::LogLogLink, mu::Array{T, 1}, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return -1 ./(mu .* log.(mu))
end
@inline function linkinv(::LogLogLink, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return exp.(-exp.(-eta))
end
# ComplementaryLogLogLink Functions
@inline function linkfun(::ComplementaryLogLogLink, mu::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return log.(-log.(1 .- mu))
end
@inline function deta_dmu(::ComplementaryLogLogLink, mu::Array{T, 1}, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return ((mu .- 1) .* log.(1 .- mu)).^-1
end
@inline function linkinv(::ComplementaryLogLogLink, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return 1 .- exp.(-exp.(eta))
end
# NegativeBinomialLink Functions - rewrite with 1/Alpha
@inline function linkfun(link::NegativeBinomialLink, mu::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  iAlpha = T(1)/link.Alpha
  return log.(mu ./(mu .+ iAlpha))
end
@inline function deta_dmu(link::NegativeBinomialLink, mu::Array{T, 1}, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  return 1 ./ (mu .+ link.Alpha .* mu .^2)
end
@inline function linkinv(link::NegativeBinomialLink, eta::Array{T, 1})::Array{T, 1} where {T <: AbstractFloat}
  tmp = exp.(eta)
  return tmp ./(link.Alpha .* (1 .- tmp))
end

# For block algorithm
function linkfun(link::AbstractLink, mu::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  return [linkfun(link, mu[i]) for i in 1:nBlocks]
end
function deta_dmu(link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  return [deta_dmu(link, mu[i], eta[i]) for i in 1:nBlocks]
end
function linkinv(link::AbstractLink, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(eta)
  return [linkinv(link, eta[i]) for i in 1:nBlocks]
end

# For parallel block algorithm
function linkfun(::Block1DParallel, link::AbstractLink, mu::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  ret::Array{Array{T, 1}, 1} = Array{Array{T, 1}, 1}(undef, nBlocks)
  @threads for i in 1:nBlocks
    ret[i] = linkfun(link, mu[i])
  end
  return ret
end
function deta_dmu(::Block1DParallel, link::AbstractLink, mu::Array{Array{T, 1}, 1}, eta::Array{Array{T}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(mu)
  ret::Array{Array{T, 1}, 1} = Array{Array{T, 1}, 1}(undef, nBlocks)
  @threads for i in 1:nBlocks
    ret[i] = deta_dmu(link, mu[i], eta[i])
  end
  return ret
end
function linkinv(::Block1DParallel, link::AbstractLink, eta::Array{Array{T, 1}, 1}) where {T <: AbstractFloat}
  nBlocks::Int64 = length(eta)
  ret::Array{Array{T, 1}, 1} = Array{Array{T, 1}, 1}(undef, nBlocks)
  @threads for i in 1:nBlocks
    ret[i] = linkinv(link, eta[i])
  end
  return ret
end


__precompile__(true)

module ModifiedHankelFunctionsOfOrderOneThird

using SpecialFunctions
using StaticArrays

export modifiedhankel

# Number of terms to sum in both `powerseries` and `asymptotic` solutions
const NUMTERMS = 20

# Series for auxiliary functions `f`, `f′`, `g`, and `g′` in `powerseries`.
const a₀ = cbrt(2)/gamma(2/3)
a(m) = m == 0 ? a₀ : -a(m-1)/((3m-1)*3m)
const avec = @SVector [a(i) for i = 0:NUMTERMS]

const b₀ = cbrt(2)/(3^(2/3)*gamma(4/3))
b(m) = m == 0 ? b₀ : -b(m-1)/((3m+1)*3m)
const bvec = @SVector [b(i) for i = 0:NUMTERMS]

const c₀ = a₀/2
c(m) = m == 0 ? c₀ : -c(m-1)/((3m+2)*3m)
const cvec = @SVector [c(i) for i = 0:NUMTERMS]

const d₀ = b₀
d(m) = m == 0 ? d₀ : -d(m-1)/((3m-2)*3m)
const dvec = @SVector [d(i) for i = 0:NUMTERMS]

Cterm(m) = 9*(2m - 1)^2 - 4
C(m) = convert(Float64, prod(Cterm, 1:m)/(2^(4m)*3^m*factorial(m)))
# `big` required for `factorial`; also, this begins with 1 to simplify evalpoly below
const Cvec = @SVector [i == 0 ? 1.0 : C(big(i)) for i = 0:NUMTERMS]
# begin with 0 to simplify evalpoly; Cvec[i+1] to offset from the 0 term of Cvec
const Cveci = @SVector [i == 0 ? 0.0 : Cvec[i+1]*i for i = 0:NUMTERMS]

"""
    modifiedhankel(z)

Return ``h₁``, ``h₂``, ``h₁'``, and ``h₂'``, the first and second modified Hankel functions
of order 1/3 and their derivatives.

These functions solve Stokes' equation ``d²u/dz² + zu = 0`` as a power series in ``z`` for
`abs2(z) < 36` and an approximate asymptotic expansion otherwise. The asymptotic solution is
necessary because the ``z³ⁱ`` in the power series blows up as ``i → ∞``.

!!! warning

    For very large arguments `z`, the exponential function called within the asymptotic
    expansion may overflow and return `Inf`, `Nan`, or `0` values. No warnings will be
    thrown from this module or Base for floating point issues.

For more information, see [^SCL1945].

See also: [`powerseries`](@ref), [`asymptotic`](@ref)

# Examples

```julia
julia> h1, h2, h1prime, h2prime = modifiedhankel(complex(2.687, -0.648));
```

# References

[^SCL1945]:

    The Staff of the Computation Library (1945), Tables of the modified Hankel function of
    order one-third and of their derivatives. Cambridge, MA: Harvard University Press.
"""
function modifiedhankel(z)
    if abs2(z) < 36
        # Power series solution
        h1, h2, h1p, h2p = powerseries(z)
    else
        # Asymptotic expansion
        h1, h2, h1p, h2p = asymptotic(z)
    end

    return h1, h2, h1p, h2p
end

"""
    powerseries(z)

Return ``h₁``, ``h₂``, ``h₁'``, and ``h₂'``, the first and second modified Hankel functions
of order 1/3 and their derivatives using a *power series* with 20 terms.

For more information, see [^SCL1945].

See also: [`modifiedhankel`](@ref), [`asymptotic`](@ref)

# References

[^SCL1945]:

    The Staff of the Computation Library (1945), Tables of the modified Hankel function of
    order one-third and of their derivatives. Cambridge, MA: Harvard University Press.
"""
function powerseries(z)
    z² = z^2
    z³ = z²*z  # z³

    # Auxiliary functions, beginning with zeroth terms (a₀, b₀, c₀, d₀)
    fval = evalpoly(z³, avec.data)
    gval = evalpoly(z³, bvec.data)
    f′val = evalpoly(z³, cvec.data)
    g′val = evalpoly(z³, dvec.data)

    gval *= z
    f′val *= -z²

    k1 = 1im*sqrt(3)/3*(gval - 2*fval)
    k2 = 1im*sqrt(3)/3*(g′val - 2*f′val)

    h1 = gval + k1
    h2 = gval - k1
    h1p = g′val + k2
    h2p = g′val - k2

    return h1, h2, h1p, h2p
end

"""
    asymptotic(z)

Return ``h₁``, ``h₂``, ``h₁'``, and ``h₂'``, the first and second modified Hankel functions
of order 1/3 and their derivatives using an *asymptotic expansion*, useful for `abs2(z) > 36`.

For more information, see [^SCL1945].

See also: [`modifiedhankel`](@ref), [`powerseries`](@ref)

# References

[^SCL1945]:

    The Staff of the Computation Library (1945), Tables of the modified Hankel function of
    order one-third and of their derivatives. Cambridge, MA: Harvard University Press.
"""
function asymptotic(z)
    α = cbrt(2)*3^(1/6)/sqrt(π)
    rootz = sqrt(z)
    rootz_cubed = rootz*z  # z^(3/2)

    zinv = inv(z)
    zterm = 1im/rootz_cubed  # im*z^(-3/2)
    negative_zterm = -zterm

    s = evalpoly(negative_zterm, Cvec.data)
    t = evalpoly(zterm, Cvec.data)
    sp = evalpoly(negative_zterm, Cveci.data)
    tp = evalpoly(zterm, Cveci.data)

    k1 = -3/2*zterm*zinv  # -3/2*im*z^(-5/2)
    sp *= k1
    tp *= k1  # yes, same for both

    tmpa = α/sqrt(rootz)  # α*z^(-1/4)
    k2 = 2im/3*rootz_cubed
    tmpb = k2 - 5im*π/12
    tmpc = k2 + 11im*π/12
    k3 = zinv/4  # 1/4/z

    e2 = exp(tmpb)  # exp(-tmpb) = 1/exp(tmpb)
    e2inv = inv(e2)
    h1 = e2*s
    h2 = t*e2inv
    h1p = e2*(s*(1im*rootz - k3) + sp) # rootz*z^(-1/4) = z^(1/4)
    h2p = (t*(-1im*rootz - k3) + tp)*e2inv

    e3 = exp(tmpc)  # exp(-tmpc) = 1/exp(tmpc)
    e3inv = inv(e3)
    argz = angle(z)
    if -4π/3 < argz < 0
        h1 += t*e3inv
        h1p += (t*(-1im*rootz - k3) + tp)*e3inv
    else
        h2 += e3*s
        h2p += e3*(s*(1im*rootz - k3) + sp)
    end

    h1 *= tmpa
    h2 *= tmpa
    h1p *= tmpa
    h2p *= tmpa

    return h1, h2, h1p, h2p
end

end # module


using Plots, Statistics, LinearAlgebra
using Flux, DiffEqFlux, StochasticDiffEq, DiffEqBase.EnsembleAnalysis
using FileIO, Serialization
using Distributions
using Optim
# don't call this function unless you really want to. Expensive.
# function generateTrueData(u0, tspan, tsteps, numTrajectories)
function generateTrueData(u0, tspan, tsteps, outputPath)

    numTrajectories = size(u0)[2];

    function trueDriftFunc(du, u, p, t)
        ρ = reshape(u[1:9],(3,3));
        M = reshape(u[10:18],(3,3));
        k = inv(ρ);
        Π = (k*k')/tr(k*k');
        placeHolder = [(M'*ρ); (-M^2 + tr(M^2)*Π)];
        du .= [reshape(placeHolder[1:3,1:3],9); reshape(placeHolder[4:6,1:3],9)] #du/dt = [ dρdt; dMdt]
    end

    function prob_func(prob, i, repeat)
        remake(prob,u0=u0[:,i]);
    end

    trueOdeProb = ODEProblem(trueDriftFunc, u0[:,1], tspan);

    ensembleProb = EnsembleProblem(trueOdeProb,
                                   prob_func = prob_func);#start each new problem with a different initial conditions

    #ensembleSol = solve(ensembleProb, SOSRI(), trajectories = numTrajectories);
    @time ensembleSol = solve(ensembleProb, Tsit5(), tstops = tsteps, trajectories = numTrajectories, saveat = tsteps);

    #return Array.(timeseries_point_meanvar(ensembleSol, tsteps));
    #save(ensembleSol,"/groups/chertkov/cmhyett/PureModelData/LearningTetrad/finalState.dat")
    #serialize("/groups/chertkov/cmhyett/PureModelData/LearningTetrad/trainingData.dat", ensembleSol)
    serialize(outputPath, ensembleSol);
    return ensembleSol;
end


include("macros.jl")

using StaticArrays

# old style call
function eval_UC_spline(a, b, orders, C::Array{Float64}, S::Matrix{Float64})
    d,N = size(S)

    n_x = size(C,1)
    dims = tuple(size(C)[2:end]...)
    out = zeros(n_x,N)
    CC = reshape(reinterpret(SVector{n_x,Float64},vec(C)),dims)
    SS = reshape(reinterpret(SVector{d,Float64},vec(S)),(N,))
    V  = reshape(reinterpret(SVector{n_x,Float64},vec(out)),(N,))
    eval_UC_spline!(a, b, orders, CC, SS, V)
    return out
end

function eval_UC_spline(a, b, orders, C::Array{T, d}, S::Vector{SVector{d,Float64}}) where T where d
    N = length(S)
    vals = zeros(T,N)
    eval_UC_spline!(a, b, orders, C, S, vals)
    return vals
end

@generated function eval_UC_spline!(a, b, orders, C::Array{T, d}, S::Vector{SVector{d, Float64}}, V::Vector{T}) where d where T
    # d = C.parameters[2]-1 # first dimension of C indexes the splines
    # the speed penalty of extrapolating points when iterating over point
    # seems very small so this is the default
    fun = (create_function(d,"natural"))
    return fun.args[2].args[2]
end


# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# Licensed under ASL - see ASL.md for terms and conditions.
# --------------------------------------------------------------------------

module SphericalHarmonics

using StaticArrays, LinearAlgebra

import ACE

import JuLIP.MLIPs: IPBasis, alloc_B, alloc_dB, fltype
import JuLIP: alloc_temp, alloc_temp_d, evaluate!, evaluate_d!,
			     write_dict, read_dict

const JVec = SVector{3}

export SHBasis, RSHBasis

# --------------------------------------------------------
#     Coordinates
# --------------------------------------------------------

"""
`struct PseudoSpherical` : a simple datatype storing spherical coordinates
of a point (x,y,z) in the format (r, cosφ, sinφ, cosθ, sinθ).

Use `spher2cart` and `cart2spher` to convert between cartesian and spherical
coordinates.
"""
struct PseudoSpherical{T}
	r::T
	cosφ::T
	sinφ::T
	cosθ::T
	sinθ::T
end

spher2cart(S::PseudoSpherical) = S.r * JVec(S.cosφ*S.sinθ, S.sinφ*S.sinθ, S.cosθ)

function cart2spher(R::AbstractVector)
	@assert length(R) == 3
	r = norm(R)
	φ = atan(R[2], R[1])
	sinφ, cosφ = sincos(φ)
	cosθ = R[3] / r
	sinθ = sqrt(R[1]^2+R[2]^2) / r
	return PseudoSpherical(r, cosφ, sinφ, cosθ, sinθ)
end

PseudoSpherical(φ, θ) = PseudoSpherical(1.0, cos(φ), sin(φ), cos(θ), sin(θ))

"""
convert a gradient with respect to spherical coordinates to a gradient
with respect to cartesian coordinates
"""
function dspher_to_dcart(S, f_φ_div_sinθ, f_θ)
	r = S.r + eps()
   return SVector( - (S.sinφ * f_φ_div_sinθ) + (S.cosφ * S.cosθ * f_θ),
			            (S.cosφ * f_φ_div_sinθ) + (S.sinφ * S.cosθ * f_θ),
			 			                                 - (   S.sinθ * f_θ) ) / r
end

# --------------------------------------------------------
#     Indexing
# --------------------------------------------------------

"""
	sizeP(maxDegree)

Return the size of the set of Associated Legendre Polynomials ``P_l^m(x)`` of
degree less than or equal to the given maximum degree
"""
sizeP(maxDegree) = div((maxDegree + 1) * (maxDegree + 2), 2)

"""
	sizeY(maxDegree)

Return the size of the set of real spherical harmonics ``Y_{l,m}(θ,φ)`` of
degree less than or equal to the given maximum degree
"""
sizeY(maxDegree) = (maxDegree + 1) * (maxDegree + 1)

"""
	index_p(l,m)

Return the index into a flat array of Associated Legendre Polynomials ``P_l^m``
for the given indices ``(l,m)``.
``P_l^m`` are stored in l-major order i.e. [P(0,0), [P(1,0), P(1,1), P(2,0), ...]
"""
# Base.@pure index_p(l,m) = m + div(l*(l+1), 2) + 1
@inline index_p(l::Int,m::Int) = m + div(l*(l+1), 2) + 1
index_p(l::Integer, m::Integer) = index_y(Int(l), Int(m))

"""
	index_y(l,m)

Return the index into a flat array of real spherical harmonics ``Y_{l,m}``
for the given indices ``(l,m)``.
``Y_{l,m}`` are stored in l-major order i.e.
[Y(0,0), [Y(1,-1), Y(1,0), Y(1,1), Y(2,-2), ...]
"""
Base.@pure index_y(l::Int, m::Int) = m + l + (l*l) + 1
index_y(l::Integer, m::Integer) = index_y(Int(l), Int(m))

# # i = m + l + l^2 + 1
# # m = i - l - l^2 - 1
# _li2m(l, i) = i - l - l^2 - 1

# --------------------------------------------------------
#     Associated Legendre Polynomials
#     TODO: rewrite within general interface?
#           - alloc_B, alloc_dB, alloc_temp, ...
# --------------------------------------------------------

"""
`ALPCoefficients` : an auxiliary datastructure for
evaluating the associated lagrange functions
used for the spherical harmonics
"""
struct ALPCoefficients{T}
	A::Vector{T}
	B::Vector{T}
end

ALPCoefficients(maxDegree::Integer, T::Type=Float64) =
	ALPCoefficients( Vector{T}(undef, sizeP(maxDegree)),
						  Vector{T}(undef, sizeP(maxDegree)) )

"""
	compute_coefficients(L)

Precompute coefficients ``a_l^m`` and ``b_l^m`` for all l <= L, m <= l
"""
function compute_coefficients(L::Integer, T::Type=Float64)
	coeff = ALPCoefficients(L, T)
	for l in 2:L
		ls = l*l
		lm1s = (l-1) * (l-1)
		for m in 0:(l-2)
			ms = m * m
			coeff.A[index_p(l, m)] = sqrt((4 * ls - 1.0) / (ls - ms))
			coeff.B[index_p(l, m)] = -sqrt((lm1s - ms) / (4 * lm1s - 1.0))
		end
	end
	return coeff
end

"""
	compute_coefficients(L)

Create an array large enough to store an entire set of Associated Legendre
Polynomials ``P_l^m(x)`` of maximum degree L.
"""
allocate_p(L::Int) = Array{Float64}(undef, sizeP(L))

"""
	compute_p(L, x, coeff, P)

Compute an entire set of Associated Legendre Polynomials ``P_l^m(x)``
using the given coefficients, and store in the array P.
"""
function compute_p!(L::Integer, S::PseudoSpherical{T}, coeff::ALPCoefficients{T},
					     P::Array{T,1}) where {T}
	@assert length(coeff.A) >= sizeP(L)
	@assert length(coeff.B) >= sizeP(L)
	@assert length(P) >= sizeP(L)

	temp = sqrt(0.5/π)
	P[index_p(0, 0)] = temp
	if L == 0; return P; end

	P[index_p(1, 0)] = S.cosθ * sqrt(3) * temp
	temp = - sqrt(1.5) * S.sinθ * temp
	P[index_p(1, 1)] = temp

	for l in 2:L
		il = ((l*(l+1)) ÷ 2) + 1
		ilm1 = il - l
		ilm2 = ilm1 - l + 1
		for m in 0:(l-2)
			# FOR DEBUGGING TURN ON THESE ASSERTS!
			# @assert il+m == index_p(l, m)
			# @assert ilm1+m == index_p(l-1,m)
			# @assert ilm2+m == index_p(l-2,m)
			@inbounds P[il+m] = coeff.A[il+m] * (     S.cosθ * P[ilm1+m]
					                           + coeff.B[il+m] * P[ilm2+m] )
		end
		@inbounds P[il+l-1] = S.cosθ * sqrt(2 * (l - 1) + 3) * temp
		temp = -sqrt(1.0 + 0.5 / l) * S.sinθ * temp
		@inbounds P[il+l] = temp
	end

	return P
end

function compute_dp!(L::Integer, S::PseudoSpherical{T}, coeff::ALPCoefficients{T},
					     P::Array{T,1}, dP::Array{T,1}) where T
	@assert length(coeff.A) >= sizeP(L)
	@assert length(coeff.B) >= sizeP(L)
	@assert length(P) >= sizeP(L)

	temp = sqrt(0.5/π)
	P[index_p(0, 0)] = temp
	temp_d = 0.0
	dP[index_p(0, 0)] = temp_d
	if L == 0; return nothing; end

	P[index_p(1, 0)] = S.cosθ * sqrt(3) * temp
	dP[index_p(1, 0)] = -S.sinθ * sqrt(3) * temp + S.cosθ * sqrt(3) * temp_d
	temp1, temp_d = ( - sqrt(1.5) * temp,
					      - sqrt(1.5) * (S.cosθ * temp + S.sinθ * temp_d) )
	P[index_p(1, 1)] = temp1
	dP[index_p(1, 1)] = temp_d

	for l in 2:L
		m = 0
		@inbounds P[index_p(l, m)] =
				coeff.A[index_p(l, m)] * (     S.cosθ * P[index_p(l - 1, m)]
				             + coeff.B[index_p(l, m)] * P[index_p(l - 2, m)] )
		@inbounds dP[index_p(l, m)] =
			coeff.A[index_p(l, m)] * (
							- S.sinθ * P[index_p(l - 1, m)]
							+ S.cosθ * dP[index_p(l - 1, m)]
			             + coeff.B[index_p(l, m)] * dP[index_p(l - 2, m)] )

		for m in 1:(l-2)
			@inbounds P[index_p(l, m)] =
					coeff.A[index_p(l, m)] * (     S.cosθ * P[index_p(l - 1, m)]
					             + coeff.B[index_p(l, m)] * P[index_p(l - 2, m)] )
			@inbounds dP[index_p(l, m)] =
				coeff.A[index_p(l, m)] * (
								- S.sinθ^2 * P[index_p(l - 1, m)]
								+ S.cosθ * dP[index_p(l - 1, m)]
				             + coeff.B[index_p(l, m)] * dP[index_p(l - 2, m)] )
		end
		@inbounds P[index_p(l, l - 1)] = sqrt(2 * (l - 1) + 3) * S.cosθ * temp1
		@inbounds dP[index_p(l, l - 1)] = sqrt(2 * (l - 1) + 3) * (
									        -S.sinθ^2 * temp1 + S.cosθ * temp_d )

      (temp1, temp_d) = (
					-sqrt(1.0+0.5/l) * S.sinθ * temp1,
		         -sqrt(1.0+0.5/l) * (S.cosθ * temp1 * S.sinθ + S.sinθ * temp_d) )
		@inbounds P[index_p(l, l)] = temp1
		@inbounds dP[index_p(l, l)] = temp_d
	end
	return nothing
end

"""
	compute_p(L, x)

Compute an entire set of Associated Legendre Polynomials ``P_l^m(x)`` where
``0 ≤ l ≤ L`` and ``0 ≤ m ≤ l``. Assumes ``|x| ≤ 1``.
"""
function compute_p(L::Integer, S::PseudoSpherical{T}) where {T}
	P = Array{T}(undef, sizeP(L))
	coeff = compute_coefficients(L)
	compute_p!(L, S, coeff, P)
	return P
end

function compute_dp(L::Integer, S::PseudoSpherical{T}) where {T}
	P = Array{T}(undef, sizeP(L))
	dP = Array{T}(undef, sizeP(L))
	coeff = compute_coefficients(L)
	compute_dp!(L, S, coeff, P, dP)
	return P, dP
end

compute_p(L::Integer, θ::Real) =
	compute_p(L, PseudoSpherical(0.0, 0.0, 0.0, cos(θ), sin(θ)))

compute_dp(L::Integer, θ::Real) =
	compute_dp(L, PseudoSpherical(0.0, 0.0, 0.0, cos(θ), sin(θ)))

# ------------------------------------------------------------------------
#                  Spherical Harmonics
# ------------------------------------------------------------------------

"""
evaluate complex spherical harmonics
"""
function cYlm!(Y, L, S::PseudoSpherical, P)
	@assert length(P) >= sizeP(L)
	@assert length(Y) >= sizeY(L)
   @assert abs(S.cosθ) <= 1.0

	ep = 1 / sqrt(2) + im * 0
	for l = 0:L
		Y[index_y(l, 0)] = P[index_p(l, 0)] * ep
	end

   sig = 1
   ep_fact = S.cosφ + im * S.sinφ
	for m in 1:L
		sig *= -1
		ep *= ep_fact            # ep =   exp(i *   m  * φ)
		em = sig * conj(ep)      # ep = ± exp(i * (-m) * φ)
		for l in m:L
			p = P[index_p(l,m)]
			@inbounds Y[index_y(l, -m)] = em * p   # (-1)^m * p * exp(-im*m*phi) / sqrt(2)
			@inbounds Y[index_y(l,  m)] = ep * p   #          p * exp( im*m*phi) / sqrt(2)
		end
	end

	return Y
end

"""
evaluate gradients of complex spherical harmonics
"""
function cYlm_d!(Y, dY, L, S::PseudoSpherical, P, dP)
	@assert length(P) >= sizeP(L)
	@assert length(Y) >= sizeY(L)
	@assert length(dY) >= sizeY(L)
   # @assert abs(S.cosθ) < 1.0

	# m = 0 case
	ep = 1 / sqrt(2)
	for l = 0:L
		Y[index_y(l, 0)] = P[index_p(l, 0)] * ep
		dY[index_y(l, 0)] = dspher_to_dcart(S, 0.0, dP[index_p(l, 0)] * ep)
	end

   sig = 1
   ep_fact = S.cosφ + im * S.sinφ

	for m in 1:L
		sig *= -1
		ep *= ep_fact            # ep =   exp(i *   m  * φ)
		em = sig * conj(ep)      # ep = ± exp(i * (-m) * φ)
		dep_dφ = im *   m  * ep
		dem_dφ = im * (-m) * em
		for l in m:L
			p_div_sinθ = P[index_p(l,m)]
			@inbounds Y[index_y(l, -m)] = em * p_div_sinθ * S.sinθ
			@inbounds Y[index_y(l,  m)] = ep * p_div_sinθ * S.sinθ

			dp_dθ = dP[index_p(l,m)]
			@inbounds dY[index_y(l, -m)] = dspher_to_dcart(S, dem_dφ * p_div_sinθ, em * dp_dθ)
			@inbounds dY[index_y(l,  m)] = dspher_to_dcart(S, dep_dφ * p_div_sinθ, ep * dp_dθ)
		end
	end
	# return Y, dY
end

"""
evaluate real spherical harmonics
"""
function rYlm!(Y::AbstractVector{T}, L, S::PseudoSpherical, P) where {T <: Real}
	@assert length(P) >= sizeP(L)
	@assert length(Y) >= sizeY(L)
   @assert abs(S.cosθ) <= 1.0

   oort2 = 1 / sqrt(2)
	for l = 0:L
		Y[index_y(l, 0)] = P[index_p(l, 0)] * oort2
	end

   sig = 1
	ec = 1.0 + 0 * im
   ec_fact = S.cosφ + im * S.sinφ
	for m in 1:L
		sig *= -1                # sig = (-1)^m
		ec *= ec_fact            # ec = exp(i * m  * φ) / sqrt(2)
		# cYlm = p * ec,    (also cYl{-m} = sig * p * conj(ec)), but not needed)
		# rYlm    =  Re(cYlm)
		# rYl{-m} = -Im(cYlm)
		for l in m:L
			p = P[index_p(l,m)]
			@inbounds Y[index_y(l, -m)] = -p * imag(ec)
			@inbounds Y[index_y(l,  m)] =  p * real(ec)
		end
	end

	return Y
end

"""
evaluate gradients of real spherical harmonics
"""
function rYlm_d!(Y::AbstractVector{T}, dY,
					  L, S::PseudoSpherical, P, dP) where {T <: Real}
	@assert length(P) >= sizeP(L)
	@assert length(Y) >= sizeY(L)
   @assert abs(S.cosθ) <= 1.0

   oort2 = 1 / sqrt(2)
	for l = 0:L
		Y[index_y(l, 0)] = P[index_p(l, 0)] * oort2
		dY[index_y(l, 0)] = dspher_to_dcart(S, 0.0, dP[index_p(l, 0)] * oort2)
	end

   sig = 1
	ec = 1.0 + 0 * im
   ec_fact = S.cosφ + im * S.sinφ
	drec_dφ = 0.0
	for m in 1:L
		sig *= -1                # sig = (-1)^m
		ec *= ec_fact            # ec = exp(i * m  * φ) / sqrt(2)
		dec_dφ = im * m * ec

		# cYlm = p * ec,    (also cYl{-m} = sig * p * conj(ec)), but not needed)
		# rYlm    =  Re(cYlm)
		# rYl{-m} = -Im(cYlm)
		for l in m:L
			p_div_sinθ = P[index_p(l,m)]
			p = p_div_sinθ * S.sinθ
			Y[index_y(l, -m)] = -p * imag(ec)
			Y[index_y(l,  m)] =  p * real(ec)

			dp_dθ = dP[index_p(l,m)]
			dY[index_y(l, -m)] = dspher_to_dcart(S, - imag(dec_dφ) * p_div_sinθ,
															    - imag(ec) * dp_dθ)
			dY[index_y(l,  m)] = dspher_to_dcart(S,   real(dec_dφ) * p_div_sinθ,
															      real(ec) * dp_dθ)
		end
	end

	return dY
end

# ---------------------------------------------
#      The nice basis interface
# ---------------------------------------------

abstract type AbstractSHBasis{T} <: IPBasis end

"""
complex spherical harmonics
"""
struct SHBasis{T} <: AbstractSHBasis{T}
	maxL::Int
	coeff::ALPCoefficients{T}
end

"""
real spherical harmonics
"""
struct RSHBasis{T} <: AbstractSHBasis{T}
	maxL::Int
	coeff::ALPCoefficients{T}
end

import Base.==
==(B1::AbstractSHBasis, B2::AbstractSHBasis) =
	( (B1.maxL == B1.maxL) &&
	  (typeof(B1) == typeof(B2)) )

write_dict(SH::SHBasis{T}) where {T} =
		Dict("__id__" => "ACE_SHBasis",
			  "T" => write_dict(T),
			  "maxL" => SH.maxL)

read_dict(::Val{:SHIPs_SHBasis}, D::Dict) =
	read_dict(Val{:ACE_SHBasis}(), D)

read_dict(::Val{:ACE_SHBasis}, D::Dict) =
		SHBasis(D["maxL"], read_dict(D["T"]))

SHBasis(maxL::Integer, T::Type=Float64) =
		SHBasis(Int(maxL), compute_coefficients(maxL, T))

RSHBasis(maxL::Integer, T::Type=Float64) =
		RSHBasis(Int(maxL), compute_coefficients(maxL, T))

rfltype(SH::AbstractSHBasis{T}) where {T} = T
Base.length(S::AbstractSHBasis) = sizeY(S.maxL)

fltype(::SHBasis{T}) where {T} = Complex{T}
fltype(::RSHBasis{T}) where {T} = T

alloc_B( S::SHBasis{T}, args...) where {T} =
		Vector{Complex{T}}(undef, length(S))

alloc_B( S::RSHBasis{T}, args...) where {T} =
		Vector{T}(undef, length(S))

alloc_dB(S::SHBasis{T}) where {T} =
		Vector{JVec{Complex{T}}}(undef, length(S))

alloc_dB(S::RSHBasis{T}) where {T} =
		Vector{JVec{T}}(undef, length(S))

alloc_dB(S::AbstractSHBasis, N::Integer) = alloc_dB(S)

alloc_temp(SH::AbstractSHBasis{T}, args...) where {T} = (
		P = Vector{T}(undef, sizeP(SH.maxL)), )

alloc_temp_d(SH::AbstractSHBasis{T}, args...) where {T} = (
		 P = Vector{T}(undef, sizeP(SH.maxL)),
		dP = Vector{T}(undef, sizeP(SH.maxL)) )

_evaluate!(Y, L, S, P, ::SHBasis) = cYlm!(Y, L, S, P)
_evaluate!(Y, L, S, P, ::RSHBasis) = rYlm!(Y, L, S, P)

_evaluate_d!(Y, dY, L, S, P, dP, ::SHBasis) = cYlm_d!(Y, dY, L, S, P, dP)
_evaluate_d!(Y, dY, L, S, P, dP, ::RSHBasis) = rYlm_d!(Y, dY, L, S, P, dP)

function evaluate!(Y, tmp, SH::AbstractSHBasis, R::JVec)
	L=SH.maxL
	@assert 0 <= L <= SH.maxL
	@assert length(Y) >= sizeY(L)
	S = cart2spher(R)
	compute_p!(L, S, SH.coeff, tmp.P)
	_evaluate!(Y, L, S, tmp.P, SH)
	return Y
end

function evaluate_d!(Y, dY, tmp, SH::AbstractSHBasis, R::JVec)
	L=SH.maxL
	@assert 0 <= L <= SH.maxL
	@assert length(Y) >= sizeY(L)
	# if R[1]^2+R[2]^2 < 1e-20 * R[3]^2
	# 	R = JVec(R[1]+1e-9, R[2], R[3])
	# end
	S = cart2spher(R)
	compute_dp!(L, S, SH.coeff, tmp.P, tmp.dP)
	_evaluate_d!(Y, dY, L, S, tmp.P, tmp.dP, SH)
	# return Y, dY
end

end


include("src/A1.jl")
using .A1

using GRUtils

a = 0
b = 1

p = Problem(
    (x) -> 2 / √(x + 1),
    (x) -> -2 * √(x^2 + 1),
    (x) -> 10 / (1 + x * (1 - x)),
    a,
    0,
    0.5,
    b,
    -1,
    0,
)

open("result", "w") do io
    println(io, solve(p))
end

for n in (10, 100)
    h = (b - a) / n
    x = [a + i * h for i in 0:n]
    y = solve(p, Options(n))
    if n == 10
        plot(x, y, "+", x, y, "b--", hold=true)
    else
        plot(x, y, hold=true)
    end
end

legend("n = 10", "n = 10", "n = 100", location = "upper left")
savefig("figures/comparison.pdf")


module SpiralExample

using Flux
using DiffEqFlux
using OrdinaryDiffEq
using Random
using Distributions

const noise_std = 0.1 # standard deviation of random noise added to samples
const noise_logvar = 2*log(noise_std)

# The elements of `SpiralSample`s are vectors with two entries for the two dimensions.
const SpiralSample = Vector{Vector{Float32}}

logp_x_z(x, z) = sum(log_normal_pdf(x, z, noise_logvar))

function spiral_clockwise(; start, stop, ntotal, a, b)
    zs = stop .+ 1. .- range(start, length = ntotal, stop=stop)
    rs = a .+ (b .* zs)
    xs = rs .* cos.(zs) .- 5.
    ys = rs .* sin.(zs)
    hcat(xs, ys)
end

function spiral_counterclockwise(; start, stop, ntotal, a, b)
    zs = range(start, length = ntotal, stop = stop)
    rw = a .+ (b .* 20 ./ zs)
    xs = rw .* cos.(zs) .+ 5.
    ys = rw .* sin.(zs)
    hcat(xs, ys)
end

function spiral_samples(;
        nspiral::Int = 1000, # no of spirals generated
        nsample::Int = 100, # no of trajectory points drawn from each spiral
        start::Float64 = 0.0, # spiral starting phi value
        stop::Float64 = 6*pi, # spiral ending phi value
        a::Float64 = 0.0, # parameters defining shape of Archimedean spiral
        b::Float64 = 1.0)
    # Parametric formula for 2d spiral is `r = a + b * phi`.

    # returns named tuple:
    # - orig_trajs: vector of `SpiralSample`s over a full spiral
    # - samp_trajs: noisy observations, vector of `SpiralSample`s
    # - orig_ts: vector of length `ntotal` containing the timepoints for `orig_trajs`
    # - orig_ts: vector of length `nsample` containing the timepoints for `samp_trajs`

    ntotal = 5*nsample # no of trajectory points for each spiral sample

    orig_traj_cc = spiral_counterclockwise(start = start, stop = stop, ntotal = ntotal, a = a, b = b)
    orig_traj_cw = spiral_clockwise(start = start, stop = stop, ntotal = ntotal, a = a, b = b)

    # sample starting timestamps
    orig_ts = range(start, length = ntotal, stop = stop)
    samp_ts = orig_ts[1:nsample]

    orig_trajs = []
    samp_trajs = []

    for _ in 1:nspiral
        # don't sample t0 very near the start or the end
        pvec = [1. ./ (ntotal.- 4. .* nsample) for _ in 1:(ntotal .- (4 .* nsample))]
        t0_idx = rand(Multinomial(1, pvec))
        t0_idx = argmax(t0_idx) + nsample

        cc = rand() > .5 # uniformly select rotation
        orig_traj = cc ? orig_traj_cc : orig_traj_cw
        push!(orig_trajs, Float32.(orig_traj))

        samp_traj = deepcopy(orig_traj[t0_idx:2:(t0_idx + 2*(nsample-1)),:])
        samp_traj += randn(size(samp_traj)) .* noise_std
        push!(samp_trajs, Float32.(samp_traj))
    end

    datatrafo(x) = [[x[i][j,:] for j in 1:size(x[1],1)] for i in 1:length(x)]

    return (orig_trajs = datatrafo(orig_trajs),
            samp_trajs = datatrafo(samp_trajs),
            orig_ts = orig_ts, samp_ts = samp_ts)
end

struct LatentTimeSeriesVAE
    rnn
    latentODEfunc
    latentODEparams
    decoder
end

function LatentTimeSeriesVAE(; latent_dim, obs_dim, rnn_nhidden, f_nhidden, dec_nhidden)
    rnn = Chain(RNN(obs_dim, rnn_nhidden), Dense(rnn_nhidden, latent_dim*2))

    latentODEfunc = Chain(Dense(latent_dim, f_nhidden, Flux.elu),
                          Dense(f_nhidden, f_nhidden, Flux.elu),
                          Dense(f_nhidden, latent_dim))
    latentODEparams, re = Flux.destructure(latentODEfunc)

    decoder = Chain(Dense(latent_dim, dec_nhidden, Flux.relu),
                    Dense(dec_nhidden, obs_dim))
    LatentTimeSeriesVAE(rnn, re, latentODEparams, decoder)
end

latentz0(μ, logσ) = μ .+ exp.(logσ) .* randn(Float32)

function n_ode(model, z0, t)
    tspan = (t[1], t[end])
    latent_dynamics(u, latentODEparams, t) = model.latentODEfunc(latentODEparams)(u)
    latentODEprob = ODEProblem(latent_dynamics, z0, tspan)
    Array(solve(latentODEprob, Tsit5(), u0=z0, p=model.latentODEparams, saveat=t))
end

function latent_mu_logsd(model, x::SpiralSample)
    latent_dim = nhiddennodes(model.latentODEfunc.m)
    [model.rnn(xi) for xi in reverse(x, dims=1)[1:(end-1)]]
    rnn_encoded = model.rnn(reverse(x, dims=1)[end])
    μ = rnn_encoded[1:latent_dim]
    logσ = rnn_encoded[(latent_dim+1):end]
    μ, logσ
end

# p(x,z)
function log_normal_pdf(x,mean,logvar)
    constant = log(2*pi)
    -0.5*(constant .+ logvar .+ ((x.-mean).^2 ./ exp.(logvar)))
end

# Kullback-Leibler-divergence
kl_q_p(μ, logσ) = 0.5 * sum(exp.(2 .* logσ) + μ.^2 .- 1 .- (2 .* logσ))

nhiddennodes(d::Dense) = size(d.W, 1)
nhiddennodes(c::Chain) = nhiddennodes(c.layers[end])

# loss function - ELBO
function elbo(model, x::SpiralSample, t)
    empmu, emplogsd = latent_mu_logsd(model, x)
    Flux.reset!(model.rnn)
    z0 = latentz0(empmu, emplogsd)
    pred_z = n_ode(model, z0, t)
    sumlogp_x_z = sum([logp_x_z(x[i], model.decoder(pred_z[:,i])) for i in 1:size(pred_z,2)])
    sumlogp_x_z - kl_q_p(empmu, emplogsd)
end

function train!(model, xs, t;
        epochs = 20, learningrate = 0.01, monitoring = (args...) -> nothing)

    opt = ADAM(learningrate)
    ps = Flux.params(model.rnn, model.latentODEparams, model.decoder)
    cumloss = 0.0f0

    function loss(x)
        -elbo(model, x, t) + 0.01 * sum(x->sum(x.^2), Flux.params(model.rnn))
    end

    for epoch in 1:epochs
        local curloss
        for x in xs
            grads = Flux.gradient(ps) do
              curloss = loss(x)
              return curloss
            end
            cumloss += curloss
            Flux.Optimise.update!(opt, ps, grads)
        end
        monitoring(epoch, cumloss)
        cumloss = 0.0f0
    end
end

function predictspiral(model, x::SpiralSample, t)
    predμ, predlogσ = latent_mu_logsd(model, x)
    predz0 = latentz0(predμ, predlogσ)
    predz = n_ode(model, predz0, t)
    predx = [model.decoder(predz[:,i]) for i in 1:size(predz,2)]
end

end # module SpiralExample


include("bm_test.jl")
include("bm_drift_test.jl")
include("brownian_bridge_test.jl")
include("first_time_test.jl")
include("poisson_test.jl")

using Base.Test
using StochasticProcesses

t=linspace(0, 1, 10)

let process = GeometricBrownianMotion(1., .1, 100)
  @test size(cumsim(process, t, 3)) == (10, 3)
  @test length(sim(process, t, 3)) == 3
  distribution(process, 0)
  distribution(process, 10);
  rand(process, t, 10)
end

let process = ItoIntegral((t, b)->t)
  @test size(cumsim(process, t, 3)) == (10, 3)
  @test length(sim(process, t, 3)) == 3
  rand(process, t, 10)
end

let process = SDE((x,y)->x, (x,y)->y, 0.)
  @test size(cumsim(process, t, 3)) == (10, 3)
  @test length(sim(process, t, 3)) == 3
  rand(process, t, 10)
end

let process = CompositeProcess(BrownianMotion(), (t, y) -> y.^2)
  @test size(cumsim(process, t, 3)) == (10, 3)
  @test length(sim(process, t, 3)) == 3
  rand(process, t, 10)
end


rawdata = []
open("data.txt") do file
    for line in eachline(file)
        push!(rawdata,parse(Int,line))
    end
end

sorteddata = sort(rawdata)

for y in 200:-1:1
    x = 1
    a = sorteddata[x]
    b = sorteddata[y]
    while (a + b != 2020)
        println("x=", x)
        println("y=", y)
        println("a=", a)
        println("b=", b)
        if (a + b < 2020)
            x += 1
            a = sorteddata[x]
        end
        if (a + b > 2020)
            break
        end
        if a + b == 2020
            println("Found it")
            println("a=", a)
            println("b=", b)
            answer = a * b
            println("answer=", answer)
            exit(0)
        end
    end
end


# MDP parameters, ValueIteration minimizes the cost => cost is positive, reward is negative
no_states = 10
_horizon = 5
actions = [:l, :r]
actionCost = 1.
actionsImpact = Base.ImmutableDict(:l => -1, :r => 1)
reward_states = [1, no_states]
reward = -10.
discount_factor = 1.
noise = .6

mdp = CustomFHExample(no_states, _horizon, actions, actionCost, actionsImpact, reward_states, reward, discount_factor, noise)

println("Verbose Check")
fhsolver = FiniteHorizonSolver(verbose=true)
FHPolicy = FiniteHorizonPOMDPs.solve(fhsolver, mdp);


nt = 1000
seed = 99
ssm = acqr(0.8,1.5,2,3,0.5,2)
m = acqr_simula(ssm,nt,seed)
y = m["y"]
a = ssm.a
c = ssm.c
q = ssm.q
r = ssm.r
x10 = ssm.m1
P10 = ssm.P1

###########################
print("Testing acqr_kfilter: ")
xtt,Ptt,xtt1,Ptt1,et,St,Kt,loglik = acqr_kfilter(y,a,c,q,r,x10,P10)
if sum(xtt.^2) > 1e-6
    println("OK")
end


module SurfaceGeometry

using LinearAlgebra
eye(n) = Matrix(1.0I, n, n)

include("Iterators.jl")
include("ComplexDS.jl")
export FaceBasedDS, ConnectivityDS
export FaceVRing, VertexVRing, DoubleVertexVRing

include("Utils.jl")
export subdivision, volume

include("Properties.jl")
export NormalVectors!, MeanCurvatures!, FitEllipsoid

### Essential from stabilistation
include("StabilisationMethods/stabilisationV1.jl")
include("StabilisationMethods/stabilisationV2.jl")
include("StabilisationMethods/stabilisationV3.jl")
include("StabilisationMethods/Stabilisation.jl")
export stabilise

### Just redirecting
#using ElTopo
#export Elparameters, improvemesh, improvemeshcol

end # module


arg_dict = Dict{String,Any}(
"physics" => "Advection",
"run_type" => 20,
"order" => 1,
"dimensions" => 2,
"smb_name" => "SRCMESHES/square_x0-2_y0-2_20x20_tri.smb",
"IC_name" => "ICsinwave",
"numBC" => 1,
"BC1" => [0, 1, 2, 3],
# "BC1" => [0],
"BC1_name" => "sinwaveBC",
"delta_t" => 0.1,
"t_max" => 1.0,
# "res_abstol" => -1.0,
# "res_reltol" => -1.0,
"use_DG" => true,
"operator_type" => "SBPGamma",
# "Flux_name" => "RoeFlux",
"Flux_name" => "LFFlux",
"output_freq" => 100,
"real_time" => true,
"vortex_x0" => 5.0,
"use_itermax" => true,
"itermax" => 1000,
"output_freq" => 1,
"jac_method" => 1,        # 1: FD
"jac_type" => 2,          # 1: dense matrix, 2: Julia sparse
"do_postproc" => true,
"exact_soln_func" => "ICsinwave",
"use_Minv" => true,
)


using Roots

function detach_indexer(repo::Repository, id::AbstractString)::Tuple{Indexer,Repository}
    # Get requested indexer
    i = findfirst(x -> x.id == id, repo.indexers)
    if isnothing(i)
        throw(UnknownIndexerError())
    end
    indexer = repo.indexers[i]

    # Remove indexer from repository
    indexers = filter(x -> x.id != id, repo.indexers)
    frepo = Repository(indexers, repo.subgraphs)
    return indexer, frepo
end

function stakes(r::Repository)
    # Loop over subgraphs, getting their names
    subgraphs = ipfshash.(r.subgraphs)

    # Match name to indexer allocations
    ω = reduce(vcat, allocation.(r.indexers))
    Ω = []
    for subgraph in subgraphs
        subgraph_allocations = filter(x -> x.ipfshash == subgraph, ω)
        subgraph_amounts = allocated_stake.(subgraph_allocations)

        # If no match, then set to 0
        if isempty(subgraph_amounts)
            subgraph_amounts = [0.0]
        end
        stake_on_subgraph = sum(subgraph_amounts)
        append!(Ω, stake_on_subgraph)
    end
    return Ω
end

function solve_dual(Ω, ψ, σ)
    lower_bound = eps(Float64)
    upper_bound = (sum(.√(ψ .* Ω)))^2 / σ
    sol = find_zero(
        x -> sum(max.(0.0, .√(ψ .* Ω / x) .- Ω)) - σ,
        (lower_bound, upper_bound),
        Roots.Brent(),
    )
    return sol
end

solve_primal(Ω, ψ, v) = max.(0.0, .√(ψ .* Ω / v) - Ω)

function optimize(indexer::Indexer, repo::Repository)
    ψ = signal.(repo.subgraphs)
    Ω = stakes(repo)
    σ = indexer.stake

    # Solve the dual and use that value to solve the primal
    v = solve_dual(Ω, ψ, σ)
    ω = solve_primal(Ω, ψ, v)

    return ω
end


# =========================================================================================
"""
    abstract type QuadraturePoints{T <: AbstractFloat} end

Abstract base type for all sets of quadrature points
"""
abstract type QuadraturePoints{T <: AbstractFloat} end

# =========================================================================================
"""
    struct QuadPts2D{T} <: QuadraturePoints{T}
        num   ::Int         # number of points
        x     ::Vector{T}   # x values
        y     ::Vector{T}   # y values
        weight::Vector{T}   # weights
    end

Quadrature points and weights for triangles
"""
struct QuadPts2D{T} <: QuadraturePoints{T}
    """Number of points"""
    num::Int
    """x values"""
    x::Vector{T}
    """y values"""
    y::Vector{T}
    """Weights"""
    weight::Vector{T}
end

# =========================================================================================
"""
    struct QuadPts3D{T} <: QuadraturePoints{T}
        num   ::Int         # number of points
        x     ::Vector{T}   # x values
        y     ::Vector{T}   # y values
        z     ::Vector{T}   # z values
        weight::Vector{T}   # weights
    end

Quadrature points and weights for tetrahedra
"""
struct QuadPts3D{T} <: QuadraturePoints{T}
    """Number of points"""
    num::Int
    """x values"""
    x::Vector{T}
    """y values"""
    y::Vector{T}
    """z values"""
    z::Vector{T}
    """Weights"""
    weight::Vector{T}
end

# =========================================================================================
@doc """
    quadraturepoints(::Type{Tetrahedron{Float64}})
    quadraturepoints(::Type{Tetrahedron{Float32}})
    quadraturepoints(::Type{Triangle{Float64}})
    quadraturepoints(::Type{Triangle{Float32}})

Generator function for quadrature points:
 * *Triangles*: 7 points per element [[Rad48]](@ref Bibliography)
 * *Tetrahedra*: 5 points per element [[Kea86]](@ref Bibliography)

## Return type
`QuadPts2D` or `QuadPts3D`
""" quadraturepoints
for T in [:Float64, :Float32]
    varname = Symbol("triquadpts_", T)
    @eval begin
        const $(varname) = QuadPts2D(7,
            $(T)[1/3, (6+√15)/21, (9-2*√15)/21, (6+√15)/21,
                (6-√15)/21, (9+2*√15)/21, (6-√15)/21],
            $(T)[1/3, (9-2*√15)/21, (6+√15)/21, (6+√15)/21,
                (9+2*√15)/21, (6-√15)/21, (6-√15)/21],
            $(T)[9/80, (155+√15)/2400, (155+√15)/2400, (155+√15)/2400,
                (155-√15)/2400, (155-√15)/2400, (155-√15)/2400]
        )
        quadraturepoints(::Type{Triangle{$(T)}}) = $(varname)
    end
end

for T in [:Float64, :Float32]
    varname = Symbol("tetraquadpts_", T)
    @eval begin
        const $(varname) = QuadPts3D(5,
            $(T)[.25, .5, 1/6, 1/6, 1/6],
            $(T)[.25, 1/6, 1/6, 1/6, .5],
            $(T)[.25, 1/6, 1/6, .5, 1/6],
            $(T)[-.8, .45, .45, .45, .45]/6
        )
        quadraturepoints(::Type{Tetrahedron{$(T)}}) = $(varname)
    end
end


# FCA.jl

"""
Module providing the free componant analysis (FCA) algorithm.
"""
module FCA

# Package requried
using Optim # for manifold optimization
using LinearAlgebra
using GenericLinearAlgebra
using Statistics
using Combinatorics

import StatsBase: sample # for sample in embed.jl

# export functions
export  freecf, # main function

        # loss function
        κ₄,
        neg_abs_sum_free_kurt,
        free_ent,
        sum_free_ent,

        # gradient function
        grad_neg_abs_sum_free_kurt,
        grad_sum_free_ent,

        # freely center and whiten
        mat_center,
        free_whiten,

        # embed
        mat_embed,

        # icf
        icf,
        kurt,
        neg_abs_sum_kurt,
        grad_neg_abs_sum_kurt,
        negent,
        ent_sum,
        grad_ent_sum,

        # errpd
        err_pd

# include sources
include("loss.jl")
include("gradient.jl")
include("fcf.jl")
include("free_whiten.jl")
include("embed.jl")
include("icf.jl")
include("orthopt.jl")
include("errpd.jl")

end


# Kaggle tutorial:
# https://www.kaggle.com/c/street-view-getting-started-with-julia/overview/julia-tutorial

# Load packages

using Images
using DataFrames
using CSV
using DecisionTree
using Statistics

;

"""
dataset_type:
    "train" or "test"
labels_info:
    IDs of each image to read
image_size:
    amount of pixels in image
dataset_dir_path:
    path to directory with "trainResized" and "testResized" directories.
"""
function read_image(dataset_type, labels_info, image_size, dataset_dir_path)
    # Initialize X matrix
    X = zeros(size(labels_info, 1), image_size)
    @show size(labels_info)
    for (idx, image_id) in enumerate(labels_info[:ID])
        # Load image
        image_full_path = "$(dataset_dir_path)/$(dataset_type)Resized/$(image_id).Bmp"
        image = load(image_full_path)
        image_greyscale = Gray.(image)
        # Transform image matrix to a vector and store it in data matrix
        X[idx, :] = reshape(image_greyscale, 1, image_size)
    end
    return X
end

;

# The images in the trainResized and testResized directories are 20x20 pixels,
# so image_size should be set to 400.

image_size = 20 * 20
dataset_dir_path = "/home/max/Documents/ai/julia/julia_kaggle/data"

# Load labels (ID and Class) for train and test

# train_set
train_labels_info = CSV.read("$(dataset_dir_path)/trainLabels.csv")
@show size(train_labels_info)
@show names(train_labels_info)

# test_set
test_labels_info = CSV.read("$(dataset_dir_path)/sampleSubmission.csv")
@show size(test_labels_info)
@show names(test_labels_info)

# Show some info about these labels
@show first(train_labels_info, 5)
@show first(test_labels_info, 5)

;

# Read training and testing sets
@time X_train = read_image("train", train_labels_info, image_size, dataset_dir_path)
@time X_test = read_image("test", test_labels_info, image_size, dataset_dir_path)

@show size(X_train)
@show size(X_test)

;

# Create Y vector from each image.
# Target variable == first character of 'Class' column

y_train = map(col_value -> col_value[1], train_labels_info[:Class])
y_train = convert(Array{Int64, 1}, y_train)

@show size(y_train)
@show typeof(y_train)
@show y_train[1:10]

;

# Create a training model

# Learn patterns in the images that identify the character in the label.
# Use Julia version of Random Forest algorithm.
# Parameters to set:
#     number of features to choose at each split
#         sqrt(number_of_features)
#     number of trees
#         larger is better, but it takes more time to train
#     ratio of sampling
#         usually 1.0

num_of_features = sqrt(size(X_train, 2))
num_of_trees = 50
ratio_of_sampling = 1.0

# build_forest(labels::Array{T,1}, features::Array{S,2}, n_subfeatures, n_trees, partial_sampling)

model = build_forest(
    y_train, X_train,  # labels, features
    num_of_features,
    num_of_trees,
    ratio_of_sampling
)

@show model

;

# Use trained model: identify the characters in the test data

test_set_predictions = apply_forest(model, X_test)

;

# Convert predictions back to characters
test_labels_info[:Class] = convert(Array{Char, 1}, test_set_predictions)

;

# Print model accuracy

n_folds = 4

# nfoldCV_forest(
#     labels::Array{T,1}, features::Array{S,2},
#     n_folds::Integer, n_subfeatures::Integer, n_trees::Integer, partial_sampling::Float64,
# )

accuracy = nfoldCV_forest(
    y_train, X_train,  # labels, features
    n_folds,
    convert(Int, num_of_features),
    num_of_trees,
    ratio_of_sampling
)

@show mean(accuracy)

;

# Save predictions

CSV.write(
    "$(dataset_dir_path)/julia_submission.csv",
    test_labels_info,
    writeheader=true
)

;


@testset "util.jl" begin

    @testset "gridmake" begin
        want = [1 4 6
                2 4 6
                3 4 6
                1 5 6
                2 5 6
                3 5 6
                1 4 7
                2 4 7
                3 4 7
                1 5 7
                2 5 7
                3 5 7]

        # test allocating version
        @test want == @inferred gridmake([1, 2, 3], [4, 5], [6, 7])

        # test non-allocating version
        out = zeros(Int, 12, 3)
        @inferred gridmake!(out, [1, 2, 3], [4, 5], [6, 7])
        @test out == want

        # test single array version returns matrix of same type
        @test gridmake([1, 2, 3]) == [1 2 3]'

        # make sure we get the error we were expecting
        @test_throws AssertionError gridmake!(zeros(Int, 2, 2), [1, 2, 3], [4, 5])
    end

    @testset "fix" begin
        @test 1 == @inferred QuantEcon.fix(1.2)
        @test [0, 2] == @inferred QuantEcon.fix([0.9, 2.1])
        @test [0, 2] == @inferred QuantEcon.fix([0, 2])

        out = [100, 100]
        @inferred QuantEcon.fix!([0, 2], out)
        @test [0, 2] == out

        @test [0 2; 2 0] == @inferred QuantEcon.fix([0.5 2.9999; 3-2eps() -0.9])
    end

    @test ([1 2; 1 2], [3 3; 4 4]) == @inferred meshgrid([1, 2], [3, 4])

    @testset "simplex_tools" begin
        grid_3_4 = [0  0  0  0  0  1  1  1  1  2  2  2  3  3  4
                    0  1  2  3  4  0  1  2  3  0  1  2  0  1  0
                    4  3  2  1  0  3  2  1  0  2  1  0  1  0  0]
        points = [0 1 4
                  0 1 0
                  4 2 0]
        idx = Array{Int}(3)

        for i in 1:3
            idx[i] = simplex_index(points[:, i], 3, 4)
        end

        @test all(simplex_grid(3, 4) .== grid_3_4)
        @test all(grid_3_4[:, idx] .== points)
        @test size(grid_3_4, 2) == num_compositions(3, 4)

        # Output from QuantEcon.py
        grid_5_4 =
            [0 0 0 0 4 0 0 0 1 3 0 0 0 2 2 0 0 0 3 1 0 0 0 4 0
             0 0 1 0 3 0 0 1 1 2 0 0 1 2 1 0 0 1 3 0 0 0 2 0 2
             0 0 2 1 1 0 0 2 2 0 0 0 3 0 1 0 0 3 1 0 0 0 4 0 0
             0 1 0 0 3 0 1 0 1 2 0 1 0 2 1 0 1 0 3 0 0 1 1 0 2
             0 1 1 1 1 0 1 1 2 0 0 1 2 0 1 0 1 2 1 0 0 1 3 0 0
             0 2 0 0 2 0 2 0 1 1 0 2 0 2 0 0 2 1 0 1 0 2 1 1 0
             0 2 2 0 0 0 3 0 0 1 0 3 0 1 0 0 3 1 0 0 0 4 0 0 0
             1 0 0 0 3 1 0 0 1 2 1 0 0 2 1 1 0 0 3 0 1 0 1 0 2
             1 0 1 1 1 1 0 1 2 0 1 0 2 0 1 1 0 2 1 0 1 0 3 0 0
             1 1 0 0 2 1 1 0 1 1 1 1 0 2 0 1 1 1 0 1 1 1 1 1 0
             1 1 2 0 0 1 2 0 0 1 1 2 0 1 0 1 2 1 0 0 1 3 0 0 0
             2 0 0 0 2 2 0 0 1 1 2 0 0 2 0 2 0 1 0 1 2 0 1 1 0
             2 0 2 0 0 2 1 0 0 1 2 1 0 1 0 2 1 1 0 0 2 2 0 0 0
             3 0 0 0 1 3 0 0 1 0 3 0 1 0 0 3 1 0 0 0 4 0 0 0 0]
        @test simplex_grid(5, 4) == reshape(grid_5_4', 5, 70)
    end

    @testset "next_k_array" begin

        k_arrays = [1  2  3;
                    1  2  4;
                    1  3  4;
                    2  3  4;
                    1  2  5;
                    1  3  5;
                    2  3  5;
                    1  4  5;
                    2  4  5;
                    3  4  5;
                    1  2  6;
                    1  3  6;
                    2  3  6;
                    1  4  6;
                    2  4  6;
                    3  4  6;
                    1  5  6;
                    2  5  6;
                    3  5  6;
                    4  5  6]
        L, k = size(k_arrays)

        k_arrays_computed = similar(k_arrays)
        k_arrays_computed[1, :] = collect(1:k)
        @test k_array_rank(k_arrays_computed[1, :]) == 1
        for i = 2:L
            k_arrays_computed[i, :] = next_k_array!(k_arrays_computed[i-1, :])
            @test k_array_rank(k_arrays_computed[i, :]) == i
        end

        @test k_arrays_computed == k_arrays

        # test InexactError when ranking is out of range
        n, k = 100, 50
        @test k_array_rank(BigInt, collect(n-k+1:n)) == binomial(BigInt(n), BigInt(k))
        @test_throws InexactError k_array_rank(collect(n-k+1:n))

        # test returning wrong value when ranking is out of range
        n, k = 68, 29
        @test k_array_rank(BigInt, collect(n-k+1:n)) == binomial(BigInt(n), BigInt(k))
        @test k_array_rank(collect(n-k+1:n)) != binomial(BigInt(n), BigInt(k))

    end

end


function mean(iterable)
    state = start(iterable)
    if done(iterable, state)
        error("mean of empty collection undefined: $(repr(iterable))")
    end
    count = 1
    total, state = next(iterable, state)
    while !done(iterable, state)
        value, state = next(iterable, state)
        total += value
        count += 1
    end
    return total/count
end
mean(v::AbstractArray) = sum(v) / length(v)

function mean!{T}(r::AbstractArray{T}, v::AbstractArray)
    sum!(r, v; init=true)
    rs = convert(T, length(v) / length(r))
    if rs != 1
        for i = 1:length(r)
            @inbounds r[i] /= rs
        end
    end
    return r
end

meantype{T}(::Type{T}) = typeof((zero(T) + zero(T)) / 2)
mean{T}(v::AbstractArray{T}, region) =
    mean!(Array(meantype(T), reduced_dims(size(v), region)), v)

function median!{T<:Real}(v::AbstractVector{T}; checknan::Bool=true)
    isempty(v) && error("median of an empty array is undefined")
    if checknan
        for x in v
            isnan(x) && return x
        end
    end
    n = length(v)
    if isodd(n)
        return select!(v,div(n+1,2))
    else
        m = select!(v, div(n,2):div(n,2)+1)
        return (m[1] + m[2])/2
    end
end
median{T<:Real}(v::AbstractArray{T}; checknan::Bool=true) =
    median!(vec(copy(v)), checknan=checknan)
median{T}(v::AbstractArray{T}, region; checknan::Bool=true) =
    mapslices( x->median(x; checknan=checknan), v, region )

## variances

function sumabs2_pairwise{T<:Base.LinAlg.BlasFloat}(A::StridedArray{T}, i1::Int, n::Int)
    if n <= 2048
        BLAS.dot(n, pointer(A, i1), stride(A, 1), pointer(A, i1), stride(A, 1))
    else
        n2 = div(n,2)
        sumabs2_pairwise(A, i1, n2) + sumabs2_pairwise(A, i1+n2, n-n2)
    end
end

function sumabs2_pairwise(A::AbstractArray, i1::Int, n::Int)
    if n < 256
        @inbounds s = abs2(A[i1])
        for i=i1+1:i1+n-1
            @inbounds s += abs2(A[i])
        end
        return s
    else
        n2 = div(n,2)
        return sumabs2_pairwise(A, i1, n2) + sumabs2_pairwise(A, i1+n2, n-n2)
    end
end

sumabs2(v::AbstractArray) = sumabs2_pairwise(v, 1, length(v))

plusabs2(x, y) = x + abs2(y)
eval(ngenerate(:N, :(typeof(R)), :(_sumabs2!{T,N}(R::AbstractArray, A::AbstractArray{T,N})), N->gen_reduction_body(N, plusabs2)))
sumabs2!{R}(r::AbstractArray{R}, A::AbstractArray; init::Bool=true) = _sumabs2!(initarray!(r, zero(R), init), A)
sumabs2{T}(A::AbstractArray{T}, region) = _sumabs2!(reduction_init(A, region, abs2(zero(T))+abs2(zero(T))), A)

function varzm(v::AbstractArray; corrected::Bool=true)
    n = length(v)
    n == 0 && return NaN
    return sumabs2(v) / (n - int(corrected))
end

function varzm(v::AbstractArray, region; corrected::Bool=true)
    cn = regionsize(v, region) - int(corrected)
    sumabs2(v, region) / cn
end

function varm_pairwise(A::AbstractArray, m::Number, i1::Int, n::Int) # see sum_pairwise
    if n < 256
        @inbounds s = abs2(A[i1] - m)
        for i = i1+1:i1+n-1
            @inbounds s += abs2(A[i] - m)
        end
        return s
    else
        n2 = div(n,2)
        return varm_pairwise(A, m, i1, n2) + varm_pairwise(A, m, i1+n2, n-n2)
    end
end

function varm(v::AbstractArray, m::Number; corrected::Bool=true)
    n = length(v)
    n == 0 && return NaN
    return varm_pairwise(v, m, 1, n) / (n - int(corrected))
end

@ngenerate N Array{typeof((abs2(zero(T))+abs2(zero(T)))/1), N} function _varm{S,T,N}(v::AbstractArray{S,N}, m::AbstractArray{T,N}, region, corrected::Bool)
    rdims = reduced_dims(v, region)
    rdims == size(m) || error(DimensionMismatch("size of mean does not match reduced dimensions"))

    R = fill!(similar(v, typeof((abs2(zero(T))+abs2(zero(T)))/1), rdims), 0)
    @nextract N sizeR d->size(R,d)
    @nloops N i v d->(j_d = sizeR_d==1 ? 1 : i_d) begin
        @inbounds (@nref N R j) += abs2((@nref N v i) - (@nref N m j))
    end
    scale!(R, 1/(regionsize(v, region) - int(corrected)))
end
varm{S,T,N}(v::AbstractArray{S,N}, m::AbstractArray{T,N}, region; corrected::Bool=true) = _varm(v, m, region, corrected)

function var(v::AbstractArray; corrected::Bool=true, mean=nothing)
    mean == 0 ? varzm(v; corrected=corrected) :
    mean == nothing ? varm(v, Base.mean(v); corrected=corrected) :
    isa(mean, Number) ? varm(v, mean; corrected=corrected) :
    error("Invalid value of mean.")
end

function var(v::AbstractArray, region; corrected::Bool=true, mean=nothing)
    mean == 0 ? varzm(v, region; corrected=corrected) :
    mean == nothing ? varm(v, Base.mean(v, region), region; corrected=corrected) :
    isa(mean, AbstractArray) ? varm(v, mean, region; corrected=corrected) :
    error("Invalid value of mean.")
end

function var(iterable; corrected::Bool=true, mean=nothing)
    state = start(iterable)
    if done(iterable, state)
        error("variance of empty collection undefined: $(repr(iterable))")
    end
    count = 1
    value, state = next(iterable, state)
    if mean == nothing
        # Use Welford algorithm as seen in (among other places)
        # Knuth's TAOCP, Vol 2, page 232, 3rd edition.
        M = value / 1
        S = zero(M)
        while !done(iterable, state)
            value, state = next(iterable, state)
            count += 1
            new_M = M + (value - M) / count
            S = S + (value - M) * (value - new_M)
            M = new_M
        end
        return S / (count - int(corrected))
    else # mean provided
        # Cannot use a compensated version, e.g. the one from
        # "Updating Formulae and a Pairwise Algorithm for Computing Sample Variances."
        # by Chan, Golub, and LeVeque, Technical Report STAN-CS-79-773,
        # Department of Computer Science, Stanford University,
        # because user can provide mean value that is different to mean(iterable)
        sum2 = (value - mean)^2
        while !done(iterable, state)
            value, state = next(iterable, state)
            count += 1
            sum2 += (value - mean)^2
        end
        return sum2 / (count - int(corrected))
    end
end

varm(iterable, m::Number; corrected::Bool=true) =
    var(iterable, corrected=corrected, mean=m)

## variances over ranges

varm(v::Range, m::Number) = var(v)

function var(v::Range)
    s = step(v)
    l = length(v)
    if l == 0 || l == 1
        return NaN
    end
    return abs2(s) * (l + 1) * l / 12
end

## standard deviation

function sqrt!(v::AbstractArray)
    for i = 1:length(v)
        v[i] = sqrt(v[i])
    end
    v
end

stdm(v::AbstractArray, m::Number; corrected::Bool=true) =
    sqrt(varm(v, m; corrected=corrected))

std(v::AbstractArray; corrected::Bool=true, mean=nothing) =
    sqrt(var(v; corrected=corrected, mean=mean))

std(v::AbstractArray, region; corrected::Bool=true, mean=nothing) =
    sqrt!(var(v, region; corrected=corrected, mean=mean))

std(iterable; corrected::Bool=true, mean=nothing) =
    sqrt(var(iterable, corrected=corrected, mean=mean))

stdm(iterable, m::Number; corrected::Bool=true) =
    std(iterable, corrected=corrected, mean=m)

## pearson covariance functions ##

# auxiliary functions

_conj{T<:Real}(x::AbstractArray{T}) = x
_conj(x::AbstractArray) = conj(x)

_getnobs(x::AbstractVector, vardim::Int) = length(x)
_getnobs(x::AbstractMatrix, vardim::Int) = size(x, vardim)

function _getnobs(x::AbstractVecOrMat, y::AbstractVecOrMat, vardim::Int)
    n = _getnobs(x, vardim)
    _getnobs(y, vardim) == n || throw(DimensionMismatch("Dimensions of x and y mismatch."))
    return n
end

_vmean(x::AbstractVector, vardim::Int) = mean(x)
_vmean(x::AbstractMatrix, vardim::Int) = mean(x, vardim)

# core functions

unscaled_covzm(x::AbstractVector) = dot(x, x)
unscaled_covzm(x::AbstractMatrix, vardim::Int) = (vardim == 1 ? _conj(x'x) : x * x')

unscaled_covzm(x::AbstractVector, y::AbstractVector) = dot(x, y)
unscaled_covzm(x::AbstractVector, y::AbstractMatrix, vardim::Int) =
    (vardim == 1 ? At_mul_B(x, _conj(y)) : At_mul_Bt(x, _conj(y)))
unscaled_covzm(x::AbstractMatrix, y::AbstractVector, vardim::Int) =
    (c = vardim == 1 ? At_mul_B(x, _conj(y)) :  x * _conj(y); reshape(c, length(c), 1))
unscaled_covzm(x::AbstractMatrix, y::AbstractMatrix, vardim::Int) =
    (vardim == 1 ? At_mul_B(x, _conj(y)) : A_mul_Bc(x, y))

# covzm (with centered data)

covzm(x::AbstractVector; corrected::Bool=true) = unscaled_covzm(x, x) / (length(x) - int(corrected))

covzm(x::AbstractMatrix; vardim::Int=1, corrected::Bool=true) =
    scale!(unscaled_covzm(x, vardim), inv(size(x,vardim) - int(corrected)))

covzm(x::AbstractVector, y::AbstractVector; corrected::Bool=true) =
    unscaled_covzm(x, y) / (length(x) - int(corrected))

covzm(x::AbstractVecOrMat, y::AbstractVecOrMat; vardim::Int=1, corrected::Bool=true) =
    scale!(unscaled_covzm(x, y, vardim), inv(_getnobs(x, y, vardim) - int(corrected)))

# covm (with provided mean)

covm(x::AbstractVector, xmean; corrected::Bool=true) =
    covzm(x .- xmean; corrected=corrected)

covm(x::AbstractMatrix, xmean; vardim::Int=1, corrected::Bool=true) =
    covzm(x .- xmean; vardim=vardim, corrected=corrected)

covm(x::AbstractVector, xmean, y::AbstractVector, ymean; corrected::Bool=true) =
    covzm(x .- xmean, y .- ymean; corrected=corrected)

covm(x::AbstractVecOrMat, xmean, y::AbstractVecOrMat, ymean; vardim::Int=1, corrected::Bool=true) =
    covzm(x .- xmean, y .- ymean; vardim=vardim, corrected=corrected)

# cov (API)

function cov(x::AbstractVector; corrected::Bool=true, mean=nothing)
    mean == 0 ? covzm(x; corrected=corrected) :
    mean == nothing ? covm(x, Base.mean(x); corrected=corrected) :
    isa(mean, Number) ? covm(x, mean; corrected=corrected) :
    error("Invalid value of mean.")
end

function cov(x::AbstractMatrix; vardim::Int=1, corrected::Bool=true, mean=nothing)
    mean == 0 ? covzm(x; vardim=vardim, corrected=corrected) :
    mean == nothing ? covm(x, _vmean(x, vardim); vardim=vardim, corrected=corrected) :
    isa(mean, AbstractArray) ? covm(x, mean; vardim=vardim, corrected=corrected) :
    error("Invalid value of mean.")
end

function cov(x::AbstractVector, y::AbstractVector; corrected::Bool=true, mean=nothing)
    mean == 0 ? covzm(x, y; corrected=corrected) :
    mean == nothing ? covm(x, Base.mean(x), y, Base.mean(y); corrected=corrected) :
    isa(mean, (Number,Number)) ? covm(x, mean[1], y, mean[2]; corrected=corrected) :
    error("Invalid value of mean.")
end

function cov(x::AbstractVecOrMat, y::AbstractVecOrMat; vardim::Int=1, corrected::Bool=true, mean=nothing)
    if mean == 0
        covzm(x, y; vardim=vardim, corrected=corrected)
    elseif mean == nothing
        covm(x, _vmean(x, vardim), y, _vmean(y, vardim); vardim=vardim, corrected=corrected)
    elseif isa(mean, (Any,Any))
        covm(x, mean[1], y, mean[2]; vardim=vardim, corrected=corrected)
    else
        error("Invalid value of mean.")
    end
end

# cov2cor!

function cov2cor!{T}(C::AbstractMatrix{T}, xsd::AbstractArray)
    nx = length(xsd)
    size(C) == (nx, nx) || throw(DimensionMismatch("Inconsistent dimensions."))
    for j = 1:nx
        for i = 1:j-1
            C[i,j] = C[j,i]
        end
        C[j,j] = one(T)
        for i = j+1:nx
            C[i,j] /= (xsd[i] * xsd[j])
        end
    end
    return C
end

function cov2cor!(C::AbstractMatrix, xsd::Number, ysd::AbstractArray)
    nx, ny = size(C)
    length(ysd) == ny || throw(DimensionMismatch("Inconsistent dimensions."))
    for j = 1:ny
        for i = 1:nx
            C[i,j] /= (xsd * ysd[j])
        end
    end
    return C
end

function cov2cor!(C::AbstractMatrix, xsd::AbstractArray, ysd::Number)
    nx, ny = size(C)
    length(xsd) == nx || throw(DimensionMismatch("Inconsistent dimensions."))
    for j = 1:ny
        for i = 1:nx
            C[i,j] /= (xsd[i] * ysd)
        end
    end
    return C
end

function cov2cor!(C::AbstractMatrix, xsd::AbstractArray, ysd::AbstractArray)
    nx, ny = size(C)
    (length(xsd) == nx && length(ysd) == ny) ||
        throw(DimensionMismatch("Inconsistent dimensions."))
    for j = 1:ny
        for i = 1:nx
            C[i,j] /= (xsd[i] * ysd[j])
        end
    end
    return C
end

# # corzm (non-exported, with centered data)

corzm{T}(x::AbstractVector{T}) = float(one(T) * one(T))

corzm(x::AbstractMatrix; vardim::Int=1) =
    (c = unscaled_covzm(x, vardim); cov2cor!(c, sqrt!(diag(c))))

function corzm(x::AbstractVector, y::AbstractVector)
    n = length(x)
    length(y) == n || throw(DimensionMismatch("Inconsistent lengths."))
    x1 = x[1]
    y1 = y[1]
    xx = abs2(x1)
    yy = abs2(y1)
    xy = x1 * conj(y1)
    i = 1
    while i < n
        i += 1
        @inbounds xi = x[i]
        @inbounds yi = y[i]
        xx += abs2(xi)
        yy += abs2(yi)
        xy += xi * conj(yi)
    end
    return xy / (sqrt(xx) * sqrt(yy))
end

corzm(x::AbstractVector, y::AbstractMatrix; vardim::Int=1) =
    cov2cor!(unscaled_covzm(x, y, vardim), sqrt(sumabs2(x)), sqrt!(sumabs2(y, vardim)))

corzm(x::AbstractMatrix, y::AbstractVector; vardim::Int=1) =
    cov2cor!(unscaled_covzm(x, y, vardim), sqrt!(sumabs2(x, vardim)), sqrt(sumabs2(y)))

corzm(x::AbstractMatrix, y::AbstractMatrix; vardim::Int=1) =
    cov2cor!(unscaled_covzm(x, y, vardim), sqrt!(sumabs2(x, vardim)), sqrt!(sumabs2(y, vardim)))

# corm

corm(x::AbstractVector, xmean) = corzm(x .- xmean)

corm(x::AbstractMatrix, xmean; vardim::Int=1) = corzm(x .- xmean; vardim=vardim)

corm(x::AbstractVector, xmean, y::AbstractVector, ymean) = corzm(x .- xmean, y .- ymean)

corm(x::AbstractVecOrMat, xmean, y::AbstractVecOrMat, ymean; vardim::Int=1) =
    corzm(x .- xmean, y .- ymean; vardim=vardim)

# cor

function cor(x::AbstractVector; mean=nothing)
    mean == 0 ? corzm(x) :
    mean == nothing ? corm(x, Base.mean(x)) :
    isa(mean, Number) ? corm(x, mean) :
    error("Invalid value of mean.")
end

function cor(x::AbstractMatrix; vardim::Int=1, mean=nothing)
    mean == 0 ? corzm(x; vardim=vardim) :
    mean == nothing ? corm(x, _vmean(x, vardim); vardim=vardim) :
    isa(mean, AbstractArray) ? corm(x, mean; vardim=vardim) :
    error("Invalid value of mean.")
end

function cor(x::AbstractVector, y::AbstractVector; mean=nothing)
    mean == 0 ? corzm(x, y) :
    mean == nothing ? corm(x, Base.mean(x), y, Base.mean(y)) :
    isa(mean, (Number,Number)) ? corm(x, mean[1], y, mean[2]) :
    error("Invalid value of mean.")
end

function cor(x::AbstractVecOrMat, y::AbstractVecOrMat; vardim::Int=1, mean=nothing)
    if mean == 0
        corzm(x, y; vardim=vardim)
    elseif mean == nothing
        corm(x, _vmean(x, vardim), y, _vmean(y, vardim); vardim=vardim)
    elseif isa(mean, (Any,Any))
        corm(x, mean[1], y, mean[2]; vardim=vardim)
    else
        error("Invalid value of mean.")
    end
end

## nice-valued ranges for histograms

function histrange{T<:FloatingPoint,N}(v::AbstractArray{T,N}, n::Integer)
    if length(v) == 0
        return 0.0:1.0:0.0
    end
    lo, hi = extrema(v)
    if hi == lo
        step = 1.0
    else
        bw = (hi - lo) / n
        e = 10.0^floor(log10(bw))
        r = bw / e
        if r <= 2
            step = 2*e
        elseif r <= 5
            step = 5*e
        else
            step = 10*e
        end
    end
    start = step*(ceil(lo/step)-1)
    nm1 = iceil((hi - start)/step)
    start:step:(start + nm1*step)
end

function histrange{T<:Integer,N}(v::AbstractArray{T,N}, n::Integer)
    if length(v) == 0
        return 0:1:0
    end
    lo, hi = extrema(v)
    if hi == lo
        step = 1
    else
        bw = (hi - lo) / n
        e = 10^max(0,ifloor(log10(bw)))
        r = bw / e
        if r <= 1
            step = e
        elseif r <= 2
            step = 2*e
        elseif r <= 5
            step = 5*e
        else
            step = 10*e
        end
    end
    start = step*(ceil(lo/step)-1)
    nm1 = iceil((hi - start)/step)
    start:step:(start + nm1*step)
end

## midpoints of intervals
midpoints(r::Range) = r[1:length(r)-1] + 0.5*step(r)
midpoints(v::AbstractVector) = [0.5*(v[i] + v[i+1]) for i in 1:length(v)-1]

## hist ##
function sturges(n)  # Sturges' formula
    n==0 && return one(n)
    iceil(log2(n))+1
end

function hist!{HT}(h::AbstractArray{HT}, v::AbstractVector, edg::AbstractVector; init::Bool=true)
    n = length(edg) - 1
    length(h) == n || error("length(h) must equal length(edg) - 1.")
    if init
        fill!(h, zero(HT))
    end
    for x in v
        i = searchsortedfirst(edg, x)-1
        if 1 <= i <= n
            h[i] += 1
        end
    end
    edg, h
end

hist(v::AbstractVector, edg::AbstractVector) = hist!(Array(Int, length(edg)-1), v, edg)
hist(v::AbstractVector, n::Integer) = hist(v,histrange(v,n))
hist(v::AbstractVector) = hist(v,sturges(length(v)))

function hist!{HT}(H::AbstractArray{HT,2}, A::AbstractMatrix, edg::AbstractVector; init::Bool=true)
    m, n = size(A)
    size(H) == (length(edg)-1, n) || error("Incorrect size of H.")
    if init
        fill!(H, zero(HT))
    end
    for j = 1:n
        hist!(sub(H, :, j), sub(A, :, j), edg)
    end
    edg, H
end

hist(A::AbstractMatrix, edg::AbstractVector) = hist!(Array(Int, length(edg)-1, size(A,2)), A, edg)
hist(A::AbstractMatrix, n::Integer) = hist(A,histrange(A,n))
hist(A::AbstractMatrix) = hist(A,sturges(size(A,1)))

## hist2d
function hist2d!{HT}(H::AbstractArray{HT,2}, v::AbstractMatrix,
                     edg1::AbstractVector, edg2::AbstractVector; init::Bool=true)
    size(v,2) == 2 || error("hist2d requires an Nx2 matrix.")
    n = length(edg1) - 1
    m = length(edg2) - 1
    size(H) == (n, m) || error("Incorrect size of H.")
    if init
        fill!(H, zero(HT))
    end
    for i = 1:size(v,1)
        x = searchsortedfirst(edg1, v[i,1]) - 1
        y = searchsortedfirst(edg2, v[i,2]) - 1
        if 1 <= x <= n && 1 <= y <= m
            @inbounds H[x,y] += 1
        end
    end
    edg1, edg2, H
end

hist2d(v::AbstractMatrix, edg1::AbstractVector, edg2::AbstractVector) =
    hist2d!(Array(Int, length(edg1)-1, length(edg2)-1), v, edg1, edg2)

hist2d(v::AbstractMatrix, edg::AbstractVector) = hist2d(v, edg, edg)

hist2d(v::AbstractMatrix, n1::Integer, n2::Integer) =
    hist2d(v, histrange(sub(v,:,1),n1), histrange(sub(v,:,2),n2))
hist2d(v::AbstractMatrix, n::Integer) = hist2d(v, n, n)
hist2d(v::AbstractMatrix) = hist2d(v, sturges(size(v,1)))

## quantiles ##

# for now, use the R/S definition of quantile; may want variants later
# see ?quantile in R -- this is type 7
function quantile!(v::AbstractVector, q::AbstractVector)
    isempty(v) && error("empty data array")
    isempty(q) && error("empty quantile array")

    # make sure the quantiles are in [0,1]
    q = bound_quantiles(q)

    lv = length(v)
    lq = length(q)

    index = 1 .+ (lv-1)*q
    lo = ifloor(index)
    hi = iceil(index)
    sort!(v)
    isnan(v[end]) && error("quantiles are undefined in presence of NaNs")
    i = find(index .> lo)
    r = float(v[lo])
    h = (index.-lo)[i]
    r[i] = (1.-h).*r[i] + h.*v[hi[i]]
    return r
end
quantile(v::AbstractVector, q::AbstractVector) = quantile!(copy(v),q)
quantile(v::AbstractVector, q::Number) = quantile(v,[q])[1]

function bound_quantiles(qs::AbstractVector)
    epsilon = 100*eps()
    if (any(qs .< -epsilon) || any(qs .> 1+epsilon))
        error("quantiles out of [0,1] range")
    end
    [min(1,max(0,q)) for q = qs]
end