# Test residual block
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

using Test, LinearAlgebra, InvertibleNetworks

# Input
nx1 = 32
nx2 = 32
nx_in = 8
n_hidden = 16   # same for x and y
batchsize = 2

ny1 = 64
ny2 = 24
ny_in = 1

# Residual blocks
RB = ConditionalResidualBlock(nx1, nx2, nx_in, ny1, ny2, ny_in, n_hidden, batchsize)

function loss(RB, X, Y)
    Zx, Zy = RB.forward(X, Y)
    f = -log_likelihood(Zx) - log_likelihood(Zy)
    ΔZx = -∇log_likelihood(Zx)
    ΔZy = -∇log_likelihood(Zy)
    ΔX, ΔY = RB.backward(ΔZx, ΔZy, X, Y)[1:2]
    return f, ΔX, ΔY, RB.W0.grad, RB.W1.grad, RB.W2.grad, RB.W3.grad,
        RB.b0.grad, RB.b1.grad, RB.b2.grad
end

###################################################################################################
# Gradient tests

# Input image and data
X0 = glorot_uniform(nx1, nx2, nx_in, batchsize)
X = glorot_uniform(nx1, nx2, nx_in, batchsize)
Y0 = glorot_uniform(ny1, ny2, ny_in, batchsize)
Y = glorot_uniform(ny1, ny2, ny_in, batchsize)
dX = X - X0
dY = Y - Y0

# Gradient test w.r.t. input
f0, ΔX, ΔY = loss(RB, X0, Y0)[1:3]
h = 0.1f0
maxiter = 5
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)

print("\nGradient test convolutions\n")
for j=1:maxiter
    f = loss(RB, X0 + h*dX, Y0 + h*dY)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX, ΔX) - h*dot(dY, ΔY))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Gradient test for weights
RB0 = ConditionalResidualBlock(nx1, nx2, nx_in, ny1, ny2, ny_in, n_hidden, batchsize)
RBini = deepcopy(RB0)
dW0 = RB.W0.data - RB0.W0.data
dW1 = RB.W1.data - RB0.W1.data
dW2 = RB.W2.data - RB0.W2.data
dW3 = RB.W3.data - RB0.W3.data
f0, ΔX, ΔY, ΔW0, ΔW1, ΔW2, ΔW3 = loss(RB0, X, Y)[1:7]
h = 0.1f0
maxiter = 5
err3 = zeros(Float32, maxiter)
err4 = zeros(Float32, maxiter)

print("\nGradient test convolutions\n")
for j=1:maxiter
    RB0.W0.data = RBini.W0.data + h*dW0
    RB0.W1.data = RBini.W1.data + h*dW1
    RB0.W2.data = RBini.W2.data + h*dW2
    RB0.W3.data = RBini.W3.data + h*dW3
    f = loss(RB0, X, Y)[1]
    err3[j] = abs(f - f0)
    err4[j] = abs(f - f0 - h*dot(dW0, ΔW0) - h*dot(dW1, ΔW1) - h*dot(dW2, ΔW2) - h*dot(dW3, ΔW3))
    print(err3[j], "; ", err4[j], "\n")
    global h = h/2f0
end

