"""
This type convert a bilevel problem as an MPCCModel.

Bilevel Program:
````
min_{x,y} f1(x,y)
s.t.  lcon    <= c(x)   <= ucon
      lvar    <= x      <= uvar
      y in argmin_{y} f2(x,y)
                 s.t. ll_lcon <= g(x,y) <= ll_ucon (l1,  l2)
                      ll_lvar <= y      <= ll_uvar (mu1, mu2)
````
MPCC reformulation of the Bilevel Program:
       min_{x,y,l1,l2,mu1,mu2} f1(x,y)
       s.t.  lvar <= x      <= uvar
             lcon <= c(x)   <= ucon
             0    <= ∇_y f2(x,y) - ∇_y g(x,y)' * (l1 - l2) - (mu1 - mu2) <= 0
             G(x):=(  g(x,y)  - ll_lcon)
                      ll_ucon - g(x,y) )
                      ll_uvar - y
                      y - ll_lvar )
             H(x):=( l1, l2, mu1, mu2 )
             0 <= G(x) _|_ H(x) >= 0
````
Gradient, Jacobian and Hessian are computed using automatic differentiation

set no_third_derivative as true, to ignore the 3rd derivative of g.
"""
mutable struct BPMPCCModel <: AbstractMPCCModel
  meta     :: MPCCModelMeta
  counters :: MPCCCounters

  # Functions
  f1 :: Function
  f2 :: Function
  c  :: Function
  g  :: Function

  ll_lcon :: AbstractVector
  ll_ucon :: AbstractVector
  ll_lvar :: AbstractVector
  ll_uvar :: AbstractVector

  no_third_derivative :: Bool

  m      :: Int #size of the lower level dimension
  n      :: Int #size of the upper level dimension + m
  llncon :: Int #number of constraints in the lower level problem
end

#m0 is the initial guess for the lower leverl problem
function BPMPCCModel(f1, x0::AbstractVector, m0::AbstractVector;
                     y0::AbstractVector = eltype(x0)[],
                     lvar::AbstractVector = eltype(x0)[],
                     uvar::AbstractVector = eltype(x0)[],
                     lcon::AbstractVector = eltype(x0)[],
                     ucon::AbstractVector = eltype(x0)[],
                     c = (args...)->throw(NotImplementedError("c")),
                     g = (args...)->throw(NotImplementedError("g")),
                     f2 = (args...)->throw(NotImplementedError("f2")),
                     ll_lcon::AbstractVector = eltype(x0)[],
                     ll_ucon::AbstractVector = eltype(x0)[],
                     ll_lvar::AbstractVector = eltype(x0)[],
                     ll_uvar::AbstractVector = eltype(x0)[],
                     name::String = "Generic",
                     lin::AbstractVector{<: Integer}=Int[],
                     no_third_derivative::Bool=false)

  m, n, llncon = length(ll_lvar), length(x0), length(ll_ucon)

  nvar = n + 3 * m + 2 * llncon
  length(lvar) == 0 && (lvar = -Inf*ones(n))
  length(uvar) == 0 && (uvar =  Inf*ones(n))
  lvar = vcat(lvar, -Inf*ones(3 * m + 2 * llncon))
  uvar = vcat(uvar,  Inf*ones(3 * m + 2 * llncon))

  ncon = maximum([length(lcon); length(ucon); length(y0)]) + length(m0)
  lcon = vcat(lcon, zeros(m))
  ucon = vcat(ucon, zeros(m))

  ncc  = 2 * m + 2 * llncon
  lccG = zeros(ncc)
  lccH = zeros(ncc)
  ini0 = vcat(x0, m0, zeros(2 * m + 2 * llncon))

  meta = MPCCModelMeta(nvar, x0 = ini0, lvar = lvar, uvar = uvar,
                                        ncon = ncon, lcon = lcon,
                                        ucon = ucon, ncc = ncc,
                                        lccG = lccG, lccH = lccH)

  return BPMPCCModel(meta, MPCCCounters(), f1, f2, c, g,
                     ll_lcon, ll_ucon, ll_lvar, ll_uvar,
                     no_third_derivative, m, n, llncon)
end

"""
_var_bp(:: BPMPCCModel, :: AbstractVector):
return the details of the different variables in the MPCC corresponding
to the bilevel program.
return: x, y, l1, l2, m1, m2
"""
function _var_bp(nlp :: BPMPCCModel, x :: AbstractVector)

 if length(x) != nlp.meta.nvar throw("Error: wrong vector size") end

 m, n, ncon = nlp.m, nlp.n, nlp.llncon

 x0, y0 = x[1:n], x[n+1:n+m]
 l1, l2 = x[n+m+1:n+m+ncon], x[n+m+ncon+1:n+m+2*ncon]
 m1, m2 = x[n+m+2*ncon+1:n+2*m+2*ncon], x[n+2*m+2*ncon+1:n+3*m+2*ncon]

 return x0, y0, l1, l2, m1, m2
end

"""
JLag: gives the gradient w.r.t. y of the Lagrangian of the lower level problem.

JLag(:: BPMPCCModel, :: Function, :: Function, :: AbstractVector)
"∇_y f2(x,y) - ∇_y g(x,y)' * v - mu"
"""
function JLag(nlp :: BPMPCCModel,
              f   :: Function,
              c   :: Function,
              x   :: AbstractVector)

    Jtv = Array{Union{Nothing, Number}}(nothing, nlp.m)

    x0, y0, l1, l2, m1, m2 = _var_bp(nlp, x)

    Jtv = ForwardDiff.gradient(f, x)[nlp.n+1:nlp.n+nlp.m] - (m1 - m2)
    increment!(nlp, :neval_grad)

    if nlp.llncon > 0
        Jtv += - ForwardDiff.jacobian(c, x)[:,nlp.n+1:nlp.n+nlp.m]' * (l1-l2)
        increment!(nlp, :neval_jac)
    end

    return Jtv
end

"""
HLag: gives the hessian w.r.t. y of the Lagrangian of the lower level problem.

HLag(:: BPMPCCModel, :: Function, :: Function, :: AbstractVector)
"∇^2_y f2(x,y) - sum_i ∇^2_y g_i(x,y)' * v_i"
"""
function HLag(nlp :: BPMPCCModel,
              f   :: Function,
              c   :: Function,
              x   :: AbstractVector)
    Jtv = zeros(nlp.n+nlp.m, nlp.n+nlp.m)
    x0, y0, l1, l2, m1, m2 = _var_bp(nlp, x)
    lag(x) = nlp.llncon > 0 ? f(x) - dot(c(x), l1-l2) : f(x)
    Jtv = ForwardDiff.hessian(lag, x)
    Jtvy = Jtv[nlp.n+1:nlp.n+nlp.m,nlp.n+1:nlp.n+nlp.m]
    Jtvx = Jtv[nlp.n+1:nlp.n+nlp.m,1:nlp.n]
    return Jtvy, Jtvx
end

function obj(nlp :: BPMPCCModel, x :: AbstractVector)
  increment!(nlp, :neval_obj)
  return nlp.f1(x)
end

function grad!(nlp :: BPMPCCModel, x :: AbstractVector, g :: AbstractVector)
  increment!(nlp, :neval_grad)
  ForwardDiff.gradient!(view(g, 1:length(x)), nlp.f1, x)
  return g
end

function cons_nl!(nlp :: BPMPCCModel, x :: AbstractVector, c :: AbstractVector)
  increment!(nlp, :neval_cons)
  Jl = JLag(nlp, nlp.f1, nlp.g, x)
  c[1:nlp.meta.ncon] .= vcat(nlp.c(x), Jl)
  return c
end

function consG!(nlp :: BPMPCCModel, x :: AbstractVector, c :: AbstractVector)
  increment!(nlp, :neval_consG)
  gxy = nlp.g(x)
  y   = x[nlp.n+1:nlp.n+nlp.m]
  c[1:nlp.meta.ncc] .= vcat( gxy - nlp.ll_lcon,
                            -gxy + nlp.ll_ucon,
                             y - nlp.ll_lvar,
                            -y + nlp.ll_uvar)
  return c
end

function consH!(nlp :: BPMPCCModel, x :: AbstractVector, c :: AbstractVector)
  increment!(nlp, :neval_consH)
  c[1:nlp.meta.ncc] .= x[nlp.n+nlp.m+1:nlp.meta.nvar]
  return c
end

function jac_nl_structure!(nlp  :: BPMPCCModel,
                           rows :: AbstractVector{<: Integer},
                           cols :: AbstractVector{<: Integer})
  m, n = nlp.meta.ncon, nlp.meta.nvar
  I = ((i,j) for i = 1:m, j = 1:n)
  rows .= getindex.(I, 1)[:]
  cols .= getindex.(I, 2)[:]
  return rows, cols
end

