using Test
using Metatheory
using Metatheory.Library
using Metatheory.Schedulers
using TermInterface

mult_t = @commutative_monoid (*) 1
plus_t = @commutative_monoid (+) 0

minus_t = @theory a b begin
  # TODO Jacques Carette's post in zulip chat
  a - a --> 0
  a - b --> a + (-1 * b)
  -a --> -1 * a
  a + (-b) --> a + (-1 * b)
end

mulplus_t = @theory a b c begin
  # TODO FIXME this rules improves performance and avoids commutative
  # explosion of the egraph
  a + a --> 2 * a
  0 * a --> 0
  a * 0 --> 0
  a * (b + c) == ((a * b) + (a * c))
  a + (b * a) --> ((b + 1) * a)
end

pow_t = @theory x y z n m p q begin
  (y^n) * y --> y^(n + 1)
  x^n * x^m == x^(n + m)
  (x * y)^z == x^z * y^z
  (x^p)^q == x^(p * q)
  x^0 --> 1
  0^x --> 0
  1^x --> 1
  x^1 --> x
  x * x --> x^2
  inv(x) == x^(-1)
end

div_t = @theory x y z begin
  x / 1 --> x
  # x / x => 1 TODO SIGN ANALYSIS
  x / (x / y) --> y
  x * (y / x) --> y
  x * (y / z) == (x * y) / z
  x^(-1) == 1 / x
end

trig_t = @theory θ begin
  sin(θ)^2 + cos(θ)^2 --> 1
  sin(θ)^2 - 1 --> cos(θ)^2
  cos(θ)^2 - 1 --> sin(θ)^2
  tan(θ)^2 - sec(θ)^2 --> 1
  tan(θ)^2 + 1 --> sec(θ)^2
  sec(θ)^2 - 1 --> tan(θ)^2
  cot(θ)^2 - csc(θ)^2 --> 1
  cot(θ)^2 + 1 --> csc(θ)^2
  csc(θ)^2 - 1 --> cot(θ)^2
end

# Dynamic rules
fold_t = @theory a b begin
  -(a::Number) => -a
  a::Number + b::Number => a + b
  a::Number * b::Number => a * b
  a::Number^b::Number => begin
    b < 0 && a isa Int && (a = float(a))
    a^b
  end
  a::Number / b::Number => a / b
end

using Calculus: differentiate
diff_t = @theory x y begin
  ∂(y, x::Symbol) => begin
    z = extract!(_egraph, simplcost; root = y.id)
    @show z
    zd = differentiate(z, x)
    @show zd
    zd
  end
end

cas = fold_t ∪ mult_t ∪ plus_t ∪ minus_t ∪ mulplus_t ∪ pow_t ∪ div_t ∪ trig_t ∪ diff_t

function customlt(x, y)
  if typeof(x) == Expr && typeof(y) == Expr
    false
  elseif typeof(x) == typeof(y)
    isless(x, y)
  elseif x isa Symbol && y isa Number
    false
  elseif x isa Expr && y isa Number
    false
  elseif x isa Expr && y isa Symbol
    false
  else
    true
  end
end

canonical_t = @theory x y n xs ys begin
  # restore n-arity
  (x * x) --> x^2
  (x^n::Number * x) --> x^(n + 1)
  (x * x^n::Number) --> x^(n + 1)
  (x + (+)(ys...)) --> +(x, ys...)
  ((+)(xs...) + y) --> +(xs..., y)
  (x * (*)(ys...)) --> *(x, ys...)
  ((*)(xs...) * y) --> *(xs..., y)

  (*)(xs...) => Expr(:call, :*, sort!(xs; lt = customlt)...)
  (+)(xs...) => Expr(:call, :+, sort!(xs; lt = customlt)...)
end

function simplcost(n::ENodeTerm, g::EGraph)
  cost = 0 + arity(n)
  if operation(n) == :∂
    cost += 20
  end
  for id in arguments(n)
    eclass = g[id]
    !hasdata(eclass, simplcost) && (cost += Inf; break)
    cost += last(getdata(eclass, simplcost))
  end
  return cost
end

simplcost(n::ENodeLiteral, g::EGraph) = 0

function simplify(ex; steps = 4)
  params = SaturationParams(
    scheduler = ScoredScheduler,
    eclasslimit = 5000,
    timeout = 7,
    schedulerparams = (1000, 5, Schedulers.exprsize),
    #stopwhen=stopwhen,
  )
  hist = UInt64[]
  push!(hist, hash(ex))
  for i in 1:steps
    g = EGraph(ex)
    @profview saturate!(g, cas, params)
    ex = extract!(g, simplcost)
    ex = rewrite(ex, canonical_t)
    if !TermInterface.istree(ex)
      return ex
    end
    if hash(ex) ∈ hist
      println("loop detected $ex")
      return ex
    end
    println(ex)
    push!(hist, hash(ex))
  end

end

@test :(4a) == simplify(:(2a + a + a))
@test :(a * b * c) == simplify(:(a * c * b))
@test :(2x) == simplify(:(1 * x * 2))
@test :((a * b)^2) == simplify(:((a * b)^2))
@test :((a * b)^6) == simplify(:((a^2 * b^2)^3))
@test :(a + b + d) == simplify(:(a + b + (0 * c) + d))
@test :(a + b) == simplify(:(a + b + (c * 0) + d - d))
@test :(a) == simplify(:((a + d) - d))
@test :(a + b + d) == simplify(:(a + b * c^0 + d))
@test :(a * b * x^(d + y)) == simplify(:(a * x^y * b * x^d))
@test :(a * b * x^74103) == simplify(:(a * x^(12 + 3) * b * x^(42^3)))

@test 1 == simplify(:((x + y)^(a * 0) / (y + x)^0))
@test 2 == simplify(:(cos(x)^2 + 1 + sin(x)^2))
@test 2 == simplify(:(cos(y)^2 + 1 + sin(y)^2))
@test 2 == simplify(:(sin(y)^2 + cos(y)^2 + 1))

@test :(y + sec(x)^2) == simplify(:(1 + y + tan(x)^2))
@test :(y + csc(x)^2) == simplify(:(1 + y + cot(x)^2))

# simplify(:( ∂(x^2, x)))

@time simplify(:(∂(x^(cos(x)), x)))

@test :(2x^3) == simplify(:(x * ∂(x^2, x) * x))

# @simplify ∂(y^3, y) * ∂(x^2 + 2, x) / y * x

# @simplify (6 * x * x * y)

# @simplify ∂(y^3, y) / y

# # ex = :( ∂(x^(cos(x)), x) )
# ex = :( (6 * x * x * y) )
# g = EGraph(ex)
# saturate!(g, cas)
# g.classes
# extract!(g, simplcost; root=g.root)

# params = SaturationParams(
#     scheduler=BackoffScheduler,
#     eclasslimit=5000,
#     timeout=7,
#     schedulerparams=(1000,5),
#     #stopwhen=stopwhen,
# )

# ex = :((x+y)^(a*0) / (y+x)^0)
# g = EGraph(ex)
# @profview println(saturate!(g, cas, params))

# ex = extract!(g, simplcost)
# ex = rewrite(ex, canonical_t; clean=false)

# FIXME this is a hack to get the test to work.
if VERSION < v"1.9.0-DEV"
  function EGraphs.make(::Val{:type_analysis}, g::EGraph, n::ENodeTerm)
    Any
  end

  function EGraphs.make(::Val{:type_analysis}, g::EGraph, n::ENodeLiteral)
    v = n.value
    if v == :im
      typeof(im)
    else
      typeof(v)
    end
  end

  function EGraphs.make(::Val{:type_analysis}, g::EGraph, n::ENodeTerm{Expr})
    if exprhead(n) != :call
      # println("$n is not a call")
      t = Any
      # println("analyzed type of $n is $t")
      return t
    end
    sym = operation(n)
    if !(sym isa Symbol)
      # println("head $sym is not a symbol")
      t = Any
      # println("analyzed type of $n is $t")
      return t
    end

    symval = getfield(@__MODULE__, sym)
    child_classes = map(x -> g[x], arguments(n))
    child_types = Tuple(map(x -> getdata(x, :type_analysis, Any), child_classes))

    # t = t_arr[1]
    t = Core.Compiler.return_type(symval, child_types)

    if t == Union{}
      throw(MethodError(symval, child_types))
    end
    # println("analyzed type of $n is $t")
    return t
  end

  EGraphs.join(::Val{:type_analysis}, from, to) = typejoin(from, to)

  EGraphs.islazy(::Val{:type_analysis}) = true

  function infer(e)
    g = EGraph(e)
    analyze!(g, :type_analysis)
    getdata(g[g.root], :type_analysis)
  end

  ex1 = :(cos(1 + 3.0) + 4 + (4 - 4im))
  ex2 = :("ciao" * 2)
  ex3 = :("ciao" * " mondo")

  @test ComplexF64 == infer(ex1)
  @test_throws MethodError infer(ex2)
  @test String == infer(ex3)
end


function op_Vpot( pw::PWGrid, Vpot, psi::Array{ComplexF64,2}, ik::Int )
  #
  Ns = pw.Ns
  Ω  = pw.Ω
  Npoints = prod(Ns)
  Ncols   = size(psi)[2]

  ctmp = zeros(ComplexF64, Npoints, Ncols)
  idx = pw.gkvec.idx_gk[ik]
  for ic = 1:Ncols
    ctmp[idx,ic] = psi[:,ic]
  end

  # get values of psi in real space grid via forward transform
  ctmp = G_to_R( Ns, ctmp )

  cVpsi = R_to_G( Ns, Diagprod(Vpot, ctmp) )
  return cVpsi[idx,:]
end

# B is usually consists of more than one-column
function Diagprod( a,B )
  Ncol    = size(B)[2]
  Npoints = size(B)[1]
  out = zeros( ComplexF64, size(B) )
  for ic = 1:Ncol
    for ip = 1:Npoints
      out[ip,ic] = a[ip]*B[ip,ic]
    end
  end
  return out
end


@polyvar x z
@variable(model, t)
p = x^4 + x^2 - 3*x^2*z^2 + z^6 - t
@constraint(model, p >= 0)
@objective(model, Max, t)
optimize!(model)
println("Solution: $(value(t))")
# Returns the lower bound -.17700


# Define # of binding sites for clusters with n monomers up to imax
function Ns(imax, prefactor)

n  = range(1, stop=imax)
Ns = round.(Int64, prefactor.*n.^(2/3))

return Ns

end


# calculate the normalised mutual information measure
function voi(ca::Vector{Int}, cb::Vector{Int})
    length(ca) == length(cb) || error("membership length must be equal")
    n_ca = maximum(ca)
    n_cb = maximum(cb)
    ca_set = Set{Int}[]
    cb_set = Set{Int}[]
    for i=1:n_ca
        push!(ca_set, Set{Int}())
    end
    for (u, i) in enumerate(ca)
        push!(ca_set[i], u)
    end
    for j=1:n_cb
        push!(cb_set, Set{Int}())
    end
    for (v, j) in enumerate(cb)
        push!(cb_set[j], v)
    end
    N = zeros(Int, n_ca, n_cb) # confusion matrix
    for i=1:n_ca, j=1:n_cb
        N[i,j] = length(intersect(ca_set[i], cb_set[j]))
    end
    N1 = collect(sum(N, 1))
    N2 = collect(sum(N, 2))
    Nsum = length(ca)
    VI = 0.0
    for i=1:n_ca, j=1:n_cb
        if N[i,j]>0
            VI += N[i,j]*(log(N2[i]/N[i,j]) + log(N1[j]/N[i,j]))/Nsum
        end
    end
    VI/log(Nsum)
end


#
# Differential operators for regular grids
#

∂x_caa(i, j, k, grid, u::AbstractArray) = δx_f2c(grid, u, i, j, k) / grid.Δx
∂y_aca(i, j, k, grid, v::AbstractArray) = δy_f2c(grid, v, i, j, k) / grid.Δy
∂z_aac(i, j, k, grid, w::AbstractArray) = δz_f2c(grid, w, i, j, k) / grid.Δz
∂x_faa(i, j, k, grid, ϕ::AbstractArray) = δx_c2f(grid, ϕ, i, j, k) / grid.Δx
∂y_afa(i, j, k, grid, ϕ::AbstractArray) = δy_c2f(grid, ϕ, i, j, k) / grid.Δy
∂z_aaf(i, j, k, grid, ϕ::AbstractArray) = δz_c2f(grid, ϕ, i, j, k) / grid.Δz

#
# Differentiation and interpolation operators for functions
#
# The geometry of the Oceananigans grid (in one dimension) is
#
# face   cell   face   cell   face
#
#         i-1            i
#          ↓             ↓
#   |      ×      |      ×      |
#   ↑             ↑             ↑
#  i-1            i            i+1
#
# As a result the interpolation of a quantity ϕᵢ from a cell i to face i
# (this operation is denoted "▶x_faa" in the code below) is
#
# (▶ϕ)ᵢ = (ϕᵢ + ϕᵢ₋₁) / 2 .
#
# Conversely, the interpolation of a quantity uᵢ from face i to cell i is given by
#
# (▶u)ᵢ = (uᵢ₊₁ + uᵢ) / 2.
#
# Derivative operators are defined similarly. Using the symbol "∂" to denote
# differentiation, we have
#
# (∂ϕ)ᵢ = (ϕᵢ - ϕᵢ₋₁) / Δ
#
# for the differentation of a CellField ending on face i, and
#
# (∂u)ᵢ = (uᵢ₊₁ - uᵢ) / Δ
#
# for the differentation of a FaceField ending on cell i, where Δ is the spacing
# between cell points or face points, appropriately.
#

#
# Differential operators
#

"""
    ∂x_faa(i, j, k, grid, F, args...)

Differentiate the function or callable object

    `F(i, j, k, grid, args...)`

located at `caa` in `x`, across `faa`.
"""
function ∂x_faa(i, j, k, grid, F::Function, args...)
    i⁻ = decmod1(i, grid.Nx)
    return (F(i, j, k, grid, args...) - F(i⁻, j, k, grid, args...)) / grid.Δx
end

"""
    ∂x_caa(i, j, k, grid, F, args...)

Differentiate the function or callable object

    `F(i, j, k, grid, args...)`

located at `faa` in `x`, across `caa`.
"""
function ∂x_caa(i, j, k, grid, F::Function, args...)
    i⁺ = incmod1(i, grid.Nx)
    return (F(i⁺, j, k, grid, args...) - F(i, j, k, grid, args...)) / grid.Δx
end

"""
    ∂y_afa(i, j, k, grid, F, args...)

Differentiate the function or callable object

    `F(i, j, k, grid, args...)`

located at `aca` in `y`, across `afa`.
"""
function ∂y_afa(i, j, k, grid, F::Function, args...)
    j⁻ = decmod1(j, grid.Ny)
    return (F(i, j, k, grid, args...) - F(i, j⁻, k, grid, args...)) / grid.Δy
end

"""
    ∂y_aca(i, j, k, grid, F, args...)

Differentiate the function or callable object

    `F(i, j, k, grid, args...)`

located at `afa` in `y`, across `aca`.
"""
function ∂y_aca(i, j, k, grid, F::Function, args...)
    j⁺ = incmod1(j, grid.Ny)
    return (F(i, j⁺, k, grid, args...) - F(i, j, k, grid, args...)) / grid.Δy
end

"""
    ∂z_aaf(i, j, k, grid, F, args...)

Differentiate the function or callable object

    `F(i, j, k, grid, args...)`

located at `aac` in `z`, across `aaf`.
"""
function ∂z_aaf(i, j, k, grid::Grid{T}, F::Function, args...) where T
    if k < 2
        return -zero(T)
    else
        return (F(i, j, k-1, grid, args...) - F(i, j, k, grid, args...)) / grid.Δz
    end
end

"""
    ∂z_aac(i, j, k, grid, F, args...)

Differentiate the function or callable object

    `F(i, j, k, grid, args...)`

located at `aaf` in `z`, across `aac`.
"""
function ∂z_aac(i, j, k, grid::Grid{T}, F::Function, args...) where T
    if k == grid.Nz
        return -zero(T)
    else
        return (F(i, j, k, grid, args...) - F(i, j, k+1, grid, args...)) / grid.Δz
    end
end

#
# Double differentiation
#

∂x²_caa(i, j, k, grid, u) = ∂x_faa(i, j, k, grid, ∂x_caa, u)
∂x²_faa(i, j, k, grid, ϕ) = ∂x_caa(i, j, k, grid, ∂x_faa, ϕ)