@test isapprox(err3[end] / (err3[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err4[end] / (err4[1]/4^(maxiter-1)), 1f0; atol=1f1)


using HermiteGF
using Test

include("trapz.jl")
include("test_interpolation1d.jl")
include("test_interpolation2d.jl")
include("test_interpolation3d.jl")
include("test_interpolation4d.jl")
include("test_interpolation5d.jl")


#############################################################################
# constant.jl
# Defines Constant, which is a subtype of AbstractExpr
#############################################################################
export Constant
export vexity, evaluate, sign, conic_form!

ispos(x::Real) = x >= 0
ispos(v::AbstractVecOrMat{<:Real}) = all(ispos, v)
isneg(x::Real) = x <= 0
isneg(v::AbstractVecOrMat{<:Real}) = all(isneg, v)

_size(x::Number) = (1, 1)
_size(x::AbstractVector) = (length(x), 1)
_size(x::AbstractMatrix) = size(x)

_sign(x::Union{Complex,AbstractVecOrMat{<:Complex}}) = ComplexSign()
function _sign(x::Value)
    if ispos(x)
        Positive()
    elseif isneg(x)
        Negative()
    else
        NoSign()
    end
end

struct Constant{T<:Value} <: AbstractExpr
    head::Symbol
    id_hash::UInt64
    value::T
    size::Tuple{Int, Int}
    sign::Sign

    Constant(x::Value, sign::Sign) = new{typeof(x)}(:constant, objectid(x), x, _size(x), sign)
    Constant(x::Value, check_sign::Bool=true) = Constant(x, check_sign ? _sign(x) : NoSign())
end

#### Constant Definition end     #####

vexity(::Constant) = ConstVexity()

evaluate(x::Constant) = x.value

sign(x::Constant) = x.sign

# `real(::Real)` is a no-op and should be optimized out for `Constant{<:Real}`
real_conic_form(x::Constant{<:Number}) = [real(x.value)]
real_conic_form(x::Constant{<:AbstractVecOrMat}) = vec(real(x.value))

# `imag(::Real)` always returns 0, so we can avoid the implicit conversion to `Complex`
# by multiplication with `im` and just use an explicit call to `zeros` with the appropriate
# length
imag_conic_form(x::Constant{T}) where {T<:Real} = zeros(Complex{T}, 1)
imag_conic_form(x::Constant{<:AbstractVecOrMat{T}}) where {T<:Real} = zeros(Complex{T}, length(x))
imag_conic_form(x::Constant{<:Complex}) = [im * imag(x.value)]
imag_conic_form(x::Constant{<:AbstractVecOrMat{<:Complex}}) = im * vec(imag(x.value))

# We can more efficiently get the length of a constant by asking for the length of its
# value, which Julia can get via Core.arraylen for arrays and knows is 1 for scalars
length(x::Constant) = length(x.value)

function conic_form!(x::Constant, unique_conic_forms::UniqueConicForms=UniqueConicForms())
    if !has_conic_form(unique_conic_forms, x)
        #real_Value = real_conic_form(x)
        #imag_Value = imag_conic_form(x)
        objective = ConicObj()
        objective[objectid(:constant)] = (real_conic_form(x), imag_conic_form(x))
        cache_conic_form!(unique_conic_forms, x, objective)
    end
    return get_conic_form(unique_conic_forms, x)
end


mutable struct Bard1 <:AbstractMPCCModel
  meta :: MPCCModelMeta
  counters :: MPCCCounters
end

function Bard1(; x0 = zeros(5))

  nvar = length(x0)
  lvar = -Inf*ones(nvar)
  uvar =  Inf*ones(nvar)
  ncon = 1
  lcon = [0.]
  ucon = [0.]
  y0   = [0.]
  ncc  = 3
  lccG = zeros(3)
  lccH = zeros(3)

  nnzh = nvar * (nvar + 1) / 2
  nnzj = nvar * ncon

  meta = MPCCModelMeta(nvar, x0 = x0, y0=y0, ncc = ncc, lccG = lccG, lccH = lccH,
                             lvar = lvar, uvar = uvar,
                             ncon = ncon, lcon = lcon, ucon = ucon,
                             nnzj=nnzj, nnzh=nnzh,
                             minimize=true, islp=false, name="Bard1")

  return Bard1(meta, MPCCCounters())
end

function MPCC.obj(nlp :: Bard1, x :: AbstractVector)
  #increment!(nlp, :neval_obj)
  return (x[1]-5)^2+(2*x[2]+1)^2
end

function MPCC.grad!(nlp :: Bard1, x :: AbstractVector, gx :: AbstractVector)
  #increment!(nlp, :neval_grad)
  gx .= vcat(2*(x[1] - 5), 4 * (2 * x[2] + 1), zeros(3))
  return gx
end
function MPCC.cons_nl!(nlp :: Bard1, x :: AbstractVector, c :: AbstractVector)
  #increment!(nlp, :neval_cons)
  c[1:nlp.meta.ncon] .= [2*(x[2]-1)-1.5*x[1]+x[3]-0.5*x[4]+x[5]]
  return c
end

function MPCC.consG!(nlp :: Bard1, x :: AbstractVector, c :: AbstractVector)
  #increment!(nlp, :neval_consG)
  c[1:nlp.meta.ncc] .= vcat(3x[1]-x[2]-3, -x[1]+0.5x[2]+4, -x[1]-x[2]+7 )
  return c
end

function MPCC.consH!(nlp :: Bard1, x :: AbstractVector, c :: AbstractVector)
  #increment!(nlp, :neval_consH)
  c[1:nlp.meta.ncc] .= vcat(x[3], x[4], x[5])
  return c
end

function MPCC.jac_nl_structure!(nlp :: Bard1, 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[1 : nlp.meta.nnzj] .= getindex.(I, 1)[:]
  cols[1 : nlp.meta.nnzj] .= getindex.(I, 2)[:]
  return rows, cols
end

function MPCC.jac_nl_coord!(nlp :: Bard1, x :: AbstractVector, vals :: AbstractVector)
  #increment!(nlp, :neval_jac)
  Jx = [-1.5 2. 1. -0.5 1.]
  vals[1 : nlp.meta.nnzj] .= Jx[:]
  return vals
end

function MPCC.jacG_structure!(nlp :: Bard1, 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 MPCC.jacG_coord!(nlp :: Bard1, x :: AbstractVector, vals :: AbstractVector)
  m, n = nlp.meta.ncc, nlp.meta.nvar
  #increment!(nlp, :neval_jacG)
  Jx = vcat([3 -1 0 0 0], [-1 0.5 0 0 0], [-1 -1 0 0 0])
  vals[1 : n*m] .= Jx[:]
  return vals
end

function MPCC.jacH_structure!(nlp :: Bard1, 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 MPCC.jacH_coord!(nlp :: Bard1, x :: AbstractVector, vals :: AbstractVector)
  m, n = nlp.meta.ncc, nlp.meta.nvar
  #increment!(nlp, :neval_jacH)
  Jx = zeros(3, 5)
  Jx[1,1] = 1.
  Jx[2,2] = 1.
  Jx[3,3] = 1.
  vals[1 : n*m] .= Jx[:]
  return vals
end

function MPCC.jnlprod!(nlp :: Bard1, x :: AbstractVector, v :: AbstractVector, Jv :: AbstractVector)
  #increment!(nlp, :neval_jprod)
  Jv[1:nlp.meta.ncon] .= [-1.5 2. 1. -0.5 1.] * v
  return Jv
end

function MPCC.jnltprod!(nlp :: Bard1, x :: AbstractVector, v :: AbstractVector, Jtv :: AbstractVector)
  #increment!(nlp, :neval_jtprod)
  Jtv[1:nlp.meta.nvar] .= [-1.5 2. 1. -0.5 1.]' * v
  return Jtv
end

function MPCC.jGprod!(nlp :: Bard1, x :: AbstractVector, v :: AbstractVector, Jv :: AbstractVector)
  #increment!(nlp, :neval_jGprod)
  Jv[1:nlp.meta.ncc] .=  vcat([3 -1 0 0 0], [-1 0.5 0 0 0], [-1 -1 0 0 0]) * v
  return Jv
end

function MPCC.jGtprod!(nlp :: Bard1, x :: AbstractVector, v :: AbstractVector, Jtv :: AbstractVector)
  #increment!(nlp, :neval_jGtprod)
  Jtv[1:nlp.meta.nvar] .=  vcat([3 -1 0 0 0], [-1 0.5 0 0 0], [-1 -1 0 0 0])' * v
  return Jtv
end

function MPCC.jHprod!(nlp :: Bard1, x :: AbstractVector, v :: AbstractVector, Jv :: AbstractVector)
  #increment!(nlp, :neval_jHprod)
  Jx = zeros(3, 5)
  Jx[1,1] = 1.
  Jx[2,2] = 1.
  Jx[3,3] = 1.
  Jv[1:nlp.meta.ncc] .= Jx * v
  return Jv
end

function MPCC.jHtprod!(nlp :: Bard1, x :: AbstractVector, v :: AbstractVector, Jtv :: AbstractVector)
  #increment!(nlp, :neval_jHtprod)
  Jx = zeros(3, 5)
  Jx[1,1] = 1.
  Jx[2,2] = 1.
  Jx[3,3] = 1.
  Jtv[1:nlp.meta.nvar] .= Jx' * v
  return Jtv
end

#function hess(nlp :: Bard1, x :: AbstractVector; obj_weight :: Real = one(eltype(x)))
#  #increment!(nlp, :neval_hess)
#  ℓ(x) = obj_weight * nlp.f(x)
#  Hx = ForwardDiff.hessian(ℓ, x)
#  return tril(Hx)
#end

#function hess(nlp :: Bard1, x :: AbstractVector, y :: AbstractVector; obj_weight :: Real = one(eltype(x)))
#  #increment!(nlp, :neval_hess)
#  ℓ(x) = obj_weight * nlp.f(x) + dot(nlp.c(x), y)
#  Hx = ForwardDiff.hessian(ℓ, x)
#  return tril(Hx)
#end

function MPCC.hess_structure!(nlp :: Bard1, 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 MPCC.hess_coord!(nlp :: Bard1, x :: AbstractVector, vals :: AbstractVector; obj_weight :: Real = one(eltype(x)))
  #increment!(nlp, :neval_hess)
  ℓ(x) = obj_weight * nlp.f(x)
  Hx = zeros(5,5)
  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 MPCC.hess_coord!(nlp :: Bard1, x :: AbstractVector, y :: AbstractVector, vals :: AbstractVector; obj_weight :: Real = one(eltype(x)))
  #increment!(nlp, :neval_hess)
  ncon, ncc = nlp.meta.ncon, nlp.meta.ncc
  # ℓ(x) = obj_weight * f(x) + λ₁ c(x) - λ₂ G(x) - λ₃ H(x)
  #ℓ(x) = obj_weight * nlp.f(x) + dot(nlp.c(x), y[1:ncon]) - dot(nlp.G(x), y[ncon+1:ncon+ncc]) - dot(nlp.H(x), y[ncon+ncc+1:ncon+2*ncc])
  Hx = zeros(5,5)
  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 MPCC.hessG(nlp :: Bard1, x :: AbstractVector, y :: AbstractVector)
  #increment!(nlp, :neval_hessG)
  #ℓ(x) = dot(nlp.G(x), y)
  Hx = zeros(5,5)
  return tril(Hx)
end

function MPCC.hessG_structure!(nlp :: Bard1, 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 MPCC.hessG_coord!(nlp :: Bard1, x :: AbstractVector, y :: AbstractVector, vals :: AbstractVector)
  #increment!(nlp, :neval_hessG)
  #ℓ(x) = dot(nlp.G(x), y)
  Hx = zeros(5,5)
  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 MPCC.hessH(nlp :: Bard1, x :: AbstractVector, y :: AbstractVector)
  #increment!(nlp, :neval_hessH)
  #ℓ(x) = dot(nlp.H(x), y)
  Hx = zeros(5,5)
  return tril(Hx)
end

function MPCC.hessH_structure!(nlp :: Bard1, 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 MPCC.hessH_coord!(nlp :: Bard1, x :: AbstractVector, y :: AbstractVector, vals :: AbstractVector)
  #increment!(nlp, :neval_hess)
  #ℓ(x) = dot(nlp.H(x), y)
  Hx = zeros(5,5)
  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 MPCC.hprod!(nlp :: Bard1, x :: AbstractVector, v :: AbstractVector, Hv :: AbstractVector; obj_weight :: Real = one(eltype(x)))
  #increment!(nlp, :neval_hprod)
  #ℓ(x) = obj_weight * nlp.f(x)
  Hv .= zeros(5)
  return Hv
end

function MPCC.hprod!(nlp :: Bard1, x :: AbstractVector, y :: AbstractVector, v :: AbstractVector, Hv :: AbstractVector; obj_weight :: Real = one(eltype(x)))
  #increment!(nlp, :neval_hprod)
  ncon, ncc = nlp.meta.ncon, nlp.meta.ncc
  #ℓ(x) = obj_weight * nlp.f(x) + dot(nlp.c(x), y[1:ncon]) - dot(nlp.G(x), y[ncon+1:ncon+ncc]) - dot(nlp.H(x), y[ncon+ncc+1:ncon+2*ncc])
  Hv .= zeros(5)
  return Hv
end

function MPCC.hGprod!(nlp :: Bard1, x :: AbstractVector, y :: AbstractVector, v :: AbstractVector, Hv :: AbstractVector)
  #increment!(nlp, :neval_hGprod)
  #ℓ(x) = dot(nlp.G(x), y)
  Hv .= zeros(5)
  return Hv
end

function MPCC.hHprod!(nlp :: Bard1, x :: AbstractVector, y :: AbstractVector, v :: AbstractVector, Hv :: AbstractVector)
  #increment!(nlp, :neval_hHprod)
  #ℓ(x) = dot(nlp.H(x), y)
  Hv .= zeros(5) #ForwardDiff.derivative(t -> ForwardDiff.gradient(ℓ, x + t * v), 0)
  return Hv
end


# addition
function (+)(f::SimpleRationalFunction, g::SimpleRationalFunction)
    top = f.num * g.den + f.den * g.num
    bot = f.den * g.den
    return SimpleRationalFunction(top, bot)
end

function (+)(f::SimpleRationalFunction, p::SimplePolynomial)
    return f + SimpleRationalFunction(p)
end
(+)(p::SimplePolynomial, f::SimpleRationalFunction) = f + p

(+)(f::SimpleRationalFunction, a::T) where {T<:CoefX} = f + SimpleRationalFunction(a)
(+)(a::T, f::SimpleRationalFunction) where {T<:CoefX} = f + a

# subtraction
(-)(f::SimpleRationalFunction) = SimpleRationalFunction(-f.num, f.den)

(-)(f::SimpleRationalFunction, g::SimpleRationalFunction) = f + (-g)

(-)(f::SimpleRationalFunction, p::SimplePolynomial) = f + (-p)
(-)(p::SimplePolynomial, f::SimpleRationalFunction) = p + (-f)

(-)(f::SimpleRationalFunction, a::T) where {T<:CoefX} = f + SimpleRationalFunction(-a)
(-)(a::T, f::SimpleRationalFunction) where {T<:CoefX} = SimpleRationalFunction(a) - f

# multiplication
(*)(f::SimpleRationalFunction, g::SimpleRationalFunction) =
    SimpleRationalFunction(f.num * g.num, f.den * g.den)

(*)(f::SimpleRationalFunction, p::SimplePolynomial) = f * SimpleRationalFunction(p)
(*)(p::SimplePolynomial, f::SimpleRationalFunction) = f * p

(*)(f::SimpleRationalFunction, a::T) where {T<:CoefX} = f * SimpleRationalFunction(a)
(*)(a::T, f::SimpleRationalFunction) where {T<:CoefX} = f * a

# division
inv(f::SimpleRationalFunction) = SimpleRationalFunction(f.den, f.num)

function inv(p::SimplePolynomial)
    T = eltype(p)
    SimpleRationalFunction(T(1), p)
end

SimpleThing = Union{SimplePolynomial,SimpleRationalFunction}

(/)(f::S, g::T) where {S<:SimpleThing,T<:SimpleThing} = f * inv(g)
(/)(f::SimpleRationalFunction, a::T) where {T<:CoefX} = f / SimplePolynomial(a)
(/)(a::T, f::S) where {T<:CoefX,S<:SimpleThing} = a * inv(f)

(//)(f::S, g::T) where {S<:SimpleThing,T<:SimpleThing} = f * inv(g)
(//)(f::SimpleRationalFunction, a::T) where {T<:CoefX} = f / SimplePolynomial(a)
(//)(a::T, f::S) where {T<:CoefX,S<:SimpleThing} = a * inv(f)

# exponentiation
function (^)(p::SimpleRationalFunction, k::S) where {S<:Integer}
    if k == 0
        return SimpleRationalFunction(1)
    end

    if k < 0
        return (inv(p))^(-k)
    end

    if k == 1
        return p
    end

    if k == 2
        return p * p
    end

    if k % 2 == 0  # even exponent
        j = k ÷ 2
        q = p^j
        return q * q
    end

    j = k ÷ 2
    q = p^j
    return q * q * p
end

# derivative
import SimplePolynomials: derivative
import Base: adjoint

function derivative(f::SimpleRationalFunction)
    p = numerator(f)
    q = denominator(f)
    return (p' * q - q' * p) / (q * q)
end
adjoint(f::SimpleRationalFunction) = derivative(f)


using SparseArrays

abstract type constraint end

# Constraint Types and settings

include("affineConstraints.jl")
include("socpConstraints.jl")
include("socpNewConstraints.jl")
include("socpAngleConstraints.jl")

#= Augmented Lagrangian Constraints
These are presented in the order:
- AffineEquality
- AffineInequality
- (Old) SOCP Constraints
- (Old) Functional SOCP Constraints with Slack Variables
- Simple SOCP Constraints
- Simple SOCP Constraints for N-constraints
- SOCP Constraints for Angles
- SOCP Constraints for Angles for N-constraints
=#


@testset "Jeffrey Problems" begin
    (p, q, r, x, ) = @variables p q r x

    #= ::Package:: =#

    #= ::Title:: =#
    #=David*Jeffrey - Rectifying*Transformations*for*Trig*Integration*(1997)=#

    #= ::Subsection::Closed:: =#
    #=Problem*(1.2)=#

    @test_int [3/(5 - 4*cos(x)), x, 2, x + 2*atan(sin(x)/(2 - cos(x)))]

    #= ::Subsection::Closed:: =#
    #=Problem*(1.4)=#

    @test_int [(cos(x) + 2*sin(x) + 1)/(cos(x)^2 - 2*sin(x)*cos(x) + 2*sin(x) + 3), x, -43, -atan((2*cos(x) - sin(x))/(2 + sin(x)))]

    #= ::Subsection::Closed:: =#
    #=Problem*(1.5)=#

    @test_int [(2 + 5*sin(x) + cos(x))/(4*cos(x) + sin(x)*cos(x) - 2*sin(x) - 2*sin(x)^2), x, -25, -log(1 - 3*cos(x) + sin(x)) + log(3 + cos(x) + sin(x))]

    #= ::Subsection::Closed:: =#
    #=Problem*(3.3)=#

    @test_int [(7*cos(x) + 2*sin(x) + 3)/(3*cos(x)^2 - sin(x)*cos(x) + 4*cos(x) - 5*sin(x) + 1), x, -32, -log(1 + cos(x) - 2*sin(x)) + log(3 + cos(x) + sin(x))]

    #= ::Subsection::Closed:: =#
    #=Problem =#

    @test_int [(5*cos(x)^2 + 4*cos(x) - 1)/(4*cos(x)^3 - 3*cos(x)^2 - 4*cos(x) - 1), x, -2, x - 2*atan(sin(x)/(3 + cos(x))) - 2*atan((3*sin(x) + 7*cos(x)*sin(x))/(1 + 2*cos(x) + 5*cos(x)^2))]

    #= ::Subsection::Closed:: =#
    #=Problem =#

    @test_int [(7*cos(x)^2 + 2*cos(x) - 5)/(4*cos(x)^3 - 9*cos(x)^2 + 2*cos(x) - 1), x, -2, x - 2*atan((2*cos(x)*sin(x))/(1 - cos(x) + 2*cos(x)^2))]

    #= ::Subsection::Closed:: =#
    #=Problem*(3.4)=#

    @test_int [3/(5 + 4*sin(x)), x, 2, x + 2*atan(cos(x)/(2 + sin(x)))]

    #= ::Subsection::Closed:: =#
    #=Problem*(3.6)=#

    @test_int [2/(1 + cos(x)^2), x, 3, sqrt(2)*x - sqrt(2)*atan((cos(x)*sin(x))/(1 + sqrt(2) + cos(x)^2))]

    #= ::Subsection::Closed:: =#
    #=Problem*(3.8)=#

    @test_int [1/(p + q*cos(x) + r*sin(x)), x, 3, (2*atan((r + (p - q)*tan(x/2))/sqrt(p^2 - q^2 - r^2)))/sqrt(p^2 - q^2 - r^2)]
    end

test/0 Independent test suites/Jeffrey Problems.jl


@testset "Pagerank" begin
    g5 = DiGraph(4)
    add_edge!(g5, 1, 2); add_edge!(g5, 2, 3); add_edge!(g5, 1, 3); add_edge!(g5, 3, 4)
    for g in testdigraphs(g5)
      @test pagerank(g)[3] ≈ 0.318 atol = 0.001
      @test length(@inferred(pagerank(g))) == nv(g)
      @test_throws ErrorException pagerank(g, 2)
      @test_throws ErrorException pagerank(g, 0.85, 2)
    end
end


using brainflow

# specify logging library to use
brainflow.enable_dev_brainflow_logger(brainflow.BOARD_CONTROLLER)

params = BrainFlowInputParams()
board_shim = brainflow.BoardShim(brainflow.SYNTHETIC_BOARD, params)

brainflow.prepare_session(board_shim)
brainflow.start_stream(board_shim)
sleep(5)
brainflow.stop_stream(board_shim)
data = brainflow.get_current_board_data(32, board_shim)
brainflow.release_session(board_shim)

eeg_channels = brainflow.get_eeg_channels(brainflow.SYNTHETIC_BOARD)
sampling_rate = brainflow.get_sampling_rate(brainflow.SYNTHETIC_BOARD)

data_first_channel = data[eeg_channels[1], :]
println("Original Data First Channel")
println(data_first_channel)
brainflow.perform_lowpass(data_first_channel, sampling_rate, 10.0, 3, brainflow.BUTTERWORTH, 0.0)
println("After LowPass Filter")
println(data_first_channel)

data_second_channel = data[eeg_channels[2], :]
println("Original Data Second Channel")
println(data_second_channel)
brainflow.perform_highpass(data_second_channel, sampling_rate, 5.0, 3, brainflow.CHEBYSHEV_TYPE_1, 1.0)
println("After HighPass Filter")
println(data_second_channel)

data_third_channel = data[eeg_channels[3], :]
println("Original Data Third Channel")
println(data_third_channel)
brainflow.perform_bandpass(data_third_channel, sampling_rate, 25.0, 20.0, 3, brainflow.BESSEL, 0.0)
println("After BandPass Filter")
println(data_third_channel)

data_fourth_channel = data[eeg_channels[4], :]
println("Original Data Fourth Channel")
println(data_fourth_channel)
brainflow.perform_bandstop(data_fourth_channel, sampling_rate, 50.0, 2.0, 3, brainflow.BESSEL, 0.0)
println("After BandStop Filter")
println(data_fourth_channel)


    # todo check if number of cores are the same, check if summing index matches

    import Base: +, -
    import Base: *
    # summation of two MPTs
    function +(mpt1::MPT,mpt2::MPT)
        N        = order(mpt1);
        cores    = Vector{Array{Float64,ndims(mpt1[1])}}(undef,N);
        cores[1] = cat(mpt1[1], mpt2[1],dims=ndims(mpt1[1]));
        cores[N] = cat(mpt1[N], mpt2[N],dims=1);
        sizes = size(mpt1);
        rnks1 = rank(mpt1);
        rnks2 = rank(mpt2);
        for i = 2:N-1
            row1     = cat(mpt1[i], zeros((rnks1[i][1],sizes[i]...,rnks2[i][2])),dims=ndims(mpt1[i]));
            row2     = cat(zeros((rnks2[i][1],sizes[i]...,rnks1[i][2])), mpt2[i],dims=ndims(mpt1[i]));
            cores[i] = cat(row1,row2,dims=1);
        end

        MPT(cores);
    end

    function +(mpo::MPO,mat::Matrix{Float64})
        return mpo2mat(mpo) + mat
    end

    function +(mat::Matrix{Float64},mpo::MPO)
        return +(mpo,mat)
    end

    function -(mpo1::MPO,mpo2::MPO)
        return +(mpo1,-1.0*mpo2)
    end

    function -(mpo::MPO,mat::Matrix{Float64})
        return mpo2mat(mpo) - mat
    end

    function -(mat::Matrix{Float64},mpo::MPO)
        return mat - mpo2mat(mpo)
    end

    # multiplying MPT by a number
    function *(mpt::MPT,a::Float64)
        # if mpt in site-k-canonical format that multiply norm core with scalar
        mpt[1] = a*mpt[1];
        return mpt
        #MPT(pushfirst!(mpt[2:ndims(mpt)],a*mpt[1]))
    end

    function *(a::Float64,mpt::MPT)
        *(mpt,a)
    end

    # matrix by vector product (returns mps)
    function *(mpo::MPT{4},mps::MPT{3})
        N = order(mpo);
        cores = Vector{Array{Float64,3}}(undef,N);
        for i = 1:N
            cores[i] = contractcores(mpo[i],mps[i]);
        end
        MPT(cores)
    end

    # matrix by matrix product (returns mpo)
    function *(mpo1::MPT{4},mpo2::MPT{4})
        N = order(mpo1);
        cores = Vector{Array{Float64,4}}(undef,N);
        for i = 1:N
            cores[i] = contractcores(mpo1[i],mpo2[i]);
        end
        MPT(cores)
    end

    # vector by matrix product !! check if it zip line
    function *(mps::MPT{3},mpo::MPT{4})
        N = order(mpo);
        cores = Vector{Array{Float64,3}}(undef,N);
        for i = 1:N
            cores[i] = contractcores(mps[i],mpo[i]);
        end
        MPT(cores)
    end

    # vector by vector product (returns scalar)
    function *(mps1::MPT{3},mps2::MPT{3})
        N = order(mps1);
        c = 1;
        for i = 1:N
            c = c*contractcores(mps1[i],mps2[i]);
        end
        return c[1]
    end

    function outerprod(tt1::MPT{3},tt2::MPT{3})
        N     = order(tt1);
        cores = Vector{Array{Float64,4}}(undef,N);
        r1 = rank(tt1);
        r2 = rank(tt2);
        s1 = size(tt1);
        s2 = size(tt2);
        for i = 1:N
            g1  = reshape(tt1[i],(length(tt1[i]), 1));
            g2  = reshape(tt2[i],(length(tt2[i]), 1));
            tmp = kron(g1,g2);
            tmp = reshape(tmp,(r1[i][1], s1[i][1], r1[i][2], r2[i][1], s2[i][1], r2[i][2]));
            tmp = permutedims(tmp,[1 4 2 5 3 6]);
            tmp = reshape(tmp,(r1[i][1]*r2[i][1], s1[i][1], s2[i][1], r1[i][2]*r2[i][2]));
            cores[i]  = tmp;
        end
        return MPT(cores);
    end

    function *(core1::Array{Float64,3},core2::Array{Float64,3})
        size1 = size(core1);
        size2 = size(core2);
        reshape(unfold(core1,[3],"right")*unfold(core2,[1],"left"),(size1[1],size1[2]*size2[2], size2[3]))
    end

    function *(cores::Vector{Array{Float64,3}})
        tmp = cores[1]*cores[2]
        for i = 3:length(cores)
            tmp = tmp*cores[i]
        end
        return tmp[1,:,:]
    end

    function *(ttm::MPT{4},A::Matrix,ind::Int) # for now work only if second index is in last core
        d = order(ttm)
        cores = Vector{Array{Float64,4}}(undef,d)
        for i = 1:d-1
            cores[i] = ttm[i]
        end
        #ttm[ind] = nmodeproduct(A,ttm[ind],3)
        tmp = ttm[ind]
        tmp = reshape(tmp,(rank(ttm)[ind][1]*size(ttm)[ind][1],size(ttm)[ind][2]))
        tmp = tmp*A
        cores[ind] = reshape(tmp,(rank(ttm)[ind][1],size(ttm)[ind][1],size(A,2),1))
        return MPT(cores)
    end

    function KathriRao(A::Matrix{Float64},B::Matrix{Float64},dims::Int64)
            if dims == 1
                C = zeros(size(A,1),size(A,2)*size(B,2));
                @inbounds @simd for i = 1:size(A,1)
                    @views kron!(C[i,:],A[i,:],B[i,:])
                end
            elseif dims == 2
                C = zeros(size(A,1)*size(B,1),size(A,2));
                @inbounds @simd for i = 1:size(A,2)
                    @views kron!(C[:,i],A[:,i],B[:,i])
                end
            end

            return C
    end


using Test
using LinearAlgebra
using IntervalArithmetic
using PolynomialBasis
# using Revise
using ImplicitDomainQuadrature

IDQ = ImplicitDomainQuadrature

function allapprox(v1, v2; tol = 1e-14)
    @assert length(v1) == length(v2)
    flags = [isapprox(v1[i], v2[i], atol = tol) for i = 1:length(v1)]
    return all(flags)
end

function circle_distance_function(coords, center, radius)
    npts = size(coords)[2]
    return [norm(coords[:, i] - center) - radius for i = 1:npts]
end

function plane_distance_function(coords, normal, x0)
    return (coords .- x0)' * normal
end

function plot_zero_levelset(poly)
    x = -1:1e-2:1
    contour(x, x, (x, y) -> poly(x, y), levels = [0.0])
end

struct FunctionAndGradient
    func::Any
    grad::Any
    perturbation::Any
end

function (f::FunctionAndGradient)(x)
    return f.func(x) + perturbation
end

function PolynomialBasis.gradient(f::FunctionAndGradient, x)
    return f.grad(x)
end

#######################################################################
# Test a curved interface and subdivision
radius = 0.5
center = [0.0, 0.0]
poly = InterpolatingPolynomial(1, 2, 3)
dpoly = InterpolatingPolynomial(2, 2, 3)
coeffs = circle_distance_function(poly.basis.points, center, radius)
update!(poly, coeffs)
update_interpolating_gradient!(dpoly,poly)
testarea = 4.0 - pi * 0.5^2

xL,xR = [-1.0,-1.0],[1.,1.]

numqp = 5
quad = IDQ.area_quadrature(poly,dpoly, +1, xL,xR, numqp)
@test isapprox(sum(quad.weights), testarea, atol = 1e-1)

squad = IDQ.surface_quadrature(poly,dpoly, xL,xR, numqp)
@test isapprox(sum(quad.weights), pi, atol = 1e-1)


include("wannierize.jl")

filename = "silicon"

read_amn = false
read_eig = false

p = read_system(filename,read_amn,read_eig)

nbeg = 1
nend = p.nwannier

#Set the seed of the pseudo-random generator
#srand(0)

println("$(p.nband) bands, wannierizing bands $nbeg:$nend, $(p.N1) x $(p.N2) x $(p.N3) grid, $(p.nntot) neighbors")

#Build the gauge that makes the Bloch frame continuous on the Brillouin Zone.
#This is equivalent to building a set of algebraic decaying Wannier functions
p, Obs_array, Uint = make_wannier(p)

# #Plot and display the results
# plot_results(p)
# print_error(p)

# figure()
# for i=1:5:p.N2
#     plot(real(Uint[:,i,2,2]),imag(Uint[:,i,2,2]),"x-")
# end
# title("real 22 / imag 22")

# figure()
# for i=1:10:p.N2
#     plot(real(Uint[:,i,1,2]),imag(Uint[:,i,1,2]),"x-")
# end
# title("real 12 / imag 12")

# figure()
# for i=1:5:p.N2
#     plot(real(Uint[:,i,:,2]))
# end

# #using Plots
# #println("Creating animation of the obstruction")
# #anim = @animate for i=1:2*p.N2-1
# #               plot(real(Uint[:,abs(p.N2-i)+1,2,2]),imag(Uint[:,abs(p.N2-i)+1,2,2]),xlims=[-1,1],ylims=[-1,1])
# #end
# #
# #gif(anim,"obs.gif",fps=25)


using BandedMatrices, Test

@testset "QR tests" begin
    for T in (Float64,ComplexF64,Float32,ComplexF32)
        A=brand(T,10,10,3,2)
        Q,R=qr(A)
        @test Matrix(Q)*Matrix(R) ≈ A
        b=rand(T,10)
        @test mul!(similar(b),Q,mul!(similar(b),Q',b)) ≈ b
        for j=1:size(A,2)
            @test Q' * A[:,j] ≈ R[:,j]
        end
        A=brand(T,14,10,3,2)

        Q,R=qr(A)

        for k=1:size(Q,1),j=1:size(Q,2)
            @test Q[k,j] ≈ Matrix(Q)[k,j]
        end

        @test Matrix(Q)*Matrix(R) ≈ A

        A=brand(T,10,14,3,2)

        Q,R=qr(A)

        for k=1:size(Q,1),j=1:size(Q,2)
            @test Q[k,j] ≈ Matrix(Q)[k,j]
        end

        @test Matrix(Q)*Matrix(R) ≈ A
        A=brand(T,100,100,3,4)
        Q,R=qr(A)
        b=rand(T,100)
        @test R\(Q'*b) ≈ Matrix(A)\b
    end

    A=brand(10,10,3,2)
    b=rand(ComplexF64,10)
    Q,R=qr(A)
    @test R\(Q'*b) ≈ Matrix(A)\b

    A=brand(ComplexF64,10,10,3,2)
    b=rand(10)
    Q,R=qr(A)
    @test R\(Q'*b) ≈ Matrix(A)\b
end


# Loads all files
module EffectiveWaves

# Here are the main exported functions and types. Note there are other exported functions and types in files such as "effective_wave/export.jl" and "discrete_wave/export.jl"
export  EffectivePlaneWaveMode, EffectiveRegularWaveMode, ScatteringCoefficientsField

export  DiscretePlaneWaveMode
export  MatchPlaneWaveMode # a combination of the Discrete and Effective modes
export  Material, Specie, Species, SetupSymmetry, number_density, volume_fraction
export  hole_correction_pair_correlation, gls_pair_radial_fun, pair_radial_fun, smooth_pair_corr_distance

# for MatchPlaneWaveMode
export  match_error, x_mesh_match

# for effective waves
export wavenumbers, wavenumber, asymptotic_monopole_wavenumbers, eigenvectors, convert_eigenvector_basis, eigenvector_length
export WaveModes, WaveMode, wavemode_wienerhopf
export solve_boundary_condition, scattering_field, material_scattering_coefficients, material_scattered_waves

export dispersion_equation, dispersion_complex, eigensystem # supplies a matrix used for the disperision equation and effective eignvectors

export scattering_amplitudes_average

export  reflection_coefficient, reflection_coefficients, reflection_transmission_coefficients, planewave_amplitudes
export  effective_medium

# List of shorthand for some materials
export  Brick, IronArmco, LeadAnnealed, RubberGum, FusedSilica, GlassPyrex,
        ClayRock, WaterDistilled, Glycerol, Hexadecane, Acetone, Benzene,
        Nitrobenzene, OliveOil, SodiumNitrate, AirDry,
        LimeStone, Clay, Calcite, SilicaQuartz

import Base.isequal, Base.(==), Base.zero
import SpecialFunctions: besselj, hankelh1
import StaticArrays: SVector
import Statistics: mean, std

import IterTools: groupby

using Reexport
@reexport using MultipleScattering

import MultipleScattering: RegularSource, Acoustic

using RecipesBase, OffsetArrays, LinearAlgebra

# Heavy package
using Optim: optimize, Optim, FixedParameters, Options, LBFGS, Fminbox, NelderMead
# Below doesn't work on Julia 1.1
# @reexport using Optim: optimize, FixedParameters, Options, LBFGS, Fminbox, NelderMead
import Optim: simplexer

# Heavy package
# using ApproxFun: ApproxFun.(..), Fun, Segment, Domain, Chebyshev, DefiniteIntegral, LowRankFun, Interval

using Interpolations: interpolate, BSpline, Cubic, Line, OnGrid, scale
using HCubature: hcubature, hquadrature
using ClassicalOrthogonalPolynomials: Legendre

using MDBM

include("material_types.jl")
include("statistics.jl")
include("wave_types.jl")
include("specialfunctions.jl")

include("effective_wave/export.jl")
include("acoustics/export.jl")

include("match_waves/match_waves.jl")
include("match_waves/match_arrays.jl")
include("match_waves/reflection.jl")

include("../plot/graphics.jl")
include("../plot/plot.jl")

end # module


immutable BilinearSpline <: AbstractInterpolation
    BilinearSpline() = BilinearSpline
end

immutable BilinearSplineInterpolation{T}
    x::Vector{T}
    y::Vector{T}
    z::Array{T,2}
end

function default_settings{D,T}(::Type{BilinearSpline},axes::NTuple{D,T})
    nx = length(axes[1])
    ny = length(axes[2])
    return Irregular{2,Val{(nx,ny)}}, Nearest
end

function BilinearSpline(x::AbstractVector, y::AbstractVector, z::AbstractArray)
  issorted(x) || throw(ArgumentError("x points must be in ascending order"))
  issorted(y) || throw(ArgumentError("y points must be in ascending order"))
  x = convert(Array{Float64,1},x)
  y = convert(Array{Float64,1},y)
  z = convert(Array{Float64,2},z)

  return BilinearSplineInterpolation{eltype(x)}(x,y,z)
end

@inline function preprocess(::Type{BilinearSpline}, axes, f)
    return BilinearSpline(axes...,f)
end

function apply_bc(::Type{BilinearSpline},BC::Type{Nearest}, I::Symbol, x::Symbol)
    xbc = apply_nearest(I,1, x, 1)
    ybc = apply_nearest(I,2, x, 2)
    quote
        $(xbc)
        $(ybc)
    end
end

function apply_bc(::Type{BilinearSpline},BC::Type{Reflect}, I::Symbol, x::Symbol)
    xbc = apply_reflect(I,1,x,1)
    ybc = apply_reflect(I,2,x,2)
    quote
        $(xbc)
        $(ybc)
    end
end

function apply_bc{V}(::Type{BilinearSpline},BC::Type{Value{V}}, I::Symbol, x::Symbol)
    xbc = apply_value(BC,I,1,x,1)
    ybc = apply_value(BC,I,2,x,2)
    quote
        $(xbc)
        $(ybc)
    end
end

function apply_bc(::Type{BilinearSpline},BC::Type{None}, I::Symbol, x::Symbol)
    apply_none()
end

function apply_bc(::Type{BilinearSpline},BC::Type{Error}, I::Symbol, x::Symbol)
    xbc = apply_error(I,1,x,1)
    ybc = apply_error(I,2,x,2)
    quote
        $(xbc)
        $(ybc)
    end
end

@inline function get_indices{S}(::Type{BilinearSpline}, G::Type{Uniform{2,S}}, I, x)
    dx = I.x[2]-I.x[1]
    ix = clamp(Int(floor((x[1] - I.x[1])/dx)) + 1, 1, S.parameters[1][1]-1)
    dy = I.y[2]-I.y[1]
    iy = clamp(Int(floor((x[2] - I.y[1])/dy)) + 1, 1, S.parameters[1][2]-1)
    return (ix, iy)
end

@inline function get_indices{S}(::Type{BilinearSpline}, G::Type{Irregular{2,S}}, I, x)
    xr = searchsortedlast(I.x,x[1],Base.Order.Forward)
    ix = clamp(xr, 1, S.parameters[1][1]-1)
    yr = searchsortedlast(I.y,x[2],Base.Order.Forward)
    iy = clamp(yr, 1, S.parameters[1][2]-1)
    return (ix, iy)
end

@inline function get_value(::Type{BilinearSpline}, I, x, inds)
    ix, iy = inds
    xp = x[1]
    yp = x[2]
    x1, x2 = I.x[ix], I.x[ix+1]
    y1, y2 = I.y[iy], I.y[iy+1]
    dxdy = (x2-x1)*(y2-y1)

    b11 = (x2 - xp)*(y2 - yp)
    b21 = (xp - x1)*(y2 - yp)
    b12 = (x2 - xp)*(yp - y1)
    b22 = (xp - x1)*(yp - y1)
    v = b11*I.z[ix,iy] + b12*I.z[ix,iy+1] + b21*I.z[ix+1,iy] + b22*I.z[ix+1,iy+1]
    return v/dxdy
end


using Sudoku

function main()
    #golden nugget
    puzzle = [0 0 0 0 0 0 0 3 9;
              0 0 0 0 1 0 0 0 5;
              0 0 3 0 0 5 8 0 0;
              0 0 8 0 0 9 0 0 6;
              0 7 0 0 2 0 0 0 0;
              1 0 0 4 0 0 0 0 0;
              0 0 9 0 0 8 0 5 0;
              0 2 0 0 0 0 6 0 0;
              4 0 0 7 0 0 0 0 0;]

    println("Solving:")
    println("Golden Nugget")
    display(puzzle)
    println("")

    @time solution = Sudoku.mc_solve(puzzle, 1500)

    println("Initial board:")
    display(puzzle)
    println("")
    println("Solution, energy=$(solution.energy)")
    display(solution.puzzle)
    println("")
end

main()


using Knet
using Knet: Net, params, axpy!
isdefined(:MNIST) || include("mnist.jl")

@knet function model01(x; winit=Gaussian(0,.1), hidden=32, f1=:relu, o...)
    h    = wf(x; out=hidden, f=f1, winit=winit)
    return wf(h; out=10, f=:soft, winit=winit)
end

function train01(; epochs=1000, lr=0.5, adam=false, batchsize=100, seed=0, target=nothing, o...)
    seed > 0 && setseed(seed)
    global f = compile(:model01; o...)
    global d = minibatch(MNIST.xtrn, MNIST.ytrn, batchsize)
    setp(f, lr=lr, adam=adam)
    global w0,w1,w2,dw,w2_w0,w2_w1
    w0 = w1 = w2 = nothing
    for epoch=1:epochs
        zloss = sloss = iter = 0
        for (x,ygold) in d
            ypred = forw(f, x)
            back(f, ygold, softloss)

            if target != nothing
                # we start at w0 and want to reach w2
                # currently we are at w1 and about to take the step dw
                # compare dw and w2-w1
                # compare w2-w1 and w2-w0
                w1 = weights(f)
                dw = gradients(f)
                if w0 == nothing
                    w0 = deepcopy(w1)
                    w2 = weights(target)
                    w2_w0 = deepcopy(w2)
                    w2_w1 = deepcopy(w2)
                    for i=1:length(w2); axpy!(-1,w0[i],w2_w0[i]); end
                end
                for i=1:length(w2); copy!(w2_w1[i],w2[i]); axpy!(-1,w1[i],w2_w1[i]); end
                @printf("w2-w0=%g,1 w2-w1=%g,%g dw=%g,%g\n",
                        vnorm(w2_w0), vnorm(w2_w1), vcos(w2_w1, w2_w0),
                        lr*vnorm(dw), -vcos(dw, w2_w1))
            end

            update!(f)
            zloss += zeroone(ypred, ygold)
            sloss += softloss(ypred, ygold)
            iter += 1
        end
        println((epoch, sloss/iter, zloss/iter))
        zloss == 0 && break
    end
    return f
end

weights(f::Net)=map(r->r.out, params(f))
gradients(f::Net)=map(r->r.dif, params(f))

function vnorm(w)
    sqnorm = 0
    for u in w; sqnorm += vecnorm(u)^2; end
    sqrt(sqnorm)
end

function vcos(u,v)
    uv = 0
    for i=1:length(u); uv += CUBLAS.dot(u[i],v[i]); end
    uv / (vnorm(u)*vnorm(v))
end

function test01(f, data, loss)
    sumloss = numloss = 0
    for (x,ygold) in data
        ypred = forw(f, x)
        sumloss += loss(ypred, ygold)
        numloss += 1
    end
    sumloss / numloss
end

f1 = train01(seed=1)
f2 = train01(seed=1, target=f1)

# Can get to zloss=0 with hidden=24 using adam=true lr=0.001


using ComponentArrays
using OrdinaryDiffEq
using LinearAlgebra
using Unitful

function newton(du, u, p, t)
    mu = 398600.4418u"km^3/s^2"
    r = norm(u.r)
    du.r = u.v
    du.v = -mu .* u.r / r^3
end

r0 = [1131.340, -2282.343, 6672.423]u"km"
v0 = [-5.64305, 4.30333, 2.42879]u"km/s"
Δt = 86400.0*365u"s"
rv0 = ComponentArray(r=r0, v=v0)

prob = ODEProblem(newton, rv0, (0.0u"s", Δt))
@time sol = solve(prob, Vern8(), dt=100u"s", adaptive=false)
#sol |> display
println()

@time sol2 = solve(prob, Tsit5(), dt=100u"s",adaptive=false)

sol2.u[1].r |> display
sol2.u[1].v |> display


if !isdefined(:ExoplanetsSysSim) using ExoplanetsSysSim end
using Distributions
using Roots #if we want to use the 'find_zero()' function
using Optim

##### This code is a translation of 'MRpredict.R' at 'https://github.com/Bo-Ning/Predicting-exoplanet-mass-and-radius-relationship/blob/master/MR-predict/MRpredict.R'

struct MR_param_Ning2018
    Mass_min::Float64 #log10 scale
    Mass_max::Float64 #log10 scale
    Radius_min::Float64 #log10 scale
    Radius_max::Float64 #log10 scale
    deg::Int64
    weights_mle::Vector{Float64}
    indices_keep::Vector{Int64}
    #indices_div::Vector{Int64}
    #indices_rem::Vector{Int64}
end

function pdfnorm_beta(x::Float64, x_obs::Float64, x_sd::Float64, x_max::Float64, x_min::Float64, shape1::Int64, shape2::Int64 ; log::Bool=true)
    if log == true
        norm_dist = Normal(10^x, x_sd)
    else
        norm_dist = Normal(x, x_sd)
    end
    #beta_dist = Beta(shape1, shape2)
    return pdf.(norm_dist, x_obs) .* pdf.(Beta(shape1, shape2), (x-x_min)/(x_max-x_min)) / (x_max-x_min)
end

function fn_for_integrate(x_obs::Float64, x_sd::Float64, deg::Int, degree::Int, x_max::Float64, x_min::Float64 ; log::Bool=false, abs_tol::Float64=1e-10)
    shape1, shape2 = degree, deg-degree+1
    return quadgk(x -> pdfnorm_beta(x, x_obs, x_sd, x_max, x_min, shape1, shape2; log=log), x_min, x_max, abstol=abs_tol)[1]
end

function conditional_density(y::Float64, y_max::Float64, y_min::Float64, x_max::Float64, x_min::Float64, deg::Int64, w_hat::Vector{Float64}, indices_keep::Vector{Int64}; y_sd::Union{Void,Float64}=nothing, qtl::Vector{Float64}=[0.16, 0.84])
    deg_vec = 1:deg

    #To compute conditional mean, variance, quantile, distribution:
    y_beta_indv = zeros(deg)
    if y_sd == nothing
        #y_stdardize = (y-y_min)/(y_max-y_min)
        #y_beta_indv = map(degree -> pdf.(Beta(degree, deg-degree+1), (y-y_min)/(y_max-y_min)), deg_vec)::Array{Float64,1} / (y_max-y_min)
        map!(degree -> pdf.(Beta(degree, deg-degree+1), (y-y_min)/(y_max-y_min)), y_beta_indv, deg_vec)
        y_beta_indv /= (y_max-y_min)
    else
        #y_beta_indv = map(degree -> fn_for_integrate(y, y_sd, deg, degree, y_max, y_min), deg_vec)::Array{Float64,1}
        map!(degree -> fn_for_integrate(y, y_sd, deg, degree, y_max, y_min), y_beta_indv, deg_vec)
    end
    y_beta_pdf = kron(ones(deg), y_beta_indv)[indices_keep]
    ###y_beta_pdf = y_beta_indv[indices_rem]
    denominator = sum(w_hat .* y_beta_pdf::Vector{Float64})

    #Mean:
    mean_beta_indv = deg_vec/(deg+1)*(x_max-x_min)+x_min
    mean_beta = kron(mean_beta_indv, y_beta_indv)[indices_keep]
    ###mean_beta = mean_beta_indv[indices_div] .* y_beta_indv[indices_rem]
    mean_numerator = sum(w_hat .* mean_beta::Vector{Float64})
    mean = mean_numerator / denominator
    #mean = sum(w_hat .* kron(deg_vec/(deg+1)*(x_max-x_min)+x_min, y_beta_indv)) / denominator

    #Variance:
    var_beta_indv = deg_vec .* (deg-deg_vec+1) / ((deg+1)^2*(deg+2))*(x_max-x_min)^2
    var_beta = kron(var_beta_indv, y_beta_indv)[indices_keep]
    ###var_beta = var_beta_indv[indices_div] .* y_beta_indv[indices_rem]
    var_numerator = sum(w_hat .* var_beta::Vector{Float64})
    var = var_numerator /denominator
    #var = sum(w_hat .* kron(deg_vec .* (deg-deg_vec+1) / ((deg+1)^2*(deg+2))*(x_max-x_min)^2, y_beta_indv)) / denominator

    #Quantile:
    function pbeta_conditional_density(x::Float64)::Float64
        function mix_density(j::Float64)::Float64
            #x_indv_cdf::Vector{Float64} = map(degree -> cdf.(Beta(degree, deg-degree+1), (j-x_min)/(x_max-x_min)), deg_vec)
            ###quantile_beta = x_indv_cdf[indices_div] .* y_beta_indv[indices_rem]
            #quantile_numerator = sum(w_hat .* kron(x_indv_cdf, y_beta_indv)::Vector{Float64})
            #return quantile_numerator/denominator
            return sum(w_hat .* kron(map(degree -> cdf.(Beta(degree, deg-degree+1), (j-x_min)/(x_max-x_min)), deg_vec), y_beta_indv)[indices_keep]::Vector{Float64}) / denominator
        end

        return mix_density(x)
    end

    function conditional_quantile(q::Float64)::Float64
        function g(x::Float64)::Float64
            return pbeta_conditional_density(x) - q
        end
        function g2(x::Float64)::Float64
            return (pbeta_conditional_density(x) - q)^2
        end

        #return find_zero(g, (x_min, x_max), FalsePosition())
        #return find_zero(g, (x_min, x_max))::Float64 ###'find_zero' is from 'using Roots' (is this a good function to use?); Eric suggested brent's method in the Optim package

        result = optimize(g2, x_min, x_max, Brent(), rel_tol=1e-3, abs_tol=1e-3) #Brent(); GoldenSection()
        if Optim.converged(result)
            return Optim.minimizer(result)::Float64
        else
            return find_zero(g, (x_min, x_max))::Float64
        end
    end

    quantile = map(conditional_quantile, qtl)

    return (mean, var, quantile, denominator, y_beta_indv)
end

function predict_mass_given_radius(Radius::Float64, param::MR_param_Ning2018 ; R_sigma::Union{Void,Float64}=nothing, posterior_sample::Bool=false, qtl::Vector{Float64}=[0.16, 0.84])
    #Convert data to log scale:
    l_radius = log10.(Radius)

    #The 'posterior_sample == false' condition can deal with two cases:
    #Case I: if input data do not have measurement errors
    #Case II: if input data have measurement error
    if posterior_sample == false
        predicted_value = conditional_density(l_radius, param.Radius_max, param.Radius_min, param.Mass_max, param.Mass_min, param.deg, param.weights_mle, param.indices_keep; y_sd=R_sigma, qtl=qtl)
        predicted_mean = predicted_value[1]
        predicted_quantiles = predicted_value[3]
        #predicted_lower_quantile = predicted_value[3][1]
        #predicted_upper_quantile = predicted_value[3][2]
    elseif posterior_sample == true
        #Case III: if the input are posterior samples
        radius_sample = log10(Radius)
        k = length(radius_sample)
        denominator_sample = zeros(k)
        mean_sample = zeros(k)
        y_beta_indv_sample = zeros(k, param.deg)

        #Calculate mean:
        for i in 1:k
            results = cond_density_estimation(y=radius_sample[i], y_max=param.Radius_max, y_min=param.Radius_min, x_max=param.Mass_max, x_min=param.Mass_min, deg=param.deg, w_hat=param.weights_mle, qtl=quantile, only_output_mean=true) ###what function is this??? Is it part of some CRAN package...
            mean_sample[i] = (results[1]) ###why are there parentheses here? It makes no difference...
            denominator_sample[i] = results[2]
            y_beta_indv_sample[i,:] = results[3:57]
        end
        predicted_mean = mean(mean_sample)

        #Calculate 16% and 84% quantiles:
        #Mixture of the CDF of k conditional densities
        function pbeta_conditional_density(x)
            function mix_density(j)
                deg_vec = 1:param.deg
                x_indv_cdf = map(degree -> cdf.(Beta(degree, param.deg-degree+1), (j-x_min)/(x_max-x_min)), deg_vec)
                quantile_numerator = zeros(k)
                p_beta_sample = zeros(k)
                for ii in 1:k
                    quantile_numerator[ii] = sum(param.weights_mle .* kron(x_indv_cdf, y_beta_indv_sample[ii,:]))
                    p_beta_sample[ii] = quantile_numerator[ii]/denominator_sample[ii]
                end
                return p_beta = mean(p_beta_sample)
            end

            return map(mix_density, x)
        end

        function mixture_conditional_quantile(q, x_min, x_max)
            function g(x)
                return pbeta_conditional_density(x) - q
            end
            return find_zero(g, (x_min, x_max)) ###see note above
            #return Optim.minimum(optimize(g, x_min, x_max))
        end

        predicted_quantiles = map(q -> mixture_conditional_quantile(q, param.Mass_min, param.Mass_max), qtl)
        #predicted_lower_quantile = predicted_quantiles[1]
        #predicted_upper_quantile = predicted_quantiles[2]
    end

    #Return the output:
return (predicted_mean, predicted_quantiles)::Tuple{Float64, Vector{Float64}}
end

function draw_planet_mass_from_radius_Ning2018(Radius::Float64, param::MR_param_Ning2018)
    #This function takes in a Radius (in solar radii) and draws a mass (returning in solar masses) probabilistically
    Radius_in_earths = Radius/ExoplanetsSysSim.earth_radius
    @assert param.Radius_min < log10.(Radius_in_earths) < param.Radius_max
    q::Float64 = rand()
    l_mass = predict_mass_given_radius(Radius_in_earths, param; qtl=[q])[2][1]
    return (10^l_mass)*ExoplanetsSysSim.earth_mass
end

function generate_planet_mass_from_radius_Ning2018(Radius::Float64, sim_param::SimParam)
    global MR_param
    return draw_planet_mass_from_radius_Ning2018(Radius, MR_param)
end

##### To initialize the M-R model parameters:

weights_mle = CSV.read("weights.mle.csv")[:x]
degrees = 55

N_keep = 25
indices_keep = sortperm(weights_mle, rev=true)[1:N_keep]
#indices_div = 1+div.(indices_keep .- 1, degrees)
#indices_rem = 1+rem.(indices_keep .- 1, degrees)
weights = weights_mle[indices_keep]
MR_param = MR_param_Ning2018(-1., 3.809597, -0.302, 1.357509, degrees, weights, indices_keep) #(-1., 3.809597, -0.3, 1.357509, 55, weights_mle)

##### Examples:

#=
#Observation without measurement errors:

Radius = 5. #original scale, not log scale
predict_result = predict_mass_given_radius(Radius, MR_param)
println(predict_result)

#Observation with a measurement error:

Radius = 5. #original scale, not log scale
R_sigma = 0.1
predict_result = predict_mass_given_radius(Radius, MR_param; R_sigma=0.1)
println(predict_result)

#If want to change 16% and 84% quantiles to 5% and 95% quantiles:

Radius = 5. #original scale, not log scale
predict_result = predict_mass_given_radius(Radius, MR_param; qtl=[0.05, 0.95])
println(predict_result)
=#

#Input are posterior samples: ###currently broken because the function 'cond_density_estimation' is undefined
#=
Radius_dist = Normal(5, 0.5) #original scale, not log scale
Radius_samples = rand(Radius_dist, 100)
predict_result = predict_mass_given_radius(Radius_samples, MR_param; R_sigma=0.1, posterior_sample=true)
println(predict_result)
=#

##### For timing the function:
#=
sim_param = setup_sim_param_model()
#Radii = ones(10000)*ExoplanetsSysSim.earth_radius
Radii = (10.^(linspace(MR_param.Radius_min+0.01, MR_param.Radius_max-0.01, 10000)))*ExoplanetsSysSim.earth_radius
tic()
Masses = map(r -> generate_planet_mass_from_radius_Ning2018(r, sim_param), Radii)
t_elapsed = toc()
=#

##### For plotting and timing the model:
#=
#Pkg.add("PyPlot")
using PyPlot

N_keep_array = [3025, 25, 15]
colors = ["blue", "green", "red"]
Radii = 10.^(linspace(MR_param.Radius_min+0.01, MR_param.Radius_max-0.01, 1000))

fig, ax = subplots()
for (i,N_keep) in enumerate(N_keep_array)
    indices_keep = sortperm(weights_mle, rev=true)[1:N_keep]
    weights = weights_mle[indices_keep]
    MR_param = MR_param_Ning2018(-1., 3.809597, -0.302, 1.357509, degrees, weights, indices_keep)

    tic()
    Masses_predict = map(r -> predict_mass_given_radius(r, MR_param; qtl=[0.16, 0.84]), Radii)
    t_elapsed = toc()

    Masses_mean = 10.^[x[1] for x in Masses_predict]
    Masses_lower = 10.^[x[2][1] for x in Masses_predict]
    Masses_upper = 10.^[x[2][2] for x in Masses_predict]

    loglog(Radii, Masses_mean, color=colors[i], linewidth=2.0, linestyle="-", label="M = "*string(N_keep))
    loglog(Radii, Masses_lower, color=colors[i], linewidth=1.0, linestyle="--")
    loglog(Radii, Masses_upper, color=colors[i], linewidth=1.0, linestyle="--")
end
xlabel(L"$R_p (R_\oplus)$")
ylabel(L"$M_p (M_\oplus)$")
legend(loc="upper left")
=#


# rounding.jl

function Base.floor(qdt::QDate, p::Dates.Year)
    value(p) < 1 && throw(DomainError())
    years = year(qdt)
    return QDate(years - mod(years, value(p)))
end

function Base.floor(qdt::QDate, p::Dates.Month)
    value(p) < 1 && throw(DomainError())
    y, m, l = yearmonthleap(qdt)
    qdt0 = QDate(y, m, l, 1)
    if value(p) == 1
        return qdt0
    end
    months = round(Int, (value(qdt0) - FIRST_VALUE) / _DAYS_IN_MONTH_F)
    return qdt0 - Dates.Month(mod(months, value(p)))
end

function Base.floor(qdt::QDate, p::Day)
    value(p) < 1 && throw(DomainError())
    days = value(qdt) - FIRST_VALUE
    return QDate(UTD(days - mod(days, value(p)) + FIRST_VALUE))
end


module SurfaceQG

export
  Problem,
  set_b!,
  updatevars!,

  kinetic_energy,
  buoyancy_variance,
  buoyancy_dissipation,
  buoyancy_work,
  buoyancy_advection,
  kinetic_energy_advection

  using
    CUDA,
    Reexport,
    DocStringExtensions

@reexport using FourierFlows

using LinearAlgebra: mul!, ldiv!
using FourierFlows: parsevalsum

nothingfunction(args...) = nothing

"""
    Problem(dev::Device=CPU();
                 nx = 256,
                 Lx = 2π,
                 ny = nx,
                 Ly = Lx,
                  ν = 0,
                 nν = 1,
                 dt = 0.01,
            stepper = "RK4",
              calcF = nothingfunction,
         stochastic = false,
  aliased_fraction = 1/3,
                  T = Float64)

Construct a surface quasi-geostrophic `problem` on device `dev`.

Keyword arguments
=================
  - `dev`: (required) `CPU()` or `GPU()`; computer architecture used to time-step `problem`.
  - `nx`: Number of grid points in ``x``-domain.
  - `ny`: Number of grid points in ``y``-domain.
  - `Lx`: Extent of the ``x``-domain.
  - `Ly`: Extent of the ``y``-domain.
  - `ν`: Small-scale (hyper)-viscosity coefficient.
  - `nν`: (Hyper)-viscosity order, `nν```≥ 1``.
  - `dt`: Time-step.
  - `stepper`: Time-stepping method.
  - `calcF`: Function that calculates the Fourier transform of the forcing, ``F̂``.
  - `stochastic`: `true` or `false`; boolean denoting whether `calcF` is temporally stochastic.
  - `aliased_fraction`: the fraction of high-wavenumbers that are zero-ed out by `dealias!()`.
  - `T`: `Float32` or `Float64`; floating point type used for `problem` data.
"""
function Problem(dev::Device=CPU();
  # Numerical parameters
                nx = 256,
                ny = nx,
                Lx = 2π,
                Ly = Lx,
  # Hyper-viscosity parameters
                 ν = 0,
                nν = 1,
  # Timestepper and equation options
                dt = 0.01,
           stepper = "RK4",
             calcF = nothingfunction,
        stochastic = false,
  # Float type and dealiasing
  aliased_fraction = 1/3,
                 T = Float64)

  grid = TwoDGrid(dev, nx, Lx, ny, Ly; aliased_fraction=aliased_fraction, T=T)

  params = Params{T}(ν, nν, calcF)

  vars = calcF == nothingfunction ? DecayingVars(dev, grid) : (stochastic ? StochasticForcedVars(dev, grid) : ForcedVars(dev, grid))

  equation = Equation(params, grid)

  return FourierFlows.Problem(equation, stepper, dt, grid, vars, params, dev)
end

# ----------
# Parameters
# ----------

"""
    Params{T}(ν, nν, calcF!)

A struct containing the parameters for Surface QG dynamics. Included are:

$(TYPEDFIELDS)
"""
struct Params{T} <: AbstractParams
    "buoyancy (hyper)-viscosity coefficient"
       ν :: T
    "buoyancy (hyper)-viscosity order"
      nν :: Int
    "function that calculates the Fourier transform of the forcing, ``F̂``"
  calcF! :: Function
end

Params(ν, nν) = Params(ν, nν, nothingfunction)

# ---------
# Equations
# ---------

"""
    Equation(params, grid)

Return the `equation` for surface QG dynamics with `params` and `grid`. The linear
opeartor ``L`` includes (hyper)-viscosity of order ``n_ν`` with coefficient ``ν``,

```math
L = - ν |𝐤|^{2 n_ν} .
```

Plain old viscocity corresponds to ``n_ν=1``.

The nonlinear term is computed via function `calcN!()`.
"""
function Equation(params::Params, grid::AbstractGrid)
  L = @. - params.ν * grid.Krsq^params.nν
  CUDA.@allowscalar L[1, 1] = 0

  return FourierFlows.Equation(L, calcN!, grid)
end

# ----
# Vars
# ----

abstract type SurfaceQGVars <: AbstractVars end

"""
    Vars{Aphys, Atrans, F, P}(b, u, v, bh, uh, vh, Fh, prevsol)

The variables for surface QG problem:

$(FIELDS)
"""
struct Vars{Aphys, Atrans, F, P} <: SurfaceQGVars
    "buoyancy"
        b :: Aphys
    "x-component of velocity"
        u :: Aphys
    "y-component of velocity"
        v :: Aphys
    "Fourier transform of buoyancy"
       bh :: Atrans
    "Fourier transform of x-component of velocity"
       uh :: Atrans
    "Fourier transform of y-component of velocity"
       vh :: Atrans
    "Fourier transform of forcing"
       Fh :: F
    "`sol` at previous time-step"
  prevsol :: P
end

const DecayingVars = Vars{<:AbstractArray, <:AbstractArray, Nothing, Nothing}
const ForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, Nothing}
const StochasticForcedVars = Vars{<:AbstractArray, <:AbstractArray, <:AbstractArray, <:AbstractArray}

"""
    DecayingVars(dev, grid)

Return the `vars` for unforced surface QG dynamics on device `dev` and with `grid`.
"""
function DecayingVars(::Dev, grid::AbstractGrid) where Dev
  T = eltype(grid)
  @devzeros Dev T (grid.nx, grid.ny) b u v
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) bh uh vh

  return Vars(b, u, v, bh, uh, vh, nothing, nothing)
end

"""
    ForcedVars(dev, grid)

Return the vars for forced surface QG dynamics on device `dev` and with `grid`.
"""
function ForcedVars(dev::Dev, grid) where Dev
  T = eltype(grid)
  @devzeros Dev T (grid.nx, grid.ny) b u v
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) bh uh vh Fh

  return Vars(b, u, v, bh, uh, vh, Fh, nothing)
end

"""
    StochasticForcedVars(dev, grid)

Return the `vars` for stochastically forced surface QG dynamics on device `dev` and with `grid`.
"""
function StochasticForcedVars(dev::Dev, grid) where Dev
  T = eltype(grid)

  @devzeros Dev T (grid.nx, grid.ny) b u v
  @devzeros Dev Complex{T} (grid.nkr, grid.nl) bh uh vh Fh prevsol

  return Vars(b, u, v, bh, uh, vh, Fh, prevsol)
end

# -------
# Solvers
# -------

"""
    calcN_advection(N, sol, t, clock, vars, params, grid)

Calculate the Fourier transform of the advection term, ``- 𝖩(ψ, b)`` in conservative
form, i.e., ``- ∂_x[(∂_y ψ)b] - ∂_y[(∂_x ψ)b]`` and store it in `N`:

```math
N = - \\widehat{𝖩(ψ, b)} = - i k_x \\widehat{u b} - i k_y \\widehat{v b} .
```
"""
function calcN_advection!(N, sol, t, clock, vars, params, grid)
  @. vars.bh = sol
  @. vars.uh =   im * grid.l  * sqrt(grid.invKrsq) * sol
  @. vars.vh = - im * grid.kr * sqrt(grid.invKrsq) * sol

  ldiv!(vars.u, grid.rfftplan, vars.uh)
  ldiv!(vars.v, grid.rfftplan, vars.vh)
  ldiv!(vars.b, grid.rfftplan, vars.bh)

  ub, ubh = vars.u, vars.uh         # use vars.u, vars.uh as scratch variables
  vb, vbh = vars.v, vars.vh         # use vars.v, vars.vh as scratch variables

  @. ub *= vars.b                   # u*b
  @. vb *= vars.b                   # v*b

  mul!(ubh, grid.rfftplan, ub)      # \hat{u*b}
  mul!(vbh, grid.rfftplan, vb)      # \hat{v*b}

  @. N = - im * grid.kr * ubh - im * grid.l * vbh

  return nothing
end

"""
    calcN!(N, sol, t, clock, vars, params, grid)

Calculate the nonlinear term, that is the advection term and the forcing,

```math
N = - \\widehat{𝖩(ψ, b)} + F̂ .
```
"""
function calcN!(N, sol, t, clock, vars, params, grid)
  dealias!(sol, grid)

  calcN_advection!(N, sol, t, clock, vars, params, grid)

  addforcing!(N, sol, t, clock, vars, params, grid)

  return nothing
end

"""
    addforcing!(N, sol, t, clock, vars, params, grid)

When the problem includes forcing, calculate the forcing term ``F̂`` and add it to the
nonlinear term ``N``.
"""
addforcing!(N, sol, t, clock, vars::Vars, params, grid) = nothing

function addforcing!(N, sol, t, clock, vars::ForcedVars, params, grid)
  params.calcF!(vars.Fh, sol, t, clock, vars, params, grid)
  @. N += vars.Fh

  return nothing
end

function addforcing!(N, sol, t, clock, vars::StochasticForcedVars, params, grid)
  if t == clock.t # not a substep
    @. vars.prevsol = sol # sol at previous time-step is needed to compute budgets for stochastic forcing
    params.calcF!(vars.Fh, sol, t, clock, vars, params, grid)
  end

  @. N += vars.Fh

  return nothing
end

# ----------------
# Helper functions
# ----------------

"""
    updatevars!(prob)

Update variables in `vars` with solution in `sol`.
"""
function updatevars!(prob)
  vars, grid, sol = prob.vars, prob.grid, prob.sol

  dealias!(sol, grid)

  @. vars.bh = sol
  @. vars.uh =   im * grid.l  * sqrt(grid.invKrsq) * sol
  @. vars.vh = - im * grid.kr * sqrt(grid.invKrsq) * sol

  ldiv!(vars.b, grid.rfftplan, deepcopy(vars.bh))
  ldiv!(vars.u, grid.rfftplan, deepcopy(vars.uh))
  ldiv!(vars.v, grid.rfftplan, deepcopy(vars.vh))

  return nothing
end

"""
    set_b!(prob, b)

Set the solution `sol` as the transform of `b` and update all variables.
"""
function set_b!(prob, b)
  mul!(prob.sol, prob.grid.rfftplan, b)
  CUDA.@allowscalar prob.sol[1, 1] = 0 # zero domain average

  updatevars!(prob)

  return nothing
end

"""
    kinetic_energy(prob)

Return the domain-averaged surface kinetic energy. Since ``u² + v² = |{\\bf ∇} ψ|²``, we get
```math
\\int \\frac1{2} |{\\bf ∇} ψ|² \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} \\frac1{2} |𝐤|² |ψ̂|² .
```
In SQG, this is identical to half the domain-averaged surface buoyancy variance.
"""
@inline function kinetic_energy(prob)
  sol, vars, grid = prob.sol, prob.vars, prob.grid

  ψh = vars.uh                     # use vars.uh as scratch variable
  kinetic_energyh = vars.bh        # use vars.bh as scratch variable

  @. ψh = sqrt(grid.invKrsq) * sol
  @. kinetic_energyh = 1 / 2 * grid.Krsq * abs2(ψh)

  return 1 / (grid.Lx * grid.Ly) * parsevalsum(kinetic_energyh, grid)
end

"""
    buoyancy_variance(prob)

Return the buoyancy variance,
```math
\\int b² \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} |b̂|² .
```
In SQG, this is identical to the velocity variance (i.e., twice the domain-averaged kinetic
energy).
"""
@inline function buoyancy_variance(prob)
  sol, grid = prob.sol, prob.grid

  return 1 / (grid.Lx * grid.Ly) * parsevalsum(abs2.(sol), grid)
end

"""
    buoyancy_dissipation(prob)

Return the domain-averaged dissipation rate of surface buoyancy variance due
to small scale (hyper)-viscosity,
```math
2 ν (-1)^{n_ν} \\int b ∇^{2n_ν} b \\frac{𝖽x 𝖽y}{L_x L_y} = - 2 ν \\sum_{𝐤} |𝐤|^{2n_ν} |b̂|² ,
```
where ``ν`` the (hyper)-viscosity coefficient ``ν`` and ``nν`` the (hyper)-viscosity order.
In SQG, this is identical to twice the rate of kinetic energy dissipation
"""
@inline function buoyancy_dissipation(prob)
  sol, vars, params, grid = prob.sol, prob.vars, prob.params, prob.grid
  buoyancy_dissipationh = vars.uh         # use vars.uh as scratch variable

  @. buoyancy_dissipationh = 2 * params.ν * grid.Krsq^params.nν * abs2(sol)
  CUDA.@allowscalar buoyancy_dissipationh[1, 1] = 0

  return 1 / (grid.Lx * grid.Ly) * parsevalsum(buoyancy_dissipationh, grid)
end

"""
    buoyancy_work(prob)
    buoyancy_work(sol, vars, grid)

Return the domain-averaged rate of work of buoyancy variance by the forcing,
```math
\\int 2 b F \\frac{𝖽x 𝖽y}{L_x L_y} = \\sum_{𝐤} 2 b̂ F̂^* .
```
"""
@inline function buoyancy_work(sol, vars::ForcedVars, grid)
  buoyancy_workh = vars.uh         # use vars.uh as scratch variable

  @. buoyancy_workh =  2 * sol * conj(vars.Fh)    # 2*b̂*conj(f̂)
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(buoyancy_workh, grid)
end

@inline function buoyancy_work(sol, vars::StochasticForcedVars, grid)
  buoyancy_workh = vars.uh         # use vars.uh as scratch variable

  @. buoyancy_workh =  (vars.prevsol + sol) * conj(vars.Fh) # Stratonovich
  return 1 / (grid.Lx * grid.Ly) * parsevalsum(buoyancy_workh, grid)
end

@inline buoyancy_work(prob) = buoyancy_work(prob.sol, prob.vars, prob.grid)

end # module


using ApproxFun, IntervalSets, SpecialFunctions, LinearAlgebra, Random, Test
    import ApproxFun: HeavisideSpace, PointSpace, DiracSpace, PiecewiseSegment

@testset "Singularities" begin
    @testset "sqrt" begin
        x=Fun(identity);
        @test sqrt(cos(π/2*x))(.1) ≈ sqrt(cos(.1π/2))

        x=Fun(identity,-2..2)
        u=sqrt(4-x^2)/(2π)

        @test u(.1) ≈ sqrt(4-0.1^2)/(2π)
        @test sum(u) ≈ 1

        #this call threw an error, which we check
        @test length(values(u)) == 1

        f = Fun(x->x*cot(π*x/2))
        x = Fun(identity)
        u = Fun(JacobiWeight(1.,1.,ChebyshevInterval()), (f/(1-x^2)).coefficients)
        @test 1/(0.1*cot(π*.1/2)) ≈ (1/u)(.1)

        @test (x/u)(.1) ≈ tan(π*.1/2)

        f=Fun(x->exp(-x^2),Line(0.,0.,-.5,-.5),400)
        @test sum(f) ≈ sqrt(π)

        f=Fun(x->exp(x)/sqrt(1-x.^2),JacobiWeight(-.5,-.5))
        @test f(.1) ≈ (x->exp(x)/sqrt(1-x.^2))(.1)

        @test norm(Fun(exp,Legendre(0..1))+sqrt(Fun(0..1))) ≈ 2.491141949903508

        # sampling Chebyshev
        x=Fun(identity)
        f = exp(x)/sqrt(1-x^2)
        g = cumsum(f)
        @test abs(g(-1)) ≤ 1E-15
        @test g'(0.1) ≈ f(0.1)
    end

    @testset "JacobiWeight Derivative" begin
        S = JacobiWeight(-1.,-1.,Chebyshev(0..1))

        # Checks bug in Derivative(S)
        @test typeof(ConstantSpace(0..1)) <: Space{ClosedInterval{Int},Float64}

        D=Derivative(S)
        f=Fun(S,Fun(exp,0..1).coefficients)
        x=0.1
        @test f(x) ≈ exp(x)*x^(-1)*(1-x)^(-1)/4
        @test (D*f)(x) ≈ -exp(x)*(1+(x-3)*x)/(4*(x-1)^2*x^2)

        S=JacobiWeight(-1.,0.,Chebyshev(0..1))
        D=Derivative(S)

        f=Fun(S,Fun(exp,0..1).coefficients)
        x=.1
        @test f(x) ≈ exp(x)*x^(-1)/2
        @test (D*f)(x) ≈ exp(x)*(x-1)/(2x^2)
    end

    @testset "Jacobi singularity" begin
        x = Fun(identity)
        f = exp(x)/(1-x.^2)

        @test f(.1) ≈ exp(.1)/(1-.1^2)
        f = exp(x)/(1-x.^2).^1
        @test f(.1) ≈ exp(.1)/(1-.1^2)
        f = exp(x)/(1-x.^2).^1.0
        @test f(.1) ≈ exp(.1)/(1-.1^2)

        ## 1/f with poles

        x=Fun(identity)
        f=sin(10x)
        g=1/f

        @test g(.123) ≈ csc(10*.123)
    end

    @testset "Ray and Line" begin
        @testset "Ray" begin
            @test Inf in Ray()   # this was a bug
            @test Space(0..Inf) == Chebyshev(Ray())
            f=Fun(x->exp(-x),0..Inf)
            @test f(0.1) ≈ exp(-0.1)
            @test f'(.1) ≈ -f(.1)

            x=Fun(identity,Ray())
            f=exp(-x)
            u=integrate(f)
            @test (u(1.)-u(0)-1) ≈ -f(1)

            x=Fun(identity,Ray())
            f=x^(-0.123)*exp(-x)
            @test integrate(f)'(1.) ≈ f(1.)

            @test ≈(sum(Fun(sech,0..Inf)),sum(Fun(sech,0..40));atol=1000000eps())
            @test Line() ∪ Ray() == Line()
            @test Line() ∩ Ray() == Ray()

            f=Fun(sech,Line())
            Fun(f,Ray())(2.0) ≈ sech(2.0)
            Fun(f,Ray(0.,π))(-2.0) ≈ sech(-2.0)
            Fun(sech,Ray(0.,π))(-2.0) ≈ sech(-2.0)
        end

        @testset "Ei (Exp Integral)" begin
            y=Fun(Ray())
            q=integrate(exp(-y)/y)
            @test (q-last(q))(2.) ≈ (-0.04890051070806113)

            ## Line
            f=Fun(x->exp(-x^2),Line())

            @test f'(0.1) ≈ -2*0.1exp(-0.1^2)
            @test (Derivative()*f)(0.1) ≈ -2*0.1exp(-0.1^2)
        end

        @testset "PeriodicLine" begin
            d=PeriodicLine()
            D=Derivative(d)

            f = Fun(x->sech(x-0.1),d,200)
            @test f(1.) ≈ sech(1-0.1)

            f=Fun(x->sech(x-0.1),d)
            @test f(1.) ≈ sech(1-0.1)

            @test ≈((D*f)(.2),-0.0991717226583897;atol=100000eps())
            @test ≈((D^2*f)(.2),-0.9752522555114987;atol=1000000eps())
        end
    end

    @testset "LogWeight" begin
        x=Fun(identity,-1..1)
        f=exp(x+1)-1
        @test log(f)(0.1) ≈ log(f(0.1))

        x=Fun(identity,0..1)
        f=exp(x)-1
        @test log(f)(0.1) ≈ log(f(0.1))

        x=Fun(identity,0..1)
        @test Fun(exp(x)/x-1/x,Chebyshev)(0.1) ≈ (exp(0.1)-1)/0.1

        x=Fun(identity,0..1)
        f=1/x
        p=integrate(f)
        @test (p-p(1.))(0.5) ≈ log(0.5)

        f=1/(1-x)
        p=integrate(f)
        @test (p-p(0.))(0.5) ≈ -log(1-0.5)
    end

    @testset "Complex domains sqrt" begin
        a=1+10*im; b=2-6*im
        d=Curve(Fun(x->1+a*x+b*x^2))

        x=Fun(d)
        w=sqrt(abs(leftendpoint(d)-x))*sqrt(abs(rightendpoint(d)-x))

        @test sum(w/(x-2.))/(2π*im) ≈ (-4.722196879007759+2.347910413861846im)
        @test linesum(w*log(abs(x-2.)))/π ≈ (88.5579588360686)

        a=Arc(0.,1.,0.,π/2)
        ζ=Fun(identity,a)
        f=Fun(exp,a)*sqrt(abs((ζ-1)*(ζ-im)))
    end

    @testset "DiracDelta and PointSpace" begin
        a,b=DiracDelta(0.),DiracDelta(1.)
        f=Fun(exp)
        g=a+0.2b+f
        @test components(g)[2](0.) ≈ 1.
        @test g(.1) ≈ exp(.1)
        @test sum(g) ≈ (sum(f)+1.2)

        #Checks prevoius bug
        δ=DiracDelta()
        x=Fun()
        w=sqrt(1-x^2)
        w+δ

        ## PointSpace

        @test eltype(domain(PointSpace([0,0.1,1])) ) == Float64

        f=Fun(x->(x-0.1),PointSpace([0,0.1,1]))
        @test roots(f) == [0.1]

        a=Fun(exp,space(f))
        @test f/a == Fun(x->(x-0.1)*exp(-x),space(f))

        f = Fun(space(f),[1.,2.,3.])

        g = f + Fun(2..3)
        @test f(0.0) ≈ g(0.0) ≈ 1.0
        @test f(0.1) ≈ g(0.1) ≈ 2.0
        @test f(1.0) ≈ g(1.0) ≈ 3.0

        @test g(2.3) ≈ 2.3

        h = a + Fun(2..3)

        # for some reason this test is broken only on Travis
        @test_skip g/h ≈ f/a + Fun(1,2..3)
    end

    @testset "DiracDelta integration and differentiation" begin
        δ = DiracDelta()
        h = integrate(δ)
        @test domain(h) == PiecewiseSegment([0,Inf])
        @test h(-2) == 0
        @test h(2) == 1

        δ = 0.3DiracDelta(0.1) + 3DiracDelta(2.3)
        h = integrate(δ)
        @test domain(h) == PiecewiseSegment([0.1,2.3,Inf])
        @test h(-2) == 0
        @test h(2) == 0.3
        @test h(3) == 3.3

        δ = (0.3+1im)DiracDelta(0.1) + 3DiracDelta(2.3)
        h = integrate(δ)
        @test domain(h) == PiecewiseSegment([0.1,2.3,Inf])
        @test h(-2) == 0
        @test h(2) == 0.3+1im
        @test h(3) == 3.3+1im
    end

    @testset "DiracDelta sampling" begin
        δ = 0.3DiracDelta(0.1) + 3DiracDelta(2.3)
        Random.seed!(0)
        for _=1:10
            @test sample(δ) ∈ [0.1, 2.3]
        end
        Random.seed!(0)
        r = sample(δ, 10_000)
        @test count(i -> i == 0.1, r)/length(r) ≈ 0.3/(3.3) atol=0.01
    end

    @testset "Multiple roots" begin
        x=Fun(identity,-1..1)

        @test (1/x^2)(0.1) ≈ 100.
        @test (1/x^2)(-0.1) ≈ 100.

        fc=x*(1+x)^2
        @test (1/fc)(0.1) ≈ 1/fc(0.1)

        fc=x*(1-x)^2
        @test (1/fc)(0.1) ≈ 1/fc(0.1)
    end

    @testset "special function singularities" begin
        x=Fun(0..1)
        @test erf(sqrt(x))(0.1) ≈ erf(sqrt(0.1))
        @test erfc(sqrt(x))(0.1) ≈ erfc(sqrt(0.1))

        ## roots of log(abs(x-y))
        x=Fun(-2..(-1))
        @test space(abs(x)) == Chebyshev(-2 .. (-1))

        @test roots(abs(x+1.2)) ≈ [-1.2]

        f = sign(x+1.2)
        @test space(f) isa PiecewiseSpace
        @test f(-1.4) == -1
        @test f(-1.1) == 1

        f=abs(x+1.2)
        @test abs(f)(-1.3) ≈ f(-1.3)
        @test abs(f)(-1.1) ≈ f(-1.1)
        @test norm(abs(f)-f)<10eps()

        @test norm(sign(f)-Fun(1,space(f)))<10eps()

        @test log(f)(-1.3) ≈ log(abs(-1.3+1.2))
        @test log(f)(-1.1) ≈ log(abs(-1.1+1.2))
    end

    @testset "#393" begin
        x = Fun(0..1)
        f = exp(x)*sqrt(x)*log(1-x)
        @test f(0.1) ≈ exp(0.1)*sqrt(0.1)*log(1-0.1)
    end

    @testset "Jacobi conversions" begin
        S1,S2=JacobiWeight(3.,1.,Jacobi(1.,1.)),JacobiWeight(1.,1.,Jacobi(0.,1.))
        f=Fun(S1,[1,2,3.])
        C=Conversion(S1,S2)
        Cf=C*f
        @test Cf(0.1) ≈ f(0.1)

        S1,S2=JacobiWeight(3.,2.,Jacobi(1.,1.)),JacobiWeight(1.,1.,Jacobi(0.,0.))
        f=Fun(S1,[1,2,3.])
        C=Conversion(S1,S2)
        Cf=C*f
        @test Cf(0.1) ≈ f(0.1)
    end

    @testset "Derivative operator for HeavisideSpace" begin
        H = HeavisideSpace([-1.0,0.0,1.0])
        @test Fun(H, [1.0])(1.0) == 0.0
        @test Fun(H, [0.0,1.0])(1.0) == 1.0

        H=HeavisideSpace([1,2,3])
        D=Derivative(H)
        @test domain(D)==PiecewiseSegment([1,2,3])
        @test D[1,1]==-1
        @test D[1,2]==1

        H=HeavisideSpace([1,2,3,Inf])
        D=Derivative(H)
        @test domain(D)==PiecewiseSegment([1,2,3,Inf])
        @test D[1,1]==-1
        @test D[2,2]==-1
        @test D[1,2]==1

        S = HeavisideSpace([-1.0,0.0,1.0])
        @test Derivative(S) === Derivative(S,1)
    end
end


using Test

using MarkovChains
@testset "poisson_trunc" begin
    λ = 20.0
    tol = 1e-6
    ltp, rtp, probs = poisson_trunc_point(λ, tol)
    @test ltp == 3
    @test rtp == 45
    @test abs(1.0 - sum(probs)) < tol
    λ = 400.0
    ltp, rtp, probs = poisson_trunc_point(λ, tol)
    @test ltp == 306
    @test rtp == 502
    @test abs(1.0 - sum(probs)) < tol
end


# module ProFit

## References ##

# written by: Péter Salvi
# original source: https://github.com/salvipeter/profile-fitter

# The general workflow is based on:
#
# P. Benkő, T. Várady:
#   Best fit translational and rotational surfaces for reverse engineering shapes.
#   In: The mathematics of surfaces IX, pp. 70-81, 2000.

# The thinning algorithm for the guiding polygon is based on:
#
# I-K. Lee:
#   Curve reconstruction from unorganized points.
#   In: Computer Aided Geometric Design, Vol. 17(2), pp. 161-177, 2000.

# Circle fitting is based on:
#
# S.J. Ahn, W. Rauh, H-J. Warnecke:
#   Least-squares orthogonal distances fitting of circle, sphere, ellipse, hyperbola, and parabola.
#   In: Pattern Recognition, Vol. 34(12), pp. 2283-2303, 2001.

## TODO ##

# Beautification:
# - delete short arcs/segments
# - unite similar arcs
# - create segments from arcs with large radii
# - ensure C0 continuity
# - ensure G1 continuity where plausible

#using LinearAlgebra
#using Plots
#using Random

# Delaunay triangulation

"""
    bounding_triangle(points)

Returns a triangle that contains all points.
"""
function bounding_triangle(points)
    minx, maxx = extrema((x->x[1]).(points))
    miny, maxy = extrema((x->x[2]).(points))
    center = [(minx + maxx) / 2, (miny + maxy) / 2]
    r = norm(center - [minx, miny])
    [center - [sqrt(3) * r, r],
     center + [0, 2r],
     center + [sqrt(3) * r, -r]]
end

"""
    circumcircle(a, b, c)

Given three points, gives back the circumcircle as the tuple (center, radius).
"""
function circumcircle(a, b, c)
    d = 2 * (a[1] * (b[2] - c[2]) + b[1] * (c[2] - a[2]) + c[1] * (a[2] - b[2]))
    if abs(d) < 1e-10
        # Degenerate case
        if norm(a - c) > norm(a - b)
            b, c = c, b
        end
        return ((a + b) / 2, norm(a - b) / 2)
    end
    x = (a[1]^2 + a[2]^2) * (b[2] - c[2]) +
        (b[1]^2 + b[2]^2) * (c[2] - a[2]) +
        (c[1]^2 + c[2]^2) * (a[2] - b[2])
    y = (a[1]^2 + a[2]^2) * (c[1] - b[1]) +
        (b[1]^2 + b[2]^2) * (a[1] - c[1]) +
        (c[1]^2 + c[2]^2) * (b[1] - a[1])
    center = [x / d, y / d]
    center, norm(center - a)
end

"""
    to_edges!(triangles)

Converts a list of triangles (index-triples) to a list of edges (index-pairs).
The indices in the individual edges are sorted.
"""
function to_edges!(triangles)
    edges(tri) = sort.([[tri[1], tri[2]], [tri[1], tri[3]], [tri[2], tri[3]]])
    mapreduce(edges, append!, triangles)
end

"""
    walk_polygon(edges)

Given a list of edges (pairs of indices) defining a closed loop,
returns a sorted list of points going around the edges.
"""
function walk_polygon(edges)
    poly = [edges[1][1], edges[1][2]]
    while length(poly) != length(edges)
        for edge in edges
            if edge[1] == poly[end] && edge[2] != poly[end-1]
                push!(poly, edge[2])
                break
            elseif edge[2] == poly[end] && edge[1] != poly[end-1]
                push!(poly, edge[1])
                break
            end
        end
    end
    poly
end

"""
    delaunay(points)

Returns a list of triangles, each consisting of 3 indices to the input list.
"""
function delaunay(points)
    n = length(points)
    p = copy(points)
    append!(p, bounding_triangle(points))
    triangles = [[n+1, n+2, n+3]]

    for i in 1:n
        bad = filter(triangles) do tri
            c, r = circumcircle(p[tri[1]], p[tri[2]], p[tri[3]])
            norm(p[i] - c) <= r
        end
        setdiff!(triangles, bad)
        #isempty(bad) && continue
        bad_edges = to_edges!(bad)
        bad_boundary = filter(bad_edges) do edge
            count(x -> x == edge, bad_edges) == 1
        end
        poly = walk_polygon(bad_boundary)
        m = length(poly)
        for j in 1:m
            push!(triangles, [i, poly[j], poly[mod1(j+1,m)]])
        end
    end

    filter(triangles) do tri
        all(x -> x <= n, tri)
    end
end

# Thinning

"""
    spanning_tree(edges, weights)

Given a list of edges (index pairs) and associated weights,
returns a list of edges defining a minimal weight spanning tree.
"""
function spanning_tree(edges, weights)
    order = sortperm(weights)
    result = Array{Array{Int64,1},1}(undef, 0)
    sets = Dict()
    for edge in edges
        for p in edge
            sets[p] = p
        end
    end
    for i in order
        a, b = edges[i]
        if sets[a] != sets[b]
            push!(result, edges[i])
            new, old = minmax(sets[a], sets[b])
            for (k, v) in sets
                if v == old
                    sets[k] = new
                end
            end
        end
    end
    result
end

"""
    thinning(points, edges, thickness)

Given an Euclidean minimum spanning tree over a set of points taken from a (noisy) curve,
this function generates an ordered list of points representing a thinned version of the curve.
The `thickness` parameter should be around the same size as the noise.
"""
function thinning(points, edges, thickness)
    lookup = [Int[] for _ in points]
    for e in edges
        push!(lookup[e[1]], e[2])
        push!(lookup[e[2]], e[1])
    end

    function bfs(start)
        seen = fill(false, length(points))
        result = []
        queue = [start]
        while !isempty(queue)
            v = popfirst!(queue)
            push!(result, v)
            for w in lookup[v]
                if !seen[w]
                    push!(queue, w)
                    seen[w] = true
                end
            end
        end
        result
    end
    function neighbors(start)
        x = points[start]
        result = []
        function rec(next)
            push!(result, next)
            for p in lookup[next]
                if !(p in result) && norm(x - points[p]) < thickness
                    rec(p)
                end
            end
        end
        rec(start)
        result
    end

    ordered = bfs(bfs(1)[end])
    thinned = Array{Array{Float64,1},1}(undef, 0)
    chunks = Array{Array{Int64,1},1}(undef, 0)
    while !isempty(ordered)
        adjacent = neighbors(ordered[1])
        push!(thinned, mapreduce(x -> points[x], +, adjacent) / length(adjacent))
        push!(chunks, adjacent)
        setdiff!(ordered, adjacent)
    end
    thinned, chunks
end

"""
    thinning(points, thickness)

Given a set of unsorted points taken from a (noisy) curve,
this function generates an ordered list of points representing a thinned version of the curve.
The `thickness` parameter should be around the same size as the noise.
"""
function thinning(points, thickness)
    @logmsg IterLow1  "Delaunay"
    tris = delaunay(points)
    all_edges = to_edges!(tris)
    @logmsg IterLow1 "Computing distance for $(length(all_edges)) edges, with $(length(points)) points"
    weights = [norm(points[e[1]] - points[e[2]]) for e in all_edges]
    @logmsg IterLow1 "Spanning tree"
    tree = spanning_tree(all_edges, weights)
    @logmsg IterLow1 "Thinning"
    thinning(points, tree, thickness)
end

function alle(l)
    ae = Vector{Vector{Int}}(undef, 0)
    for i in 1:l
        for j in i+1:l
            #i == j && continue
            push!(ae, [i,j])
        end
    end
    ae
end

"""
    thinning_slow(points, thickness)

Given a set of unsorted points taken from a (noisy) curve,
this function generates an ordered list of points representing a thinned version of the curve.
The `thickness` parameter should be around the same size as the noise.
Delaunay triangulation is not used to accelerate the code.
"""
function thinning_slow(points, thickness)
    @logmsg IterLow1 "Collect edges for $(length(points)) points."
    all_edges = alle(size(points,1))
    @logmsg IterLow1 "Computing distance"
    weights = [norm(points[e[1]] - points[e[2]]) for e in all_edges]
    @logmsg IterLow1 "Spanning tree"
    tree = spanning_tree(all_edges, weights)
    @logmsg IterLow1 "Thinning"
    thinning(points, tree, thickness)
end

function thinning_deldir(points, thickness)
    miv, mav = findAABB(points)
    rect = [miv[1]; mav[1]; miv[2]; mav[2]]
    pix = (x->x[1]).(points)
    piy = (x->x[2]).(points)
    @logmsg IterLow1  "Delaunay"
    deli1, deli2 = deldir(pix, piy, rect)
    all_edges = [ [deli1[i], deli2[i]] for i in 1:size(deli1,1)]
    #tris = delaunay(points)
    #all_edges = to_edges!(tris)
    @logmsg IterLow1 "Computing distance for $(length(all_edges)) edges, with $(length(points)) points"
    weights = [norm(points[e[1]] - points[e[2]]) for e in all_edges]
    @logmsg IterLow1 "Spanning tree"
    tree = spanning_tree(all_edges, weights)
    @logmsg IterLow1 "Thinning"
    thinning(points, tree, thickness)
end

# end # module


module PlasmaParameters

using Unitful
using ..Units
using ..Constants: q, me, ϵ0

function electron_cyclotron_frequency(B::Unitful.BField)
    return uconvert(rad/s, q*B/me)
end

function electron_cyclotron_frequency(B::U) where U
    B_units = isa(B, Unitful.Quantity) ? B : B*T
    return ustrip(electron_cyclotron_frequency(B_units))
end

function ion_cyclotron_frequency(B::Unitful.BField, Zi::Int, mi::Unitful.Mass)
    return uconvert(rad/s, Zi*q*B/mi)
end

function ion_cyclotron_frequency(B, Zi, mi)
    B_units = isa(B, Unitful.Quantity) ? B : B*T
    mi_units = isa(mi, Unitful.Quantity) ? mi : mi*u
    return ustrip(ion_cyclotron_frequency(B_units, Zi, mi_units))
end

export electron_cyclotron_frequency, ion_cyclotron_frequency

function electron_plasma_frequency(ne::NumberDensity)
    return uconvert(rad/s,sqrt((ne*q^2)/(me*ϵ0)))
end

function electron_plasma_frequency(ne)
    ne_units = isa(ne, Unitful.Quantity) ? ne : ne*cm^-3
    return ustrip(electron_plasma_frequency(ne_units))
end

function ion_plasma_frequency(ni::NumberDensity, Zi::Int, mi::Unitful.Mass)
    return uconvert(rad/s,sqrt((ni*(Zi*q)^2)/(mi*ϵ0)))
end

function ion_plasma_frequency(ni, Zi, mi)
    ni_units = isa(ni, Unitful.Quantity) ? ni : ni*cm^-3
    mi_units = isa(mi, Unitful.Quantity) ? mi : mi*u
    return ustrip(ion_plasma_frequency(ni_units, Zi, mi_units))
end

export electron_plasma_frequency, ion_plasma_frequency

function electron_thermal_speed(Te::Unitful.Energy)
    return uconvert(m/s,sqrt(Te/me))
end

function electron_thermal_speed(Te)
    Te_units = isa(Te, Unitful.Quantity) ? Te : Te*keV
    return ustrip(electron_thermal_speed(Te_units))
end

function ion_thermal_speed(Ti::Unitful.Energy, mi::Unitful.Mass)
    return uconvert(m/s,sqrt(Ti/mi))
end

function ion_thermal_speed(Ti, mi)
    Ti_units = isa(Ti, Unitful.Quantity) ? Ti : Ti*keV
    mi_units = isa(mi, Unitful.Quantity) ? mi : mi*u
    return ustrip(ion_thermal_speed(Ti_units, mi_units))
end

export electron_thermal_speed, ion_thermal_speed

function electron_gyro_radius(Te::Unitful.Energy, B::Unitful.BField)
    return uconvert(m, sqrt(2*me*Te)/(q*B))
end

function electron_gyro_radius(Te, B)
    Te_units = isa(Te, Unitful.Quantity) ? Te : Te*keV
    B_units = isa(B, Unitful.Quantity) ? B : B*T
    return ustrip(electron_gyro_radius(Te_units,B_units))
end

function ion_gyro_radius(Ti::Unitful.Energy, B::Unitful.BField, Zi::Int, mi::Unitful.Mass)
    return uconvert(m, sqrt(2*mi*Ti)/(Zi*q*B))
end

function ion_gyro_radius(Ti, B, Zi, mi)
    Ti_units = isa(Ti, Unitful.Quantity) ? Ti : Ti*keV
    B_units = isa(B, Unitful.Quantity) ? B : B*T
    mi_units = isa(mi, Unitful.Quantity) ? mi : mi*u
    return ustrip(ion_gyro_radius(Ti_units,B_units, Zi, mi_units))
end

export electron_gyro_radius, ion_gyro_radius

end #module


@testset "LocalField" begin

  @testset "Creation" begin
    Qx, x = FlintQQ["x"]
    f = x^2-2*x+2
    K, a = local_field(f, 2, 10, "a", Hecke.EisensteinLocalField, check = false)
    @test precision(K) == 20
    @test characteristic(K) == 0
    @test prime(K) == 2

    Kt, t = PolynomialRing(K, "t")
    g = t^2+2
    L, b = local_field(g, "b", Hecke.EisensteinLocalField, check = false)
    @test precision(L) == 2*precision(g)
    @test precision(b) == precision(L)
    @test degree(L) == 2
    @test absolute_degree(L) == 4
    @test prime(L) == 2

    Q2 = PadicField(2, 10)
    Q2s, s = PolynomialRing(Q2, "s")
    f = s^2+s+1
    Ku, c = local_field(f, "s", Hecke.UnramifiedLocalField, check = false)
    @test precision(Ku) == precision(f)
    @test precision(f) == precision(c)
    @test degree(Ku) == 2
    @test absolute_degree(Ku) == 2
    @test inertia_degree(Ku) == 2
    @test ramification_index(Ku) == 1
    @test absolute_ramification_index(Ku) == 1
    @test absolute_inertia_degree(Ku) == 2

    Lu, u = PolynomialRing(L, "u")
    Lu, d = local_field(u^2+u+1, "s", Hecke.UnramifiedLocalField, check = false)
    @test absolute_degree(Lu) == 8
    @test ramification_index(Lu, K) == 2
    @test inertia_degree(Lu, K) == 2

  end

  @testset "Norm" begin
    K = QadicField(3, 4, 10)[1]
    Kx, x = PolynomialRing(K, "x")
    L = eisenstein_extension(x^20+3)[1]
    b = @inferred basis(L)
    for i = 1:10
      r = 1+2*uniformizer(L)^i * sum([rand(1:10)*b[i] for i in 1:5])
      M = @inferred representation_matrix(r)
      n = @inferred norm(r)
      @test n == det(M)
      @test trace(r) == trace(M)
    end
  end

  @testset "Completions" begin
    K, gK = cyclotomic_field(15)
    OK = maximal_order(K)
    lP = prime_decomposition(OK, 3)
    P = lP[1][1]
    Kp, mKp = @inferred Hecke.generic_completion(K, P, 100)
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test valuation(mKp(elem_in_nf(uniformizer(P)))) == fmpq(1, 2)
    @test valuation(mKp(gen(K)))*2 == valuation(gen(K), P)
    @test iszero(defining_polynomial(K)(mKp(gen(K))))

    P = prime_decomposition(OK, 7)[1][1]
    Kp, mKp = @inferred Hecke.unramified_completion(K, P, 100)
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test isone(valuation(mKp(elem_in_nf(uniformizer(P)))))
    @test valuation(mKp(gen(K))) == valuation(gen(K), P)
    @test iszero(defining_polynomial(K)(mKp(gen(K))))

    K, gK = cyclotomic_field(9)
    OK = maximal_order(K)
    lP = prime_decomposition(OK, 3)
    P = lP[1][1]
    Kp, mKp = @inferred Hecke.totally_ramified_completion(K, P, 100)
    @test valuation(mKp(elem_in_nf(uniformizer(P)))) == fmpq(1, 6)
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test valuation(mKp(gen(K)))*2 == valuation(gen(K), P)
    @test iszero(defining_polynomial(K)(mKp(gen(K))))

    K, gK = cyclotomic_field(55)
    OK = maximal_order(K)
    lP = prime_decomposition(OK, 11)
    P = lP[1][1]
    Kp, mKp = @inferred Hecke.generic_completion(K, P, 100)
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test valuation(mKp(elem_in_nf(uniformizer(P)))) == fmpq(1, 10)
    @test valuation(mKp(gen(K)))*10 == valuation(gen(K), P)
    @test iszero(defining_polynomial(K)(mKp(gen(K))))

    lP = prime_decomposition(OK, 7)
    P = lP[1][1]
    Kp, mKp = @inferred Hecke.unramified_completion(K, P, 100)
    @test isone(valuation(mKp(elem_in_nf(uniformizer(P)))))
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test valuation(mKp(gen(K))) == valuation(gen(K), P)
    @test iszero(defining_polynomial(K)(mKp(gen(K))))

    Qx, x = PolynomialRing(FlintQQ, "x", cached = false)
    f = x^10 + 5*x^9 - 633*x^8 + 38157*x^7 + 1360601*x^6 - 39808345*x^5 - 464252491*x^4 - 17622187401*x^3 + 2826886632714*x^2 - 96335539805104*x + 1313743135204448
    K, a = number_field(f, check = false, cached = false)
    OK = maximal_order(K)
    P = prime_decomposition(OK, 2)[1][1]
    Kp, mKp = @inferred Hecke.unramified_completion(K, P, 100)
    @test isone(valuation(mKp(elem_in_nf(uniformizer(P)))))
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test valuation(mKp(gen(K))) == valuation(gen(K), P)
    @test iszero(defining_polynomial(K)(mKp(gen(K))))

    lP = prime_decomposition(OK, 11)
    if lP[1][2] == 5
      P = lP[1][1]
    else
      P = lP[2][1]
    end
    Kp, mKp = @inferred Hecke.totally_ramified_completion(K, P, 100)
    @test valuation(mKp(elem_in_nf(uniformizer(P)))) == fmpq(1, 5)
    @test isone(valuation(mKp\(uniformizer(Kp)), P))
    @test valuation(mKp(gen(K)))*5 == valuation(gen(K), P)
    #@test iszero(defining_polynomial(K)(mKp(gen(K))))

  end

  @testset "Exp and Log" begin
    K = PadicField(2, 100)
    Kx, x = PolynomialRing(K, "x", cached = false)
    L, b = Hecke.eisenstein_extension(x^7+2, "a")
    pi = uniformizer(L)
    @test iszero(log(pi))
    B = basis(L)
    for i = 15:20
      el = sum([rand(FlintZZ, 0:10)*B[j] for j = 1:7])*pi^i
      explog = exp(log(1+el))
      logexp = log(exp(el))
      @test iszero(setprecision(explog, precision(explog)-35) - 1 - el)
      @test iszero(setprecision(logexp, precision(logexp)-35)-el)
    end
  end

  @testset "Maps" begin
    # FlintQadicField -> FlintQadicField
    Qq, a = QadicField(2, 3, 100)
    rt = roots(map_coefficients(Qq, defining_polynomial(Qq)))

    i = findfirst(x -> x == a, rt)
    img = rt[mod(i-2, 3)+1]
    invimg = rt[mod(i, 3)+1]

    f = @inferred hom(Qq, Qq, img)
    @test f(a) == img
    @test f\(f(a)) == a
    @test f(f(a)) == invimg

    k, mk = ResidueField(Qq)
    for i in 1:10
      z = mk\(rand(k))
      @test z == f\(f(z))
      fl, w = @inferred haspreimage(f, z)
      @test fl
      @test f(w) == z
    end

    @test_throws ErrorException hom(Qq, Qq, Qq(2))

    f = @inferred hom(Qq, Qq, img, inverse = (invimg))
    @test f(a) == img
    @test f\a == invimg
    @test_throws ErrorException hom(Qq, Qq, Qq(2), inverse = a + 1)
    @test_throws ErrorException hom(Qq, Qq, a, inverse = Qq(2))

    g = @inferred inv(f)

    for i in 1:10
      z = mk\(rand(k))
      @test f(g(z)) == z
      @test g(f(z)) == z
    end

    # FlintQadicField -> LocalField
    Qqt, t = Qq["t"]
    L, b = eisenstein_extension(t^3 + 2, "b")
    f = @inferred hom(Qq, L, L(gen(Qq)))

    @test f(a) == L(gen(Qq))
    @test_throws ErrorException hom(Qq, L, b)
    @test f\(L(gen(Qq))) == gen(Qq)
    fl, z = @inferred haspreimage(f, b^3)
    @test fl
    @test f(z) == L(-2)

    # LocalField -> FlintQadicField
    Qp = PadicField(2, 100)
    Qpx, x = PolynomialRing(Qp)
    K, a = Hecke.unramified_extension(x^2+x+1)
    Qq, gQq = QadicField(2, 2, 100)
    rt = roots(map_coefficients(Qq, defining_polynomial(K)))

    f = @inferred hom(K, Qq, rt[1])
    g = @inferred inv(f)
    k, mk = ResidueField(Qq)
    for i in 1:10
      w = mk\(rand(k))
      @test f(g(w)) == w
    end
  end

  @testset "Automorphisms" begin
    K = PadicField(2, 200)
    Kt, t = PolynomialRing(K)
    L, b = Hecke.eisenstein_extension(t^2+2, "a")
    @test length(automorphisms(L)) == 2
    Qq, a = QadicField(2, 2, 100)
    @test length(automorphisms(Qq)) == 2
    Qqx, x = PolynomialRing(Qq)
    L, b = Hecke.eisenstein_extension(x^3+2, "a")
    @test length(automorphisms(L)) == 3
    @test length(absolute_automorphisms(L)) == 6
    G, mG = absolute_automorphism_group(L)
    @test order(G) == 6
    @test !isabelian(G)
  end

  @testset "Unramified extension" begin
    Qx,x = FlintQQ["x"]
    f = Qx([ 1, 8, -40, -46, 110, 71, -113, -43, 54, 11, -12, -1, 1 ])
    L = number_field(f)[1];
    P = prime_decomposition(maximal_order(L),7)[1][1];
    lp, mp = Hecke.generic_completion(L,P);
    Qy,y = PolynomialRing(lp,"y")
    f, mf = ResidueField(lp)
    N = Hecke.unramified_extension(y^3+preimage(mf,(gen(f))) +4 )[1]
    F, mF = ResidueField(N)
    @test order(F) == 7^6
    G, mG = automorphism_group(N)
    @test order(G) == 3
    @test mG(G[1]^2) == mG(G[1])^2
  end

  @testset "Ali-Inverse" begin
    Qx,x = FlintQQ["x"]
    L = number_field(x^6-6*x^4+9*x^2+23)[1]
    P = prime_decomposition(maximal_order(L),23)[1][1]
    lp,mp = Hecke.generic_completion(L,P)
    a = uniformizer(lp)
    Qy,y = PolynomialRing(lp,"y")
    N = Hecke.unramified_extension(y^2+13*y+4)[1]
    r = N(53)*basis(N)[1] + N(165)*basis(N)[2]
    @test isone(r*r^-1)
  end
end


x, xi  = 2.0, 0.5
y, yi  = 4.0, 0.25
z, zi  = x + y, 1.0 / ( x + y )

multiplier = (n1, n2) -> (m) -> n1 * n2 * m

numlist  = [x ,  y,  z]
numlisti = [xi, yi, zi]

@show collect(multiplier(n, invn)(0.5) for (n, invn) in zip(numlist, numlisti))


using MinimaxCalibratedEBayes
using Test
using Plots

mhist =  MCEBHistogram(-6:0.02:6)
marginal_target = MarginalDensityTarget(StandardNormalSample(0.0))

marginal_target_cb = DiscretizedAffineEstimator(mhist, ButuceaComte(marginal_target, 10_000))

#test for symmetry
@test marginal_target_cb(-5.0001) == marginal_target_cb(+5.0001)

pgfplotsx()
plot(cb_disc)

prior_target = PriorDensityTarget(0.0)

prior_target_cb = DiscretizedAffineEstimator(mhist, ButuceaComte(prior_target, 10_000))

@test prior_target_cb(-5.0001) == prior_target_cb(+5.0001)

plot(prior_target_cb)

gmix_tmp = GaussianMixturePriorClass(0.1, [0.0], Gurobi.Optimizer)

Z = DiscretizedStandardNormalSample(1.0, mhist)

res = worst_case_bias(marginal_target_cb,
                 Z,
                 gmix_tmp,
                 marginal_target)

tmp = JuMP.value.(res[1][:πs])

findmax(tmp)

res2 = worst_case_bias(prior_target_cb,
                 Z,
                 gmix_tmp,
                 prior_target) #point is, I can check bias...

sum(prior_target_cb.Q .* prob_mass) - pdf(Normal(0.0, 0.1), 0.0)
prob_mass = pdf(marginalize(Normal(0.0, 0.1),Z))

sum(prob_mass)

integrate()

# check bias now.


module DiscreteMarkovChains
import LinearAlgebra

include("Utils.jl")

include("MC.jl")
include("DMC.jl")
include("CMC.jl")

export DiscreteMarkovChain, state_space, transition_matrix, probability_matrix,
communication_classes, periodicities, decompose, canonical_form,
is_regular, is_ergodic, is_absorbing, is_reversible,
stationary_distribution, fundamental_matrix,
mean_time_to_absorption, exit_probabilities, first_passage_probabilities,
mean_recurrence_time, mean_first_passage_time

export ContinuousMarkovChain, embedded

end  # Module


# Simple model objects
parameters_simple  = @with_kw (μ = 0.0048,
                                υ = 0.02,
                                θ = 2.1,
                                r = t -> 0.05,
                                ζ = 14.5,
                                ξ = 1.0, π = (t, x) -> 1,
                                x = t -> ζ)

settings_simple = @with_kw (z_ex = unique([range(0., 0.1, length = 20)' range(0.1, 1., length = 20)' range(1., 5., length = 20)']),
                            T = 100.0, iterations = 1000,
                            ts = range(0., 100., length = length(z_ex)),
                            stationary_x0 = default_simple_stationary_x0,
                            ode_solve_algorithm = IDA())

default_simple_stationary_x0(parameters, settings) = [parameters.r(settings.T)/2] # or something like that


struct WeightedLeastSquares
    estimate::Array{Float64,1}
    rmse::Float64
    mae::Float64
    wrss::Float64
end

struct ErrorMetric
    rmse::Float64
    mae::Float64
    wrss::Float64
end

struct ErrorMetricWiden
    rmse::Float64
    mae::Float64
    wrss::Float64
    rmseGBPWLS::Float64
    maeGBPWLS::Float64
end

########## Compute the GBP error metrics ##########
function errorMetric(gbp::Union{ContinuousModel, ContinuousTreeModel})
    observationGBP = gbp.system.jacobian * gbp.inference.mean
    NactiveRows = activeRows(gbp)

    rmse = 0.0; mae = 0.0; wrss = 0.0
    @inbounds for i = 1:length(gbp.system.observation)
        rmse += (gbp.system.observation[i] - observationGBP[i])^2
        mae += abs(gbp.system.observation[i] - observationGBP[i])
        wrss += (gbp.system.observation[i] - observationGBP[i])^2 / gbp.system.variance[i]
    end
    rmse = (rmse / NactiveRows)^(1/2)
    mae = mae / NactiveRows

    return ErrorMetric(rmse, mae, wrss)
end

########## Compute the GBP error metrics and error according to the WLS ##########
@inline function errorMetric(gbp::Union{ContinuousModel, ContinuousTreeModel}, wls::WeightedLeastSquares)
    observationGBP = gbp.system.jacobian * gbp.inference.mean
    NactiveRows = activeRows(gbp)

    rmse = 0.0; mae = 0.0; wrss = 0.0
    @inbounds for i = 1:length(gbp.system.observation)
        rmse += (gbp.system.observation[i] - observationGBP[i])^2
        mae += abs(gbp.system.observation[i] - observationGBP[i])
        wrss += (gbp.system.observation[i] - observationGBP[i])^2 / gbp.system.variance[i]
    end
    rmse = (rmse / NactiveRows)^(1/2)
    mae = mae / NactiveRows

    rmseGBPWLS = 0.0; maeGBPWLS = 0.0
    @inbounds for i = 1:gbp.graph.Nvariable
        rmseGBPWLS += (wls.estimate[i] - gbp.inference.mean[i])^2
        maeGBPWLS += abs(wls.estimate[i] - gbp.inference.mean[i])
    end
    rmseGBPWLS = (rmseGBPWLS / gbp.graph.Nvariable)^(1/2)
    maeGBPWLS = maeGBPWLS / gbp.graph.Nvariable

    return ErrorMetricWiden(rmse, mae, wrss, rmseGBPWLS, maeGBPWLS)
end

########## Compute WLS solution and error metrics ##########
function wls(gbp::Union{ContinuousModel, ContinuousTreeModel})
    W = spdiagm(0 =>  @. 1.0 / sqrt(gbp.system.variance))
    A = W * gbp.system.jacobian
    G = A' * A
    b = A' * W * gbp.system.observation
    x = G \ b

    observationWLS = gbp.system.jacobian * x
    NactiveRows = activeRows(gbp)

    rmse = 0.0; mae = 0.0; wrss = 0.0
    @inbounds for i = 1:length(gbp.system.observation)
        rmse += (gbp.system.observation[i] - observationWLS[i])^2
        mae += abs(gbp.system.observation[i] - observationWLS[i])
        wrss += (gbp.system.observation[i] - observationWLS[i])^2 / gbp.system.variance[i]
    end
    rmse = (rmse / NactiveRows)^(1/2)
    mae = mae / NactiveRows

    return WeightedLeastSquares(x, rmse, mae, wrss)
end

########## Find number of nonzeros rows ##########
@inline function activeRows(gbp)
    NactiveRows = length(gbp.system.observation)
    @inbounds for (k, i) in enumerate(gbp.system.observation)
        if i == 0.0
            if all(gbp.system.jacobianTranspose[:, k] .== 0.0)
                NactiveRows -= 1
            end
        end
    end

    return NactiveRows
end

########## Check keyword arguments ##########
function checkKeywords(prob, alpha, variance)
    #### Check convergence parameters
    if prob <= 0.0 || prob >= 1.0
        error("Invalid prob value.")
    end
    if alpha <= 0.0 || alpha >= 1.0
        error("Invalid alpha value.")
    end

    #### Check initial variance
    if variance < 0.0
        error("Invalid variance value.")
    end
end


using Multisets
using Test

A = Multiset(1:5)
B = A+A
@test length(B)==2*length(A)
@test A<B
@test A<=B
@test B>A
@test B>=A
B[1]=0
@test Set(B)==Set(2:5)
@test length(A|B) == 9
@test length(A&B) == 4
@test length(A*B) == 40
@test length(A+B) == 13
@test collect(A) == collect(1:5)
@test hash(A) != hash(B)