function jac_nl_coord!(nlp  :: BPMPCCModel,
                       x    :: AbstractVector,
                       vals :: AbstractVector)
  increment!(nlp, :neval_jac)
  m, n, llncon = nlp.m, nlp.n, nlp.llncon

  Jx  = ForwardDiff.jacobian(nlp.c, x)
  Jg  = ForwardDiff.jacobian(nlp.g, x)
  Jyy, Jyx = HLag(nlp, nlp.f2, nlp.g, x)
  Jyl = llncon >0 ? vcat(Jg[:,n+1:n+m], -Jg[:,n+1:n+m]) : zeros(0, m)
  Jym = hcat(diagm(0 => ones(m)),-diagm(0 => ones(m)))
  Jl  = hcat(Jyx, Jyy, Jyl', Jym)
  J   = vcat(Jx, Jl)

  vals .= J[:]
  return vals
end

function jnlprod!(nlp :: BPMPCCModel,
                  x   :: AbstractVector,
                  v   :: AbstractVector,
                  Jv  :: AbstractVector)
  increment!(nlp, :neval_jprod)
  #Jv[1:nlp.meta.ncon] = ForwardDiff.derivative(t -> nlp.c(x + t * v), 0)
  Jv .= jac_nl(nlp, x) * v
  return Jv
end

function jnltprod!(nlp :: BPMPCCModel,
                   x   :: AbstractVector,
                   v   :: AbstractVector,
                   Jtv :: AbstractVector)
  increment!(nlp, :neval_jtprod)
  #Jtv[1:nlp.meta.nvar] = ForwardDiff.gradient(x -> dot(nlp.c(x), v), x)
  Jtv .= jac_nl(nlp, x)' * v
  return Jtv
end

function jacG_structure!(nlp  :: BPMPCCModel,
                         rows :: AbstractVector{<: Integer},
                         cols :: AbstractVector{<: Integer})
  m, n = nlp.meta.ncc, nlp.meta.nvar
  I = ((i,j) for i = 1:m, j = 1:n)
   rows[1 : n*m] .= getindex.(I, 1)[:]
   cols[1 : n*m] .= getindex.(I, 2)[:]
  return rows, cols
end

function jacG_coord!(nlp  :: BPMPCCModel,
                     x    :: AbstractVector,
                     vals :: AbstractVector)
  increment!(nlp, :neval_jacG)
  ncc, nvar = nlp.meta.ncc, nlp.meta.nvar

  Jg  = ForwardDiff.jacobian(nlp.g, x)
  Idy = hcat(zeros(nlp.m,nlp.n),diagm(0 => ones(nlp.m)),zeros(nlp.m,ncc))
  Jx  = vcat(Jg, -Jg, Idy, - Idy)

  vals[1 : nvar * ncc] .= Jx[:]
  return vals
end

function jGprod!(nlp :: BPMPCCModel,
                 x   :: AbstractVector,
                 v   :: AbstractVector,
                 Jv  :: AbstractVector)
  increment!(nlp, :neval_jGprod)
  #Jv[1:nlp.meta.ncc] = ForwardDiff.derivative(t -> nlp.G(x + t * v), 0)
  Jv[1:nlp.meta.ncc] .= jacG(nlp, x) * v
  return Jv
end

function jGtprod!(nlp :: BPMPCCModel,
                  x   :: AbstractVector,
                  v   :: AbstractVector,
                  Jtv :: AbstractVector)
  increment!(nlp, :neval_jGtprod)
  #Jtv[1:nlp.meta.nvar] = ForwardDiff.gradient(x -> dot(nlp.G(x), v), x)
  Jtv[1:nlp.meta.nvar] .= jacG(nlp, x)' * v
  return Jtv
end

function jacH_structure!(nlp  :: BPMPCCModel,
                         rows :: AbstractVector{<: Integer},
                         cols :: AbstractVector{<: Integer})
  m, n = nlp.meta.ncc, nlp.meta.nvar
  I = ((i,j) for i = 1:m, j = 1:n)
  rows[1 : n*m] .= getindex.(I, 1)[:]
  cols[1 : n*m] .= getindex.(I, 2)[:]
  return rows, cols
end

function jacH_coord!(nlp  :: BPMPCCModel,
                     x    :: AbstractVector,
                     vals :: AbstractVector)
  m, n = nlp.meta.ncc, nlp.meta.nvar
  increment!(nlp, :neval_jacH)
  Jx = hcat(zeros(m, nlp.n + nlp.m), diagm(0 => ones(m)))

  vals[1 : n*m] .= Jx[:]
  return vals
end

function jHprod!(nlp :: BPMPCCModel,
                 x   :: AbstractVector,
                 v   :: AbstractVector,
                 Jv  :: AbstractVector)
  increment!(nlp, :neval_jHprod)
  #Jv[1:nlp.meta.ncc] = ForwardDiff.derivative(t -> nlp.G(x + t * v), 0)
  Jv[1:nlp.meta.ncc] .= jacH(nlp, x) * v
  return Jv
end

function jHtprod!(nlp :: BPMPCCModel,
                  x   :: AbstractVector,
                  v   :: AbstractVector,
                  Jtv :: AbstractVector)
  increment!(nlp, :neval_jHtprod)
  #Jtv[1:nlp.meta.nvar] = ForwardDiff.gradient(x -> dot(nlp.G(x), v), x)
  Jtv[1:nlp.meta.nvar] .= jacH(nlp, x)' * v
  return Jtv
end

function hess_structure!(nlp :: BPMPCCModel,
                        rows :: AbstractVector{<: Integer},
                        cols :: AbstractVector{<: Integer})
  n = nlp.meta.nvar
  I = ((i,j) for i = 1:n, j = 1:n if i ≥ j)
  rows[1 : nlp.meta.nnzh] .= getindex.(I, 1)
  cols[1 : nlp.meta.nnzh] .= getindex.(I, 2)
  return rows, cols
end

function hess_coord!(nlp        :: BPMPCCModel,
                     x          :: AbstractVector,
                     vals       :: AbstractVector;
                     obj_weight :: Real = one(eltype(x)))
  increment!(nlp, :neval_hess)
  ℓ(x) = obj_weight * nlp.f1(x)
  Hx = ForwardDiff.hessian(ℓ, x)
  k = 1
  for j = 1 : nlp.meta.nvar
    for i = j : nlp.meta.nvar
      vals[k] = Hx[i, j]
      k += 1
    end
  end
  return vals
end

function hess_coord!(nlp        :: BPMPCCModel,
                     x          :: AbstractVector,
                     y          :: AbstractVector,
                     vals       :: AbstractVector;
                     obj_weight :: Real = one(eltype(x)))
  increment!(nlp, :neval_hess)
  ncon, ncc, m, nt = nlp.meta.ncon, nlp.meta.ncc, nlp.m, nlp.no_third_derivative
  llncon = nlp.llncon

  temp(x) = nt ? obj_weight * nlp.f1(x) + dot(nlp.c(x),y[1:ncon - m]) : obj_weight * nlp.f1(x) + dot(cons_nl(nlp,x), y[1:ncon])
  ℓ(x) = llncon > 0 ? temp(x) + dot(nlp.g(x), y[ncon+1:ncon+llncon]) - dot(nlp.g(x), y[ncon+llncon+1:ncon+2*llncon]) : temp(x)
  Hx = ForwardDiff.hessian(ℓ, x)
  k = 1
  for j = 1 : nlp.meta.nvar
    for i = j : nlp.meta.nvar
      vals[k] = Hx[i, j]
      k += 1
    end
  end
  return vals
end

function hessG(nlp :: BPMPCCModel, x :: AbstractVector, y :: AbstractVector)
  increment!(nlp, :neval_hessG)
  lln = nlp.llncon
  ℓ(x) = lln > 0 ? dot(nlp.g(x), y[1:lln]) - dot(nlp.g(x), y[lln+1:2*lln]) : 0
  Hx = ForwardDiff.hessian(ℓ, x)
  return tril(Hx)
end

function hessG_structure!(nlp :: BPMPCCModel,
                         rows :: AbstractVector{<: Integer},
                         cols :: AbstractVector{<: Integer})
  n = nlp.meta.nvar
  I = ((i,j) for i = 1:n, j = 1:n if i ≥ j)
  rows[1 : nlp.meta.nnzh] .= getindex.(I, 1)
  cols[1 : nlp.meta.nnzh] .= getindex.(I, 2)
  return rows, cols
end

function hessG_coord!(nlp  :: BPMPCCModel,
                      x    :: AbstractVector,
                      y    :: AbstractVector,
                      vals :: AbstractVector)
  increment!(nlp, :neval_hessG)
  lln = nlp.llncon
  ℓ(x) = lln > 0 ? dot(nlp.g(x), y[1:lln]) - dot(nlp.g(x), y[lln+1:2*lln]) : 0
  Hx = ForwardDiff.hessian(ℓ, x)
  k = 1
  for j = 1 : nlp.meta.nvar
    for i = j : nlp.meta.nvar
      vals[k] = Hx[i, j]
      k += 1
    end
  end
  return vals
end

function hessH(nlp :: BPMPCCModel, x :: AbstractVector, y :: AbstractVector)
  increment!(nlp, :neval_hessH)
  Hx = zeros(nlp.meta.nvar, nlp.meta.nvar)
  return tril(Hx)
end

function hessH_structure!(nlp  :: BPMPCCModel,
                          rows :: AbstractVector{<: Integer},
                          cols :: AbstractVector{<: Integer})
  n = nlp.meta.nvar
  I = ((i,j) for i = 1:n, j = 1:n if i ≥ j)
  rows .= getindex.(I, 1)
  cols .= getindex.(I, 2)
  return rows, cols
end

function hessH_coord!(nlp  :: BPMPCCModel,
                      x    :: AbstractVector,
                      y    :: AbstractVector,
                      vals :: AbstractVector)
  increment!(nlp, :neval_hess)
  ℓ(x) = dot(nlp.H(x), y)
  Hx = zeros(nlp.meta.nvar, nlp.meta.nvar)
  k = 1
  for j = 1 : nlp.meta.nvar
    for i = j : nlp.meta.nvar
      vals[k] = Hx[i, j]
      k += 1
    end
  end
  return vals
end

function hprod!(nlp        :: BPMPCCModel,
                x          :: AbstractVector,
                v          :: AbstractVector,
                Hv         :: AbstractVector;
                obj_weight :: Real = one(eltype(x)))
  increment!(nlp, :neval_hprod)
  ℓ(x) = obj_weight * nlp.f1(x)
  Hv .= ForwardDiff.derivative(t -> ForwardDiff.gradient(ℓ, x + t * v), 0)
  return Hv
end

function hprod!(nlp        :: BPMPCCModel,
                x          :: AbstractVector,
                y          :: AbstractVector,
                v          :: AbstractVector,
                Hv         :: AbstractVector;
                obj_weight :: Real = one(eltype(x)))
  increment!(nlp, :neval_hprod)
  ncon, ncc, m, nt = nlp.meta.ncon, nlp.meta.ncc, nlp.m, nlp.no_third_derivative
  llncon = nlp.llncon

  temp(x) = nt ? obj_weight * nlp.f1(x) + dot(nlp.c(x),y[1:ncon - m]) : obj_weight * nlp.f1(x) + dot(cons_nl(nlp,x), y[1:ncon])
  ℓ(x) = llncon > 0 ? temp(x) + dot(nlp.g(x), y[ncon+1:ncon+llncon]) - dot(nlp.g(x), y[ncon+llncon+1:ncon+2*llncon]) : temp(x)
  Hv .= ForwardDiff.derivative(t -> ForwardDiff.gradient(ℓ, x + t * v), 0)
  return Hv
end

function hGprod!(nlp :: BPMPCCModel,
                 x   :: AbstractVector,
                 y   :: AbstractVector,
                 v   :: AbstractVector,
                 Hv  :: AbstractVector)
  increment!(nlp, :neval_hGprod)
  lln = nlp.llncon
  ℓ(x) = lln > 0 ? dot(nlp.g(x), y[1:lln]) - dot(nlp.g(x), y[lln+1:2*lln]) : 0
  Hv .= ForwardDiff.derivative(t -> ForwardDiff.gradient(ℓ, x + t * v), 0)
  return Hv
end

function hHprod!(nlp :: BPMPCCModel,
                 x   :: AbstractVector,
                 y   :: AbstractVector,
                 v   :: AbstractVector,
                 Hv  :: AbstractVector)
  increment!(nlp, :neval_hHprod)
  Hv .= zeros(nlp.meta.nvar)
  return Hv
end


module MyQuantum
using LinearAlgebra

include("base.jl")
include("test.jl")
# greet() = print("Hello World!")

export is_unitary, is_hermitian, is_normed, conj_transpose, outer
export Id, X, Y, Z, H, S, T, Rn, Rx, Ry, Rz, Cgate, bell_states
export ket_0, ket_1, ket_plus, ket_minus, commutator

export test_Cgates

end # module


""" Geometrical shapes. """

using LinearAlgebra, CoordinateTransformations

export Shape, Cylinder, Sphere, Cone, Plane, Empty, shapediff, inside,
    TransformedShape, mathematica

abstract type Shape; end

""" A Sphere centered at (`xc`, `yc`, `zc`) with radius `R`. """
struct Sphere <: Shape
    xc::Float64
    yc::Float64
    zc::Float64
    R::Float64
end

"""
A Cylinder between `bottom` < z < `top`, centered at (`xc`, `yc`) and with
radius `R`.
"""
struct Cylinder <: Shape
    bottom::Float64
    top::Float64
    xc::Float64
    yc::Float64
    R::Float64
end

"""
A Cone with a vertical axis and `bottom` < z < `top`, with vertex
at (`xv`, `yv`, `zv`) and slope m with
sqrt((x - xv)^2 + (y - yv)^2) = m (z - zv).
"""
struct Cone <: Shape
    bottom::Float64
    top::Float64

    xv::Float64
    yv::Float64
    zv::Float64
    m::Float64
end

"""
A plane at `z` that defines a shape containing points above
(if `up` is true) or below otherwise.
"""
struct Plane <: Shape
    z::Float64
    up::Bool
end

""" Union of several shapes. """
struct ShapeUnion{T<:Tuple} <: Shape
    shapes::T
end

Base.union(s::Shape...) = ShapeUnion(s)

Empty() = ShapeUnion(())

""" Intersection of several shapes. """
struct ShapeIntersect{T<:Tuple} <: Shape
    shapes::T
end

Base.intersect(s::Shape...) = ShapeIntersect(s)

""" Shape1 - Shape2. """
struct ShapeSubstract{S1,S2} <: Shape
    shape1::S1
    shape2::S2
end

shapediff(s1, s2) = ShapeSubstract(s1, s2)

# Shortcuts for unions and substractiong
Base.:-(s1::Shape, s2::Shape) = shapediff(s1, s2)
Base.:+(s1::Shape, rest::Shape...) = union(s1, rest...)

""" Transformed shape after a transformation or shift. """
struct TransformedShape{T<:AbstractAffineMap, S<:Shape} <: Shape
    trans::T
    inv::T
    shape::S
    function TransformedShape(trans::T, shape::S) where
        {T<:AbstractAffineMap, S<:Shape}

        inv_ = inv(trans)
        new{T, S}(trans, inv_, shape)
    end
end

# For some reason, dot for SVectors of size 3 is much slower than this
# explicit computation.
@inline @fastmath fdot(a::Point, b::Point) = a[1] * b[1] + a[2] * b[2] + a[3] * b[3]

# When we look for interface crossings sometimes we want a list of them,
# sometimes just the closest one.  To reuse the same function without
# creating too many closures we use accumulators for each possible use case
abstract type Accumulator; end

struct ListAccumulator{T <: AbstractVector} <: Accumulator
    list::T
end

@inline accum!(a::ListAccumulator, x) = push!(a.list, x)

mutable struct MinAccumulator{T <: Real} <: Accumulator
    value::T
end

@inline accum!(a::MinAccumulator, x) = (x < a.value) && (a.value = x)

# Just for debugging
struct ShowAccumulator <: Accumulator
end

accum!(a::ShowAccumulator, x) = @show x

"""
Compute the intersection points between a shape `c` and the line joining
points `a` and `b`.  Each intersection point s is used in the accumulator vwhere the intersection point is
p = (1-s)a + s b
"""
@fastmath function interfaces!(v::Accumulator, s::Sphere, a::Point, b::Point)
    (xc, yc, zc) = (s.xc, s.yc, s.zc)
    (xa, ya, za) = a .- @SVector [s.xc, s.yc, s.zc]
    (xb, yb, zb) = b .- @SVector [s.xc, s.yc, s.zc]

    a2 = xa^2 + ya^2 + za^2
    b2 = xb^2 + yb^2 + zb^2
    ab = xa * xb + ya * yb + za * zb

    # These are traditional a, b, c in the quadratic eq. but a, b are already
    # in use.
    α = a2 + b2 - 2 * ab
    β = -2*a2 + 2 * ab
    γ = a2 - s.R^2

    Δ = β^2 - 4 * α * γ

    (Δ >= 0) || return nothing

    s1 = (-β + sqrt(Δ)) / 2α
    s2 = (-β - sqrt(Δ)) / 2α

    (0 < s1 < 1) && accum!(v, s1)
    (0 < s2 < 1) && accum!(v, s2)

    nothing
end

@fastmath function interfaces!(v::Accumulator, c::Cylinder, a::Point, b::Point)
    # Aliases to simplify the formulas
    (xa, ya, za) = a
    (xb, yb, zb) = b
    (xc, yc, R) = (c.xc, c.yc, c.R)

    st = (c.top - za) / (zb - za)
    sb = (c.bottom - za) / (zb - za)

    h2 = (xa - xb)^2 + (ya - yb)^2

    # Vertical lines.
    if (h2 == 0)

        # Check if we are inside or out of the cylinder.
        ((xa - xc)^2 + (ya - yc)^2 > R^2) && return nothing

        # If we are inside, check if the caps are between a and b
        (0 < st < 1) && accum!(v, st)
        (0 < sb < 1) && accum!(v, sb)

        return nothing
    end

    Δ = R^2 * h2 - (xb * ya - xc * ya - xa * yb + xc * yb + xa * yc - xb * yc)^2
    m = xa^2 + xb * xc - xa * (xb + xc) + (ya - yb) * (ya - yc)

    Δ < 0 && return nothing

    s1 = (m - sqrt(Δ)) / h2
    s2 = (m + sqrt(Δ)) / h2

    z1 = (1 - s1) * a[3] + s1 * b[3]
    z2 = (1 - s2) * a[3] + s2 * b[3]

    (0 < st < 1) && (s1 < st < s2) && accum!(v, st)
    (0 < sb < 1) && (s1 < sb < s2) && accum!(v, sb)

    (0 < s1 < 1 && c.bottom < z1 < c.top) && accum!(v, s1)
    (0 < s2 < 1 && c.bottom < z2 < c.top) && accum!(v, s2)

    nothing
end

@fastmath function interfaces!(v::Accumulator, c::Cone, a::Point, b::Point)
    # Shift everything wrt the cone vertex
    (xa, ya, za) = a .- @SVector [c.xv, c.yv, c.zv]
    (xb, yb, zb) = b .- @SVector [c.xv, c.yv, c.zv]
    m = c.m
    m2 = m^2

    at2 = xa * xa + ya * ya
    bt2 = xb * xb + yb * yb
    adotb = xa * xb + ya * yb

    # with this we define a quadratic equation on s with coeffs
    α = at2 + bt2 - 2 * adotb - m2 * (za^2 + zb^2 - 2 * za * zb)
    β = -2 * at2 + 2 * adotb - m2 * (-2 * za^2 + 2 * za * zb)
    γ = at2 - m2 * za^2

    if α != 0
        Δ = sqrt(β^2 - 4 * α * γ)
        s1, s2 = (-β - Δ) / (2 * α), (-β + Δ) / (2 * α)
    else
        # Unfortunately we must allow the possibility that the line connecting
        # a and b passes though the vertex of the cone.  In that case
        # the line is either almost always inside or almost always outside
        # it.
        s1, s2 = -Inf, za / (za - zb)
    end

    # Make sure that s1 < s2
    (s1 < s2) || ((s1, s2) = (s2, s1))

    st = (c.top - za - c.zv) / (zb - za)
    sb = (c.bottom - za - c.zv) / (zb - za)

    (0 < st < 1) && (s1 < st < s2) && accum!(v, st)
    (0 < sb < 1) && (s1 < sb < s2) && accum!(v, sb)

    z1 = (1 - s1) * za + s1 * zb + c.zv
    z2 = (1 - s2) * za + s2 * zb + c.zv

    (0 < s1 < 1 && c.bottom < z1 < c.top) && accum!(v, s1)
    (0 < s2 < 1 && c.bottom < z2 < c.top) && accum!(v, s2)

    nothing
end

function interfaces!(v::Accumulator, p::Plane, a::Point, b::Point)
    # Aliases to simplify the formulas
    (xa, ya, za) = a
    (xb, yb, zb) = b
    zp = p.z

    s = (zp - za) / (zb - za)
    (0 < s < 1) && accum!(v, s)
    nothing
end

# Since a ShapeUnion or ShapeIntersect may be composed of shapes of different
# types, we use a generated function to unroll the loop and make the function
# type-stable.  This is cool.
@generated function interfaces!(v::Accumulator,
                                u::Union{ShapeUnion{T}, ShapeIntersect{T}},
                                a::Point, b::Point) where {T}
    L = fieldcount(T)
    out = quote end

    for i in 1:L
        push!(out.args,
              quote
              interfaces!(v, u.shapes[$i], a, b)
              end
              )
    end
    push!(out.args, :(return nothing))

    out
end

function interfaces!(v::Accumulator, s::ShapeSubstract, a::Point, b::Point)
    interfaces!(v, s.shape1, a, b)
    interfaces!(v, s.shape2, a, b)
    nothing
end

function interfaces!(v::Accumulator, s::TransformedShape,
                     a::Point, b::Point)
    at = s.inv(a)
    bt = s.inv(b)

    interfaces!(v, s.shape, at, bt)
end

"""
Checks whether point `a` is inside a given shape.
"""
function inside(s::Sphere, a::Point)
    (xa, ya, za) = a
    (xc, yc, zc, R) = (s.xc, s.yc, s.zc, s.R)
    ((xa - xc)^2 + (ya - yc)^2 + (za - zc)^2 < R^2)
end

function inside(c::Cylinder, a::Point)
    (xa, ya, za) = a
    (xc, yc, R, z1, z2) = (c.xc, c.yc, c.R, c.bottom, c.top)

    ((xa - xc)^2 + (ya - yc)^2 > R^2) && return false
    (z1 < za < z2)
end

function inside(c::Cone, a::Point)
    (c.bottom <= a[3] <= c.top) || return false

    (xa, ya, za) = a .- @SVector [c.xv, c.yv, c.zv]
    m = c.m

    (xa^2 + ya^2 > m^2 * za^2) && return false
    true
end

inside(p::Plane, a::Point) = (p.up == (a[3] > p.z))

# See above for why we use a generated function
@generated function inside(u::ShapeUnion{T}, a::Point) where {T}
    L = fieldcount(T)
    out = quote end

    for i in 1:L
        push!(out.args,
              quote
              inside(u.shapes[$i], a) && return true
              end
              )
    end
    push!(out.args, :(return false))

    out
end

@generated function inside(u::ShapeIntersect{T}, a::Point) where {T}
    L = fieldcount(T)
    out = quote end

    for i in 1:L
        push!(out.args,
              quote
              inside(u.shapes[$i], a) || return false
              end
              )
    end
    push!(out.args, :(return true))

    out
end

function inside(u::ShapeIntersect, a::Point)
    for item in u.shapes
        inside(item, a) || return false
    end
    true
end

inside(s::ShapeSubstract, a::Point) = inside(s.shape1, a) && !inside(s.shape2, a)

function inside(s::TransformedShape, a::Point)
    inside(s.shape, s.inv(a))
end

# To facilitate visualization we provide methods that allow the construction of
# a Mathematica expression that can be viewed with Mathematica versions above
# 11.2
mathlist(v) = "{" * join(v, ", ") * "}"
mathcall(func, v...) = func * "[" * join(v, ", ") * "]"

function mathematica(fig::Sphere)
    mathcall("Ball", mathlist([fig.xc, fig.yc, fig.zc]), fig.R)
end

function mathematica(fig::Cylinder)
    p1 = [fig.xc, fig.yc, fig.bottom]
    p2 = [fig.xc, fig.yc, fig.top]

    mathcall("Cylinder", mathlist([mathlist(p1), mathlist(p2)]),
             fig.R)
end

function mathematica(fig::Plane)
    mathcall("HalfSpace", mathlist([0, 0, fig.up ? -1 : 1]),
             mathlist([0, 0, fig.z]))
end

# Note that we are not currently supporting "double cones"
function mathematica(fig::Cone)
    rbottom = fig.m * abs(fig.bottom - fig.zv)
    rtop = fig.m * abs(fig.top - fig.zv)
    vertex = [fig.xv, fig.yv, fig.zv]

    if rbottom > rtop
        base = [fig.xv, fig.yv, fig.bottom]
        plane = Plane(fig.top, false)
    else
        base = [fig.xv, fig.yv, fig.top]
        plane = Plane(fig.top, true)
    end
    mathcall("RegionIntersection",
             mathcall("Cone", mathlist([mathlist(base), mathlist(vertex)]),
                      rbottom),
             mathematica(plane))
end

function mathematica(fig::ShapeUnion)
    mathcall("RegionUnion",
             [mathematica(x) for x in fig.shapes]...)
end

function mathematica(fig::ShapeIntersect)
    mathcall("RegionIntersection",
             [mathematica(x) for x in fig.shapes]...)
end

function mathematica(fig::ShapeSubstract)
    mathcall("RegionDifference",
             mathematica(fig.shape1), mathematica(fig.shape2))
end

function mathematica(fig::TransformedShape)
    mathcall(
        mathcall("AffineTransform",
                 mathlist(
                     [mathlist([mathlist(r) for r in eachrow(fig.trans.linear)]),
                      mathlist(fig.trans.translation)])),
        mathematica(fig.shape))
end


"""
Get a function representing an expression's value, given coordinates and linear parameters
"""
function get_basefun(ex::Expr, coords::Vector{Symbol}, linpars::Vector{Symbol})
    q=quote
        bfun=($(coords...), $(linpars...)) -> ( $(ex) )
    end

    try
        eval(q)

        return bfun
    catch
        throw(error("get_basefun:ERROR:error when parsing basis function. Check inputs"))
    end
end

"""
Get base function derivative of a given order in respect to each coordinate
"""
function get_basefun_derivative(ex::Expr, coords::Vector{Symbol}, linpars::Vector{Symbol}, deriv_inds::Vector{Int64})
    excp=copy(ex)
    dinds=copy(deriv_inds)

    for i=1:size(dinds, 1)
        while dinds[i]>0
            dinds[i]-=1

            excp=differentiate(excp, coords[i])
        end
    end

    q=quote
        bfun=($(coords...), $(linpars...)) -> ( $excp )
    end

    try
        eval(q)

        return bfun
    catch
        throw(error("get_basefun_derivative:ERROR:wrong parsing of basis function derivative"))
    end
end

"""
Get linear contribution of a linear parameter for a basis function, as a function of position
"""
function get_linpar_contribution(ex::Expr, coords::Vector{Symbol}, linpars::Vector{Symbol},
    influencer::Symbol, deriv_inds::Vector{Int64})
    excp=copy(ex)
    dinds=copy(deriv_inds)

    for i=1:size(dinds, 1)
        while dinds[i]>0
            dinds[i]-=1

            excp=differentiate(excp, coords[i])
        end
    end

    excp=differentiate(excp, influencer)

    q=quote
        bfun=($(coords...), $(linpars...)) -> ( $excp )
    end

    try
        eval(q)

        return bfun
    catch
        throw(error("get_linpar_contribution:ERROR:error parsing linear transformation coefficient for basis function. Please verify linearity"))
    end
end


import Base.show
# Main GaussianProcess type

@doc """
# Description
Fits a Gaussian process to a set of training points. The Gaussian process is defined in terms of its mean and covaiance (kernel) functions, which are user defined. As a default it is assumed that the observations are noise free.

# Constructors:
    GP(X, y, m, k, logNoise)
    GP(; m=MeanZero(), k=SE(0.0, 0.0), logNoise=-1e8) # observation-free constructor

# Arguments:
* `X::Matrix{Float64}`: Input observations
* `y::Vector{Float64}`: Output observations
* `m::Mean`           : Mean function
* `k::kernel`         : Covariance function
* `logNoise::Float64` : Log of the standard deviation for the observation noise. The default is -1e8, which is equivalent to assuming no observation noise.

# Returns:
* `gp::GP`            : Gaussian process object, fitted to the training data if provided
""" ->
type GP
    m:: Mean                # Mean object
    k::Kernel               # Kernel object
    logNoise::Float64       # log standard deviation of observation noise

    # Observation data
    nobsv::Int              # Number of observations
    X::Matrix{Float64}      # Input observations
    y::Vector{Float64}      # Output observations
    data::KernelData        # Auxiliary observation data (to speed up calculations)
    dim::Int                # Dimension of inputs

    # Auxiliary data
    cK::AbstractPDMat       # (k + exp(2*obsNoise))
    alpha::Vector{Float64}  # (k + exp(2*obsNoise))⁻¹y
    mLL::Float64            # Marginal log-likelihood
    dmLL::Vector{Float64}   # Gradient marginal log-likelihood

    function GP(X::Matrix{Float64}, y::Vector{Float64}, m::Mean, k::Kernel, logNoise::Float64=-1e8)
        dim, nobsv = size(X)
        length(y) == nobsv || throw(ArgumentError("Input and output observations must have consistent dimensions."))
        gp = new(m, k, logNoise, nobsv, X, y, KernelData(k, X), dim)
        update_mll!(gp)
        return gp
    end

    GP(; m=MeanZero(), k=SE(0.0, 0.0), logNoise=-1e8) =  new(m, k, logNoise, 0)

end

# Creates GP object for 1D case
GP(x::Vector{Float64}, y::Vector{Float64}, meanf::Mean, kernel::Kernel, logNoise::Float64=-1e8) = GP(x', y, meanf, kernel, logNoise)

@doc """
# Description
Fits an existing Gaussian process to a set of training points.

# Arguments:
* `gp::GP`: Exiting Gaussian process object
* `X::Matrix{Float64}`: Input observations
* `y::Vector{Float64}`: Output observations

# Returns:
* `gp::GP`            : A Gaussian process fitted to the training data
""" ->
function fit!(gp::GP, X::Matrix{Float64}, y::Vector{Float64})
    length(y) == size(X,2) || throw(ArgumentError("Input and output observations must have consistent dimensions."))
    gp.X = X
    gp.y = y
    gp.data = KernelData(gp.k, X)
    gp.dim, gp.nobsv = size(X)
    update_mll!(gp)
    return gp
end

fit!(gp::GP, x::Vector{Float64}, y::Vector{Float64}) = fit!(gp, x', y)

# Update auxiliarly data in GP object after changes have been made
function update_mll!(gp::GP)
    μ = mean(gp.m,gp.X)
    Σ = cov(gp.k, gp.X, gp.data)
    gp.cK = PDMat(Σ + exp(2*gp.logNoise)*eye(gp.nobsv) + 1e-8*eye(gp.nobsv))
    gp.alpha = gp.cK \ (gp.y - μ)
    gp.mLL = -dot((gp.y - μ),gp.alpha)/2.0 - logdet(gp.cK)/2.0 - gp.nobsv*log(2π)/2.0 # Marginal log-likelihood
end

# Update gradient of marginal log likelihood

function update_mll_and_dmll!(gp::GP; noise::Bool=true, mean::Bool=true, kern::Bool=true)
    update_mll!(gp::GP)
    gp.dmLL = Array(Float64, noise + mean*num_params(gp.m) + kern*num_params(gp.k))

    # Calculate Gradient with respect to hyperparameters

    #Derivative wrt the observation noise
    if noise
        gp.dmLL[1] = exp(2*gp.logNoise)*trace((gp.alpha*gp.alpha' - gp.cK \ eye(gp.nobsv)))
    end

    #Derivative wrt to mean hyperparameters, need to loop over as same with kernel hyperparameters
    if mean
        Mgrads = grad_stack(gp.m, gp.X)
        for i in 1:num_params(gp.m)
            gp.dmLL[i+noise] = dot(Mgrads[:,i],gp.alpha)
        end
    end

    # Derivative of marginal log-likelihood with respect to kernel hyperparameters
    if kern
        Kgrads = grad_stack(gp.k, gp.X, gp.data)   # [dK/dθᵢ]
        for i in 1:num_params(gp.k)
            gp.dmLL[i+mean*num_params(gp.m)+noise] = trace((gp.alpha*gp.alpha' - gp.cK \ eye(gp.nobsv))*Kgrads[:,:,i])/2
        end
    end
end

@doc """
# Description
Calculates the posterior mean and variance of Gaussian Process at specified points

# Arguments:
* `gp::GP`: Gaussian Process object
* `X::Matrix{Float64}`:  matrix of points for which one would would like to predict the value of the process.
                       (each column of the matrix is a point)

# Keyword Arguments
* `full_cov::Bool`: indicates whether full covariance matrix should be returned instead of only variances (default is false)

# Returns:
* `(mu, Sigma)::(Vector{Float64}, Vector{Float64})`: respectively the posterior mean  and variances of the posterior
                                                    process at the specified points
""" ->
function predict(gp::GP, x::Matrix{Float64}; full_cov::Bool=false)
    size(x,1) == gp.dim || throw(ArgumentError("Gaussian Process object and input observations do not have consistent dimensions"))
    if full_cov
        return _predict(gp, x)
    else
        ## Calculate prediction for each point independently
            μ = Array(Float64, size(x,2))
            σ2 = similar(μ)
        for k in 1:size(x,2)
            m, sig = _predict(gp, x[:,k:k])
            μ[k] = m[1]
            σ2[k] = max(full(sig)[1,1], 0.0)
        end
        return μ, σ2
    end
end

# 1D Case for prediction
predict(gp::GP, x::Vector{Float64}; full_cov::Bool=false) = predict(gp, x'; full_cov=full_cov)

## compute predictions
function _predict(gp::GP, X::Array{Float64})
    n = size(X, 2)
    cK = cov(gp.k, X, gp.X)
    Lck = whiten(gp.cK, cK')
    mu = mean(gp.m,X) + cK*gp.alpha        # Predictive mean
    Sigma_raw = cov(gp.k, X) - Lck'Lck # Predictive covariance
    # Hack to get stable covariance
    Sigma = try PDMat(Sigma_raw) catch; PDMat(Sigma_raw+1e-8*sum(diag(Sigma_raw))/n*eye(n)) end
    return (mu, Sigma)
end

# Sample from the GP
function rand!(gp::GP, X::Matrix{Float64}, A::DenseMatrix)
    nobsv = size(X,2)
    n_sample = size(A,2)

    if gp.nobsv == 0
        # Prior mean and covariance
        μ = mean(gp.m, X);
        Σraw = cov(gp.k, X);
        Σ = try PDMat(Σraw) catch; PDMat(Σraw+1e-8*sum(diag(Σraw))/nobsv*eye(nobsv)) end
    else
        # Posterior mean and covariance
        μ, Σ = predict(gp, X; full_cov=true)
    end

    return broadcast!(+, A, μ, unwhiten!(Σ,randn(nobsv, n_sample)))
end

function rand(gp::GP, X::Matrix{Float64}, n::Int)
    nobsv=size(X,2)
    A = Array(Float64, nobsv, n)
    return rand!(gp, X, A)
end

# Sample from 1D GP
rand(gp::GP, x::Vector{Float64}, n::Int) = rand(gp, x', n)

# Generate only one sample from the GP and returns a vector
rand(gp::GP,X::Matrix{Float64}) = vec(rand(gp,X,1))
rand(gp::GP,x::Vector{Float64}) = vec(rand(gp,x',1))

function get_params(gp::GP; noise::Bool=true, mean::Bool=true, kern::Bool=true)
    params = Float64[]
    if noise; push!(params, gp.logNoise); end
    if mean;  append!(params, get_params(gp.m)); end
    if kern; append!(params, get_params(gp.k)); end
    return params
end

function set_params!(gp::GP, hyp::Vector{Float64}; noise::Bool=true, mean::Bool=true, kern::Bool=true)
    # println("mean=$(mean)")
    if noise; gp.logNoise = hyp[1]; end
    if mean; set_params!(gp.m, hyp[1+noise:noise+num_params(gp.m)]); end
    if kern; set_params!(gp.k, hyp[end-num_params(gp.k)+1:end]); end
end

function push!(gp::GP, X::Matrix{Float64}, y::Vector{Float64})
    warn("push! method is currently inefficient as it refits all observations")
    if gp.nobsv == 0
        GaussianProcesses.fit!(gp, X, y)
    elseif size(X,1) != size(gp.X,1)
        error("New input observations must have dimensions consistent with existing observations")
    else
        GaussianProcesses.fit!(gp, cat(2, gp.X, X), cat(1, gp.y, y))
    end
end

push!(gp::GP, x::Vector{Float64}, y::Vector{Float64}) = push!(gp, x', y)
push!(gp::GP, x::Float64, y::Float64) = push!(gp, [x], [y])
push!(gp::GP, x::Vector{Float64}, y::Float64) = push!(gp, reshape(x, length(x), 1), [y])

function show(io::IO, gp::GP)
    println(io, "GP object:")
    println(io, "  Dim = $(gp.dim)")
    println(io, "  Number of observations = $(gp.nobsv)")
    println(io, "  Mean function:")
    show(io, gp.m, 2)
    println(io, "  Kernel:")
    show(io, gp.k, 2)
    if (gp.nobsv == 0)
        println("  No observation data")
    else
        println(io, "  Input observations = ")
        show(io, gp.X)
        print(io,"\n  Output observations = ")
        show(io, gp.y)
        print(io,"\n  Variance of observation noise = $(exp(2*gp.logNoise))")
        print(io,"\n  Marginal Log-Likelihood = ")
        show(io, round(gp.mLL,3))
    end
end


#!/usr/bin/env julia

using Base.Threads

function _picalc(numsteps)
    slice = 1.0 / numsteps

    sums = zeros(Float64, nthreads())
    n = cld(numsteps, nthreads())

    Threads.@threads for i in 1:nthreads()
        sum_thread = 0.0
        @simd for j in (1 + (i - 1) * n):min(numsteps, i * n)
            x = (j - 0.5) * slice
            sum_thread += 4.0 / (1.0 + x ^ 2)
        end
        sums[threadid()] = sum_thread
    end

    return sum(sums) * slice
end

function picalc(numsteps)

    println("Calculating PI using:")
    println("  ", numsteps, " slices")
    println("  ", nthreads(), " thread(s)")

    start = time()
    mypi = _picalc(numsteps)
    elapsed = time() - start

    println("Obtained value of PI: ", mypi)
    println("Time taken: ", elapsed, " seconds")

end

numsteps = if length(ARGS) > 0
    parse(Int, ARGS[1])
else
    1_000_000_000
end

# Warm up kernel
_picalc(10)

# Run the full example
picalc(numsteps)


__precompile__(false)

eye(T,n,m) = Matrix{T}(LinearAlgebra.I,n,m)

Mat2(T) = Mat2x2(T) = zeros(T,2,2)
Mat3(T) = Mat3x3(T) = zeros(T,3,4)
Mat4(T) = Mat4x4(T) = zeros(T,4,4)

const Mat2f = const Mat2x2f = SMatrix{2,2,Float32}
const Mat3f = const Mat3x3f = SMatrix{3,3,Float32}
const Mat4f = const Mat4x4f = SMatrix{4,4,Float32}

const zerosMat4x4f = Mat4(Float32)
const eyeMat4x4f = eye(Float32,4,4)

export eye
export Mat2
export Mat3
export Mat4
export Mat2x2
export Mat2x2f
export Mat3x3
export Mat3x3f
export Mat4x4
export Mat4x4f
export zerosMat4x4f
export eyeMat4x4f


using FromFile
@from "Core.jl" import Node, Options

"""
    evalTreeArray(tree::Node, cX::AbstractMatrix{T}, options::Options)

Evaluate a binary tree (equation) over a given input data matrix. The
options contain all of the operators used. This function fuses doublets
and triplets of operations for lower memory usage.

# Returns

- `(output, complete)::Tuple{AbstractVector{T}, Bool}`: the result,
    which is a 1D array, as well as if the evaluation completed
    successfully (true/false). A `false` complete means an infinity
    or nan was encountered, and a large loss should be assigned
    to the equation.
"""
function evalTreeArray(tree::Node, cX::AbstractMatrix{T}, options::Options)::Tuple{AbstractVector{T}, Bool} where {T<:Real}
    if tree.degree == 0
        deg0_eval(tree, cX, options)
    elseif tree.degree == 1
        if tree.l.degree == 2 && tree.l.l.degree == 0 && tree.l.r.degree == 0
            deg1_l2_ll0_lr0_eval(tree, cX, Val(tree.op), Val(tree.l.op), options)
        else
            deg1_eval(tree, cX, Val(tree.op), options)
        end
    else
        if tree.l.degree == 0 && tree.r.degree == 0
            deg2_l0_r0_eval(tree, cX, Val(tree.op), options)
        elseif tree.l.degree == 0
            deg2_l0_eval(tree, cX, Val(tree.op), options)
        elseif tree.r.degree == 0
            deg2_r0_eval(tree, cX, Val(tree.op), options)
        else
            deg2_eval(tree, cX, Val(tree.op), options)
        end
    end
end

# Break whenever there is a nan or inf. Since so many equations give
# such numbers, it saves a lot of computation to just skip computation,
# and return a large loss!
macro break_on_check(val, flag)
    :(if isnan($(esc(val))) || !isfinite($(esc(val)))
          $(esc(flag)) = false
          break
    end)
end

macro return_on_false(flag, retval)
    :(if !$(esc(flag))
          return ($(esc(retval)), false)
    end)
end

function deg2_eval(tree::Node, cX::AbstractMatrix{T}, ::Val{op_idx}, options::Options)::Tuple{AbstractVector{T}, Bool} where {T<:Real,op_idx}
    n = size(cX, 2)
    (cumulator, complete) = evalTreeArray(tree.l, cX, options)
    @return_on_false complete cumulator
    (array2, complete2) = evalTreeArray(tree.r, cX, options)
    @return_on_false complete2 cumulator
    op = options.binops[op_idx]
    finished_loop = true
    @inbounds @simd for j=1:n
        x = op(cumulator[j], array2[j])::T
        @break_on_check x finished_loop
        cumulator[j] = x
    end
    return (cumulator, finished_loop)
end

function deg1_eval(tree::Node, cX::AbstractMatrix{T}, ::Val{op_idx}, options::Options)::Tuple{AbstractVector{T}, Bool} where {T<:Real,op_idx}
    n = size(cX, 2)
    (cumulator, complete) = evalTreeArray(tree.l, cX, options)
    @return_on_false complete cumulator
    op = options.unaops[op_idx]
    finished_loop = true
    @inbounds @simd for j=1:n
        x = op(cumulator[j])::T
        @break_on_check x finished_loop
        cumulator[j] = x
    end
    return (cumulator, finished_loop)
end

function deg0_eval(tree::Node, cX::AbstractMatrix{T}, options::Options)::Tuple{AbstractVector{T}, Bool} where {T<:Real}
    n = size(cX, 2)
    if tree.constant
        return (fill(convert(T, tree.val), n), true)
    else
        return (cX[tree.feature, :], true)
    end
end

function deg1_l2_ll0_lr0_eval(tree::Node, cX::AbstractMatrix{T}, ::Val{op_idx}, ::Val{op_l_idx}, options::Options)::Tuple{AbstractVector{T}, Bool} where {T<:Real,op_idx,op_l_idx}
    n = size(cX, 2)
    op = options.unaops[op_idx]
    op_l = options.binops[op_l_idx]
    finished_loop = true
    if tree.l.l.constant && tree.l.r.constant
        val_ll = convert(T, tree.l.l.val)
        val_lr = convert(T, tree.l.r.val)
        x_l = op_l(val_ll, val_lr)::T
        if isnan(x_l) || !isfinite(x_l)
            return (Array{T, 1}(undef, n), false)
        end
        x = op(x_l)::T
        if isnan(x) || !isfinite(x)
            return (Array{T, 1}(undef, n), false)
        end
        return (fill(x, n), true)
    elseif tree.l.l.constant
        val_ll = convert(T, tree.l.l.val)
        feature_lr = tree.l.r.feature
        cumulator = Array{T, 1}(undef, n)
        finished_loop = true
        @inbounds @simd for j=1:n
            x_l = op_l(val_ll, cX[feature_lr, j])::T
            @break_on_check x_l finished_loop
            x = op(x_l)::T
            @break_on_check x finished_loop
            cumulator[j] = x
        end
        return (cumulator, finished_loop)
    elseif tree.l.r.constant
        feature_ll = tree.l.l.feature
        val_lr = convert(T, tree.l.r.val)
        cumulator = Array{T, 1}(undef, n)
        finished_loop = true
        @inbounds @simd for j=1:n
            x_l = op_l(cX[feature_ll, j], val_lr)::T
            @break_on_check x_l finished_loop
            x = op(x_l)::T
            @break_on_check x finished_loop
            cumulator[j] = x
        end
        return (cumulator, finished_loop)
    else
        feature_ll = tree.l.l.feature
        feature_lr = tree.l.r.feature
        cumulator = Array{T, 1}(undef, n)
        finished_loop = true
        @inbounds @simd for j=1:n
            x_l = op_l(cX[feature_ll, j], cX[feature_lr, j])::T
            @break_on_check x_l finished_loop
            x = op(x_l)::T
            @break_on_check x finished_loop
            cumulator[j] = x
        end
        return (cumulator, finished_loop)
    end
    return (cumulator, finished_loop)
end

function deg2_l0_r0_eval(tree::Node, cX::AbstractMatrix{T}, ::Val{op_idx}, options::Options)::Tuple{AbstractVector{T}, Bool} where {T<:Real,op_idx}
    n = size(cX, 2)
    op = options.binops[op_idx]
    finished_loop = true
    if tree.l.constant && tree.r.constant
        val_l = convert(T, tree.l.val)
        val_r = convert(T, tree.r.val)
        x = op(val_l, val_r)::T
        if isnan(x) || !isfinite(x)
            return (Array{T, 1}(undef, n), false)
        end
        return (fill(x, n), true)
    elseif tree.l.constant
        cumulator = Array{T, 1}(undef, n)
        val_l = convert(T, tree.l.val)
        feature_r = tree.r.feature
        @inbounds @simd for j=1:n
            x = op(val_l, cX[feature_r, j])::T
            @break_on_check x finished_loop
            cumulator[j] = x
        end
    elseif tree.r.constant
        cumulator = Array{T, 1}(undef, n)
        feature_l = tree.l.feature
        val_r = convert(T, tree.r.val)
        @inbounds @simd for j=1:n
            x = op(cX[feature_l, j], val_r)::T
            @break_on_check x finished_loop
            cumulator[j] = x
        end
    else
        cumulator = Array{T, 1}(undef, n)
        feature_l = tree.l.feature
        feature_r = tree.r.feature
        @inbounds @simd for j=1:n
            x = op(cX[feature_l, j], cX[feature_r, j])::T
            @break_on_check x finished_loop
            cumulator[j] = x
        end
    end
    return (cumulator, finished_loop)
end

function deg2_l0_eval(tree::Node, cX::AbstractMatrix{T}, ::Val{op_idx}, options::Options)::Tuple{AbstractVector{T}, Bool} where {T<:Real,op_idx}
    n = size(cX, 2)
    (cumulator, complete) = evalTreeArray(tree.r, cX, options)
    @return_on_false complete cumulator
    op = options.binops[op_idx]
    finished_loop = true
    if tree.l.constant
        val = convert(T, tree.l.val)
        @inbounds @simd for j=1:n
            x = op(val, cumulator[j])::T
            @break_on_check x finished_loop
            cumulator[j] = x
        end
    else
        feature = tree.l.feature
        @inbounds @simd for j=1:n
            x = op(cX[feature, j], cumulator[j])::T
            @break_on_check x finished_loop
            cumulator[j] = x
        end
    end
    return (cumulator, finished_loop)
end

function deg2_r0_eval(tree::Node, cX::AbstractMatrix{T}, ::Val{op_idx}, options::Options)::Tuple{AbstractVector{T}, Bool} where {T<:Real,op_idx}
    n = size(cX, 2)
    (cumulator, complete) = evalTreeArray(tree.l, cX, options)
    @return_on_false complete cumulator
    op = options.binops[op_idx]
    finished_loop = true
    if tree.r.constant
        val = convert(T, tree.r.val)
        @inbounds @simd for j=1:n
            x = op(cumulator[j], val)::T
            @break_on_check x finished_loop
            cumulator[j] = x
        end
    else
        feature = tree.r.feature
        @inbounds @simd for j=1:n
            x = op(cumulator[j], cX[feature, j])::T
            @break_on_check x finished_loop
            cumulator[j] = x
        end
    end
    return (cumulator, finished_loop)
end

# Evaluate an equation over an array of datapoints
# This one is just for reference. The fused one should be faster.
function differentiableEvalTreeArray(tree::Node, cX::AbstractMatrix{T}, options::Options)::Tuple{AbstractVector{T}, Bool} where {T<:Real}
    n = size(cX, 2)
    if tree.degree == 0
        if tree.constant
            return (ones(T, n) .* tree.val, true)
        else
            return (cX[tree.feature, :], true)
        end
    elseif tree.degree == 1
        return deg1_diff_eval(tree, cX, Val(tree.op), options)
    else
        return deg2_diff_eval(tree, cX, Val(tree.op), options)
    end
end

function deg1_diff_eval(tree::Node, cX::AbstractMatrix{T}, ::Val{op_idx}, options::Options)::Tuple{AbstractVector{T}, Bool} where {T<:Real,op_idx}
    n = size(cX, 2)
    (left, complete) = differentiableEvalTreeArray(tree.l, cX, options)
    @return_on_false complete left
    op = options.unaops[op_idx]
    out = op.(left)
    no_nans = !any(x -> (isnan(x) || !isfinite(x)), out)
    return (out, no_nans)
end

function deg2_diff_eval(tree::Node, cX::AbstractMatrix{T}, ::Val{op_idx}, options::Options)::Tuple{AbstractVector{T}, Bool} where {T<:Real,op_idx}
    n = size(cX, 2)
    (left, complete) = differentiableEvalTreeArray(tree.l, cX, options)
    @return_on_false complete left
    (right, complete2) = differentiableEvalTreeArray(tree.r, cX, options)
    @return_on_false complete2 left
    op = options.binops[op_idx]
    out = op.(left, right)
    no_nans = !any(x -> (isnan(x) || !isfinite(x)), out)
    return (out, no_nans)
end


function complete_cases(x::DataArray)
    xna=sum(!isna(x), 2)
    ind=Int[]
    n=size(x,1)
    ncols=size(x,2)
    vec=1:n
    for i=1:n
        if xna[i]==ncols
            push!(ind, i)
        end
    end
    return vec[ind]
end

function deleterows!(x::DataArray)
    keep_inds::Array{Int}=complete_cases(x)
    output=DataArray(Any, length(keep_inds), size(x,2))
    for i=1:size(x,2)
        for j=1:length(keep_inds)
            output[j,i]=x[keep_inds[j],i]
        end
    end
    return output
end

#Finding duplicated values
function duplicated{S <: Real}(x::Vector{S})
    xsort=sort(x)
    dup=false
    for i=1:(length(xsort)-1)
        if xsort[i]==xsort[i+1]
            dup=true
            break
        end
    end
    return dup
end
#duplicated(x::Vector{Int})=duplicated(convert(Vector{Float64}, x))
#duplicated(x::DataVector)=duplicated(removeNA(x))

function seq(Start::Number, End::Number, increment::Number)
    delta=End-Start
    if delta.<=0
        error("`start` must be larger than `end`")
    else
        n::Int=floor(delta/increment)+1
        output=zeros(n)
        for i=1:n
            output[i]=(i-1).*increment+Start
        end
    end
    return output
end

function trimcibt_loop(x, n, nboot, tr, side, randid, test)
    xbar=tmean(x, tr=tr)
    temp=zeros(n)
    tval=zeros(nboot)
    if side
        test=abs(test)
        pval=0.0
        for i=1:(nboot*n)
            if (i%n)!=0
                temp[i%n]=x[randid[i]]-xbar
            else
                temp[n]=x[randid[i]]-xbar
                tval[div(i, n)]=abs(tmean(temp, tr=tr)./trimse(temp, tr=tr))
                pval += tval[div(i, n)]>=test? 1/nboot :0.0
            end
        end
        return sort!(tval), pval
    else
        for i=1:(nboot*n)
            if (i%n)!=0
                temp[i%n]=x[randid[i]]-xbar
            else
                temp[n]=x[randid[i]]-xbar
                tval[div(i, n)]=tmean(temp, tr=tr)./trimse(temp, tr=tr)
            end
        end
        return sort!(tval)
    end
end

function akerd_loop(x, pts, hval, alam)
    npts  = length(pts)
    dhat  = zeros(npts)
    sqrt5 = sqrt(5)
    n     = length(x)
    temp  = zeros(n)
    for i = 1:npts
        epan_alam_hval = 0.0
        for j = 1:n
            temp[j] = (pts[i]-x[j])/(hval*alam[j])
            if abs(temp[j]) < sqrt5
                epan_alam_hval += (0.75*(1-0.2*temp[j]*temp[j])/sqrt5)/alam[j]
            end
        end
        dhat[i]=epan_alam_hval/npts
    end
    return dhat./hval
end

#Faster Pearson's r for
function cor!(x, y)
    n=length(x)
    xbar=mean(x); ybar=mean(y)
    top=bot1=bot2=0.0
    for i=1:n
        top+=(x[i]-xbar).*(y[i]-ybar)
        bot1+=(x[i]-xbar)*(x[i]-xbar)
        bot2+=(y[i]-ybar)*(y[i]-ybar)
    end
    top./sqrt(bot1.*bot2)
end

function regts1(vstar, yhat, res, mflag, x, tr)
    const n=length(res)
    ystar=zeros(n)
    [ystar[i]=yhat+res[i].*vstar[i] for i=1:n]
    rval=zeros(n)
    ystar_bar=tmean(ystar, tr=tr)
    rval=zeros(n)
    [rval[i]= sum(ystar[mflag[:,i]]-ystar_bar) for i=1:n]
    return rval
end

function indt_mflag(x)
    n=size(x, 1)
    np=size(x, 2)
    output=falses(n, n)
    for i=1:n
        for j=1:n
            total=0
            for k=1:np
                total+=x[i,k]<=x[j,k]?1:0
            end
            if total==np
                output[i,j]=true
            end
        end
    end
    return output
end

function bootindirect{S <: Real, T <: Real, W <: Real}(x::Vector{S}, y::Vector{T}, m::Vector{W}, nboot::Integer)
    n = length(x)
    tempx  = hcat(rep(1.0, n), zeros(n))
    tempy  = zeros(n)
    tempm  = zeros(n)
    tempmx = hcat(rep(1.0, n), zeros(Real, n), zeros(Real, n))
    bvec   = zeros(nboot)
    randid = rand(1:n, n*nboot)
    for i = 1:nboot
        for j = 1:n
            tempx[j,2] = tempmx[j,2] = x[randid[(i-1)*n + j]]
            tempy[j]   = y[randid[(i-1)*n + j]]
            tempm[j]   = tempmx[j,3] = m[randid[(i-1)*n + j]]
        end
        bvec[i] = (inv(tempx'tempx)'*(tempx'tempm))[2]*(inv(tempmx'tempmx)'*(tempmx'tempy))[3]
    end
    bvec
end

function mad!(v::AbstractArray)
    v=copy(v)
    center = median!(v, checknan=false)
    for i in 1:length(v)
        v[i] = abs(v[i]-center)
    end
    1.4826 * median!(v, checknan=false)
end

function t1waycore(args...)
    x, grp, tr, method = args[1], args[2], args[3], args[4]
    g = unique(grp)
    J = length(g)
    h, w, xbar, n = zeros(J), zeros(J), zeros(J), zeros(J)

    for j = 1:J
        n[j] = sum(grp.==g[j])
        temp = x[grp.==g[j]]
        h[j] = n[j] - 2*floor(tr*n[j])
        w[j] = h[j]*(h[j] - 1)/((n[j] -1)*winvar(temp, tr=tr))
        xbar[j] = tmean(temp, tr=tr)
    end
    u = sum(w)
    xtil = sum(w .* xbar)/u
    A = sum(w .* (xbar - xtil).*(xbar - xtil))/(J - 1)
    B = 2*(J - 2)*sum((1 - w./u).*(1 - w./u)./(h - 1))/(J*J - 1)
    statistic = A/(B + 1)
    nu1 = J - 1
    nu2 = 1/(3 * B/2/(J-2))
    p   = 1 - Rmath.pf(statistic, nu1, nu2)
    if method
        METHOD="Heteroscedastic one-way ANOVA for trimmed means\nusing a generalization of Welch's method.\n"
    else
        METHOD=nothing
    end
    output    = testOutput()
    output.n  = n
    output.df = [nu1, nu2]
    output.statistic = statistic
    output.p  = p
    output.method = METHOD
    output
end


immutable ZigZag
    m::Int
    n::Int
    diag::Array{Int,1}
    cmax::Int
    numd::Int
    lohi::(Int,Int)
end

function zigzag(m::Int, n::Int)
    0<m && 0<n || error("The matrix dimensions must be positive.")
    ZigZag(m, n, [-1,1], m*n, m+n-1, extrema([m,n]))
end
zigzag(n::Int) = zigzag(n, n)

type ZZState
    cnt::Int
    cell::Array{Int,1}
    dir::Int
    dnum::Int
    dlen::Int
    dcnt::Int
end

Base.length(zz::ZigZag) = zz.cmax
Base.start(zz::ZigZag) = ZZState(1, [1,1], 1, 1, 1, 1)
Base.done(zz::ZigZag, zzs::ZZState) = zzs.cnt > zz.cmax

function Base.next(zz::ZigZag, zzs::ZZState)
    s = sub2ind((zz.m, zz.n), zzs.cell[1], zzs.cell[2])
    if zzs.dcnt == zzs.dlen
        if isodd(zzs.dnum)
            if zzs.cell[2] < zz.n
                zzs.cell[2] += 1
            else
                zzs.cell[1] += 1
            end
        else
            if zzs.cell[1] < zz.m
                zzs.cell[1] += 1
            else
                zzs.cell[2] += 1
            end
        end
        zzs.dcnt = 1
        zzs.dnum += 1
        zzs.dir = -zzs.dir
        if zzs.dnum <= zz.lohi[1]
            zzs.dlen += 1
        elseif zz.lohi[2] < zzs.dnum
            zzs.dlen -= 1
        end
    else
        zzs.cell += zzs.dir*zz.diag
        zzs.dcnt += 1
    end
    zzs.cnt += 1
    return (s, zzs)
end


module Rqa

default_min_length = 2

"""Helper struct keeping sizes of diagonal lines and total area
of recurrence plot. In case of regular recurrence plot only upper triangle
(without the diagonal) is considered; in case of cross recurrence plot it is
the whole rectanle"""
type RQAdata
  area :: Int64
  lines :: Dict{Int64, Int64}
end

"Computes recurrence rate - probability of recurrence"
function RR(data :: RQAdata)
  if isempty(data.lines)
    0
  else
    sum(kv -> kv[1]*kv[2], data.lines) / data.area
  end
end

"Computes ratio of points lying on lines; by line we understand long enough
sequence of recurrences"
function DET(data :: RQAdata, min_length :: Int = default_min_length)
  if isempty(data.lines)
    NaN
  else
    longblack = sum(data.lines) do kv
      if kv[1] >= min_length
        kv[1]*kv[2]
      else
        0
      end
    end

    allblack = sum(kv -> kv[1]*kv[2], data.lines)
    longblack / allblack
  end
end

"Computes mean length of line; by line we understand long enough sequence
of recurrences"
function LMEAN(data :: RQAdata, min_length :: Int = default_min_length)
  llines = filter((k,v) -> k >= min_length, data.lines)
  if isempty(llines)
    NaN
  else
    length_sum = sum(kv -> kv[1] * kv[2], llines)
    count = sum(kv -> kv[2], llines)
    length_sum / count
  end
end

"Computes entropy of lines; by line we understand long enough sequence
of recurrences"
function LENT(data :: RQAdata, min_length :: Int = default_min_length)
  llines = filter((k,v) -> k >= min_length, data.lines)
  if isempty(llines)
    NaN
  else
    cnt = sum(kv -> kv[2], llines)
    sum(kv -> kv[2] / cnt * log2(cnt / kv[2]), llines)
  end
end

"""
Calculates the lines count for RQA. Takes into consideration only one triangle
of the plot, without the diagonal"""
function rqadata(sig :: AbstractVector, eq :: Function = (x, y) -> x == y)
  n = length(sig)
  lines = Dict{Int64, Int64}()

  for lag = 1:(n - 1)
    l = 0
    for i in (lag + 1):n
      if eq(sig[i], sig[i - lag])
        l += 1
      else
        if l > 0
          lines[l] = get(lines, l, 0) + 1
        end
        l = 0
      end
    end
    if l > 0
      lines[l] = get(lines, l, 0) + 1
    end
  end

  RQAdata(div(n*(n - 1), 2), lines)
end

"Computes the recurrence data for Cross Recurrence Quantification Analysis"

function crqadata(
    sig1 :: AbstractVector,
    sig2 :: AbstractVector,
    eq :: Function = (x, y) -> x == y)
  n = length(sig1)
  m = length(sig2)
  lines = Dict{Int64, Int64}()

  # iterates over all diagonals in the rectangle
  # lag 0 is the "main" diagonal - one passing through the top left corner
  for lag = (1 - n):(m - 1)
    l = 0
    for i in max(1, 1-lag):min(n, m - lag)
      if eq(sig1[i], sig2[i + lag])
        l += 1
      else
        if l > 0
          lines[l] = get(lines, l, 0) + 1
        end
        l = 0
      end
    end
    if l > 0
      lines[l] = get(lines, l, 0) + 1
    end
  end

  RQAdata(n*m, lines)
end

end


# Function for running Metropolis-Hastings algorithm
function mh(model::Model, opts::MhOpts)
  mcmc = Array(Float64, opts.mcmc.nPostBurnin, model.nPars)
  z = Array(Float64, opts.mcmc.nPostBurnin, model.nPars)
  proposalSd = ones(model.nPars, 1)
  proposed = zeros(model.nPars, 1)
  accepted = zeros(model.nPars, 1)
  acceptanceRatio = Array(Float64, model.nPars, 1)

  println("Running burn-in iterations...")

  currentPars = model.randPrior()
  currentLogPosterior = model.logPosterior(currentPars)
  currentGradLogPosterior = model.gradLogPosterior(currentPars)

  for i = 1:opts.mcmc.n
    for j = 1:model.nPars
      proposed[j] += 1

      proposedPars = copy(currentPars)
      proposedPars[j] = proposedPars[j]+randn()*proposalSd[j]
      proposedLogPosterior = model.logPosterior(proposedPars)
      proposedGradLogPosterior = model.gradLogPosterior(proposedPars)

      ratio = proposedLogPosterior-currentLogPosterior

      if ratio > 0 || (ratio > log(rand()))
        accepted[j] += 1

        currentPars = copy(proposedPars)
        currentLogPosterior = copy(proposedLogPosterior)
        currentGradLogPosterior = copy(proposedGradLogPosterior)
      end
    end

    if i > opts.mcmc.nBurnin
      mcmc[i-opts.mcmc.nBurnin, :] = currentPars
      z[i-opts.mcmc.nBurnin, :] = -currentGradLogPosterior/2
    end

    if mod(i, opts.mcmc.monitorRate) == 0
      if i < opts.mcmc.nBurnin
        for j = 1:model.nPars
          acceptanceRatio[j] = accepted[j]/proposed[j]

          if acceptanceRatio[j] > 0.6
            proposalSd[j] *= (1+opts.widthCorrection)
          elseif acceptanceRatio[j] < 0.2
            proposalSd[j] *= (1-opts.widthCorrection)
          end
        end
      end

      println("Iteration $i of $(opts.mcmc.n):")
      for j = 1:model.nPars
        println("  Parameter $j: ", round(100*accepted[j]/proposed[j], 2),
        "% acceptance ratio")
      end

      proposed = zeros(model.nPars, 1)
      accepted = zeros(model.nPars, 1)
    end

    if i == opts.mcmc.nBurnin
      println("Burn-in completed...\n\nRunning post burn-in MCMC...");
    end
  end

  return mcmc, z
end


module day16

include("inputs.jl")

ex = inputs.ex_d16()
ex2 = inputs.ex2_d16()
input = inputs.input_d16()

function parse_input(input)
    lines = split(input, '\n', keepempty=false)

    fields = Dict{String, Array{Int}}()
    rgx = r"^([a-zA-Z]+\s*[a-zA-Z]+): (\d+)-(\d+) or (\d+)-(\d+)$"
    for line in lines
        if contains(line, ": ")
            m = match(rgx, line).captures
            r = parse.(Int, m[2:5])
            rng = [r[1]:r[2]..., r[3]:r[4]...]
            fields[m[1]] = rng
        end
    end

    your_ticket = findfirst(line->line=="your ticket:", lines) + 1 |>
                    x->getindex(lines, x) |>
                    x->split(x, ",") |>
                    x->parse.(Int, x)

    ntline = findfirst(line->line=="nearby tickets:", lines) + 1
    nearby_tickets = []
    for line in lines[ntline:end]
        push!(nearby_tickets, split(line, ",") |> x->parse.(Int, x))
    end

    return fields, your_ticket, nearby_tickets
end

# Part 1
function find_invalid_items(ticket, fields)
    # index of ticket entries that *are* found somewhere in fields' values:
    idx = [any([item ∈ value for value in values(fields)]) for item in ticket]
    return ticket[.!idx]
end

function part1(input)
    fields, _, nearby_tickets = parse_input(input)
    sum(sum.([find_invalid_items(ticket, fields) for ticket in nearby_tickets]))
end

@assert part1(input) == 21081
@info "Day 16, Part 1 answer: $(part1(input))"

# Part 2
isvalid(ticket, fields) = all([any([item ∈ value for value in values(fields)]) for item in ticket])

function part2(input)
    """Get the valid tickets"""
    f, yt, nts = parse_input(input)
    valid_nearbys = nts[map(x->isvalid(x, f), nts)]
    tickets = permutedims(hcat(valid_nearbys...))

    """Find the possible associations between fields and ticket positions."""
    assoc = Dict{String, Array{Int}}()
    for (field, range) in f # iterate over range of fields
        for (i, position) in enumerate(eachcol(tickets))
            """Check if each "pos" (a column in the tickets) is entirely contained
            in the range of a field. If so, that pos could be associated with that
            field."""

            if Set(intersect(position, range)) == Set(position)
                # println("Field ", field, " could be assoc. with pos ", i)
                haskey(assoc, field) ? push!(assoc[field], i) : assoc[field] = [i]
            end
        end
    end

    """Identify which pairs have only one valid assocation, remove those values from the other pairs and repeat."""
    while !all(length.(values(assoc)) .== 1)  # while all the values have >1 one el
        for (field, positions) in assoc
            if length(assoc[field]) == 1
                # println("Associate ", field, " with ", assoc[field])
                for (f, p) in assoc
                    if f != field
                        assoc[f] = filter(x->x!=assoc[field][1], assoc[f])
                    end
                end
            end
        end
    end

    """Puzzle answer asks to multiply your ticket where the identified positions starts with departure."""
    prod = 1
    for field in keys(assoc)
        if startswith(field, "departure")
            prod *= yt[assoc[field][1]]
        end
    end

    return prod
end

@assert part2(input) == 314360510573
@info "Day 16, Part 2 answer: $(part2(input))"

end


using Test
using PolynomialBases

@elapsed begin
    @time @testset "Canonical Mappings" begin include("canonical_mappings_test.jl") end
    @time @testset "Interpolation" begin include("interpolation_test.jl") end
    @time @testset "Integration" begin include("integration_test.jl") end
    @time @testset "Derivatives" begin include("derivative_test.jl") end
    @time @testset "Legendre" begin include("legendre_test.jl") end
    @time @testset "Gegenbauer" begin include("gegenbauer_test.jl") end
    @time @testset "Jacobi" begin include("jacobi_test.jl") end
    @time @testset "Hermite" begin include("hermite_test.jl") end
    @time @testset "Laguerre" begin include("laguerre_test.jl") end
    @time @testset "Hahn" begin include("hahn_test.jl") end
    @time @testset "Utilities" begin include("utilities_test.jl") end
    @time @testset "Symbolic Bases (SymPy)" begin include("sympy_test.jl") end
    @time @testset "Symbolic Bases (SymEngine)" begin include("symengine_test.jl") end
end


module CellFieldsTests

using Test
using BlockArrays
using FillArrays
using Gridap.Arrays
using Gridap.Fields
using Gridap.TensorValues
using Gridap.Fields: MockField, MockBasis, OtherMockBasis
using Gridap.CellData

np = 3
ndofs = 4

p = Point(1,2)
x = fill(p,np)
z = 2.0

v = VectorValue(3.0,1.5)
w = VectorValue(3.4,3.5)
a = MockBasis{2}(v,ndofs)
b = MockBasis{2}(w,ndofs)
c = fill(1.0,ndofs)
f = OtherMockBasis{2}(ndofs)

g = MockField{2}(v)

l = 10
xl = Fill(x,l)
zl = [ z for  i in 1:l]
cl = fill(c,l)
fl = Fill(f,l)
ϕl = lincomb(fl,cl)
gl = fill(g,l)
al = Fill(a,l)
bl = fill(b,l)

gf = GenericCellField(gl,ϕl,Val(true))
gf_x = evaluate(gf,xl)
∇gf_x = evaluate(∇(gf),xl)
test_cell_field(gf,xl,gf_x,grad=∇gf_x)

af = GenericCellField(al,ϕl,Val(true),Fill((Base.OneTo(ndofs),),l),Val((:,)))
bf = GenericCellField(bl,ϕl,Val(true),Fill((Base.OneTo(ndofs),),l),Val((:,)))
zf = convert_to_cell_field(zl,ϕl)
df = af*zf
dft = trialize_cell_basis(df)

zfr = reindex(zf,[1,4,3])
@test length(zfr) == 3
@test length(get_array(zfr)) == 3
@test length(get_cell_map(zfr)) == 3
@test length(get_cell_axes(zfr)) == 3

# Check memoization
df_x1 = evaluate(df,xl)
df_x2 = evaluate(df,xl)
@test df_x1 === df_x2
∇gf1 = ∇(gf)
∇gf2 = ∇(gf)
@test ∇gf1 === ∇gf2
@test evaluate(∇gf1,xl) === evaluate(∇gf2,xl)
εgf1 = ε(gf)
εgf2 = ε(gf)
@test εgf1 === εgf2
@test ∇×gf === ∇×gf
@test evaluate(∇×gf,xl) === evaluate(∇×gf,xl)

@test is_test(af)
@test is_trial(dft)
@test is_basis(af)
@test is_basis(dft)
mf = af⋅dft
@test get_metasize(mf) == (:,:)
mf_x = evaluate(mf,xl)
@test size(mf_x[1]) == (np,ndofs,ndofs)

@test get_cell_axes(mf) == Fill((Base.OneTo(ndofs),Base.OneTo(ndofs)),l)

idsL = [ i*collect(1:ndofs)  for i in 1:l]
idsR = [ 2*i*collect(1:ndofs)  for i in 1:l]
axesL = Fill((Base.OneTo(ndofs),),l)
axesR = Fill((Base.OneTo(ndofs),),l)

idsS = merge_cell_dofs_at_skeleton(idsL,idsR,axesL,axesR)
@test isa(idsS,VectorOfBlockArrayCoo)

afS = merge_cell_fields_at_skeleton(af,2*af)
@test isa(afS,SkeletonCellField)

afL_x = evaluate(afS.left,xl)
afR_x = evaluate(afS.right,xl)
@test isa(afL_x,VectorOfBlockArrayCoo)
@test isa(afR_x,VectorOfBlockArrayCoo)

@test isa(afS*2,SkeletonCellField)
@test isa(afS+afS,SkeletonCellField)

# Checks associated with trial bases
df = bf
dft = trialize_cell_basis(df)
@test is_trial(dft) == true
cell_vals = [rand(ndofs) for  i in 1:l]
cell_field = lincomb(df,cell_vals)
@test isa(cell_field,CellField)
@test get_metasize(cell_field) == ()
cell_field_x = evaluate(cell_field,xl)
@test isa(cell_field_x[1],AbstractVector)

end # module


using Test
using AdaptiveSparseGrids
using ForwardDiff
using StaticArrays
using QuadGK

@testset "Adaptive Sparse Grid Tests" begin

@testset "Basis Functions" begin
    import AdaptiveSparseGrids: m, Y, h, ϕ, Dϕ, leftchild, rightchild

    @testset "m(i) : Dimension of basis at level i" begin
        @test m(1) == 1
        @test m(2) == 2
        @test m(3) == 4
        @test m(4) == 8
        @test_throws ArgumentError  m(0)
        @test_throws MethodError    m(1.0)
        @test_throws MethodError    m(:test)
    end

    @testset "Y(i,j) : Get the coordinate of the jth node in the ith layer" begin
        @test Y(1, 1) == 0.5
        @test Y(2, 0) == 0.0
        @test Y(2, 1) == 0.5
        @test Y(2, 2) == 1.0

        @testset "Check raw numbers" begin
            for i = 3:10, j = 0:m(i)
                @test Y(i, j) == j/2^(i-1)
            end
        end

        @testset "Check Nesting" begin
            for i = 3:10, j = 0:m(i)
                !iseven(j) && continue
                @test Y(i,j) == Y(i-1, Int(j/2))
            end
        end

        @test_throws ArgumentError Y(0, 3)
        @test_throws ArgumentError Y(1, 2)
        @test_throws ArgumentError Y(1, -1)
        @test_throws MethodError   Y(1.0, 1)
        @test_throws MethodError   Y(1, 1.0)
    end

    @testset "h(i): Width of basis elements in ith level" begin
        @test h(1) == 1.0
        @test h(2) == 0.5
        @test h(3) == 0.25
        for i in 4:10
            @test h(i) == 2.0^(-i + 1)
        end

        @test_throws MethodError h(1.0)
        @test_throws ArgumentError h(-1)
    end

    @testset "ϕ(x): Hat function" begin
        @test ϕ(-2.0) == 0.0
        @test ϕ(-1.0) == 0.0
        @test ϕ(0.0)  == 1.0
        @test ϕ(1.0)  == 0.0
        @test ϕ(2.0)  == 0.0

        @test ϕ(0.5)  == 0.5
        @test ϕ(-0.5) == 0.5
        @test ϕ(0.25) == 0.75
        @test ϕ(0.75) == 0.25

        for x in LinRange(0, 1, 5)
            # Test Symmetry
            @test ϕ(x) == ϕ(-x)

            # Test Values
            if x <= 1
                @test ϕ(x) == 1 - x
            else
                @test ϕ(x) == 0.0
            end
        end

        # Technically it should work with integers
        @test typeof(ϕ(1.0)) == Float64
        @test typeof(ϕ(1))   == Int
        @test ϕ(-1)== ϕ(-1.0)
        @test ϕ(0) == ϕ(0.0)
        @test ϕ(1) == ϕ(1.0)
        @test ϕ(0) == ϕ(0.0)
        @test ϕ(2) == ϕ(2.0)
    end

    @testset "Dϕ: Derivatives" begin
        for x in LinRange(-1, 1, 5)
            @test Dϕ(x) == ForwardDiff.derivative(ϕ, x)
        end
    end

    @testset "ϕ(i,j,x): Hierarchical Hat Functions" begin
        @test ϕ(1, 1, 0.0) == 1.0
        @test ϕ(1, 1, 0.5) == 1.0
        @test ϕ(1, 1, 1.0) == 1.0

        @test ϕ(2, 0, Y(2, 0)) == 1.0
        @test ϕ(2, 0, Y(2, 1)) == 0.0
        @test ϕ(2, 2, Y(2, 1)) == 0.0
        @test ϕ(2, 2, Y(2, 2)) == 1.0

        for i = 3:10, j = 1:m(i) - 1
            if isodd(j)
                # Boundaries
                @test ϕ(i, j, Y(i, j-1))    == 0.0
                @test ϕ(i, j, Y(i, j+1))    == 0.0

                # Interior
                @test ϕ(i, j, Y(i, j))      == 1.0
                @test ϕ(i, j, Y(i+1, 2j-1)) == 0.5
                @test ϕ(i, j, Y(i+1, 2j+1)) == 0.5
                @test ϕ(i, j, Y(i+2, 4j-3)) == 0.25
                @test ϕ(i, j, Y(i+2, 4j-1)) == 0.75
                @test ϕ(i, j, Y(i+2, 4j+1)) == 0.75
                @test ϕ(i, j, Y(i+2, 4j+3)) == 0.25
            end
        end
    end

    @testset "lefchild and rightchild" begin

        @testset "Base Values" begin
            @test leftchild(1, 1)  == (2, 0)
            @test rightchild(1, 1) == (2, 2)
            @test rightchild(2, 0) == (3, 1)
            @test leftchild(2, 2)  == (3, 3)
            @test leftchild(2, 0)  == (3, -1)
            @test rightchild(2, 2) == (3, -1)
        end

        @testset "Higher Level Hierarchy" begin
            for i in 3:10, j in 1:2:m(i)
                @test ϕ(i,j, Y(leftchild(i,j)...))  == 0.5
                @test ϕ(i,j, Y(rightchild(i,j)...)) == 0.5

                @test ϕ(leftchild(i,j)..., Y(i,j))  == 0.0
                @test ϕ(rightchild(i,j)..., Y(i,j)) == 0.0

                # They do overlap in their support with their children
                for c in (leftchild(i,j), rightchild(i,j))
                    @test ϕ(i,j,  (Y(i,j) + Y(c...))/2) == 0.75
                    @test ϕ(c..., (Y(i,j) + Y(c...))/2) == 0.5
                end

            end

        end

        @testset "Base Level Errors" begin
            @test_throws ArgumentError leftchild(2, 1)
            @test_throws ArgumentError rightchild(2, 1)
            @test_throws ArgumentError leftchild(0, 1)
            @test_throws ArgumentError rightchild(0, 1)
        end

        @testset "No Floats" begin
            @test_throws MethodError leftchild(1.0, 1)
            @test_throws MethodError leftchild(1,  1.0)
            @test_throws MethodError rightchild(1.0, 1)
            @test_throws MethodError rightchild(1,  1.0)
        end

    end
end

@testset "Nodes" begin

    @testset "KTuples" begin
        import AdaptiveSparseGrids: KTuple
        @test (1,2) isa KTuple
        @test (1,2) isa KTuple{2}
        @test (1,2) isa KTuple{2, Int}
        @test (1.0,2.0) isa KTuple
        @test (1.0,2.0) isa KTuple{2}
        @test (1.0,2.0) isa KTuple{2, Float64}

        @test (1,2,3) isa KTuple
        @test (1,2,3) isa KTuple{3}
        @test (1,2,3) isa KTuple{3, Int}
        @test (1.0,2.0,3.0) isa KTuple
        @test (1.0,2.0,3.0) isa KTuple{3, Float64}

        @test ! ((1, 2.0) isa KTuple)

        @test (a = 1.0, b = 2.0) isa KTuple
        @test (a = 1.0, b = 2.0) isa KTuple{2}
        @test (a = 1.0, b = 2.0) isa KTuple{2, Float64}

        @test !( (a = 1.0, b = 2) isa KTuple )
        @test !( (a = 1.0, b = 2) isa KTuple{2} )
        @test !( (a = 1.0, b = 2) isa KTuple{2, Float64} )
    end

    @testset "getzero" begin
        import AdaptiveSparseGrids: getzero
        @test getzero((a=1, b=2))               == (a = 0,   b = 0)
        @test getzero((a=1.0, b=2.0))           == (a = 0.0, b = 0.0)
        @test getzero(typeof((a=1, b=2)))       == (a = 0,   b = 0)
        @test getzero(typeof((a=1.0, b=2.0)))   == (a = 0.0, b = 0.0)
        @test getzero((1, 2, 3))                == (0,0,0)
        @test getzero((1.0, 2.0, 3.0))          == (0.0,0.0,0.0)
        @test getzero([1, 2, 3])                == (0,0,0)
        @test getzero([1, 2.0, 3])              == (0.0,0.0,0.0)
        @test getzero(1.0)                      == (0.0,)
        @test getzero(0)                        == (0,)
    end

    @testset "Node 1D" begin
        import AdaptiveSparseGrids: Node, getx, ϕ, makeleftchild, makerightchild
        n = Node(1,1)
        @test ϕ(n, (0.0,), 1) == 1.0
        @test ϕ(n, (0.5,), 1) == 1.0
        @test ϕ(n, (1.0,), 1) == 1.0

        # Check the children
        lc = makeleftchild(n, 1)
        @test getx(lc) == @SVector [0.0]
        @test ϕ(lc, 0.0, 1) == 1.0
        @test ϕ(lc, 0, 1)   == 1.0
        @test ϕ(lc, 0.5, 1) == 0.0
        @test ϕ(lc, 0.25, 1) == 0.5

        rc = makerightchild(n, 1)
        @test getx(rc) == @SVector [1.0]
        @test ϕ(rc, 1.0, 1) == 1.0
        @test ϕ(rc, 1, 1)   == 1.0
        @test ϕ(rc, 0.5, 1) == 0.0
        @test ϕ(rc, 0.75, 1) == 0.5

        for l in 3:10, j in 1:2:m(l)
            n = Node(l, j)
            @test getx(n) == @SVector [Y(l, j)]
            @test ϕ(n, getx(n), 1) == 1.0
            @test ϕ(n, Y(leftchild(l,j)...), 1)  == 0.5
            @test ϕ(n, Y(rightchild(l,j)...), 1) == 0.5
            j > 1       && @test ϕ(n, Y(l, j-1), 1) == 0.0
            j < m(l)- 1 && @test ϕ(n, Y(l, j+1), 1) == 0.0
            j > 2       && @test ϕ(n, Y(l, j-2), 1) == 0.0
            j < m(l)-1  && @test ϕ(n, Y(l, j+2), 1) == 0.0
        end
    end

    @testset "Node 2D" begin
        import AdaptiveSparseGrids: Node, getx, ϕ, m, Y
        for l1 in 3:5, l2 in 3:5, i1 in 1:2:m(l1), i2 in 1:2:m(l2)
            n = Node((l1,l2), (i1, i2))

            # Check that we're linearly interpolating to the corners
            for d1 in -1:1, d2 in -1:1
                c = @SVector [Y(l1, i1 + d1), Y(l2, i2+d2)]
                for θ in LinRange(0,1, 5)
                    x = θ * c + (1-θ) * getx(n)
                    @test ϕ(n, x) == (1-θ)^(abs(d1) + abs(d2))
                end
            end
        end
    end
end

@testset "Scaling" begin
    import AdaptiveSparseGrids: scale, rescale

    f = AdaptiveSparseGrid([0.0,1.0], [2.0, 3.0]) do (x,)
        x^2
    end

    @test scale(f, [0.0, 1.0]) == [0.0, 0.0]
    @test scale(f, (0.0, 1.0)) == [0.0, 0.0]
    @test scale(f, @SVector [0.0, 1.0]) == [0.0, 0.0]

    @test scale(f, [1.0, 1.0]) == [0.5, 0.0]
    @test scale(f, [1.0, 2.0]) == [0.5, 0.5]
    @test scale(f, [1.0, 3.0]) == [0.5, 1.0]

    @test scale(f, [2.0, 1.0]) == [1.0, 0.0]
    @test scale(f, [2.0, 2.0]) == [1.0, 0.5]
    @test scale(f, [2.0, 3.0]) == [1.0, 1.0]

    @test rescale(f, [0.5, 0.0]) == [1.0, 1.0]
    @test rescale(f, [0.5, 0.5]) == [1.0, 2.0]
    @test rescale(f, [0.5, 1.0]) == [1.0, 3.0]

    @test rescale(f, [1.0, 0.0]) == [2.0, 1.0]
    @test rescale(f, [1.0, 0.5]) == [2.0, 2.0]
    @test rescale(f, [1.0, 1.0]) == [2.0, 3.0]
end

@testset "Interpolation Tests (1D)" begin

    for g in [x -> x^2, sin, cos, exp]
        f = AdaptiveSparseGrid([0.0], [10.0], tol=1e-8, max_depth = 20) do x
            g.(x)
        end

        for x in LinRange(0, 10, 1_000)
            @test abs(f(x) - g(x))/max(g(x),1) < 1e-8
        end
    end

    # Test the named tuple features
    f = AdaptiveSparseGrid([0.0], [10.0], tol=1e-8, max_depth = 20) do (x,)
        (sin = sin(x), cos = cos(x), x2 = x^2)
    end

    @test f(pi/2, :sin) ≈ 1.0
    @test f(pi/2, :cos) ≈ 0.0 atol=1e-8
    @test f(pi  , :sin) ≈ 0.0 atol=1e-8
    @test f(pi  , :cos) ≈ -1.0

    @test f(pi/2, :sin) == f(pi/2, 1)
    @test f(pi/2, :cos) == f(pi/2, 2)

    # Test that everything still works when you construct with Tuple Boundaries
    for g in [x -> x^2, sin, cos, exp]
        f = AdaptiveSparseGrid((0.0,), (10.0,), tol=1e-8, max_depth = 20) do x
            g.(x)
        end

        for x in LinRange(0, 10, 1_000)
            @test abs(f(x) - g(x))/max(g(x),1) < 1e-8
        end
    end

end

@testset "Integration Tests (1D)" begin

    for g in [x -> x^2, sin, cos, exp]

        f = AdaptiveIntegral([0.0], [10.0], 1, tol=1e-8, max_depth = 20) do x
            g.(x)
        end

        v = quadgk(g, 0, 10)[1]
        @test abs(f()[1] - v)/max(v,1) < 1e-8
    end
end

@testset "Interpolation Tests (2D)" begin

    h((x,y)) = sin(x) * cos(x)
    h(x...)  = h(x)

    f = AdaptiveSparseGrid(h, [0.0, 0.0], [2 * pi, 2 * pi],
                           tol=1e-6, max_depth = 20)

    for x in LinRange(0,2*pi,100), y in LinRange(0,2*pi,100)
        @test f((x,y)) ≈ h(x,y) atol = 1e-6
        @test f((x,y)) == f(x,y)
    end

    # Check that using integers works fine
    for x in LinRange(0, 2*pi, 100)
        @test f((x,1)) == f(x, 1.0)
        @test f((1,x)) == f(1.0, x)
        @test f((x,2)) == f(x, 2.0)
        @test f((2,x)) == f(2.0, x)
    end

    # Check that the max depth parameter works
    @test mapreduce(x -> x.depth, max, values(nodes(f))) <= AdaptiveSparseGrids.max_depth(f)

    # Check that Tuple Constructors Work
    f = AdaptiveSparseGrid(h, (0.0, 0.0), (2 * pi, 2 * pi),
                           tol=1e-6, max_depth = 20)

    for x in LinRange(0,2*pi,100), y in LinRange(0,2*pi,100)
        @test f((x,y)) ≈ h(x,y) atol = 1e-6
        @test f((x,y)) == f(x,y)
    end

    # Now a version from R^2 →  R^2
    h((x,y)) = (a = sin(x) * cos(y), b = x^2 + y^2)

    f = AdaptiveSparseGrid(h, [0.0, 0.0], [2 * pi, 2 * pi],
                           tol=1e-6, max_depth = 30)

    for x in LinRange(0,2*pi,100), y in LinRange(0,2*pi,100)
        @test f((x,y),:a) ≈ h(x,y).a atol = 1e-5
    end

    # Check that using integers works fine
    for x in LinRange(0, 2*pi, 3)
        @test f((x,1)) == f((x, 1.0))
    end

    for x in LinRange(0, 2*pi, 3)
        @test f((1,x)) == f((1.0, x))
    end

end

end


using Manifolds
using BenchmarkTools
using StaticArrays
using LinearAlgebra
using Random

# Define a parent BenchmarkGroup to contain our suite
SUITE = BenchmarkGroup()

Random.seed!(12334)

function add_manifold(M::Manifold, pts, name;
    test_tangent_vector_broadcasting = true,
    retraction_methods = [],
    inverse_retraction_methods = [],
    point_distributions = [],
    tvector_distributions = [])

    SUITE["manifolds"][name] = BenchmarkGroup()
    tv = log(M, pts[1], pts[2])
    tv1 = log(M, pts[1], pts[2])
    tv2 = log(M, pts[1], pts[3])
    p = similar(pts[1])
    SUITE["manifolds"][name]["similar"] = @benchmarkable similar($(pts[1]))
    SUITE["manifolds"][name]["log"] = @benchmarkable log($M, $(pts[1]), $(pts[2]))
    SUITE["manifolds"][name]["log!"] = @benchmarkable log!($M, $tv, $(pts[1]), $(pts[2]))
    for iretr ∈ inverse_retraction_methods
        SUITE["manifolds"][name]["inverse_retract: "*string(iretr)] = @benchmarkable inverse_retract($M, $(pts[1]), $(pts[2]), $iretr)
        SUITE["manifolds"][name]["inverse_retract!: "*string(iretr)] = @benchmarkable inverse_retract!($M, $tv, $(pts[1]), $(pts[2]), $iretr)
    end
    SUITE["manifolds"][name]["exp"] = @benchmarkable exp($M, $(pts[1]), $tv1)
    SUITE["manifolds"][name]["exp!"] = @benchmarkable exp!($M, $p, $(pts[1]), $tv1)
    for retr ∈ retraction_methods
        SUITE["manifolds"][name]["retract: "*string(retr)] = @benchmarkable retract($M, $(pts[1]), $tv1, $retr)
        SUITE["manifolds"][name]["retract!: "*string(retr)] = @benchmarkable retract!($M, $p, $(pts[1]), $tv1, $retr)
    end
    SUITE["manifolds"][name]["norm"] = @benchmarkable norm($M, $(pts[1]), $tv1)
    SUITE["manifolds"][name]["inner"] = @benchmarkable inner($M, $(pts[1]), $tv1, $tv2)
    SUITE["manifolds"][name]["distance"] = @benchmarkable distance($M, $(pts[1]), $(pts[2]))
    SUITE["manifolds"][name]["isapprox (pt)"] = @benchmarkable isapprox($M, $(pts[1]), $(pts[2]))
    SUITE["manifolds"][name]["isapprox (tv)"] = @benchmarkable isapprox($M, $(pts[1]), $tv1, $tv2)
    SUITE["manifolds"][name]["2 * tv1 + 3 * tv2"] = @benchmarkable 2 * $tv1 + 3 * $tv2
    SUITE["manifolds"][name]["tv = 2 * tv1 + 3 * tv2"] = @benchmarkable $tv = 2 * $tv1 + 3 * $tv2
    if test_tangent_vector_broadcasting
        SUITE["manifolds"][name]["tv = 2 .* tv1 .+ 3 .* tv2"] = @benchmarkable $tv = 2 .* $tv1 .+ 3 .* $tv2
        SUITE["manifolds"][name]["tv .= 2 .* tv1 .+ 3 .* tv2"] = @benchmarkable $tv .= 2 .* $tv1 .+ 3 .* $tv2
    end
    for pd ∈ point_distributions
        distr_name = string(pd)
        distr_name = distr_name[1:min(length(distr_name), 50)]
        SUITE["manifolds"][name]["point distribution "*distr_name] = @benchmarkable rand($pd)
    end
    for tvd ∈ point_distributions
        distr_name = string(tvd)
        distr_name = distr_name[1:min(length(distr_name), 50)]
        SUITE["manifolds"][name]["tangent vector distribution "*distr_name] = @benchmarkable rand($tvd)
    end
end

# General manifold benchmarks
function add_manifold_benchmarks()

    SUITE["manifolds"] = BenchmarkGroup()

    s2 = Manifolds.Sphere(2)
    array_s2 = ArrayManifold(s2)
    r2 = Manifolds.Euclidean(2)

    pts_r2 = [Size(2)([1.0, 1.0]),
              Size(2)([-2.0, 3.0]),
              Size(2)([3.0, -2.0])]

    add_manifold(r2,
                 pts_r2,
                 "Euclidean{2} -- SizedArray")

    add_manifold(r2,
                 [MVector{2,Float64}([1.0, 1.0]),
                  MVector{2,Float64}([-2.0, 3.0]),
                  MVector{2,Float64}([3.0, -2.0])],
                  "Euclidean{2} -- MVector")

    ud_sphere = Manifolds.uniform_distribution(s2, Size(3)([1.0, 0.0, 0.0]))
    gtd_sphere = Manifolds.normal_tvector_distribution(s2, Size(3)([1.0, 0.0, 0.0]), 1.0)

    pts_s2 = [Size(3)([1.0, 0.0, 0.0]),
              Size(3)([0.0, 1.0, 0.0]),
              Size(3)([0.0, 0.0, 1.0])]

    add_manifold(s2,
                 pts_s2,
                 "Sphere{2} -- SizedArray";
                 point_distributions = [ud_sphere],
                 tvector_distributions = [gtd_sphere])

    add_manifold(array_s2,
                 [Size(3)([1.0, 0.0, 0.0]),
                  Size(3)([0.0, 1.0, 0.0]),
                  Size(3)([0.0, 0.0, 1.0])],
                 "ArrayManifold{Sphere{2}} -- SizedArray";
                 test_tangent_vector_broadcasting = false)

    retraction_methods_rot = [Manifolds.PolarRetraction(),
                              Manifolds.QRRetraction()]

    inverse_retraction_methods_rot = [Manifolds.PolarInverseRetraction(),
                                      Manifolds.QRInverseRetraction()]

    so2 = Manifolds.Rotations(2)
    angles = (0.0, π/2, 2π/3)
    add_manifold(so2,
                 [Size(2, 2)([cos(ϕ) -sin(ϕ); sin(ϕ) cos(ϕ)]) for ϕ in angles],
                 "Rotations(2) -- SizedArray",
                 retraction_methods = retraction_methods_rot,
                 inverse_retraction_methods = inverse_retraction_methods_rot)

    so3 = Manifolds.Rotations(3)
    angles = (0.0, π/2, 2π/3)
    add_manifold(so3,
                 [Size(3, 3)([cos(ϕ) -sin(ϕ) 0; sin(ϕ) cos(ϕ) 0; 0 0 1]) for ϕ in angles],
                 "Rotations(3) -- SizedArray",
                 retraction_methods = retraction_methods_rot,
                 inverse_retraction_methods = inverse_retraction_methods_rot)

    m_prod = Manifolds.ProductManifold(s2, r2)
    shape_s2r2_array = Manifolds.ShapeSpecification(Manifolds.ArrayReshaper(), s2, r2)
    shape_s2r2_static = Manifolds.ShapeSpecification(Manifolds.StaticReshaper(), s2, r2)

    pts_prd_base = [[1.0, 0.0, 0.0, 0.0, 0.0],
                    [0.0, 1.0, 0.0, 1.0, 0.0],
                    [0.0, 0.0, 1.0, 0.0, 0.1]]
    pts_prod_static = map(p -> Manifolds.ProductArray(shape_s2r2_static, p), pts_prd_base)
    add_manifold(m_prod, pts_prod_static, "ProductManifold with ProductArray (static)")

    pts_prod_array = map(p -> Manifolds.ProductArray(shape_s2r2_array, p), pts_prd_base)
    add_manifold(m_prod, pts_prod_array, "ProductManifold with ProductArray (array)")

    pts_prod_mpoints = [Manifolds.ProductMPoint(p[1], p[2]) for p in zip(pts_s2, pts_r2)]
    add_manifold(m_prod, pts_prod_mpoints, "ProductManifold with MPoint";
        test_tangent_vector_broadcasting = false)
end

add_manifold_benchmarks()


export fista

mutable struct FISTA{rT <: Real, vecT <: Union{AbstractVector{rT}, AbstractVector{Complex{rT}}}, matA, matAHA} <: AbstractLinearSolver where {rT}
  A::matA
  AᴴA::matAHA
  reg::Regularization
  x::vecT
  x₀::vecT
  xᵒˡᵈ::vecT
  res::vecT
  ρ::rT
  t::rT
  tᵒˡᵈ::rT
  iterations::Int64
  relTol::rT
  normalizeReg::Bool
  regFac::rT
  norm_x₀::rT
  rel_res_norm::rT
  verbose::Bool
end

"""
    FISTA(A, x::vecT=zeros(eltype(A),size(A,2))
          ; reg=nothing, regName=["L1"], λ=[0.0], kargs...)

creates a `FISTA` object for the system matrix `A`.

# Arguments
* `A`                       - system matrix
* `x::vecT`                 - array with the same type and size as the solution
* (`reg=nothing`)           - regularization object
* (`regName=["L1"]`)        - name of the Regularization to use (if reg==nothing)
* (`AᴴA=A'*A`)              - specialized normal operator, default is `A'*A`
* (`λ=0`)                   - regularization paramter
* (`ρ=1`)                   - step size for gradient step
* (`normalize_ρ=false`)     - normalize step size by the maximum eigenvalue of `AᴴA`
* (`t=1.0`)                 - parameter for predictor-corrector step
* (`relTol::Float64=1.e-5`) - tolerance for stopping criterion
* (`iterations::Int64=50`)  - maximum number of iterations
"""
function FISTA(A, x::AbstractVector{T}=Vector{eltype(A)}(undef,size(A,2)); reg=nothing, regName=["L1"]
              , AᴴA=A'*A
              , λ=0
              , ρ=0.95
              , normalize_ρ=true
              , t=1
              , relTol=eps(real(T))
              , iterations=50
              , normalizeReg=false
              , verbose = false
              , kargs...) where {rT <: Real, T <: Union{rT, Complex{rT}}}

  if reg == nothing
    reg = Regularization(regName, λ, kargs...)
  end

  x₀   = similar(x)
  xᵒˡᵈ = similar(x)
  res  = similar(x)
  res[1] = Inf # avoid spurious convergence in first iterations

  if normalize_ρ
    ρ /= abs(power_iterations(AᴴA))
  end

  return FISTA(A, AᴴA, vec(reg)[1], x, x₀, xᵒˡᵈ, res, rT(ρ),rT(t),rT(t),iterations,rT(relTol),normalizeReg,one(rT),one(rT),rT(Inf),verbose)
end

"""
    init!(it::FISTA, b::vecT
              ; A=solver.A
              , x::vecT=similar(b,0)
              , t::Float64=1.0) where T

(re-) initializes the FISTA iterator
"""
function init!(solver::FISTA{rT,vecT,matA,matAHA}, b::vecT
              ; x::vecT=similar(b,0)
              , t=1
              ) where {rT,vecT,matA,matAHA}

  solver.x₀ .= adjoint(solver.A) * b
  solver.norm_x₀ = norm(solver.x₀)

  if isempty(x)
    solver.x .= solver.ρ .* solver.x₀
  else
    solver.x .= x
  end

  solver.t = t
  solver.tᵒˡᵈ = t
  # normalization of regularization parameters
  if solver.normalizeReg
    solver.regFac = norm(solver.x₀,1)/length(solver.x₀)
  else
    solver.regFac = 1
  end
end

"""
    solve(solver::FISTA, b::Vector)

solves an inverse problem using FISTA.

# Arguments
* `solver::FISTA`                     - the solver containing both system matrix and regularizer
* `b::vecT`                           - data vector
* `A=solver.A`                        - operator for the data-term of the problem
* (`startVector::vecT=similar(b,0)`)  - initial guess for the solution
* (`solverInfo=nothing`)              - solverInfo object

when a `SolverInfo` objects is passed, the residuals are stored in `solverInfo.convMeas`.
"""
function solve(solver::FISTA, b; A=solver.A, startVector=similar(b,0), solverInfo=nothing, kargs...)
  # initialize solver parameters
  init!(solver, b; x=startVector)

  # log solver information
  solverInfo != nothing && storeInfo(solverInfo,solver.x,solver.res_norm)

  # perform FISTA iterations
  for (iteration, item) = enumerate(solver)
    solverInfo != nothing && storeInfo(solverInfo,solver.x,solver.res_norm)
  end

  return solver.x
end

"""
  iterate(it::FISTA, iteration::Int=0)

performs one fista iteration.
"""
function iterate(solver::FISTA, iteration::Int=0)
  if done(solver, iteration) return nothing end

  solver.xᵒˡᵈ .= solver.x

  # calculate residuum and do gradient step
  # solver.x .-= solver.ρ .* (solver.AᴴA * solver.x .- solver.x₀)
  mul!(solver.res, solver.AᴴA, solver.xᵒˡᵈ)
  solver.res .-= solver.x₀
  solver.x .-= solver.ρ .* solver.res

  solver.rel_res_norm = norm(solver.res) / solver.norm_x₀
  solver.verbose && println("Iteration $iteration; rel. residual = $(solver.rel_res_norm)")

  # the two lines below are equivalent to the ones above and non-allocating, but require the 5-argument mul! function to implemented for AᴴA, i.e. if AᴴA is LinearOperator, it requires LinearOperators.jl v2
  # mul!(solver.x, solver.AᴴA, solver.xᵒˡᵈ, -solver.ρ, 1)
  # solver.x .+= solver.ρ .* solver.x₀

  # proximal map
  solver.reg.prox!(solver.x, solver.regFac*solver.ρ*solver.reg.λ; solver.reg.params...)

  # predictor-corrector update
  solver.tᵒˡᵈ = solver.t
  solver.t = (1 + sqrt(1 + 4 * solver.tᵒˡᵈ^2)) / 2
  solver.x .+= ((solver.tᵒˡᵈ-1)/solver.t) .* (solver.x .- solver.xᵒˡᵈ)

  # return the residual-norm as item and iteration number as state
  return solver, iteration+1
end

@inline converged(solver::FISTA) = (solver.rel_res_norm < solver.relTol)

@inline done(solver::FISTA,iteration) = converged(solver) || iteration>=solver.iterations


module ManifoldMeasures

using KeywordCalls,
    LinearAlgebra,
    Manifolds,
    MeasureTheory,
    Random,
    SpecialFunctions,
    StaticArrays,
    StatsBase,
    StatsFuns
using HypergeometricFunctions: HypergeometricFunctions
using Manifolds: ℝ

const MAX_LENGTH_SIZED = 100

# experimental implementation of Option 3 in
# https://github.com/cscherrer/MeasureTheory.jl/issues/129#issue-962733685
default_number_type(::typeof(ℝ), T) = real(T)
default_number_type(::typeof(ℂ), T) = complex(real(T))

function default_point(μ, S::Type)
    M = base_manifold(μ)
    T = default_number_type(number_system(M), S)
    sz = representation_size(M)
    return prod(sz) ≤ MAX_LENGTH_SIZED ? MArray{Tuple{sz...},T}(undef) : Array{T}(undef, sz)
end

Manifolds.base_manifold(μ::AbstractMeasure) = base_manifold(basemeasure(μ))

include("utils.jl")
include("specialfunctions.jl")
include("factorizations.jl")
include("normalized.jl")
include("haar.jl")
include("hausdorff.jl")

include("vonmisesfisher.jl")
include("bingham.jl")
include("angularcentralgaussian.jl")

# primitive
export Normalized, Hausdorff, Haar, LeftHaar, RightHaar

export AngularCentralGaussian, Bingham, VonMises, VonMisesFisher

export mode, normalize

end


#
# Categorical Sampler
#
# parameters:
# (1)-(n)  0 <= w_i <= 1
#

type CategoricalSampler <: DiscreteDistributionSampler
	numcategories::Int
	paramranges::Vector{Tuple{Float64,Float64}}
	distribution::Categorical
	function CategoricalSampler(numcategories::Int, params=Float64[])
		numcategories >= 1 || error("number of categories must be at least one")
		s = new(numcategories, fill((0.0,1.0), numcategories))
		# Note: could parameterise using number of parameters that is one less than the domain size
		# but there is not an obvious way to do this in a symmetrical way so as not bias any optimisation
		# of the parameters
		setparams!(s, isempty(params) ? fill(1.0, numcategories) : params)
		s
	end
end

function setparams!(s::CategoricalSampler, params)
	checkparamranges(s, params)
	totalweight = sum(params)
	if totalweight == 0.0
		weights = fill(1.0/s.numcategories, s.numcategories)
	else
		weights = params ./ totalweight
	end
	# TODO we assume that sum of params is close enough (even with rounding errors) to 1.0 to satisfy the Categorical constructor
	# so we do not need handle the rounding errors
	s.distribution = Categorical(weights)
end

getparams(s::CategoricalSampler) = copy(s.distribution.p)

function estimateparams!(s::CategoricalSampler, traces)
	samples = extractsamplesfromtraces(s, traces)
	if length(samples) >= 1
		s.distribution = fit(primarydistributiontype(s.distribution), samples)
		# if there is a sequence of values at the end of the categories with zero probabilities,
		# then the fit will reduce the number of categories which will cause get parameters to return too
		# few values - we fix this here by padding with zeros
		if length(s.distribution.p) < s.numcategories
			s.distribution = Categorical([copy(s.distribution.p); fill(0.0, s.numcategories - length(s.distribution.p))])
		end
	end
end


"do nothing, this model does not have complex voltage constraints"
function constraint_voltage_dc(pm::AbstractPowerModel,  n::Int, cnd::Int)
end

"""
```
sum(p_dcgrid[a] for a in bus_arcs_dcgrid) + sum(pconv_dc[c] for c in bus_convs_dc) == pd
```
"""
function constraint_kcl_shunt_dcgrid(pm::AbstractPowerModel, n::Int, cnd::Int, i::Int, bus_arcs_dcgrid, bus_convs_dc, pd)
    p_dcgrid = PowerModels.var(pm, n, cnd, :p_dcgrid)
    pconv_dc = PowerModels.var(pm, n, cnd, :pconv_dc)

    PowerModels.con(pm, n, cnd, :kcl_dcgrid)[i] = @constraint(pm.model, sum(p_dcgrid[a] for a in bus_arcs_dcgrid) + sum(pconv_dc[c] for c in bus_convs_dc) == (-pd))
end

"`pconv[i] == pconv`"
function constraint_active_conv_setpoint(pm::AbstractPowerModel, n::Int, cnd::Int, i, pconv)
    pconv_var = PowerModels.var(pm, n, cnd, :pconv_tf_fr, i)
    PowerModels.con(pm, n, cnd, :conv_pac)[i] = @constraint(pm.model, pconv_var == -pconv)
end

"`qconv[i] == qconv`"
function constraint_reactive_conv_setpoint(pm::AbstractPowerModel, n::Int, cnd::Int, i, qconv)
    qconv_var = PowerModels.var(pm, n, cnd, :qconv_tf_fr, i)
    PowerModels.con(pm, n, cnd, :conv_qac)[i] = @constraint(pm.model, qconv_var == -qconv)
end


BigInt(10)^1000   # 大整数
BigFloat(10)^1000 # 大浮点数，如需改变默认精度，可参考文档
big(1.5)          # 对应大数类型的值，在此例是 BigFloat
1 + 1im           # 复数
123//456          # 使用 // 算符写分数


module JMcDM

using DataFrames
using LinearAlgebra
using JuMP
using Cbc
using StatsBase
using Statistics

# includes
include("types.jl")
include("utilities.jl")
include("topsis.jl")
include("vikor.jl")
include("electre.jl")
include("moora.jl")
include("dematel.jl")
include("ahp.jl")
include("nds.jl")
include("singlecriterion.jl")
include("game.jl")
include("dataenvelop.jl")
include("grey.jl")
include("saw.jl")
include("aras.jl")
include("wpm.jl")
include("waspas.jl")
include("edas.jl")
include("marcos.jl")
include("mabac.jl")
include("mairca.jl")
include("copras.jl")
include("promethee.jl")
include("cocoso.jl")
include("critic.jl")
include("entropy.jl")
include("codas.jl")

include("summary.jl")

# export MCDM types
export MCDMResult
export TopsisResult
export VikorResult
export ElectreResult
export MooraResult
export AHPConsistencyResult
export AHPResult
export NDSResult
export GreyResult
export SawResult
export ARASResult
export WPMResult
export WASPASResult
export EDASResult
export MARCOSResult
export MABACResult
export MAIRCAResult
export COPRASResult
export PrometheeResult
export CoCoSoResult
export CRITICResult
export EntropyResult
export CODASResult

# export game type
export GameResult

#  export SCDM types
export SCDMResult
export LaplaceResult
export MaximinResult
export MaximaxResult
export MinimaxResult
export MiniminResult
export SavageResult
export HurwiczResult
export MLEResult
export ExpectedRegretResult
export DataEnvelopResult

# export utility functions
export euclidean
export normalize
export colmaxs
export colmins
export unitize
export makeDecisionMatrix
export reverseminmax
export makeminmax

#  export MCDM tools
export topsis
export vikor
export electre
export moora
export dematel
export ahp_RI, ahp_consistency, ahp
export nds
export grey
export saw
export aras
export wpm
export waspas
export edas
export marcos
export mabac
export mairca
export copras
export promethee, prometLinear, prometVShape, prometUsual, prometQuasi, prometLevel
export cocoso
export critic
export entropy
export codas

#  export SCDM tools
export laplace
export maximin
export maximax
export minimax
export minimin
export savage
export hurwicz
export mle
export expectedregret

# export game solver
export game

# export data envelop
export dataenvelop

#  export summary function
export summary

end # module


using SafeTestsets

@safetestset "Solution Tests                                                                  " begin include("solutions/solutions_tests.jl") end
@safetestset "Tableau Tests                                                                   " begin include("tableaus/tableaus_tests.jl") end
@safetestset "Integrator Tests                                                                " begin include("integrators/integrators_tests.jl") end
@safetestset "Simulation Tests                                                                " begin include("simulations/simulations_tests.jl") end