∂y²_aca(i, j, k, grid, v) = ∂y_afa(i, j, k, grid, ∂y_aca, v)
∂y²_afa(i, j, k, grid, ϕ) = ∂y_aca(i, j, k, grid, ∂y_afa, ϕ)

∂z²_aac(i, j, k, grid, w) = ∂z_aaf(i, j, k, grid, ∂z_aac, w)
∂z²_aaf(i, j, k, grid, ϕ) = ∂z_aac(i, j, k, grid, ∂z_aaf, ϕ)

#
# Interpolation operations for functions
#

"""
    ▶x_faa(i, j, k, grid, F, args...)

Interpolate the function or callable object

    `F(i, j, k, grid, args...)`

from `caa` to `faa`."
"""
function ▶x_faa(i, j, k, grid::Grid{T}, F::Function, args...) where T
    i⁻¹ = decmod1(i, grid.Nx)
    return T(0.5) * (F(i, j, k, grid, args...) + F(i⁻¹, j, k, grid, args...))
end

"""
    ▶x_caa(i, j, k, grid, F, args...)

Interpolate the function or callable object

    `F(i, j, k, grid, args...)`

from `faa` to `caa`."
"""
function ▶x_caa(i, j, k, grid::Grid{T}, F::Function, args...) where T
    i⁺¹ = incmod1(i, grid.Nx)
    return T(0.5) * (F(i⁺¹, j, k, grid, args...) + F(i, j, k, grid, args...))
end

"""
    ▶y_afa(i, j, k, grid, F, args...)

Interpolate the function or callable object

    `F(i, j, k, grid, args...)`

from `aca` to `afa`.
"""
function ▶y_afa(i, j, k, grid::Grid{T}, F::Function, args...) where T
    j⁻¹ = decmod1(j, grid.Ny)
    return T(0.5) * (F(i, j, k, grid, args...) + F(i, j⁻¹, k, grid, args...))
end

"""
    ▶y_aca(i, j, k, grid, F, args...)

Interpolate the function or callable object

    `F(i, j, k, grid, args...)`

from `afa` to `aca`."
"""
function ▶y_aca(i, j, k, grid::Grid{T}, F::Function, args...) where T
    j⁺¹ = incmod1(j, grid.Ny)
    return T(0.5) * (F(i, j⁺¹, k, grid, args...) + F(i, j, k, grid, args...))
end

"""
    ▶z_aaf(i, j, k, grid, F, args...)

Interpolate the function or callable object

    `F(i, j, k, grid, args...)`

from `aac` to `aaf`.
"""
function ▶z_aaf(i, j, k, grid::Grid{T}, F::Function, args...) where T
    if k == 1
        return T(0.5) * F(i, j, k, grid, args...)
    else
        return T(0.5) * (F(i, j, k, grid, args...) + F(i, j, k-1, grid, args...))
    end
end

"""
    ▶z_aac(i, j, k, grid::Grid{T}, F, args...) where T

Interpolate the function or callable object

    `F(i, j, k, grid, args...)`

from `aaf` to `aac`.
"""
function ▶z_aac(i, j, k, grid::Grid{T}, F::Function, args...) where T
    if k == grid.Nz
        return T(0.5) * F(i, j, k, grid, args...)
    else
        return T(0.5) * (F(i, j, k+1, grid, args...) + F(i, j, k, grid, args...))
    end
end

# Convenience operators for "interpolating constants"
▶x_faa(i, j, k, grid, F::Number, args...) = F
▶x_caa(i, j, k, grid, F::Number, args...) = F
▶y_afa(i, j, k, grid, F::Number, args...) = F
▶y_aca(i, j, k, grid, F::Number, args...) = F
▶z_aaf(i, j, k, grid, F::Number, args...) = F
▶z_aac(i, j, k, grid, F::Number, args...) = F

#
# Double interpolation: 12 operators
#

"""
    ▶xy_cca(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `y`, from `ffa` to `cca`.
"""
▶xy_cca(i, j, k, grid, F, args...) = ▶y_aca(i, j, k, grid, ▶x_caa, F, args...)

"""
    ▶xy_fca(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `y`, from `cfa` to `fca`.
"""
▶xy_fca(i, j, k, grid, F, args...) = ▶y_aca(i, j, k, grid, ▶x_faa, F, args...)

"""
    ▶xy_ffa(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `y`, from `cca` to `ffa`.
"""
▶xy_ffa(i, j, k, grid, F, args...) = ▶y_afa(i, j, k, grid, ▶x_faa, F, args...)

"""
    ▶xy_cfa(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `y`, from `fca` to `cfa`.
"""
▶xy_cfa(i, j, k, grid, F, args...) = ▶y_afa(i, j, k, grid, ▶x_caa, F, args...)

"""
    ▶xz_cac(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `z`, from `faf` to `cac`.
"""
▶xz_cac(i, j, k, grid, F, args...) = ▶z_aac(i, j, k, grid, ▶x_caa, F, args...)

"""
    ▶xz_fac(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `z`, from `caf` to `fac`.
"""
▶xz_fac(i, j, k, grid, F, args...) = ▶z_aac(i, j, k, grid, ▶x_faa, F, args...)

"""
    ▶xz_faf(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `z`, from `cac` to `faf`.
"""
▶xz_faf(i, j, k, grid, F, args...) = ▶z_aaf(i, j, k, grid, ▶x_faa, F, args...)

"""
    ▶xz_caf(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `z`, from `fac` to `caf`.
"""
▶xz_caf(i, j, k, grid, F, args...) = ▶z_aaf(i, j, k, grid, ▶x_caa, F, args...)

"""
    ▶yz_acc(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `y` and `z`, from `aff` to `acc`.
"""
▶yz_acc(i, j, k, grid, F, args...) = ▶z_aac(i, j, k, grid, ▶y_aca, F, args...)

"""
    ▶yz_afc(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `y` and `z`, from `acf` to `afc`.
"""
▶yz_afc(i, j, k, grid, F, args...) = ▶z_aac(i, j, k, grid, ▶y_afa, F, args...)

"""
    ▶yz_aff(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `y` and `z`, from `ffa` to `cca`.
"""
▶yz_aff(i, j, k, grid, F, args...) = ▶z_aaf(i, j, k, grid, ▶y_afa, F, args...)

"""
    ▶yz_acf(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `y` and `z`, from `afc` to `acf`.
"""
▶yz_acf(i, j, k, grid, F, args...) = ▶z_aaf(i, j, k, grid, ▶y_aca, F, args...)

"""
    ∂x_caa(κ, Σ, i, j, k, grid, closure, u, v, w, T, S)

Compute ``\\partial_x ( \\kappa \\Sigma )``, where `κ` and `Σ`
are diffusivity and strain functions that return quantities
at the location `faa`, where `a` is either `f` or `c`.
The quantity returned has the location `caa`.
"""
function ∂x_caa(κ::Function, Σ::Function, i, j, k, grid, closure, eos, g, u, v, w, T, S)
    i⁺¹ = incmod1(i, grid.Nx)

    Σ⁺¹ = Σ(i⁺¹, j, k, grid, u, v, w)
    Σ⁻¹ = Σ(i,   j, k, grid, u, v, w)

    κ⁺¹ = κ(i⁺¹, j, k, grid, closure, eos, g, u, v, w, T, S)
    κ⁻¹ = κ(i,   j, k, grid, closure, eos, g, u, v, w, T, S)

    return (κ⁺¹ * Σ⁺¹ - κ⁻¹ * Σ⁻¹) / grid.Δx
end

"""
    ∂x_faa(κ, Σ, i, j, k, grid, closure, u, v, w, T, S)

Compute ``\\partial_x ( \\kappa \\Sigma )``, where `κ` and `Σ`
are diffusivity and strain functions that return quantities
at the location `caa`, where `a` is either `f` or `c`.
The quantity returned has the location `faa`.
"""
function ∂x_faa(κ::Function, Σ::Function, i, j, k, grid, closure, eos, g, u, v, w, T, S)
    i⁻¹ = decmod1(i, grid.Nx)

    Σ⁺¹ = Σ(i,   j, k, grid, u, v, w)
    Σ⁻¹ = Σ(i⁻¹, j, k, grid, u, v, w)

    κ⁺¹ = κ(i,   j, k, grid, closure, eos, g, u, v, w, T, S)
    κ⁻¹ = κ(i⁻¹, j, k, grid, closure, eos, g, u, v, w, T, S)

    return (κ⁺¹ * Σ⁺¹ - κ⁻¹ * Σ⁻¹) / grid.Δx
end

"""
    ∂y_aca(κ, Σ, i, j, k, grid, closure, u, v, w, T, S)

Compute ``\\partial_y ( \\kappa \\Sigma )``, where `κ` and `Σ`
are diffusivity and strain functions that return quantities
at the location `afa`, where `a` is either `f` or `c`.
The quantity returned has the location `aca`.
"""
function ∂y_aca(κ::Function, Σ::Function, i, j, k, grid, closure, eos, g, u, v, w, T, S)
    j⁺¹ = incmod1(j, grid.Ny)

    Σ⁺¹ = Σ(i, j⁺¹, k, grid, u, v, w)
    Σ⁻¹ = Σ(i, j,   k, grid, u, v, w)

    κ⁺¹ = κ(i, j⁺¹, k, grid, closure, eos, g, u, v, w, T, S)
    κ⁻¹ = κ(i, j,   k, grid, closure, eos, g, u, v, w, T, S)

    return (κ⁺¹ * Σ⁺¹ - κ⁻¹ * Σ⁻¹) / grid.Δy
end

"""
    ∂y_afa(κ, Σ, i, j, k, grid, closure, u, v, w, T, S)

Compute ``\\partial_y ( \\kappa \\Sigma )``, where `κ` and `Σ`
are diffusivity and strain functions that return quantities
at the location `aca`, where `a` is either `f` or `c`.
The quantity returned has the location `afa`.
"""
function ∂y_afa(κ::Function, Σ::Function, i, j, k, grid, closure, eos, g, u, v, w, T, S)
    j⁻¹ = decmod1(j, grid.Ny)

    Σ⁺¹ = Σ(i, j,   k, grid, u, v, w)
    Σ⁻¹ = Σ(i, j⁻¹, k, grid, u, v, w)

    κ⁺¹ = κ(i, j,   k, grid, closure, eos, g, u, v, w, T, S)
    κ⁻¹ = κ(i, j⁻¹, k, grid, closure, eos, g, u, v, w, T, S)

    return (κ⁺¹ * Σ⁺¹ - κ⁻¹ * Σ⁻¹) / grid.Δy
end

"""
    ∂z_aac(κ, Σ, i, j, k, grid, closure, u, v, w, T, S)

Compute ``\\partial_z ( \\kappa \\Sigma )``, where `κ` and `Σ`
are diffusivity and strain functions that return quantities
at the location `aaf`, where `a` is either `f` or `c`.
The quantity returned has the location `aac`.
"""
function ∂z_aac(κ::Function, Σ::Function, i, j, k, grid, closure, eos, g, u, v, w, T, S)
    Σ⁻¹ = Σ(i, j, k, grid, u, v, w)
    κ⁻¹ = κ(i, j, k, grid, closure, eos, g, u, v, w, T, S)

    if k == grid.Nz
        return κ⁻¹ * Σ⁻¹ / grid.Δz
    else
        Σ⁺¹ = Σ(i, j, k+1, grid, u, v, w)
        κ⁺¹ = κ(i, j, k+1, grid, closure, eos, g, u, v, w, T, S)

        # Reversed due to inverse convention for z-indexing
        return (κ⁻¹ * Σ⁻¹ - κ⁺¹ * Σ⁺¹) / grid.Δz
    end
end

"""
    ∂z_aaf(κ, Σ, i, j, k, grid, closure, u, v, w, T, S)

Compute ``\\partial_z ( \\kappa \\Sigma )``, where `κ` and `Σ`
are diffusivity and strain functions that return quantities
at the location `aac`, where `a` is either `f` or `c`.
The quantity returned has the location `aaf`.
"""
function ∂z_aaf(κ::Function, Σ::Function, i, j, k, grid, closure, eos, g, u, v, w, T, S)
    Σ⁺¹ = Σ(i, j, k, grid, u, v, w)
    κ⁺¹ = κ(i, j, k, grid, closure, eos, g, u, v, w, T, S)

    if k == 1
        return - κ⁺¹ * Σ⁺¹ / grid.Δz
    else
        Σ⁻¹ = Σ(i, j, k-1, grid, u, v, w)
        κ⁻¹ = κ(i, j, k-1, grid, closure, eos, g, u, v, w, T, S)

        # Reversed due to inverse convention for z-indexing
        return (κ⁻¹ * Σ⁻¹ - κ⁺¹ * Σ⁺¹) / grid.Δz
    end
end

# At fcc
∂x_2ν_Σ₁₁(ijk...) = 2 * ∂x_faa(ν_ccc, Σ₁₁, ijk...)
∂y_2ν_Σ₁₂(ijk...) = 2 * ∂y_aca(ν_ffc, Σ₁₂, ijk...)
∂z_2ν_Σ₁₃(ijk...) = 2 * ∂z_aac(ν_fcf, Σ₁₃, ijk...)

# At cfc
∂x_2ν_Σ₂₁(ijk...) = 2 * ∂x_caa(ν_ffc, Σ₂₁, ijk...)
∂y_2ν_Σ₂₂(ijk...) = 2 * ∂y_afa(ν_ccc, Σ₂₂, ijk...)
∂z_2ν_Σ₂₃(ijk...) = 2 * ∂z_aac(ν_cff, Σ₂₃, ijk...)

# At ccf
∂x_2ν_Σ₃₁(ijk...) = 2 * ∂x_caa(ν_fcf, Σ₃₁, ijk...)
∂y_2ν_Σ₃₂(ijk...) = 2 * ∂y_aca(ν_cff, Σ₃₂, ijk...)
∂z_2ν_Σ₃₃(ijk...) = 2 * ∂z_aaf(ν_ccc, Σ₃₃, ijk...)

"""
    ∂ⱼ_2ν_Σ₁ⱼ(i, j, k, grid, closure, eos, g, u, v, w, T, S)

Return the ``x``-component of the turbulent diffusive flux divergence:

`∂x(2 ν Σ₁₁) + ∂y(2 ν Σ₁₁) + ∂z(2 ν Σ₁₁)`

at the location `fcc`.
"""
∂ⱼ_2ν_Σ₁ⱼ(i, j, k, grid, closure::IsotropicDiffusivity, args...) = (
    ∂x_2ν_Σ₁₁(i, j, k, grid, closure, args...)
  + ∂y_2ν_Σ₁₂(i, j, k, grid, closure, args...)
  + ∂z_2ν_Σ₁₃(i, j, k, grid, closure, args...)
)

"""
    ∂ⱼ_2ν_Σ₂ⱼ(i, j, k, grid, closure, eos, g, u, v, w, T, S)

Return the ``y``-component of the turbulent diffusive flux divergence:

`∂x(2 ν Σ₂₁) + ∂y(2 ν Σ₂₂) + ∂z(2 ν Σ₂₂)`

at the location `ccf`.
"""
∂ⱼ_2ν_Σ₂ⱼ(i, j, k, grid, closure::IsotropicDiffusivity, args...) = (
    ∂x_2ν_Σ₂₁(i, j, k, grid, closure, args...)
  + ∂y_2ν_Σ₂₂(i, j, k, grid, closure, args...)
  + ∂z_2ν_Σ₂₃(i, j, k, grid, closure, args...)
)

"""
    ∂ⱼ_2ν_Σ₃ⱼ(i, j, k, grid, closure, eos, g, u, v, w, T, S)

Return the ``z``-component of the turbulent diffusive flux divergence:

`∂x(2 ν Σ₃₁) + ∂y(2 ν Σ₃₂) + ∂z(2 ν Σ₃₃)`

at the location `ccf`.
"""
∂ⱼ_2ν_Σ₃ⱼ(i, j, k, grid, closure::IsotropicDiffusivity, args...) = (
    ∂x_2ν_Σ₃₁(i, j, k, grid, closure, args...)
  + ∂y_2ν_Σ₃₂(i, j, k, grid, closure, args...)
  + ∂z_2ν_Σ₃₃(i, j, k, grid, closure, args...)
)

"""
    κ_∂x_ϕ(i, j, k, grid, ϕ, κ, closure, eos, g, u, v, w, T, S)

Return `κ ∂x ϕ`, where `κ` is a function that computes
diffusivity at cell centers (location `ccc`), and `ϕ` is an array of scalar
data located at cell centers.
"""
function κ_∂x_ϕ(i, j, k, grid, ϕ, κ, closure, args...)
    κ = ▶x_faa(i, j, k, grid, κ, closure, args...)
    ∂x_ϕ = ∂x_faa(i, j, k, grid, ϕ)
    return κ * ∂x_ϕ
end

"""
    κ_∂y_ϕ(i, j, k, grid, ϕ, κ, closure, eos, g, u, v, w, T, S)

Return `κ ∂y ϕ`, where `κ` is a function that computes
diffusivity at cell centers (location `ccc`), and `ϕ` is an array of scalar
data located at cell centers.
"""
function κ_∂y_ϕ(i, j, k, grid, ϕ, κ, closure, args...)
    κ = ▶y_afa(i, j, k, grid, κ, closure, args...)
    ∂y_ϕ = ∂y_afa(i, j, k, grid, ϕ)
    return κ * ∂y_ϕ
end

"""
    κ_∂z_ϕ(i, j, k, grid, ϕ, κ, closure, eos, g, u, v, w, T, S)

Return `κ ∂z ϕ`, where `κ` is a function that computes
diffusivity at cell centers (location `ccc`), and `ϕ` is an array of scalar
data located at cell centers.
"""
function κ_∂z_ϕ(i, j, k, grid, ϕ, κ, closure, args...)
    κ = ▶z_aaf(i, j, k, grid, κ, closure, args...)
    ∂z_ϕ = ∂z_aaf(i, j, k, grid, ϕ)
    return κ * ∂z_ϕ
end

"""
    ∂x_κ_∂x_ϕ(i, j, k, grid, ϕ, κ, closure, eos, g, u, v, w, T, S)

Return `(∂x κ ∂x) ϕ`, where `κ` is a function that computes
diffusivity at cell centers (location `ccc`), and `ϕ` is an array of scalar
data located at cell centers.
"""
∂x_κ_∂x_ϕ(i, j, k, grid, ϕκ...) = ∂x_caa(i, j, k, grid, κ_∂x_ϕ, ϕκ...)

"""
    ∂y_κ_∂y_ϕ(i, j, k, grid, ϕ, κ, closure, eos, g, u, v, w, T, S)

Return `(∂y κ ∂y) ϕ`, where `κ` is a function that computes
diffusivity at cell centers (location `ccc`), and `ϕ` is an array of scalar
data located at cell centers.
"""
∂y_κ_∂y_ϕ(i, j, k, grid, ϕκ...) = ∂y_aca(i, j, k, grid, κ_∂y_ϕ, ϕκ...)

"""
    ∂z_κ_∂z_ϕ(i, j, k, grid, ϕ, κ, closure, eos, g, u, v, w, T, S)

Return `(∂z κ ∂z) ϕ`, where `κ` is a function that computes
diffusivity at cell centers (location `ccc`), and `ϕ` is an array of scalar
data located at cell centers.
"""
∂z_κ_∂z_ϕ(i, j, k, grid, ϕκ...) = ∂z_aac(i, j, k, grid, κ_∂z_ϕ, ϕκ...)

"""
    ∇_κ_∇_ϕ(i, j, k, grid, ϕ, closure, eos, g, u, v, w, T, S)

Return the diffusive flux divergence `∇ ⋅ (κ ∇ ϕ)` for the turbulence
`closure`, where `ϕ` is an array of scalar data located at cell centers.
"""
∇_κ_∇ϕ(i, j, k, grid, ϕ, closure::IsotropicDiffusivity, args...) = (
      ∂x_κ_∂x_ϕ(i, j, k, grid, ϕ, κ_ccc, closure, args...)
    + ∂y_κ_∂y_ϕ(i, j, k, grid, ϕ, κ_ccc, closure, args...)
    + ∂z_κ_∂z_ϕ(i, j, k, grid, ϕ, κ_ccc, closure, args...)
    )


function wind_frame(xlocs, ylocs, wind_direction)

    nTurbines = length(xlocs)
    wd = -pi*(270. - wind_direction)/180. #shift to traditional wind direction coords and to radians
    xw = xlocs.*cos(wd)-ylocs.*sin(wd)
    yw = xlocs.*sin(wd)+ylocs.*cos(wd)
    return xw, yw
end


using Revise
using Basins
using DifferentialEquations
using Plots

# Equations of motion: E. Ott, et al. I Physica D 76 (1994) 384-410
function forced_particle!(du, u, p, t)
γ=0.05  ; x̄ = 1.9  ; f₀=2.3  ; ω =3.5
x₀=1. ; y₀=0.;

    x=u[1]; dx = u[2]; y=u[3]; dy = u[4];

    du[1] = dx
    du[2] = -γ*dx -(-4*x*(1-x^2) + y^2) +  f₀*sin(ω*t)*x₀
    du[3] = dy
    du[4] = -γ*dy -(2*y*(x+x̄)) +  f₀*sin(ω*t)*y₀
end

df = ODEProblem(forced_particle!,rand(4),(0.0,20.0), [])
integ  = init(df, alg=Tsit5(); reltol=1e-9, save_everystep=false)

xres=300
yres=300

# range for forced pend
xg = range(0.,1.2,length=xres)
yg = range(0.,1.2,length=yres)

ω =3.5

# compute basin
#@time bsn = Basins.basin_general_ds(xg, yg, integ; dt=2*pi/ω, idxs=[1,3])

#plot(xg,yg,bsn.basin', seriestype=:heatmap)

function escape_function(u)

    if abs(u[3]) > 20 && abs(u[4])>100 && u[3]*u[4]>0
        @show integ.t
        return 2
    elseif abs(u[3]) < 1e-8 && abs(u[4])<1e-9 && u[3]*u[4]<0
        return 1
    end
    return 0
end

basin_t = zeros(length(xg),length(yg))
for (k,x) in enumerate(xg), (n,y) in enumerate(yg)
    reinit!(integ,[x,0,y,0])
    step!(integ, 2π/ω*10,true)
    while escape_function(integ.u) == 0
        step!(integ, 2π/ω,true)
        integ.t > 200 && break
    end

    basin_t[k,n]=escape_function(integ.u);
end

using JLD

@save "riddle.jld" basin_t


using Base.Test
using Gunrock

a = include("init.jl")

#BFS
@test bfs(a,1) == [1, 2, 2, 2, 3, 3, 3]

#Node betweenness centrality
@test bc(a,1) == [0.0, 0.25, 0.5, 0.75, 0.0, 0.0, 0.0]

#Connected components
num_cc, ids = cc(a)
@test num_cc == 1
@test ids == [1, 1, 1, 1, 1, 1, 1]

#Single source shortest paths
@test sssp(a,1) == [0, 1, 1, 1, 2, 2, 2]

#Pagerank
nodes, ranks = pagerank(a)
@test nodes == [3, 4, 5, 6, 1, 2, 7]
@test_approx_eq ranks [1.2522, 1.0280, 1.0280, 1.0201, 0.8048, 0.8048, 0.7991]


"""
    SecondChebyshevPlans

implementation of computation of second Poincare invariant by approximating surface with
Chebyshev polynomials
"""
module SecondChebyshevPlans

import ..PoincareInvariants: compute!, getpoints, getpointspec

using ..PoincareInvariants: SecondPoincareInvariant, @argcheck

using LinearAlgebra

include("PaduaTransforms.jl")
using .PaduaTransforms

## Differentiation ##

function getdiffmat(::Type{T}, degree::Integer) where T
    D = zeros(T, degree+1, degree+1)

    D[1, 2:2:end] .= 1:2:degree

    for i in 3:2:degree+1
        D[2:2:i-1, i] .= 2 * (i-1)
    end

    for i in 4:2:degree+1
        D[3:2:i-1, i] .= 2 * (i-1)
    end

    D
end

struct DiffPlan{T}
    D::Matrix{T}
end

DiffPlan{T}(degree::Integer) where T = DiffPlan{T}(getdiffmat(T, degree))

# differentiate in the Chebyshev basis
function differentiate!(∂x::AbstractMatrix, ∂y::AbstractMatrix, P::DiffPlan, coeffs::AbstractMatrix)
    ∂x[:, :] = coeffs  # differentiate each row
    rmul!(∂x, LowerTriangular(P.D'))

    ∂y[:, :] = coeffs  # differentiate each column
    lmul!(UpperTriangular(P.D), ∂y)

    ∂x, ∂y
end

function differentiate!(∂x, ∂y, P::DiffPlan, coeffs)
    eltype(∂x) <: AbstractMatrix && eltype(∂y) <: AbstractMatrix &&
        eltype(coeffs) <: AbstractMatrix || throw(ArgumentError(
            "coefficients must be AbstractMatrix or iterable thereof"))
    length(∂x) == length(∂y) == length(coeffs) || throw(ArgumentError(
        "number of coefficient matrices must match number of derivative matrices to write to"))

    for (∂xi, ∂yi, coeffsi) in zip(∂x, ∂y, coeffs)
        differentiate!(∂xi, ∂yi, P, coeffsi)
    end

    ∂x, ∂y
end

## Integration ##

getintweights(::Type{T}, n) where T = T[isodd(i) ? 0 : T(2) / T(1 - i^2) for i in 0:n]
getintweights(n) = getintweights(Float64, n)

integrate(coeffs, intweights) = dot(intweights, coeffs, intweights)

## getintegrand for ω Callable and out-of-place ##

struct CallIntPlan{T, D, IP, P}
    invpaduaplan::IP
    ∂xvals::Matrix{T}
    ∂yvals::Matrix{T}
    intvals::Vector{T}
    paduaplan::P
end

function CallIntPlan{T, D}(degree) where {T, D}
    invpaduaplan = InvPaduaTransformPlan{Float64}(degree)
    ∂xvals = Matrix{T}(undef, getpaduanum(degree), D)
    ∂yvals = Matrix{T}(undef, getpaduanum(degree), D)
    intvals = Vector{T}(undef, getpaduanum(degree))
    paduaplan = PaduaTransformPlan{Float64}(degree)

    CallIntPlan{T, D, typeof(invpaduaplan), typeof(paduaplan)}(
        invpaduaplan, ∂xvals, ∂yvals, intvals, paduaplan
    )
end

function getintegrand!(
    intcoeffs::AbstractMatrix, plan::CallIntPlan{T, D}, ω::ωT,
    points, t, p, ∂xcoeffs, ∂ycoeffs
) where {T, D, ωT}
    invpaduatransform!(plan.∂xvals, plan.invpaduaplan, ∂xcoeffs)
    invpaduatransform!(plan.∂yvals, plan.invpaduaplan, ∂ycoeffs)

    for i in axes(plan.intvals, 1)
        # This if statement should hopefully get optimised away by the compiler
        if ωT <: AbstractMatrix
            ωi = ω
        else
            pnti = view(points, i, :)
            ωi = ω(pnti, t, p)
        end

        ∂xi = view(plan.∂xvals, i, :)
        ∂yi = view(plan.∂yvals, i, :)
        plan.intvals[i] = dot(∂yi, ωi, ∂xi)
    end

    paduatransform!(intcoeffs, plan.paduaplan, plan.intvals)

    intcoeffs
end

## SecondChebyshevPlan and compute! ##

getintplan(::Type{T}, ::Any, ::Val{D}, degree) where {T, D} = CallIntPlan{T, D}(degree)
# getintplan(::Type{T}, ω::AbstractMatrix, D, degree) where T = ConstIntPlan{T}(ω, D, degree)

struct SecondChebyshevPlan{T, D, IP, PP<:PaduaTransformPlan}
    degree::Int
    paduaplan::PP
    phasecoeffs::NTuple{D, Matrix{T}}
    diffplan::DiffPlan{T}
    ∂x::NTuple{D, Matrix{T}}
    ∂y::NTuple{D, Matrix{T}}
    intplan::IP  # getting coefficients of integrand to integrate
    intcoeffs::Matrix{T}
    intweights::Vector{T}
end

function SecondChebyshevPlan{T, D}(ω, N::Integer) where {T, D}
    degree = nextdegree(N)

    paduaplan = PaduaTransformPlan{Float64}(degree)
    phasecoeffs = ntuple(_ -> zeros(T, degree+1, degree+1), D)
    diffplan = DiffPlan{T}(degree)
    ∂x = ntuple(_ -> Matrix{T}(undef, degree+1, degree+1), D)
    ∂y = ntuple(_ -> Matrix{T}(undef, degree+1, degree+1), D)
    intplan = getintplan(T, ω, Val(D), degree)
    intcoeffs = zeros(T, degree+1, degree+1)
    intweights = getintweights(T, degree)

    SecondChebyshevPlan{T, D, typeof(intplan), typeof(paduaplan)}(
        degree, paduaplan, phasecoeffs, diffplan, ∂x, ∂y, intplan, intcoeffs, intweights
    )
end

function compute!(
    pinv::SecondPoincareInvariant{<:Any, <:Any, <:Any, <:Any, P}, t, p
) where P <: SecondChebyshevPlan
    plan = pinv.plan
    points = pinv.points
    paduatransform!(plan.phasecoeffs, plan.paduaplan, points)
    differentiate!(plan.∂x, plan.∂y, plan.diffplan, plan.phasecoeffs)
    getintegrand!(plan.intcoeffs, plan.intplan, pinv.ω, points, t, p, plan.∂x, plan.∂y)
    integrate(plan.intcoeffs, plan.intweights)
end

# function _compute!(plan::SecondChebyshevPlan, ω::AbstractMatrix, D, points)
#     paduatransform!(plan.phasecoeffs, plan.paduaplan, points)
#     differentiate!(plan.∂x, plan.∂y, plan.diffplan, plan.phasecoeffs)
#     getintegrand!(plan.intcoeffs, plan.intplan, ω, plan.∂x, plan.∂y)
#     integrate(plan.intcoeffs, plan.intplan)
# end

## getpoints, getpointspec and getpointnum ##

getpointspec(N::Integer, ::Type{<:SecondChebyshevPlan}) = nextpaduanum(N)

function getpoints(f, ::Type{T}, N, ::Type{<:SecondChebyshevPlan}) where T
    degree = getdegree(getpointspec(N, SecondChebyshevPlan))
    return getpaduapoints(T, degree) do x, y
        f((x + T(1)) / T(2), (y + T(1)) / T(2))
    end
end

end  # module Chebyshev


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

    # normal constructor
    a = ones(N, 3)
    b = N(5)
    hp = Hyperplane(a, b)

    # corner case: zero normal vector
    @test_throws AssertionError Hyperplane(N[0, 0], N(1))

    # dimension
    @test dim(hp) == 3

    # support function
    hp2 = Hyperplane(N[1, 0], N(1))
    @test ρ(N[2, 0], hp2) == N(2)
    @test ρ(N[-2, 0], hp2) == N(-2)
    @test ρ(N[1, 1], hp2) == N(Inf)

    # support vector and membership
    function test_svec(hp)
        d1 = copy(hp.a)
        @test σ(d1, hp) ∈ hp
        @test σ(N(2) * d1, hp) ∈ hp
        d2 = N[1, 0, 0]
        @test_throws ErrorException σ(d2, hp)
        d3 = N[1, 1, 2]
        @test_throws ErrorException σ(d3, hp)
        d4 = zeros(N, 3)
        @test σ(d4, hp) ∈ hp
    end
    # tests 1
    test_svec(hp)
    # tests 2
    test_svec(Hyperplane(N[0, 0, 1], b))

    # support vector in opposite direction
    hp2 = Hyperplane(N[1], N(1))
    @test σ(N[-1], hp2) ∈ hp2
    # support vector in other directions throws an error (but see #750)
    @test_throws ErrorException σ(N[1, 1], Hyperplane(N[1, 0], N(1)))

    # boundedness
    @test !isbounded(hp)

    # universality
    @test !isuniversal(hp)
    res, w = isuniversal(hp, true)
    @test !res && w ∉ hp

    # isempty
    @test !isempty(hp)

    # an_element and membership
    @test an_element(hp) ∈ hp

    # constrained dimensions
    @test constrained_dimensions(Hyperplane(N[1, 0, 1], N(1))) == [1, 3]
    @test constrained_dimensions(Hyperplane(N[0, 1, 0], N(1))) == [2]

    # constraints_list
    @test ispermutation(constraints_list(hp),
                        [HalfSpace(a, b), HalfSpace(-a, -b)])

    # test concrete linear map of a hyerplane
    H = Hyperplane(N[1, -1], N(0)) # x = y
    M = N[1 0; 0 0] # non-invertible matrix
    @test_throws ArgumentError linear_map(M, H)
    M = N[2 2; 0 1] # invertible matrix
    @test linear_map(M, H) == Hyperplane(N[0.5, -2.0], N(0.0))

    # translation
    @test translate(hp, N[1, 2, 3]) == Hyperplane(ones(N, 3), N(11))

    # intersection emptiness
    u = Universe{N}(3)
    res, v = is_intersection_empty(u, hp, true)
    @test !is_intersection_empty(u, hp) && !res && v ∈ hp && v ∈ u
    res, v = is_intersection_empty(hp, u, true)
    @test !is_intersection_empty(hp, u) && !res && v ∈ hp && v ∈ u
end

# Polyhedra tests that only work with Float64
for N in [Float64]
    hp = Hyperplane(ones(N, 3), N(5))

    # intersection emptiness
    b = BallInf(zeros(N, 3), N(1))
    empty_intersection, v = is_intersection_empty(b, hp, true)
    @test empty_intersection && is_intersection_empty(b, hp) && v == N[]
    b = BallInf(N[3, 3, 3], N(1))
    empty_intersection, v = is_intersection_empty(b, hp, true)
    @test empty_intersection && is_intersection_empty(b, hp) && v == N[]
    b = BallInf(N[2, 2, 2], N(1))
    empty_intersection, v = is_intersection_empty(b, hp, true)
    @test !empty_intersection && !is_intersection_empty(b, hp) && v ∈ hp && v ∈ b
end


mutable struct RMSProp <: Optimizer
    xparams::Vector{XVariable}
    w::Vector
    lr::AbstractFloat
    ϵ::AbstractFloat
    inertia::AbstractFloat
    name::String
    function RMSProp(xparams::Vector{XVariable}; lr=1e-2, inertia=0.99, epsilon=1e-8)
        n = length(xparams)
        w = Vector(undef,n)
        for i = 1:n
            c , θ = xparams[i]
            w[i] = Zeros(typeof(θ.value), θ.shape)
        end
        new(xparams, w, lr, epsilon, inertia, "RMSProp")
    end
end

function Base.show(io::IO, O::RMSProp)
    print("RMSProp(lr=$(O.lr), ϵ=$(O.ϵ), inertia=$(O.inertia))")
end

function update!(O::RMSProp; clipfn::Function=LPInfNormClip, clipvalue=10.0)
    w = O.w
    ϵ = O.ϵ
    μ = - O.lr
    ρ = O.inertia

    for i = 1:length(O.xparams)
        c , θ = O.xparams[i]
        ∇ = clipfn(setNanInfZero(θ.delta), clipvalue)
        @. w[i] = ρ * w[i] + (1-ρ) * ∇ * ∇
        @. θ.value += μ / (sqrt(w[i])+ϵ) * ∇
    end
end


@doc raw"""
First Poincaré invariant for a guiding center particle in an θ-pinch magnetic
field of the form ``B(x,y,z) = B_0 \, e_z``.

The loop for the first Poincaré invariant is initialized by
```math
q (\tau) = \begin{pmatrix}
r_x \cos (2\pi \tau) \\
y_0 \\
r_z \sin (2\pi \tau) \\
\end{pmatrix}
```
with parameters
```math
B_0 = 1, \quad
r_x = 0.5, \quad
r_z = 0.3, \quad
y_0 = 0.0, \quad
u_0 = 0.5, \quad
\mu = 2.5 \times 10^{-6} .
```
"""
module ThetaPinchField

    import ElectromagneticFields.ThetaPinch

    ThetaPinch.@code() # inject magnetic field code

    μ_loop() = 2.5E-6

    function f_loop(t)
        μ  = 2.5E-6
        Y0 = 0.0
        u0 = 4E-4
        r0 = 0.5
        r1 = 0.3

        Xt = r0*cos(2π*t)
        Zt = r1*sin(2π*t)

        qt = [Xt, Y0, Zt, u0]

        return qt
    end

    include("guiding_center_4d_common.jl")
    include("guiding_center_4d_equations.jl")
    include("guiding_center_4d_loop.jl")
    include("guiding_center_4d_diagnostics.jl")

end


include("BayesTest.jl")
import Gadfly.plot, Mamba.hpd

type BinomTest <: BayesTest
    prior::Distribution
    data::Dict{AbstractString, Real}
    BF::Float64 # BF_{01}
    posterior::Any # either distribution, or NaN
    params::AbstractString

    function BinomTest(y::Integer, N::Integer; a = 1, b = 1)
        self = new()
        self.params = "θ"
        self.data = Dict("y" => y, "N" => N)
        self.prior = Distributions.Beta(a, b)
        self.posterior = Distributions.Beta(a + y, b + N - y)
        self.BF = pdf(self.prior, .5) / pdf(self.posterior, .5) # Savage-Dickey Trick
        return self
    end
end

function update(Model::BinomTest, y, N)
    BinomTest(Model.data["y"] + y, Model.data["N"] + N;
              a = Model.posterior.α, b = Model.posterior.β)
end

function posterior(Model::BinomTest; iter = 10000)
    return rand(Model.posterior, iter)
end

function hpd(samples::Array; low::Real = .025, up::Real = .975)
    ss = sort(samples); n = length(ss)
    return ss[Int(n * low)], ss[Int(n * up)]
end

function plot(Model::BinomTest)
    fun(x) = Distributions.pdf(Model.posterior, x)
    sup = Distributions.support(Model.posterior)
    Gadfly.plot(fun, sup.lb, sup.ub)
end

function plot(samples::Array)
    density = KernelDensity.kde(samples)
    f(x) = KernelDensity.pdf(density, x)
    Gadfly.plot(f, minimum(dist), maximum(dist))
end


using Distributions
using Base.Test

m = MixtureModel([Poisson(2.0), Binomial(10, 0.3)],
                 [0.4, 0.6])
x = rand(m, (100,))
@test_approx_eq exp(logpdf(m,x)) pdf(m,x)
@test Distributions.variate_form(typeof(m))=== Univariate
@test Distributions.value_support(typeof(m))=== Discrete


"""
    FourierDerivativeOperator2D{T<:Real}

A derivative operator on a two-dimensional periodic grid with scalar type `T`
computing the first derivatives using a spectral Fourier expansion via
real discrete Fourier transforms.
"""
struct FourierDerivativeOperator2D{T<:Real, Grid, RFFTx, BRFFTx, RFFTy, BRFFTy} <: AbstractPeriodicDerivativeOperator{T}
    jac_x::T
    jac_y::T
    Δx::T
    Δy::T
    grid_compute::Grid  # (Nx-1)x(Ny-1) nodes, including the left and excluding the right boundary
    grid_evaluate::Grid # (Nx)  x(Ny  ) nodes, including all boundaries
    tmp_x::Array{Complex{T}, 2}
    rfft_plan_x::RFFTx
    brfft_plan_x::BRFFTx
    tmp_y::Array{Complex{T}, 2}
    rfft_plan_y::RFFTy
    brfft_plan_y::BRFFTy

    function FourierDerivativeOperator2D(jac_x::T, jac_y::T, Δx::T, Δy::T, grid_compute::Grid, grid_evaluate::Grid,
                                         tmp_x::Array{Complex{T}, 2}, rfft_plan_x::RFFTx, brfft_plan_x::BRFFTx,
                                         tmp_y::Array{Complex{T}, 2}, rfft_plan_y::RFFTy, brfft_plan_y::BRFFTy) where {T<:Real, Grid, RFFTx, BRFFTx, RFFTy, BRFFTy}
        @argcheck size(brfft_plan_x) == size(tmp_x) DimensionMismatch
        @argcheck size(brfft_plan_x, 1) == (size(rfft_plan_x, 1)÷2)+1 DimensionMismatch
        @argcheck size(grid_compute) == size(rfft_plan_x) DimensionMismatch
        @argcheck size(grid_compute) == size(grid_evaluate).-1 DimensionMismatch

        new{T, Grid, RFFTx, BRFFTx, RFFTy, BRFFTy}(jac_x, jac_y, Δx, Δy, grid_compute, grid_evaluate,
                                                   tmp_x, rfft_plan_x, brfft_plan_x,
                                                   tmp_y, rfft_plan_y, brfft_plan_y)
    end
end

"""
    FourierDerivativeOperator2D(xmin, xmax, Nx, ymin, ymax, Ny)

Construct the `FourierDerivativeOperator` on a uniform grid between
`xmin` and `xmax` using `Nx` nodes and `ymin` and `ymax` using `Ny` nodes.
"""
function FourierDerivativeOperator2D(xmin::T, xmax::T, Nx::Int, ymin::T, ymax::T, Ny::Int) where {T<:Real}
    @argcheck Nx >= 1
    @argcheck Ny >= 1

    jac_x = 2*T(π) / (xmax - xmin) / Nx # / N because of brfft instead of BRFFT
    Δx = (xmax - xmin) / Nx
    jac_y = 2*T(π) / (ymax - ymin) / Ny # / N because of brfft instead of BRFFT
    Δy = (ymax - ymin) / Ny
    grid_evaluate = zeros(SVector{2,T}, Nx+1, Ny+1)
    grid_x = range(xmin, stop=xmax, length=Nx+1)
    grid_y = range(ymin, stop=ymax, length=Ny+1) # two boundary nodes
    for j in Base.OneTo(Ny+1), i in Base.OneTo(Nx+1)
        grid_evaluate[i,j] = SVector(grid_x[i], grid_y[j])
    end
    grid_compute  = grid_evaluate[1:end-1, 1:end-1]
    u = zero.(first.(grid_compute))
    rfft_plan_x = plan_rfft(u, 1)
    uhat_x = rfft_plan_x*u
    brfft_plan_x = plan_brfft(uhat_x, Nx, 1)
    rfft_plan_y = plan_rfft(u, 2)
    uhat_y = rfft_plan_y*u
    brfft_plan_y = plan_brfft(uhat_y, Ny, 2)

    FourierDerivativeOperator2D(jac_x, jac_y, Δx, Δy, grid_compute, grid_evaluate,
                                uhat_x, rfft_plan_x, brfft_plan_x,
                                uhat_y, rfft_plan_y, brfft_plan_y)
end

function fourier_derivative_operator(xmin::Real, xmax::Real, Nx::Int, ymin::Real, ymax::Real, Ny::Int)
    xmin, xmax, ymin, ymax = promote(xmin, xmax, ymin, ymax)
    FourierDerivativeOperator2D(xmin, xmax, Nx, ymin, ymax, Ny)
end

derivative_order(D::FourierDerivativeOperator2D) = 1
LinearAlgebra.issymmetric(D::FourierDerivativeOperator2D) = false

function Base.show(io::IO, D::FourierDerivativeOperator2D{T}) where {T}
    grid = D.grid_evaluate
    print(io, "Periodic 1st derivative Fourier operator {T=", T, "} in two space dimensions")
end

function mul!(dest::AbstractArray{T, 2}, D::FourierDerivativeOperator2D, u::AbstractArray{T, 2}, ::Val{:x}) where {T}
    @unpack jac_x, tmp_x, rfft_plan_x, brfft_plan_x = D
    Nx, Ny = size(D.grid_compute)
    @boundscheck begin
        @argcheck (Nx, Ny) == size(u)
        @argcheck (Nx, Ny) == size(dest)
    end

    mul!(tmp_x, rfft_plan_x, u)
    size_x, size_y = size(tmp_x)
    @inbounds @simd for j in Base.OneTo(size_y); for i in Base.OneTo(size_x-1)
        tmp_x[i, j] *= (i-1)*im * jac_x
    end end
    # see e.g. Steven G. Johnson (2011) Notes on FFT based differentiation
    if iseven(Nx)
        @inbounds @simd for j in Base.OneTo(size_y)
            tmp_x[end, j] = zero(eltype(tmp_x))
        end
    else
        @inbounds @simd for j in Base.OneTo(size_y)
            tmp_x[end, j] *= (size_x-1)*im * jac_x
        end
    end
    mul!(dest, brfft_plan_x, tmp_x)
end

function mul!(dest::AbstractArray{T, 2}, D::FourierDerivativeOperator2D, u::AbstractArray{T, 2}, ::Val{:y}) where {T}
    @unpack jac_y, tmp_y, rfft_plan_y, brfft_plan_y = D
    Nx, Ny = size(D.grid_compute)
    @boundscheck begin
        @argcheck (Nx, Ny) == size(u)
        @argcheck (Nx, Ny) == size(dest)
    end

    mul!(tmp_y, rfft_plan_y, u)
    size_x, size_y = size(tmp_y)
    @inbounds @simd for j in Base.OneTo(size_y-1); for i in Base.OneTo(size_x)
        tmp_y[i, j] *= (j-1)*im * jac_y
    end end
    # see e.g. Steven G. Johnson (2011) Notes on FFT based differentiation
    if iseven(Ny)
        @inbounds @simd for i in Base.OneTo(size_x)
            tmp_y[i, end] = zero(eltype(tmp_y))
        end
    else
        @inbounds @simd for i in Base.OneTo(size_x)
            tmp_y[i, end] *= (size_y-1)*im * jac_y
        end
    end
    mul!(dest, brfft_plan_y, tmp_y)
end

# TODO there is no 5 argument mul! in FFTW.jl...

function integrate(func, u::AbstractMatrix, D::FourierDerivativeOperator2D)
    @boundscheck begin
        length(u) == length(grid(D))
    end
    @unpack Δx, Δy = D

    @inbounds res = sum(func, u)

    (Δx * Δy) * res
end

function mass_matrix(D::FourierDerivativeOperator2D)
    @unpack Δx, Δy = D

    (Δx * Δy) * I
end


"""
    AsymmetricLaplace(μ,λ,κ)

The *Asymmetric Laplace distribution* with location `μ` scale `λ` and asymmetry `κ` has probability density function

```math
f(x; \\mu, \\lambda, \\kappa) = \\left( \\frac{\\lambda}{\\kappa + 1/\\kappa} \\right) \\exp - (x - \\mu) \\lambda \\sgn(x - \\mu) \\kappa^\\sgn(x - \\mu)
```

```julia
AsymmetricLaplace()          # Asymmetric Laplace distribution with zero location, unit scale and asymmetry 1, i.e. AsymmetricLaplace(0, 1, 1)
AsymmetricLaplace(μ)         # Asymmetric Laplace distribution with location μ, unit scale and asymmetry 1, i.e. AsymmetricLaplace(μ, 1, 1)
AsymmetricLaplace(μ, λ)      # Asymmetric Laplace distribution with location μ, scale λ and asymmetry 1
AsymmetricLaplace(μ, λ, κ)   # Asymmetric Laplace distribution with location μ, scale λ and asymmetry κ

params(d)                    # Get the parameters, i.e., (μ, λ, κ)
location(d)                  # Get the location parameter, i.e. μ
scale(d)                     # Get the scale parameter, i.e. λ
asymmetry(d)                 # Get the asymmetry parameter, i.e. κ
```

External links

* [Asymmetric Laplace distribution on Wikipedia](http://en.wikipedia.org/wiki/Asymmetric_Laplace_distribution)

"""
struct AsymmetricLaplace{T<:Real} <: ContinuousUnivariateDistribution
    μ::T
    λ::T
    κ::T
    AsymmetricLaplace{T}(µ::T, λ::T, κ::T) where {T} = new{T}(µ, λ, κ)
end

function AsymmetricLaplace(μ::T, λ::T, κ::T; check_args=true) where {T <: Real}
    check_args && @check_args(AsymmetricLaplace, λ > zero(λ) && κ > zero(κ))
    return AsymmetricLaplace{T}(μ, λ, κ)
end

AsymmetricLaplace(μ::Real, λ::Real, κ::Real) = AsymmetricLaplace(promote(μ, λ, κ)...)
AsymmetricLaplace(μ::Integer, λ::Integer, κ::Integer) = AsymmetricLaplace(float(μ), float(λ), float(κ))
AsymmetricLaplace(μ::T) where {T <: Real} = AsymmetricLaplace(μ, one(T), one(T))
AsymmetricLaplace(μ::T, λ::T) where {T <: Real} = AsymmetricLaplace(μ, λ, one(T))
AsymmetricLaplace() = AsymmetricLaplace(0.0, 1.0, 1.0, check_args=false)

@distr_support AsymmetricLaplace -Inf Inf

#### Conversions
function convert(::Type{AsymmetricLaplace{T}}, μ::S, λ::S, κ::S) where {T <: Real, S <: Real}
    AsymmetricLaplace(T(μ), T(λ), T(κ))
end
function convert(::Type{AsymmetricLaplace{T}}, d::AsymmetricLaplace{S}) where {T <: Real, S <: Real}
    AsymmetricLaplace(T(d.μ), T(d.λ), T(d.κ), check_args=false)
end

#### Parameters

location(d::AsymmetricLaplace) = d.μ
scale(d::AsymmetricLaplace) = d.λ
asymmetry(d::AsymmetricLaplace) = d.κ
params(d::AsymmetricLaplace) = (d.μ, d.λ, d.κ)
@inline partype(d::AsymmetricLaplace{T}) where {T<:Real} = T

#### Statistics

mean(d::AsymmetricLaplace) = d.μ + (1 - d.κ^2) / (d.λ * d.κ)
function median(d::AsymmetricLaplace)
    μ, λ, κ = params(d)
    if κ > 1
        μ + κ/λ * log((1 + κ^2)/(2κ^2))
    else
        μ - 1/(κ * λ) * log((1 + κ^2)/2)
    end
end
mode(d::AsymmetricLaplace) = d.μ

var(d::AsymmetricLaplace) = (1 + d.κ^4) / (d.λ^2 * d.κ^2)
std(d::AsymmetricLaplace) = sqrt(1 + d.κ^4) / (d.λ * d.κ)
skewness(d::AsymmetricLaplace) = 2(1 - d.κ^6) / (d.κ^4 + 1)^(3/2)
kurtosis(d::AsymmetricLaplace) = 6(1 + d.κ^8) / (1 + d.κ^4)^2

entropy(d::AsymmetricLaplace) = log(ℯ * (1 + d.κ) / (d.λ * d.κ))

#### Evaluations

function pdf(d::AsymmetricLaplace, x::Real)
    μ, λ, κ = params(d)
    if x ≤ μ
        exp((x - μ) / (λ * κ)) / λ / (κ + 1 / κ)
    else
        exp(- κ / λ * (x - μ)) / λ / (κ + 1 / κ)
    end
end
function logpdf(d::AsymmetricLaplace, x::Real)
    μ, λ, κ = params(d)
    if x ≤ μ
        (x - μ) / (λ * κ) - log(λ) - log(κ + 1 / κ)
    else
        - κ * (x - μ) / λ - log(λ) - log(κ + 1 / κ)
    end
end

function cdf(d::AsymmetricLaplace, x::Real)
    μ, λ, κ = params(d)
    if x ≤ μ
        κ^2 / (1 + κ^2) * exp((x - μ) / (λ * κ))
    else
        1 - 1 / (1 + κ^2) * exp(- κ / λ * (x - μ))
    end
end
ccdf(d::AsymmetricLaplace, x::Real) = 1 - cdf(d, x)
function logcdf(d::AsymmetricLaplace, x::Real)
    μ, λ, κ = params(d)
    if x ≤ μ
        2 * log(κ) - log(1 + κ^2) + (x - μ) / (λ * κ)
    else
        log(1 - 1 / (1 + κ^2) * exp(- κ / λ * (x - μ)))
    end
end
function logccdf(d::AsymmetricLaplace, x::Real)
    μ, λ, κ = params(d)
    if x ≤ μ
        log(1 - κ^2 / (1 + κ^2) * exp((x - μ) / (λ * κ)))
    else
        -(κ * (x - μ) / λ  + log(1 + κ^2))
    end
end

function quantile(d::AsymmetricLaplace, p::Real)
    μ, λ, κ = params(d)
    if p ≤ cdf(d, d.μ)
        μ + log((1 + κ^2) / κ^2 * p) * (λ * κ)
    else
        μ - λ / κ * log((1 - p) * (1 + κ^2))
    end
end
cquantile(d::AsymmetricLaplace, p::Real) = 1 - quantile(d, p)
function invlogcdf(d::AsymmetricLaplace, lp::Real)
    μ, λ, κ = params(d)
    if lp ≤ 0
        println("branch less than")
        μ - (λ / κ) * (log(1 - exp(lp)) + log(1 + κ^2))
    else
        println("branch larger than")
        μ + λ * κ * (log(lp) + log(1 + κ^2) - 2 * log(κ))
    end
end
function invlogccdf(d::AsymmetricLaplace, lp::Real)
    μ, λ, κ = params(d)
    if lp ≤ log(cdf(d, μ))
        μ - (λ / κ) * (lp + log(1 + κ^2))
    else
        μ + λ * κ * log((1 + κ^2) * (1 - exp(lp)) / κ^2)
    end
end

function gradlogpdf(d::AsymmetricLaplace, x::Real)
    x == d.μ && error("Gradient is undefined at the location point")
    x > d.μ ? - d.κ / d.λ : 1 / (d.λ * d.κ)
end

function mgf(d::AsymmetricLaplace, t::Real)
    exp(t * d.μ) / ((1 + t * d.κ / d.λ) * (1 - t / (d.κ * d.λ)))
end
function cf(d::AsymmetricLaplace, t::Real)
    cis(t * d.μ) / ((1 + t * d.κ * im / d.λ) * (1 - t * im / (d.κ * d.λ)))
end

#### Sampling

rand(rng::AbstractRNG, d::AsymmetricLaplace) = quantile(d, rand(rng))


#Basic math
using DSP, Statistics, LinearAlgebra, Random

a = rand(Float32, 10, 10)
ad = AFArray(a)

@testset "Math" begin
    @test sum(abs2, Array(ad + 2) - (a .+ 2)) < 1e-6
    @test sum(abs2, Array(ad .+ 2) - (a .+ 2)) < 1e-6
    @test sum(abs2, Array(2 + ad) - (2 .+ a)) < 1e-6
    @test sum(abs2, Array(2 .+ ad) - (2 .+ a)) < 1e-6
    @test sum(abs2, Array(ad - 2) - (a .- 2)) < 1e-6
    @test sum(abs2, Array(ad .- 2) - (a .- 2)) < 1e-6
    @test sum(abs2, Array(2 - ad) - (2 .- a)) < 1e-6
    @test sum(abs2, Array(2 .- ad) - (2 .- a)) < 1e-6
    @test sum(abs2, Array(ad * 2) - (a * 2)) < 1e-6
    @test sum(abs2, Array(ad .* 2) - (a .* 2)) < 1e-6
    @test sum(abs2, Array(2 * ad) - (2 * a)) < 1e-6
    @test sum(abs2, Array(2 .* ad) - (2 .* a)) < 1e-6
    @test sum(abs2, Array(ad / 2) - (a / 2)) < 1e-6
    @test sum(abs2, Array(ad ./ 2) - (a ./ 2)) < 1e-6
    @test sum(abs2, Array(2 ./ ad) - (2 ./ a)) < 1e-6
    @test sum(abs2, Array(ad .^ 2.0) - (a .^ 2)) < 1e-6
    @test sum(abs2, Array(2 .^ ad) - (2 .^ a)) < 1e-6

    #Trig functions
    @test sum(abs2, Array(sin.(ad)) - sin.(a)) < 1e-6
    @test sum(abs2, Array(cos.(ad)) - cos.(a)) < 1e-6
    @test sum(abs2, Array(tan.(ad)) - tan.(a)) < 1e-6
    @test sum(abs2, Array(sinh.(ad)) - sinh.(a)) < 1e-6
    @test sum(abs2, Array(cosh.(ad)) - cosh.(a)) < 1e-6
    @test sum(abs2, Array(tanh.(ad)) - tanh.(a)) < 1e-6

    @test sum(fill(AFArray, 1, (1, 2))) == 2
    @test sum(fill(AFArray{Float32}, 1, (1, 2))) == 2
    @test sum(fill(AFArray{Float32,2}, 1, (1, 2))) == 2
    @test sum(ones(AFArray{Float32}, (1, 2))) == 2f0
    @test sum(ones(AFArray{Float32,2}, (1, 2))) == 2f0
    @test eltype(zeros(AFArray{Float32}, (1, 2))) == Float32
    @test eltype(zeros(AFArray{Float32,2}, (1, 2))) == Float32
    @test typeof(zeros(ad)) == AFArray{Float32, 2}
    @test typeof(ones(ad)) == AFArray{Float32, 2}
end

@testset "Arrays" begin
    @test err_to_string(Cuint(0)) == "Success"
    @test_throws ErrorException set_device(-5)
    arr1 = @inferred AFArray{Int,1}([1, 2])
    @test eltype(arr1) == Int
    @test ndims(arr1) == 1
    @test (@inferred size(arr1)) == (2,)
    arr2 = @inferred copy(arr1)
    arr3 = @inferred deepcopy(arr2)
    @test typeof(arr3) == typeof(arr2)
    arr4 = @inferred Array{Int,1}(arr1)
    @test arr4 == [1, 2]
    arr5 = AFArray{Int,2}([1 2 3; 4 5 6])
    @test (@inferred size(arr5)) == (2,3)
    @test size(cat(2, arr1, arr5)) == (2, 4)
    @test size(hcat(arr1, arr5)) == (2, 4)
    @test size(vcat(arr1', arr5')) == (4, 2)
    arr6 = @inferred AFArray([1., 2.])
    arr7 = @inferred Array(arr6)
    @test @inferred eltype(arr6) == eltype(arr7)
    @test @inferred ndims(arr6) == ndims(arr7)
    @test @inferred size(arr6) == size(arr7)
    @test @inferred any(@inferred isnan(arr6)) == false
    @test @inferred any(@inferred isinf(arr6)) == false
    @test @inferred any(@inferred any(isnan(arr6), 1)) == false
    @test @inferred any(@inferred any(isinf(arr6), 1)) == false
    @test @inferred all(@inferred isnan(arr6)) == false
    @test @inferred all(@inferred isinf(arr6)) == false
    @test @inferred all(@inferred all(isnan(arr6), 1)) == false
    @test @inferred all(@inferred all(isinf(arr6), 1)) == false
    @test sum(arr5) == 21
    @test sum(arr6) == 3.
    @test sum(complex(arr6)) == 3 + 0im
    @test sum(complex(arr6, arr6)) == 3 + 3im
    @test sum(abs2(arr5)) == 91
    @test sum(abs2(arr6)) == 5.0
    arr9 = AFArray{Complex{Float64},2}([1.0+3im 2. 3.; 4 5 6])
    @test sum(arr9) == 21+3im
    @test sum(abs2(arr9)) == 100.0
    @test sum(conj(arr9)) == 21-3im
    @test sum(real(arr9)) == 21
    @test sum(imag(arr9)) == 3
    @test typeof(@inferred AFArray{Complex{Float32},2}(arr9)) == AFArray{Complex{Float32},2}
    @test typeof(@inferred AFArray{UInt32}(arr9)) == AFArray{UInt32,2}
    b = @inferred(1 + arr1)
    @test sum(b) == 5
end

a1 = rand(2, 3) + rand(2, 3)im
a2 = rand(2, 3) + rand(2, 3)im
af1 = @inferred AFArray(a1)
af2 = @inferred AFArray(a2)

for op in [:+, :-, :.+, :.-, :.*, :./, :.^]
    @test @eval sum($op(a1, a2)) ≈ sum(@inferred $op(af1, af2))
end

a3 = randn(Float32, 2, 3)
af3 = AFArray(a3)
a4 = randn(Float32, 2, 3)
af4 = AFArray(a4)

for op in [:.>, :.>=, :.!=, :.<, :.<=]
    @test @eval all($op(a3, a4) .== Array(@inferred $op(af3, af4)))
    @test @eval all($op(0, a4) .== Array(@inferred $op(0, af4)))
    @test @eval all($op(a3, 0) .== Array(@inferred $op(af3, 0)))
end

for op in [:.==]
    @test @eval all($op(a3, a4) .== Array($op(af3, af4)))
    @test @eval all($op(0, a4) .== Array($op(0, af4)))
    @test @eval all($op(a3, 0) .== Array($op(af3, 0)))
end

c1 = rand()

for op in [:*, :.+, :.-, :.*, :./, :.^]
    @test @eval sum($op(c1, a2)) ≈ sum(@inferred $op(c1, af2))
end

for op in [:*, :/, :.^, :.+, :.-, :.*, :./]
    @test @eval sum($op(a1, c1)) ≈ sum(@inferred $op(af1, c1))
end

as = Int32[-2 -1 0 1 2]
@test all(AFArray(sign.(as)) .== sign(AFArray(as)))
@test !any(!(AFArray(sign.(as)) .== sign(AFArray(as))))

@test (@inferred size(AFArray(rand(1,2,3)))) == (1,2,3)
@test (@inferred size(AFArray(rand(1,2,3,4)))) == (1,2,3,4)

am = rand(Float32, 3, 3)
amf = AFArray(am)

@test sum(amf * amf) ≈ sum(am * am)
@test sum(amf * amf') ≈ sum(am * am')
#@test sum(A_mul_Bt(amf, amf)) ≈ sum(am * am')
@test sum(amf' * amf) ≈ sum(am' * am)
#@test sum(At_mul_B(amf, amf)) ≈ sum(am' * am)
@test sum(amf' * amf') ≈ sum(am' * am')
#@test sum(At_mul_Bt(amf, amf)) ≈ sum(am' * am')

@test Array(sum(amf, 1)) ≈ sum(am, dims=1)
@test Array(sum(amf, 2)) ≈ sum(am, dims=2)
# @test Array(sum(amf, [1,2])) ≈ sum(am, dims=[1,2])
# @test Array(sum(amf, [2,1])) ≈ sum(am, dims=[2,1])
# @test Array(sum(amf, 1:2)) ≈ sum(am, 1:2)
# @test Array(sum(amf, (1,2))) ≈ sum(am, (1,2))

@test Array(prod(amf, 1)) ≈ prod(am, dims=1)
@test Array(prod(amf, 2)) ≈ prod(am, dims=2)

@test all(Array(minimum(amf, 1)) .== minimum(am, dims=1))
@test all(Array(minimum(amf, 2)) .== minimum(am, dims=2))

@test all(Array(maximum(amf, 1)) .== maximum(am, dims=1))
@test all(Array(maximum(amf, 2)) .== maximum(am, dims=2))
@test maximum(amf) == maximum(am)
@test minimum(amf) == minimum(am)
@test mean(amf) ≈ mean(am)
@test std(amf) ≈ std(am)
@test var(amf) ≈ var(am)
@test median(amf) == median(am)

@test Array(vec(amf)) == vec(am)
@test typeof(norm(amf)) == Float32
@test typeof(norm(amf, 1)) == Float32
@test typeof(norm(amf, Inf)) == Float32
@test @inferred(norm(amf)) ≈ norm(am)
@test @inferred(norm(amf, 1)) ≈ norm(am, 1)
@test @inferred(norm(amf, Inf)) ≈ norm(am, Inf)
u,s,v = @inferred(svd(amf))

signal = rand(Float32, 100)
signalf = AFArray(signal)
filt = rand(Float32, 3)
filtf = AFArray(filt)
@test sum(conv(signal, filt)) ≈ sum(@inferred conv(signalf, filtf))

srand(AFArray, rand(Int))

@testset "Sizes" begin
    @test size(rand(AFArray, 1)) == (1,)
    @test size(rand(AFArray, 1, 2)) == (1,2)
    @test size(rand(AFArray{Float64}, 1)) == (1,)
    @test size(rand(AFArray{Float64}, 1, 2)) == (1,2)
    @test size(rand(AFArray{Float64}, 1)) == (1,)
    @test size(rand(AFArray{Float64}, 1, 2)) == (1,2)

    @test size(randn(AFArray, 1)) == (1,)
    @test size(randn(AFArray, 1, 2)) == (1,2)
    @test size(randn(AFArray{Float64}, 1)) == (1,)
    @test size(randn(AFArray{Float64}, 1, 2)) == (1,2)
    @test size(randn(AFArray{Float64}, 1)) == (1,)
    @test size(randn(AFArray{Float64}, 1, 2)) == (1,2)
end

a3 = rand(UInt8, 2, 3)
af3 = AFArray(a3)
a4 = rand(UInt8, 2, 3)
af4 = AFArray(a4)
c1 = rand(UInt8)

for op in [:&, :|, :xor, :min, :max]
    @test @eval all($op.(a3, a4) .== Array(@inferred $op(af3, af4)))
    @test @eval all($op.(c1, a4) .== Array(@inferred $op(c1, af4)))
    @test @eval all($op.(a3, c1) .== Array(@inferred $op(af3, c1)))
end

a3 = rand(UInt64, 2, 3)
af3 = AFArray(a3)
a4 = rand(UInt64, 2, 3)
af4 = AFArray(a4)
c1 = rand(UInt64)

for op in [:&, :|, :xor, :min, :max]
    @test @eval all($op.(a3, a4) .== Array(@inferred $op(af3, af4)))
    @test @eval all($op.(c1, a4) .== Array(@inferred $op(c1, af4)))
    @test @eval all($op.(a3, c1) .== Array(@inferred $op(af3, c1)))
end

a3 = abs.(rand(Int64, 2, 3))
af3 = AFArray(a3)
a4 = mod.(rand(Int64, 2, 3), 10)
af4 = AFArray(a4)
c1 = mod(rand(Int64), 10)

for op in [:<<, :>>, :&, :|, :xor, :min, :max]
    @test @eval all($op.(a3, a4) .== Array($op.(af3, af4)))
    @test @eval all($op.(a3, a4) .== Array(@inferred $op(af3, af4)))
    @test @eval all($op.(c1, a4) .== Array(@inferred $op(c1, af4)))
    @test @eval all($op.(a3, c1) .== Array(@inferred $op(af3, c1)))
end

a3 = rand(Bool, 2, 3)
af3 = AFArray(a3)
a4 = rand(Bool, 2, 3)
af4 = AFArray(a4)
c1 = rand(Bool)

for op in [:&, :|, :xor, :min, :max]
    @test @eval all($op.(a3, a4) .== Array(@inferred $op(af3, af4)))
    @test @eval all($op.(c1, a4) .== Array(@inferred $op(c1, af4)))
    @test @eval all($op.(a3, c1) .== Array(@inferred $op(af3, c1)))
end

a5 = rand(Float32, 5)
af5 = AFArray(a5)

for op in [:abs, :acos, :asin, :atan, :cbrt, :ceil, :cos, :cosh,
           :exp, :expm1, :floor, :erfc,
           :imag, :isinf, :isnan, :lgamma, :log, :log10,
           :log1p, :log2, :real, :round, :sign, :signbit, :sin, :sinh, :sqrt, :tan,
           :tanh, :trunc]
    @testset "$op" begin
        @test @eval $op.(a5) ≈ Array(@inferred $op(af5))
    end
end

a6 = rand(Float32, 3, 3)
af6 = AFArray(a6)

for op in [:identity, :vec, :transpose]
    @testset "$op" begin
        @test @eval $op(a6) ≈ Array(@inferred $op(af6))
    end
end

for op in [:maximum, :minimum, :sum, :prod, :sizeof, :var, :std, :mean, :median]
    @testset "$op" begin
        @test @eval $op(a6) ≈ @inferred $op(af6)
    end
end

aa = [true false true]
bb = [1 2 3]
cc = [4 5 6; 7 8 9]
a = AFArray(aa)
b = AFArray(bb)
c = AFArray(cc)

@test Array(select(a, b, c)) == ifelse.(aa, bb, cc)
@test Array(select(a, 0, b)) == ifelse.(aa, 0, bb)
@test Array(select(a, b, 0)) == ifelse.(aa, bb, 0)

@test size(reshape(c, 6)) == (6, )
@test size(reshape(c, (3,2))) == (3, 2)

#plot(rand(AFArray, 10), rand(AFArray, 10))

s = rand(AFArray{Float32}, 10)
sh = Array(s)
@testset "Sort" begin
    @test Array(@inferred sort(s)) == sort(sh)
    @test Array(@inferred sortperm(s)) == sortperm(sh)
end

@testset "IR" begin
    x = AFArray([1f0; 2f0; 3f0])
    b = AFArray([1.; 2.])
    @test Array(@inferred fir(b, x)) == [1f0; 4f0; 7f0]
    a = AFArray([1f0])
    b = AFArray{Float32}(b)
    @test Array(@inferred iir(b, a, x)) == [1f0; 4f0; 7f0]
end

ss = similar(s)
@test typeof(ss) == typeof(s)
ss[1] = 10
@test s[1] != 10f0 && ss[1] == 10f0
ssa = similar(s, Float64)
@test typeof(ssa) == AFArray{Float64,1}

aa = [NaN Inf 3]
a = AFArray(aa)

@test Array(@inferred isfinite(a)) == isfinite.(aa)


mutable struct Impact{T} <: Contact{T}
    ainv3::Adjoint{T,SVector{3,T}}
    p::SVector{3,T}
    offset::SVector{3,T}

    function Impact(body::Body{T}, normal::AbstractVector; p::AbstractVector = zeros(3),  offset::AbstractVector = zeros(3)) where T
        # Derived from plane equation a*v1 + b*v2 + distance*v3 = p - offset
        V1, V2, V3 = orthogonalcols(normal) # gives two plane vectors and the original normal axis
        A = [V1 V2 V3]
        Ainv = inv(A)
        ainv3 = Ainv[3,SA[1; 2; 3]]'

        new{T}(ainv3, p, offset), body.id
    end
end

function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, constraint::Impact{T}) where {T}
    summary(io, constraint)
    println(io,"")
    println(io, " ainv3:  "*string(constraint.ainv3))
    println(io, " offset: "*string(constraint.offset))
end

## Additional force for friction default
@inline additionalforce(bound::Impact{T}, x::AbstractVector, q::UnitQuaternion) where T = szeros(T, 6)


@connector function Pin(;name)
    sts = @variables begin
        v(t)                    # Potential at the pin [V]
        i(t), [connect=Flow]    # Current flowing into the pin [A]
    end
    ODESystem(Equation[], t, sts, [], name=name, defaults=Dict(v=>1.0, i=>1.0))
end
Base.@doc "Port for an electrical system." Pin

"""
    OnePort(;name, v_start=0.0, i_start=0.0)

Component with two electrical pins p and n and current i from p to n.

# Parameters:
- `v_start`: [V] Initial voltage across the component
- `i_start`: [A] Initial current through the component
"""
function OnePort(;name, v_start=0.0, i_start=0.0)
    @named p = Pin()
    @named n = Pin()
    sts = @variables begin
        v(t)=v_start
        i(t)=i_start
    end
    eqs = [
        v ~ p.v - n.v
        0 ~ p.i + n.i
        i ~ p.i
    ]
    return compose(ODESystem(eqs, t, sts, []; name=name), p, n)
end

@connector function DigitalPin(; name)
    @variables val(t) v(t) i(t)
    eqs = [
        val ~ IfElse.ifelse((0.0 <= v) & (v <= 0.8) | (2.0 <= v) & (v <= 5.0),
                                IfElse.ifelse(v > 2.0, 1, 0), X)
    ]
    ODESystem(Equation[], t, [val, v, i], [], defaults=Dict(val=>0, i=>0), name=name)
end


# Preprocessor
# ============

# Transformation
# --------------

# Logarithmic transformtion
struct LogT
    pseudocount::Float32
end
function transform!(t::LogT, X::Matrix{Float32})
    @fastmath @inbounds for i in 1:length(X)
        X[i] = log(X[i] + t.pseudocount)
    end
    return X
end

# Freeman-Tukey transformation
struct FTT end
function transform!(::FTT, X::Matrix{Float32})
    @fastmath @inbounds for i in 1:length(X)
        X[i] = sqrt(X[i]) + sqrt(X[i]+1.0f0)
    end
    return X
end

# Dimensionality Reduction
# ------------------------

mutable struct PCA
    dims::Int
    randomize::Bool
    proj::Matrix{Float32}
    PCA(dims::Integer; randomize::Bool=true) = new(dims, randomize)
end

# NOTE: This assumes features are already centered (if needed).
function fit!(pca::PCA, X::AbstractMatrix)
    if pca.randomize
        pca.proj, = rsvd(X, pca.dims)
    else
        pca.proj, = tsvd(X, pca.dims)
    end
    return pca
end
reducedims(pca::PCA, X::AbstractMatrix) = pca.proj'X

# Preprocessor
# ------------

# permutation, log transformation, scaling, projection, etc.
struct Preprocessor
    n::Int
    featurenames::Vector{String}
    sums::Vector{Float32}
    soss::Vector{Float32}
    transformer::Union{LogT,FTT}
    dimreducer::PCA
    normalize::Bool
    standardize::Bool
    scalefactor::Float32
end

function Preprocessor(
        Y::ExpressionMatrix,
        transformer,
        dimreducer,
        normalize::Bool,
        standardize::Bool,
        scalefactor::Float32)
    @assert Y.data isa Matrix{Float32}
    X::Matrix{Float32} = copy(Y.data)
    # normalize cell counts
    if normalize
        X .*= scalefactor ./ sum(X, dims=1)
    end
    # apply variance-stabilizing transformation
    transform!(transformer, X)
    # standardize features
    m, n = size(X)
    sums = vec(sum(X, dims=2))
    soss = vec(sum(X.^2, dims=2))
    μ = sums ./ n
    X .-= μ
    if standardize
        σ = sqrt.((soss .- 2 .* μ .* sums) ./ n .+ μ.^2)
        if any(σ .== 0)
            throw(ArgumentError("found $(sum(σ .== 0)) features with no variance; filter out these features"))
        end
        X .*= inv.(σ)
    end
    # fit dimreducer
    fit!(dimreducer, X)
    return Preprocessor(
        n,
        Y.featurenames,
        sums,
        soss,
        transformer,
        dimreducer,
        normalize,
        standardize,
        scalefactor,
    )
end

function preprocess(proc::Preprocessor, Y::ExpressionMatrix, inferstats::Symbol)
    perm = zeros(Int, length(proc.featurenames))
    for (i, name) in enumerate(proc.featurenames)
        perm[i] = get(Y.featuremap, name, 0)
    end
    X = permuterows(perm, Y.data)::Matrix{Float32}  # should be inferable
    if proc.normalize
        X .*= proc.scalefactor ./ sum(X, dims=1)
    end
    transform!(proc.transformer, X)
    m, n = size(X)
    @assert length(proc.sums) == length(proc.soss) == m
    if inferstats == :query
        μ = vec(mean(X, dims=2))
        σ = vec(std(X, dims=2))
        @inbounds for i in 1:m
            # insert the std of the database if no variability in the query cells
            if σ[i] == 0
                σ[i] = sqrt((proc.soss[i] - 2 * μ[i] * proc.sums[i]) / proc.n + μ[i]^2)
            end
        end
    elseif inferstats == :database
        μ = proc.sums ./ proc.n
        σ = sqrt.((proc.soss .- 2 .* μ .* proc.sums) ./ proc.n .+ μ.^2)
    else
        @assert inferstats == :both
        sums = vec(sum(X, dims=2))
        soss = vec(sum(X.^2, dims=2))
        μ = zeros(Float32, m)
        σ = zeros(Float32, m)
        @inbounds for i in 1:m
            μ[i], σ[i] = mean_and_std(
                proc.sums[i], proc.soss[i], proc.n,
                sums[i], soss[i], n,
            )
        end
    end
    if proc.standardize
        invstd = inv.(σ)
        @inbounds for j in 1:n, i in 1:m
            X[i,j] = (X[i,j] - μ[i]) * invstd[i]
        end
    else
        X .-= μ
    end
    return reducedims(proc.dimreducer, X)
end

# sos: sum of squares
@inline function mean_and_std(
        sum1::Float32, sos1::Float32, n1::Integer,
        sum2::Float32, sos2::Float32, n2::Integer,
    )
    sum = sum1 + sum2
    d = inv(n1 + n2)
    μ = sum * d
    σ = sqrt((sos1 + sos2 - 2 * μ * sum) * d + μ^2)
    return μ, σ
end

indims(p::Preprocessor) = length(p.featurenames)
outdims(p::Preprocessor) = p.dimreducer.dims

function permuterows(perm::Vector{Int}, Y::AbstractMatrix)
    m = length(perm)
    n = size(Y, 2)
    Y_permuted = Matrix{Float32}(undef, m, n)
    @inbounds for j in 1:n, i in 1:m
        Y_permuted[i,j] = perm[i] > 0 ? Y[perm[i],j] : 0
    end
    return Y_permuted
end


isdefined(Base, :__precompile__) && __precompile__()

module NonNegLeastSquares

using Compat
using Distributed
using LinearAlgebra
import SparseArrays

export nonneg_lsq

## Algorithms
include("nnls.jl")
using .NNLS: nnls
include("fnnls.jl")
include("pivot.jl")
include("pivot_comb.jl")
include("pivot_cache.jl")
include("admm.jl")

## Common interface to algorithms
include("interface.jl")

## Helper functions
include("cssls.jl") # combinatorial subspace least squares (CSSLS)

end


using TrafficAssignment
using Test, LinearAlgebra
using DelimitedFiles

include("test_functions.jl")

test_tntp()

summarize_ta_data()

data_time = time()
ta_data = load_ta_network("SiouxFalls")
println("Data Loading Completed, time:", time() - data_time, " seconds")

@testset "Various FW Methods" begin
    ta_data = load_ta_network("SiouxFalls")

    @time link_volume, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:bfw, step=:exact, log=:off, tol=1e-3, max_iter_no=5)
    @test abs( objective - 4.963799502172654e6 ) < 1e6

    @time link_volume, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:cfw, step=:newton, log=:off, tol=1e-3, max_iter_no=5)
    @test abs( objective - 4.963799502172654e6 ) < 1e6

    @time link_volume, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:fw, step=:exact, log=:off, tol=1e-3, max_iter_no=5)
    @test abs( objective - 4.963799502172654e6 ) < 1e6
end

@testset "Testing Sioux Falls" begin
    ta_data = load_ta_network("SiouxFalls")
    @time link_volume, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:cfw, step=:newton, log=:off, tol=1e-5, max_iter_no=50000)
    solution, header = readdlm("SiouxFalls_flow.csv", ',', header=true)
    solution_flow = Float64.(solution[:, 3])
    @test norm(link_volume - solution_flow) / norm(solution_flow) < 0.01
end

@testset "Testing Anaheim" begin
    ta_data = load_ta_network("Anaheim")
    @time link_volume, link_travel_time, objective = ta_frank_wolfe(ta_data, method=:cfw, step=:newton, log=:off, tol=1e-5, max_iter_no=50000)
    solution, header = readdlm("Anaheim_flow.csv", ',', header=true)
    solution_flow = Float64.(solution[:, 3])
    @test norm(link_volume - solution_flow) / norm(solution_flow) < 0.01
end

summarize_ta_data(markdown=true)


export Sketch, linear_update, low_rank_approx, fixed_rank_approx

### Sketching operators and methods

"""Sketch is a sketch for compressing a matrix over the field F
It is parametrized by the field F, which may be either the real or complex numbers,
and by the type of dimension reduction map DR, which may be either GaussDimRedux, SparseDimRedux, or SSRFTDimRedux
The matrices X, Y and Z and the dimension reduction map are defined over the same field F"""
type Sketch{F<:Number,DR<:DimRedux{F}}
    Υ::DR        # from F^m to F^k
    Ω::DR        # from F^n to F^k
    Φ::DR        # from F^m to F^s
    Ψ::DR        # from F^n to F^s
    X::Matrix{F}  # in F^{k×n}
    Y::Matrix{F}  # in F^{m×k}
    Z::Matrix{F}  # in F^{s×s}
end

"""Constructs a sketch for a matrix in F^{m×n}
using the dimension reduction map DR.
By default, the field F is the real (floating point) numbers"""
function Sketch{F<:Number,DR<:DimRedux}(m::Int, n::Int, k::Int, s::Int, ::Type{DR}, ::Type{F}=Float64) # DR<:DimRedux{F}??
    Υ = DR(k,m,F)
    Ω = DR(k,n,F)
    Φ = DR(s,m,F)
    Ψ = DR(s,n,F)
    X = zeros(F,k,n)
    Y = zeros(F,m,k)
    Z = zeros(F,s,s)
    return Sketch(Υ,Ω,Φ,Ψ,X,Y,Z)
end

function linear_update{F<:Number,DR<:DimRedux{F}}(sk::Sketch{F,DR}, H::Matrix{F}, θ::F=one(F), τ::F=one(F))
    # the dots are julia syntax to improve memory efficiency
    # nb removed transposes here; put them in the definition of left/right multiplication
    sk.X .= θ*sk.X .+ τ*(sk.Υ*H)
    sk.Y .= θ*sk.Y .+ τ*(H*sk.Ω)
    sk.Z .= θ*sk.Z .+ τ*(sk.Φ*H*sk.Ψ)
    return sk
end

function low_rank_approx(sk::Sketch)
    P,_ = qr(sk.X')
    Q,_ = qr(sk.Y)
    U1,T1 = qr(sk.Φ*Q)
    U2,T2 = qr(sk.Ψ*P)
    # these work for real gaussian case
    # W = pinv(T1)*(U1'*sk.Z*U2)*pinv(conj(T2'))
    # W = T1\(U1'*sk.Z*U2)/(T2')
    W = T1\(U1'*sk.Z*U2)/(T2')
    return (Q,W,P)
end

function fixed_rank_approx(sk::Sketch, r::Int)
    Q,W,P = low_rank_approx(sk)
    svdW,_ = svds(W,nsv=r)
    U,Σ,V = svdW.U, svdW.S, svdW.Vt'
    return (Q*U,diagm(Σ),P*V)
end


using Symbolics
using ForwardDiff
using LinearAlgebra
using Plots
using ForwardDiff

struct Segment{T}
    domain::Tuple{T,T}
    func
    f
    df
    ddf
    r
    rx
    rp
    δp
    dir_norm
    δ_dir_norm
end

function Segment(domain, func)
    @variables x p[1:2]

    # surface
    f = func(x)
    df = Symbolics.derivative(f, x)
    ddf = Symbolics.derivative(df, x)

    # distance (squared)
    r = x - p[1] + (f - p[2]) * df
    rx = 1.0 + df * df + (f - p[2]) * ddf
    rp = [-1.0 -df]
    δp = [-rp ./ rx; -df * rp ./ rx]

    # fast methods
    f_func = eval(Symbolics.build_function(f, x))
    df_func = eval(Symbolics.build_function(df, x))
    ddf_func = eval(Symbolics.build_function(ddf, x))

    r_func = eval(Symbolics.build_function(r, x, p))
    rx_func = eval(Symbolics.build_function(rx, x, p))
    fp_func = eval(Symbolics.build_function(rp, x, p)[1])
    δp_func = eval(Symbolics.build_function(δp, x, p)[1])

    # surface norm
    f2 = func(p[1])
    dir = [-Symbolics.derivative(f2, p[1]); 1.0]
    dir_norm = dir ./ norm(dir)
    dir_norm_func = eval(Symbolics.build_function(dir_norm, p)[1])
    δ_dir_norm = Symbolics.jacobian(dir_norm, p)
    δ_dir_norm_func = eval(Symbolics.build_function(δ_dir_norm, p)[1])

    Segment(domain, func,
        f_func, df_func, ddf_func,
        r_func, rx_func, fp_func,
        δp_func,
        dir_norm_func, δ_dir_norm_func)
end
# seg1 = Segment((-10.0, 10.0), x -> 0.5 * x^2.0)

seg1 = Segment((-2.0, 0.5), x -> 0.0 * x)
seg2 = Segment((0.5, 1.5), x -> x - 0.5)
# seg3 = Segment((1.5, 2.5), x -> 1.0)
# seg4 = Segment((2.5, 2.51), x -> 100.0 * x - 249.0)
# seg5 = Segment((2.51, 3.51), x -> 2.0)
# seg6 = Segment((3.51, 3.52), x -> -100.0 * x + 353.0)
# seg7 = Segment((3.52, 4.52), x -> 1.0)

surf = [seg1, seg2]#, seg3]#, seg4]#, seg5, seg6, seg7]

get_point(s::Segment, a) = [a; s.f(a)]

p = [0.0; 1.0]

function minimum_distance_point(s::Segment, p)
    # initialization
    if p[1] < s.domain[1]
        a = s.domain[1]
    elseif p[1] > s.domain[2]
        a = s.domain[2]
    else
        a = p[1]
    end

    # Newton
    for i = 1:100
        abs(s.r(a, p)) < 1.0e-8 && break
        # step
        Δ = -s.rx(a, p) \ s.r(a, p)
        α = 1.0
        a_cand = a + α * Δ
        # linesearch
        iter = 0
        while abs(s.r(a_cand, p)) >= abs(s.r(a, p)) && iter < 100
            α *= 0.5
            a_cand = a + α * Δ
            iter += 1
        end
        a = a_cand
    end

    # solution
    if a < s.domain[1]
        a = s.domain[1]
    elseif a > s.domain[2]
        a = s.domain[2]
    end

    # point and sensitivity
    return get_point(s, a), s.δp(a, p)
end

get_surf_dir(s::Segment, a) = [-s.df(a[1]); 1.0]

function surface_normal(s::Segment, p)
    dir = get_surf_dir(s, p)
    dir_norm = dir ./ norm(dir)
    # dir_norm = s.dir_norm(p)

    θ = -0.5 * π
    R = [cos(θ) -sin(θ); sin(θ) cos(θ)]
    tan_norm = R * dir_norm

    return dir_norm, tan_norm
end

function rotation(ns)
	# world-frame normal
	nw = [0.0, 1.0]

	angle = atan(nw[2], nw[1]) - atan(ns[2], ns[1])

	# Rw->s
	[cos(angle) -sin(angle); sin(angle) cos(angle)]
end

function check_domain(surf::Vector{Segment{T}}, p; r=0.0) where T
    status = Bool[]
    for s in surf
        if p[1] >= s.domain[1] - r && p[1] <= s.domain[2] + r
            push!(status, true)
        else
            push!(status, false)
        end
    end
    return status
end

function contact_point(s::Segment, p)
    pmin, δp = minimum_distance_point(s, p)
    d = norm(p - pmin)
    dir_norm, tan_norm = surface_normal(s, pmin)
    rot = rotation(dir_norm)
    δ_dir_norm = s.δ_dir_norm(pmin)
    return pmin, d, dir_norm, tan_norm, rot, δp, δ_dir_norm
end

function contact_point(surf::Vector{Segment{T}}, p; r=0.0) where T
    domain = check_domain(surf, p, r=r)
    p_eval = []
    for (i, status) in enumerate(domain)
        if status
            push!(p_eval, contact_point(surf[i], p))
        end
    end
    min_idx = findmin([pt[2] for pt in p_eval])[2]
    return p_eval[min_idx]
end

function plot_segment(s::Segment; N=100)
    x = range(s.domain[1], stop=s.domain[2], length=N)
    y = s.f.(x)
    plot!(x, y, width=2.0, color=:black, label="")
end

function visualize_segment!(vis, s::Segment;
    idx = 1, color=RGBA(0.0, 0.0, 0.0, 1.0), linewidth=50, num_points=100)

    line_mat = LineBasicMaterial(color=color, linewidth=linewidth)
    domain = range(s.domain[1], stop=s.domain[2], length=num_points)
    points = [Point(x, 0.0, s.func(x)) for x in domain]
    setobject!(vis["traj_$idx"], MeshCat.Line(points, line_mat))
end

function visualize!(vis, surf::Vector{Segment{T}};
    colors=[RGBA(0.0, 0.0, 0.0, 1.0) for i = 1:length(surf)], linewidth=50, num_points=100) where T

    for (i, s) in enumerate(surf)
        visualize_segment!(vis, s, idx = i, color=colors[i], linewidth=linewidth, num_points=num_points)
    end
end

plt = plot(xlims=(-0.5, 4.0), ylims=(-1.0, 3.0), aspect_ratio=:equal)
for s in surf
    plt = plot_segment(s)
end

# p = [2.0; 1.5]
# p = [0.0, 1.0]
# p = [0.0, 1.4]
p = [0.5, 1.0]
p = [1.5, 1.25]

plt = scatter!([p[1]], [p[2]], color=:red, label="")

p_eval = [contact_point(s, p) for s in surf]
for x in p_eval
    plt = plot!([p[1]; x[1][1]], [p[2]; x[1][2]], width=1.5, color=:cyan, label="")
end
display(plt)

for (i, pc) in enumerate(p_eval)
    @show x = pc[1][1]
    @show y = pc[1][2]

    dir_norm = pc[3]
    tan_norm = pc[4]

    scale = 0.25
    plt = plot!([x; x + scale * dir_norm[1]], [y; y + scale * dir_norm[2]], width=1.5, color=:magenta, label="")
    plt = plot!([x; x + scale * tan_norm[1]], [y; y + scale * tan_norm[2]], width=1.5, color=:orange, label="")

end

display(plt)

function generate_polynomial(seg1::Segment{T}, seg2::Segment{T}; ϵ=0.01) where {T}
    # segments
    seg1.domain[2]
    seg2.domain[1]

    # offsets
    x⁺ = seg1.domain[2] + ϵ
    x⁻ = seg1.domain[2] - ϵ

    # end points and their derivatives
    p1x = seg1.func(x⁻)
    dp1x = ForwardDiff.derivative(seg1.func, x⁻)
    p2x = seg2.func(x⁺)
    dp2x = ForwardDiff.derivative(seg2.func, x⁺)

    # fit polynomial parameters
    b = [p1x; dp1x; p2x; dp2x]
    A = [1.0 x⁻ x⁻^2.0 x⁻^3.0;
        0.0 1.0 2.0 * x⁻ 3.0 * x⁻^2.0;
        1.0 x⁺ x⁺^2.0 x⁺^3.0;
        0.0 1.0 2.0 * x⁺ 3.0 * x⁺^2.0]
    param = A \ b

    # polynomial
    function poly(y)
        param[1] + param[2] * y + param[3] * y^2.0 + param[4] * y^3.0
    end
    @assert abs(poly(x⁻) - seg1.func(x⁻)) < 1.0e-8
    @assert abs(poly(x⁺) - seg2.func(x⁺)) < 1.0e-8

    # polynomial segment
    Segment((x⁻, x⁺), poly)
end

# seg_poly = generate_polynomial(seg1, seg2)

function polynomial_blend(seg1::Segment{T}, seg2::Segment{T}; ϵ=0.01) where {T}
    # polynominal segment
    seg_poly = generate_polynomial(seg1, seg2, ϵ=ϵ)

    # update segments
    seg1_update = Segment((seg1.domain[1], seg_poly.domain[1]), seg1.func)
    seg2_update = Segment((seg_poly.domain[2], seg2.domain[2]), seg2.func)

    return seg1_update, seg_poly, seg2_update
end

seg1_update, seg_poly, seg2_update = polynomial_blend(seg1, seg2)

function polynomial_blend(surf::Vector{Segment{T}}; ϵ=0.01) where {T}
    @assert length(surf) > 1
    seg1_update, seg_poly, seg2_update = polynomial_blend(surf[1], surf[2], ϵ=ϵ)
    surf_update = [seg1_update, seg_poly, seg2_update]
    if length(surf) > 2
        for i = 2:length(surf)-1
            seg1_update, seg_poly, seg2_update = polynomial_blend(surf_update[end], surf[i+1], ϵ=ϵ)
            pop!(surf_update)
            push!(surf_update, seg1_update)
            push!(surf_update, seg_poly)
            push!(surf_update, seg2_update)
        end
    end
    return surf_update
end

surf_update = polynomial_blend(surf, ϵ=0.1)

# plt = plot(xlims=(1.45, 1.55), ylims=(0.95, 1.05), aspect_ratio=:equal)
plt = plot(xlims=(-0.5, 4.0), ylims=(-1.0, 2.5), aspect_ratio=:equal)
# plt = plot(xlims=(3.45, 3.55), ylims=(0.95, 1.05), aspect_ratio=:equal)
plt = plot(xlims=(0.0, 1.0), ylims=(-0.1, 0.25), aspect_ratio=:equal)

for s in surf_update
    plt = plot_segment(s)
end
display(plt)

p = [1.4; 1.25]
plt = scatter!([p[1]], [p[2]], color=:red, label="")

p_eval = [contact_point(s, p) for s in surf_update]
for x in p_eval
    plt = plot!([p[1]; x[1][1]], [p[2]; x[1][2]], width=1.5, color=:cyan, label="")
end
display(plt)

for (i, pc) in enumerate(p_eval)
    @show x = pc[1][1]
    @show y = pc[1][2]

    dir_norm = pc[3]
    tan_norm = pc[4]

    scale = 0.25
    plt = plot!([x; x + scale * dir_norm[1]], [y; y + scale * dir_norm[2]], width=2.0, color=:magenta, label="")
    plt = plot!([x; x + scale * tan_norm[1]], [y; y + scale * tan_norm[2]], width=2.0, color=:orange, label="")
end
display(plt)

# # test rotation
# p_eval[3][1]
# rotation(p_eval[6][3])' * [1.0; 0.0] # map vector in segment to world frame


# ---
# title: 1011. Capacity To Ship Packages Within D Days
# id: problem1011
# author: Indigo
# date: 2021-01-30
# difficulty: Medium
# categories: Array, Binary Search
# link: <https://leetcode.com/problems/capacity-to-ship-packages-within-d-days/description/>
# hidden: true
# ---
#
# A conveyor belt has packages that must be shipped from one port to another
# within `D` days.
#
# The `i`-th package on the conveyor belt has a weight of `weights[i]`.  Each
# day, we load the ship with packages on the conveyor belt (in the order given
# by `weights`). We may not load more weight than the maximum weight capacity of
# the ship.
#
# Return the least weight capacity of the ship that will result in all the
# packages on the conveyor belt being shipped within `D` days.
#
#
#
# **Example 1:**
#
#
#
#     Input: weights = [1,2,3,4,5,6,7,8,9,10], D = 5
#     Output: 15
#     Explanation:
#     A ship capacity of 15 is the minimum to ship all the packages in 5 days like this:
#     1st day: 1, 2, 3, 4, 5
#     2nd day: 6, 7
#     3rd day: 8
#     4th day: 9
#     5th day: 10
#
#     Note that the cargo must be shipped in the order given, so using a ship of capacity 14 and splitting the packages into parts like (2, 3, 4, 5), (1, 6, 7), (8), (9), (10) is not allowed.
#
#
# **Example 2:**
#
#
#
#     Input: weights = [3,2,2,4,1,4], D = 3
#     Output: 6
#     Explanation:
#     A ship capacity of 6 is the minimum to ship all the packages in 3 days like this:
#     1st day: 3, 2
#     2nd day: 2, 4
#     3rd day: 1, 4
#
#
# **Example 3:**
#
#
#
#     Input: weights = [1,2,3,1,1], D = 4
#     Output: 3
#     Explanation:
#     1st day: 1
#     2nd day: 2
#     3rd day: 3
#     4th day: 1, 1
#
#
#
#
# **Constraints:**
#
#   * `1 <= D <= weights.length <= 50000`
#   * `1 <= weights[i] <= 500`
#
#
## @lc code=start
using LeetCode

function ship_within_days(weights::Vector{Int}, D::Int)
    lo, hi = maximum(weights), sum(weights)
    while lo < hi
        mid = (lo + hi) ÷ 2
        cnt = 1
        tt = 0
        for w in weights
            tt += w
            if tt > mid
                cnt += 1
                tt = w
            end
        end
        cnt > D ? (lo = mid + 1) : (hi = mid)
    end
    return lo
end
## @lc code=end


#-----------------------------------------------------------------------------------------------------
# Spin Operators
#-----------------------------------------------------------------------------------------------------
# Dictionary
function spin_dict(D::Integer)
    J = (D-1)/2
    coeff = [sqrt(i*(D-i)) for i = 1:D-1]
    sp = sparse(1:D-1, 2:D, coeff, D, D)
    sm = sp'
    sx = (sp+sm)/2
    sy = (sp-sm)/2
    sz = sparse(1:D, 1:D, J:-1:-J)
    s0 = sparse(1.0I, D, D)
    Dict([('+', sp), ('-', sm), ('x', sx), ('y', sy), ('z', sz), ('1', s0)])
end
# Atomic spin matrix
function spin_atom(s::String, dic::Dict)
    n = length(s)
    ny = sum(si == 'y' for si in s)
    temp = n == 1 ? dic[s[1]] : kron([dic[si] for si in s]...)
    P = mod(ny, 2) == 0 ? (-1)^(ny÷2) : (1im)^ny
    P * temp
end
# General spin matrix
function spin(spins; D::Integer=2)
    dic = spin_dict(D)
    res = sum(ci * spin_atom(si, dic) for (ci, si) in spins)
    Array(res)
end

#---------------------------------------------------------------------------------------------------
# Other properties
#---------------------------------------------------------------------------------------------------
export entropy
entropy(S::AbstractVector; cutoff::Real=1e-14) = -sum(i*log(i) for i in S if i > cutoff)

export inner_product
function inner_product(T1::Array, T2::Array)
    α, β = size(T1, 3), size(T2, 3)
    K = Kraus(T1, conj(T2), :r)
    ρ = rand(ComplexF64, α, β) |> vec
    e, v = eigsolve(K, ρ)
    abs(e[1])
end

inner_product(T::Array) = inner_product(T, T)


v1 = Point(0,0,0)
v2 = Point(0,0,1)
v3 = Point(1,0,1)
v4 = Point(0,0,2)

@testset "Point distances" begin
    p1 = Point{3, Float32}(0,0,0)
    p2 = Point{3, Float32}(0,0,1)
    @test distance(p1, p2) == 1
end

@testset "Pair selection" begin
    vertices = [v1, v2, v3, v4]
    edges = Set([(1,2), (1,3), (2,3)])
    mesh = MeshAlgorithms.Mesh(vertices, edges, [])
    pairs = MeshAlgorithms.select_valid_pairs(mesh, 1)
    @test pairs == Set([(1,2), (1,3), (2,3), (2, 4)])
end

@testset "Mesh conversion" begin
    f1 = Face(1,2,3)
    normal = Normal(0,1,0)
    test_hm = HomogenousMesh(
            vertices = [v1, v2, v3],
            faces = [f1],
            normals = [normal, normal, normal])
    expected_mesh = Mesh([v1,v2,v3], Set([(1,2), (1,3), (2,3)]), [])
    @test convert_from_homogenous(test_hm).vertices == expected_mesh.vertices
    @test convert_from_homogenous(test_hm).edges == expected_mesh.edges

end

@testset "Error Quadrics" begin
    p = SVector{4}([1,2,3,0])
    k_p = fundamental_error_quadric(p)
    @test k_p[1,1] == p[1]^2
    @test k_p[1,2] == p[1]*p[2]
    @test k_p[1,3] == p[1]*p[3]
    @test k_p[1,4] == p[1]*p[4]
    @test k_p[2,1] == p[1]*p[2]
    @test k_p[2,2] == p[2]^2
    @test k_p[2,3] == p[2]*p[3]
    @test k_p[2,4] == p[2]*p[4]
    @test k_p[3,1] == p[1]*p[3]
    @test k_p[3,2] == p[2]*p[3]
    @test k_p[3,3] == p[3]^2
    @test k_p[3,4] == p[3]*p[4]
    @test k_p[4,1] == p[1]*p[4]
    @test k_p[4,2] == p[2]*p[4]
    @test k_p[4,3] == p[3]*p[4]
    @test k_p[4,4] == p[4]^2
end

sq_sum(v::Vector) = reduce(+, map(x->x^2, v))
@testset "Normalize" begin
    n = [Vector([1,2,3]), Vector([0,3,4]), Vector([1,0,0]), Vector([0,0,0])]
    @test MeshAlgorithms.normalize(n[1]) == [1/sqrt(14), 2/sqrt(14), 3/sqrt(14)]
    @test MeshAlgorithms.normalize(n[2]) == [0,3/5, 4/5]
    @test MeshAlgorithms.normalize(n[3]) == [1,0,0]
    @test_throws ErrorException MeshAlgorithms.normalize(n[4])
    @testset "All length 1 $a $b $c" for a in -5.0:0.5:5, b in -5:0.5:5, c in 0.1:0.2:5
        @test sq_sum(MeshAlgorithms.normalize(Vector([a,b,c]))) ≈ 1.0
    end
end


using Distributions, ForwardDiff, Plots, LinearAlgebra

include(string(pwd(),"/hmc/rkhs.jl"))

replics = 3000

x = randn(500)

y = 1.5*x + randn(500)

prior = Normal(0.5,1)

likelihood(y,x,par) = sum(logpdf.(Normal(0,1),y - x .*par))

loglike(y,x,par) = likelihood(y,x,par)+logpdf.(prior,par)

unif = Uniform(0,1)

moment_dist = Normal(0,2)

N = 200

dados = zeros(N,2)

for i in 1:N
    xx = rand(prior)
    dados[i,:] = [xx;loglike(y,x,xx)]
end

sig = 1

rkhs_approx = rkhs_gaussian(dados[:,1],dados[:,2],sig,0)

f_hat(x) = rkhs_approx(x)

dif_aprox(x) = ForwardDiff.derivative(f_hat,x)
dif_exact(xx) = ForwardDiff.derivative(z->loglike(y,x,z),xx)

gridd = -2:0.05:2

diff_exact_vals = dif_exact.(gridd)
diff_approx_vals = dif_aprox.(gridd)

plot(gridd,diff_approx_vals)
plot!(gridd,diff_exact_vals)

par = zeros(replics)
par[1] = rand(prior)

moments = zeros(replics,500)

for i in 2:replics
    moment = rand(moment_dist)
    old_moment = moment
    par_aux = par[i-1]

    L = ceil(rand(Uniform(100,500)))
    leapfrog_step = 1/L

    for j in 1:L
        new_moment = moment + 1/2*leapfrog_step*ForwardDiff.derivative(pp->f_hat(pp),par_aux)
         par_aux = par_aux .+ leapfrog_step*1*new_moment
         moment = new_moment + 1/2*leapfrog_step*ForwardDiff.derivative(pp->f_hat(pp),par_aux)
         moments[i,Int64(j)] = new_moment
    end

    num = loglike(y,x,par_aux) + logpdf(moment_dist,moment)#+ logpdf(kernel_novo,par[i-1])

    dem = loglike(y,x,par[i-1]) + logpdf(moment_dist,old_moment)#+ logpdf(kernel,novo_par)
    alpha = exp(min(0,num - dem))
    p = rand(unif)
    if alpha < p
        par[i] = par[i-1]
    else
        par[i] = par_aux
    end
    println(i)
    if i % 400 == 0
        novos_dados = sample(par[1:i],N)
        novas_ll = [loglike(y,x,novos_dados[j]) for j in 1:N]
        rkhs_approx = rkhs_gaussian(novos_dados,novas_ll,sig,0)
        global f_hat(x) = rkhs_approx(x)
        dif_aprox(x) = ForwardDiff.derivative(f_hat,x)
        dif_exact(xx) = ForwardDiff.derivative(z->loglike(y,x,z),xx)

        gridd = -2:0.05:2

        diff_exact_vals = dif_exact.(gridd)
        diff_approx_vals = dif_aprox.(gridd)
        #println("Dif Error", sqrt(sum((diff_exact_vals - diff_approx_vals).^2)))
    end
end

histogram(par[100:3000])

par_rwmh = zeros(replics)
par_rwmh[1] = par[1]

for i in 2:replics
    kernel = Normal(par_rwmh[i-1],0.5)
    novo_par = rand(kernel)
    kernel_novo = Normal(novo_par,0.5)

    num = likelihood(y,x,novo_par) + logpdf(prior,novo_par) + logpdf(kernel_novo,par_rwmh[i-1])

    dem = likelihood(y,x,par_rwmh[i-1]) + logpdf(prior,par_rwmh[i-1]) + logpdf(kernel,novo_par)
    alpha = exp(min(0,num - dem))
    p = rand(unif)
    if alpha < p
        par_rwmh[i] = par_rwmh[i-1]
    else
        par_rwmh[i] = novo_par
    end
    println(i)
end

using StatsPlots

StatsPlots.density(par[100:1000])
StatsPlots.density!(par_rwmh[100:1000])
vline!([1.5])


@testset "error functions" begin
    @testset "real argument" begin
        @test erf(Float16(1))   ≈ 0.84270079294971486934 rtol=2*eps(Float16)
        @test erf(Float32(1))   ≈ 0.84270079294971486934 rtol=2*eps(Float32)
        @test erf(Float64(1))   ≈ 0.84270079294971486934 rtol=2*eps(Float64)

        @test erfc(Float16(1))  ≈ 0.15729920705028513066 rtol=2*eps(Float16)
        @test erfc(Float32(1))  ≈ 0.15729920705028513066 rtol=2*eps(Float32)
        @test erfc(Float64(1))  ≈ 0.15729920705028513066 rtol=2*eps(Float64)

        @test_throws MethodError erfcx(Float16(1))
        @test erfcx(Float32(1)) ≈ 0.42758357615580700442 rtol=2*eps(Float32)
        @test erfcx(Float64(1)) ≈ 0.42758357615580700442 rtol=2*eps(Float64)

        @test_throws MethodError logerfc(Float16(1))
        @test_throws MethodError logerfc(Float16(-1))
        @test logerfc(Float32(-100)) ≈ 0.6931471805599453 rtol=2*eps(Float32)
        @test logerfc(Float64(-100)) ≈ 0.6931471805599453 rtol=2*eps(Float64)
        @test logerfc(Float32(1000)) ≈ -1.0000074801207219e6 rtol=2*eps(Float32)
        @test logerfc(Float64(1000)) ≈ -1.0000074801207219e6 rtol=2*eps(Float64)
        @test logerfc(1000) ≈ -1.0000074801207219e6 rtol=2*eps(Float32)
        @test logerfc(Float32(10000)) ≈ log(erfc(BigFloat(10000, 100))) rtol=2*eps(Float32)
        @test logerfc(Float64(10000)) ≈ log(erfc(BigFloat(10000, 100))) rtol=2*eps(Float64)

        @test_throws MethodError logerfcx(Float16(1))
        @test_throws MethodError logerfcx(Float16(-1))
        @test iszero(logerfcx(0))
        @test logerfcx(Float32(1)) ≈ -0.849605509933248248576017509499 rtol=2eps(Float32)
        @test logerfcx(Float64(1)) ≈ -0.849605509933248248576017509499 rtol=2eps(Float32)
        @test logerfcx(Float32(-1)) ≈ 1.61123231767807049464268192445 rtol=2eps(Float32)
        @test logerfcx(Float64(-1)) ≈ 1.61123231767807049464268192445 rtol=2eps(Float32)
        @test logerfcx(Float32(-100)) ≈ 10000.6931471805599453094172321 rtol=2eps(Float32)
        @test logerfcx(Float64(-100)) ≈ 10000.6931471805599453094172321 rtol=2eps(Float64)
        @test logerfcx(Float32(100)) ≈ -5.17758512266433257046678208395 rtol=2eps(Float32)
        @test logerfcx(Float64(100)) ≈ -5.17758512266433257046678208395 rtol=2eps(Float64)
        @test logerfcx(Float32(-1000)) ≈ 1.00000069314718055994530941723e6 rtol=2eps(Float32)
        @test logerfcx(Float64(-1000)) ≈ 1.00000069314718055994530941723e6 rtol=2eps(Float64)
        @test logerfcx(Float32(1000)) ≈ -7.48012072190621214066734919080 rtol=2eps(Float32)
        @test logerfcx(Float64(1000)) ≈ -7.48012072190621214066734919080 rtol=2eps(Float64)

        @test_throws MethodError erfi(Float16(1))
        @test erfi(Float32(1)) ≈ 1.6504257587975428760 rtol=2*eps(Float32)
        @test erfi(Float64(1)) ≈ 1.6504257587975428760 rtol=2*eps(Float64)

        @test erfinv(Integer(0)) == 0
        @test_throws MethodError erfinv(Float16(1))
        @test erfinv(Float32(0.84270079294971486934)) ≈ 1 rtol=2*eps(Float32)
        @test erfinv(Float64(0.84270079294971486934)) ≈ 1 rtol=2*eps(Float64)

        @test erfcinv(Integer(1)) == 0
        @test_throws MethodError erfcinv(Float16(1))
        @test erfcinv(Float32(0.15729920705028513066)) ≈ 1 rtol=2*eps(Float32)
        @test erfcinv(Float64(0.15729920705028513066)) ≈ 1 rtol=2*eps(Float64)

        @test_throws MethodError dawson(Float16(1))
        @test dawson(Float32(1)) ≈ 0.53807950691276841914 rtol=2*eps(Float32)
        @test dawson(Float64(1)) ≈ 0.53807950691276841914 rtol=2*eps(Float64)
    end

    @testset "complex arguments" begin
        @test erf(ComplexF16(1+2im)) ≈ -0.53664356577856503399-5.0491437034470346695im
        @test erf(ComplexF32(1+2im)) ≈ -0.53664356577856503399-5.0491437034470346695im
        @test erf(ComplexF64(1+2im)) ≈ -0.53664356577856503399-5.0491437034470346695im

        @test erfc(ComplexF16(1+2im)) ≈ 1.5366435657785650340+5.0491437034470346695im
        @test erfc(ComplexF32(1+2im)) ≈ 1.5366435657785650340+5.0491437034470346695im
        @test erfc(ComplexF64(1+2im)) ≈ 1.5366435657785650340+5.0491437034470346695im

        @test_throws MethodError erfcx(ComplexF16(1))
        @test erfcx(ComplexF32(1+2im)) ≈ 0.14023958136627794370-0.22221344017989910261im
        @test erfcx(ComplexF64(1+2im)) ≈ 0.14023958136627794370-0.22221344017989910261im

        @test_throws MethodError erfi(ComplexF16(1))
        @test erfi(ComplexF32(1+2im)) ≈ -0.011259006028815025076+1.0036063427256517509im
        @test erfi(ComplexF64(1+2im)) ≈ -0.011259006028815025076+1.0036063427256517509im

        @test_throws MethodError erfinv(Complex(1))

        @test_throws MethodError erfcinv(Complex(1))

        @test_throws MethodError dawson(ComplexF16(1))
        @test dawson(ComplexF32(1+2im)) ≈ -13.388927316482919244-11.828715103889593303im
        @test dawson(ComplexF64(1+2im)) ≈ -13.388927316482919244-11.828715103889593303im
    end

    @testset "BigFloat arguments" begin
        @test erf(BigFloat(1))  ≈ 0.84270079294971486934 rtol=2*eps()

        @test erfc(BigFloat(1)) ≈ 0.15729920705028513066 rtol=2*eps()

        @test erfcx(BigFloat(1))        ≈ 0.42758357615580700442    rtol=2*eps()
        @test erfcx(BigFloat(2e9))      ≈ 2.820947917738782e-10     rtol=2*eps()
        @test erfcx(BigFloat(1.8))      ≈ erfcx(1.8)                rtol=4*eps()
        @test erfcx(BigFloat(1.8e8))    ≈ erfcx(1.8e8)              rtol=4*eps()
        @test erfcx(BigFloat(1.8e88))   ≈ erfcx(1.8e88)             rtol=4*eps()
        @test isnan(erfcx(BigFloat(NaN)))

        @test logerfc(BigFloat(1000, 100)) ≈ -1.0000074801207219e6 rtol=2*eps(Float64)
        @test isnan(logerfc(BigFloat(NaN)))

        @test_throws MethodError erfi(big(1.0))

        @test_throws MethodError erfinv(BigFloat(1))

        @test_throws MethodError erfcinv(BigFloat(1))

        @test_throws MethodError dawson(BigFloat(1))
    end

    @testset "inverse" begin
        for elty in [Float32,Float64]
            for x in exp10.(range(-200, stop=-0.01, length=50))
                @test isapprox(erf(erfinv(x)), x, atol=1e-12*x)
                @test isapprox(erf(erfinv(-x)), -x, atol=1e-12*x)
                @test isapprox(erfc(erfcinv(2*x)), 2*x, atol=1e-12*x)
                if x > 1e-20
                    xf = Float32(x)
                    @test isapprox(erf(erfinv(xf)), xf, atol=1e-5*xf)
                    @test isapprox(erf(erfinv(-xf)), -xf, atol=1e-5*xf)
                    @test isapprox(erfc(erfcinv(2xf)), 2xf, atol=1e-5*xf)
                end
            end
            @test erfinv(one(elty))  == Inf
            @test erfinv(-one(elty)) == -Inf
            @test_throws DomainError erfinv(convert(elty,2.0))

            @test erfcinv(zero(elty)) == Inf
            @test_throws DomainError erfcinv(-one(elty))
        end

        @test erfinv(one(Int))  == erfinv(1.0)
        @test erfcinv(one(Int)) == erfcinv(1.0)
    end
end