## Imported math functions

## need to import these
math_sympy_methods_base = (:sin, :cos, :tan, :sinh, :cosh, :tanh, :asin, :acos, :atan,
                           :asinh, :acosh, :atanh, :sec, :csc, :cot, :asec, :acsc, :acot,
                           :sech, :csch,
                           :coth, :acoth,
                           :log2, :log10, :log1p, :exponent, :exp, :exp2, :expm1,
                           :sqrt, :dawson,
                           :ceiling, :floor,
                           :factorial,
                           :gcd, :lcm,
                           :isqrt
)

math_object_methods_base = (:round,
                            )

## Export SymPy math functions and vectorize them
math_sympy_methods = (:radians2degrees, :degrees2radians,
                      :factorial2,
                      )

# hypot and atan2
hypot(x::Sym, y::Number) = sqrt(x^2 + y^2)
atan2(y::Sym, x::Number) = sympy_meth(:atan2, y, x)

## Log function handles arguments differently
log(x::Sym) = sympy_meth(:log, x)
log(b::Sym, x::Sym) = sympy_meth(:log, x, b)

### Trigonometry
Base.rad2deg(x::Sym) = radians2degrees(x)
Base.deg2rad(x::Sym) = degrees2radians(x)

## degree functions
for fn in (:cosd, :cotd, :cscd, :secd, :sind, :tand,
          :acosd, :acotd, :acscd, :asecd, :asind, :atand)

    rad_fn = string(fn)[1:end-1]
    @eval ($fn)(x::Sym) = sympy[Symbol($rad_fn)](x * Sym(sympy["pi"])/180)
end

for fn in (:cospi, :sinpi)
    rad_fn = string(fn)[1:end-2]
    @eval ($fn)(x::Sym) = sympy[Symbol($rad_fn)](x * Sym(sympy["pi"]))
end

## :asech, :acsch, :sinc, :cosc,
## These fail, so define from definitions
## http://mathworld.wolfram.com/InverseHyperbolicSecant.html
asech(z::Sym) = log(sqrt(1/z-1)*sqrt(1/z+1) + 1/z)
## http://mathworld.wolfram.com/InverseHyperbolicCosecant.html
acsch(z::Sym) = log(sqrt(1+1/z^2) + 1/z)
# Julia's sinc is defined to be zero at x=1
sinc(x::Sym) = piecewise((Sym(1), Eq(x, 0)), (sin(PI*x)/(PI*x), Gt(abs(x), 0)))
cosc(x::Sym) = diff(sinc(x))

# deprecate these when v0.4 support dropped in favor of `asech.(...)` form
asech(as::AbstractArray{Sym}) = map(asech, as)
acsch(as::AbstractArray{Sym}) = map(acsch, as)

## in Julia, not SymPy
cbrt(x::Sym) = x^(1//3)
Base.ceil(x::Sym) = ceiling(x)

functions_sympy_methods = (
                           :arg,
                           :conjugate,
                           :re,
                           :sign
                           )

## map Abs->abs, Max->max, Min->min
abs(ex::Sym, args...; kwargs...) = sympy_meth(:Abs, ex, args...; kwargs...)
Base.abs2(x::Sym) = re(x*conj(x))
Base.copysign(x::Sym, y::Sym) = abs(x)*sign(y)
Base.signbit(x::Sym) = x < 0
Base.flipsign(x::Sym, y) = signbit(y) ? -x : x
Base.real(x::Sym) = sympy_meth(:re, x)
Base.imag(x::Sym) = sympy_meth(:im, x)
Base.eps(::Type{Sym}) = zero(Sym)

## use SymPy Names here...
## XXX This pattern may be of general usage. XXX
## It allows fn(Any, Sym) or f(Sym, Any) to match

import Base: min, max
min(x::Sym, a) = sympy_meth(:Min, x, a)
min(a, x::Union{SA, Real}) where {SA <: Sym} = min(x,a)

max(x::Sym, a) = sympy_meth(:Max, x, a)
max(a, x::Union{SA, Real}) where {SA <: Sym} = max(x,a)

# SymPy names
Min(ex::Sym, ex1::Sym) = sympy_meth(:Min, ex, ex1)
Max(ex::Sym, ex1::Sym) = sympy_meth(:Max, ex, ex1)

## Calculus functions

## bring over for univariate function calls as well (not just expressions)

"""

Compute a symbolic limit lim_{x->c+} f(x)

By default, *right* limits are returned. The keyword argument `dir="-"` needs to be
specified for *left* limits.

The function `f(x)` can be expressed as a symbolic expression, or a univariate function

The values `x` and `c` can be expressed by two arguments `x, c` or a pair `x=>c`.

```
@vars x y
ex = 1/(x^2 - y^2)
limit(ex, x=>3)
limit(ex, x, 3)
limit(ex, x=>3, y=>2)  # first x then y not (x,y) -> (3,2)

fn(x) = x^x
limit(fn, 0)    # symbol not needed
```
"""
limit(ex::Sym, args...; kwargs...) = sympy_meth(:limit, ex, args...; kwargs...)
limit(ex::Sym, d::Pair; kwargs...) = limit(ex, d.first, d.second;kwargs...)
limit(ex::Sym, ds::Pair...) = reduce(limit, ex, ds)
limit(f::Function, c::Number=0; kwargs...) = (z = (symbols(gensym())); limit(f(z),z=>c;kwargs...))

export limit

function Base.diff(ex::Sym, args...; kwargs...)
    if funcname(ex) in map(string, relational_sympy_values)
        Eq(diff(lhs(ex), args...; kwargs...), diff(rhs(ex), args...; kwargs...))
    else
        sympy_meth(:diff, ex, args...; kwargs...)
    end
end

##

## diff for matrix doesn't handle vectors well, so we vectorize here
diff(exs::AbstractVector{Sym}, args...; kwargs...) = map(ex -> diff(ex, args...;kwargs...), exs)
diff(exs::AbstractMatrix{Sym}, args...; kwargs...) = map(ex -> diff(ex, args...;kwargs...), exs)

## find symbolic derivatives from a function
function diff(f::Function, k::Int=1; kwargs...)
    x = Sym("x")
    diff(f(x), x, k; kwargs...)
end

# set up derivative, call doit to implement
Derivative(ex::Sym, args...) = sympy_meth(:Derivative, ex, args...)
export(Derivative)
#### Piecewise functions

"""

Create a piecewise defined function.

To create conditions on the variable, the functions `Lt`, `Le`, `Eq`, `Ge`, and `Gt` can be used. For infix notation,
unicode operators can be used: `\\ll<tab>`, `\\leqq<tab>`, `\\Equal<tab>`, `\\geqq<tab>`, and `\\gg<tab>` (but *not* `\\ge<tab>` or `\\le<tab>`).

To combine terms, the unicode `\\vee<tab>` (for "or"), `\\wedge<tab>` (for "and") can be used

Examples:
```
x,a = symbols("x,a")
p = piecewise((1, x ≪ 1), (2, (Lt(1,x)) ∨ Lt(x,2)), (3, x ≫ 2)) ## using ∨ and ∧ for & and or
subs(p, x, 2) ## 2
p = piecewise((1, Lt(x, a)), (2, Ge(x,a)))  # same as piecewise((1,  x ≪ a), (2, x ≧ a))
subs(p, x, a - 1)
```

[Note: there is also an alias `Piecewise` for copy-n-pasting from python code, but despite the capital letter, this is not a constructor for a type.]
"""
function piecewise(args...)
    sympy_meth(:Piecewise, args...)
end
const Piecewise = piecewise

piecewise_fold(ex::Sym) = sympy_meth(:piecewise_fold, ex)

Base.ifelse(ex::Sym, a, b) = piecewise((a, ex), (b, true))

"""
Indicator expression: (Either `\\Chi[tab](x,a,b)` or `Indicator(x,a,b)`)

`Χ(x, a, b)` is `1` on `[a,b]` and 0 otherwise.

This is not a function taking `x`, but a symbolic expression of `x`.

"""
Χ(x::Sym, a=-oo, b=oo) = piecewise((1, Gt(x, a) ∧ Le(x, b)), (0,true))
Indicator(x::Sym, a=-oo, b=oo) = Χ(x, a, b)

import Base: &, |
(&)(a::Bool, b::Sym) = a & (b == SympyTRUE)
(|)(a::Bool, b::Sym) = a | (b == SympyTRUE)

export Indicator, Χ

### generic programming interface

## Handle ininf, and isnan by coercion to float
Base.isfinite(x::Sym) = isfinite(convert(Float64, x))
Base.isinf(x::Sym) = try isinf(convert(Float64, x)) catch e false end
Base.isnan(x::Sym) = try isnan(convert(Float64, x)) catch e false end

## we rename sympy.div -> polydiv
Base.div(x::Sym, y::SymOrNumber) = convert(Sym, sympy["floor"](x/convert(Sym,y)))
Base.rem(x::Sym, y::SymOrNumber) = x-Sym(y)*Sym(sympy["floor"](x/y))

## zero and one (zeros?)
Base.zero(x::Sym) = Sym(0)
Base.zero(::Type{Sym}) = Sym(0)

Base.one(x::Sym) = Sym(1)
Base.one(::Type{Sym}) = Sym(1)

## useful at times
Base.typemax(::Type{Sym}) = oo
Base.typemin(::Type{Sym}) = -oo

##################################################
## special numbers are initialized after compilation
if isdefined(PyCall,:PyNULL)
    pynull() = PyCall.PyNULL()
else
    pynull() = PyCall.PyObject()
end
global PI = Sym(pynull())
global E = Sym(pynull())
global IM = Sym(pynull())
global oo = Sym(pynull())

if isdefined(Base, :MathConstants)
    Base.convert(::Type{Sym}, x::Irrational{:π}) = PI
    Base.convert(::Type{Sym}, x::Irrational{:e}) = E
    Base.convert(::Type{Sym}, x::Irrational{:γ}) = Sym(sympy["EulerGamma"])
    Base.convert(::Type{Sym}, x::Irrational{:catalan}) = Sym(sympy["Catalan"])
    Base.convert(::Type{Sym}, x::Irrational{:φ}) = (1 + Sym(5)^(1//2))/2
else
    Base.convert(::Type{Sym}, x::Irrational{:π}) = PI
    Base.convert(::Type{Sym}, x::Irrational{:e}) = E
    Base.convert(::Type{Sym}, x::Irrational{:γ}) = Sym(sympy["EulerGamma"])
    Base.convert(::Type{Sym}, x::Irrational{:catalan}) = Sym(sympy["Catalan"])
    Base.convert(::Type{Sym}, x::Irrational{:φ}) = (1 + Sym(5)^(1//2))/2
end
function init_math()
    "PI is a symbolic  π. Using `julia`'s `pi` will give round off errors."
    copy!(PI.x,  sympy["pi"])

    "E is a symbolic  `e`. Using `julia`'s `e` will give round off errors."
    copy!(E.x, Sym(sympy["exp"](1)).x)

    "IM is a symbolic `im`"
    copy!(IM.x, sympy["I"])

    "oo is a symbolic infinity. Example: `integrate(exp(-x), x, 0, oo)`."
    copy!(oo.x, sympy["oo"])
end


struct LTInitialBelief
    robot_init::Union{Coord, Nothing}
    floor::Floor
end

sampletype(::Type{LTInitialBelief}) = LTState
function iterator(b::LTInitialBelief)
    states = LTState[]
    for (x,y) in product(1:b.floor.n_cols, 1:b.floor.n_rows)
        if isnull(b.robot_init)
            for (rx, ry) in product(1:b.floor.n_cols, 1:b.floor.n_rows)
                push!(states, LTState(Coord(rx,ry), Coord(x,y), false))
            end
        else
            push!(states, LTState(get(b.robot_init), Coord(x,y), false))
        end
    end
    return states
end

function Random.rand(rng::AbstractRNG, b::LTInitialBelief)
    opp = Coord(rand(rng, 1:b.floor.n_cols), rand(rng, 1:b.floor.n_rows))
    if b.robot_init === nothing
        rob = Coord(rand(rng, 1:b.floor.n_cols), rand(rng, 1:b.floor.n_rows))
    else
        rob = b.robot_init
    end
    return LTState(rob, opp, false)
end

function Distributions.pdf(b::LTInitialBelief, s::LTState)
    if s.terminal
        return 0.0
    else
        if b.robot_init === nothing
            return (1/n_pos(b.floor))^2
        else
            if s.robot == b.robot_init
                return 1/n_pos(b.floor)
            else
                return 0.0
            end
        end
    end
end

POMDPs.initialstate(p::LaserTagPOMDP) = LTInitialBelief(p.robot_init, p.floor)


export crossEntropy,∂crossEntropy

function crossEntropy(obj,pre)
  cost=-sum(obj.*log(pre))/size(pre)[end]
end
function ∂crossEntropy(obj,pre)
  ∂pre_L=-obj./pre
end


# The file calculate the uniform induced velocity
# 给出均匀入流模型

function uniforminflow(ψ,θcp,θlat,θlon,betanow,betaxnow,Mbeta_aero=0.0)
  vind_s = [0.0,0.0,-T/(2*ρ*A*norm(v_air))] # 固定坐标系诱导速度初值
  vall_s = v_air + vind_s # 固定坐标系总的滑流速度初值
  # lmdaui = -T/(2*ρ*A*norm(vall_s))/(Ω*R) # 前飞均匀诱导速度系数
  vindui_s = [0.0,0.0,-T/(2*ρ*A*norm(vall_s))]  # 第一步迭代的固定坐标系均匀入流值

  while true  # 迭代求解均匀入流
    if norm(vindui_s-vind_s)<=1e-3
      break
    else
      vind_s = vindui_s
      vall_s = v_air + vind_s
      vindui_s = [0.0,0.0,-T/(2*ρ*A*norm(vall_s))]
    end
  end
  lmdaui = vindui_s[3]/(Ω*R) # 收敛前飞均匀诱导速度系数

  vinduitmp_r = Array{Vector}(NR,Nbe)
  valluitmp_r = Array{Vector}(NR,Nbe)

  # 准静态挥舞解方法
  β = sflapre(ψ,vall_s[3]/(Ω*R),θcp,θlat,θlon)[1]
  dβ = sflapre(ψ,vall_s[3]/(Ω*R),θcp,θlat,θlon)[2]
  vβ_β = [0.0,0.0,-dβ] # 挥舞坐标系下挥舞引起的气流速度分量
  vβ_r = betatoro(vβ_β,β)

  # if Mbeta_aero==0.0 # 入流场无挥舞初始化
  #   vβ_r = [0.0,0.0,0.0]
  #   β = 0.0
  #   dβ = 0.0
  # else
  #   # 时间步进挥舞解方法
  #   betatmp = dflapre(betanow,betaxnow,Mbeta_aero)
  #   β = betatmp[1]
  #   dβ = betatmp[2]
  #   vβ_β = [0.0,0.0,-dβ] # 挥舞坐标系下挥舞引起的气流速度分量
  #   vβ_r = betatoro(vβ_β,β)
  # end

  for k in 1:NR
    ψk = ψ+(k-1)*2*π/NR
    for i in 1:Nbe
      # 将速度从固定坐标系转为旋转坐标系
      vinduitmp_r[k,i] = systoro(vindui_s,ψk)
      vOmega_r = [0,-Ω*rb[k,i],0] # 旋转坐标系下旋转导致的气流分量
      valluitmp_r[k,i] = systoro(v_air,ψk)+vinduitmp_r[k,i]+vOmega_r+vβ_r*(rb[k,i]-ecut*R)
    end
  end

  return vinduitmp_r,valluitmp_r,lmdaui,β,dβ
end


# -*- coding: utf-8 -*-
adparam = SNN.ADEXParameter(;a = 6.0,
    b = 7.3,
    cm = 803.1,
    v_rest= -63.2,
    tau_m=19.7,
    tau_w=351.0,
    v_thresh=-39.2,
    delta_T=6.4,
    v_spike = -36.6,
    v_reset = -59.2,
    spike_delta = 16.3)

E = SNN.AD(;N = 1, param=adparam)
E.I = [795.57128906]
SNN.monitor(E, [:v,:I])
SNN.sim!([E],[], dt=0.1ms, duration=2000ms)
# # Julia SNN Implementation of AdExp Neuron.
# [Adaptive_exponential_integrate and fire neuron](http://www.scholarpedia.org/article/Adaptive_exponential_integrate-and-fire_model)
# Dr. Wulfram Gerstner
# Romain Brette, Ecole Normale Supérieure, Paris, France


function d2tduftau(t,
		   fft_u,  fft_v,
		   fft_du, fft_dv,
		   A1, A2,
		   matr, conjmatr,
		   sigma, llambda, Ntaumm)

    sigma = 1

    u     = ifft(exp(1im * t * A1) * fft_u * matr)
    v     = ifft(exp(1im * t * A1) * fft_v * matr)
    dtu   = ifft(exp(1im * t * A1) * (1im * A1) * fft_u) * matr
    dtv   = ifft(exp(1im * t * A1) * (1im * A1) * fft_v) * matr
    d2tu  = ifft(exp(1im * t * A1) * (-A1 .^ 2) * fft_u) * matr
    d2tv  = ifft(exp(1im * t * A1) * (-A1 .^ 2) * fft_v) * matr

    du    = ifft(exp(1im * t * A1) * fft_du * matr)
    dv    = ifft(exp(1im * t * A1) * fft_dv * matr)
    dtdu  = ifft(exp(1im * t * A1) * (1im * A1) * fft_du) * matr
    dtdv  = ifft(exp(1im * t * A1) * (1im * A1) * fft_dv) * matr
    d2tdu = ifft(exp(1im * t * A1) * (-A1 .^ 2) * fft_du) * matr
    d2tdv = ifft(exp(1im * t * A1) * (-A1 .^ 2) * fft_dv) * matr

    z     = (u + conj(v)) / 2
    dtz   = (dtu + conj(dtv)) / 2
    d2tz  = (d2tu + conj(d2tv)) / 2
    dz    = (du + conj(dv)) / 2
    dtdz  = (dtdu + conj(dtdv)) / 2
    d2tdz = (d2tdu + conj(d2tdv)) / 2

    fz1 = (  2 * dz  * conj(z)  * d2tz
	   + 2 * z   * conj(dz) * d2tz
	   + 2 * z   * conj(z)  * d2tdz
	   + 4 * dtz * dtdz     * conj(z)
	   + 2 * dtz .^ 2 * conj(dz)
	   + 4 * dtdz * conj(dtz)  * z
	   + 4 * dtz  * conj(dtdz) * z
	   + 4 * dtz  * conj(dtz)  * dz
	   + 2 * z * dz * conj(d2tz)
	   +     z .^ 2 * conj(d2tdz))

    champu1 = -1im * llambda * A2 .* exp(-1im * t * A1) .* fft(conjmatr .* fz1)
    champv1 = -1im * llambda * A2 .* exp(-1im * t * A1) .* fft(conjmatr .* conj(fz1))

    fz1 = (  2 * dtz * conj(z) * dz
	   + 2 * z * conj(dtz) * dz
	   + 2 * abs(z) .^ 2 * dtdz
	   + 2 * z * dtz * conj(dz)
	   + z .^ 2 * conj(dtdz))

    champu2 = -1im * llambda .* A2 .* exp.(-1im * t * A1) .* (-1im * A1) .* fft(conjmatr .* fz1)
    champv2 = -1im * llambda .* A2 .* exp.(-1im * t * A1) .* (-1im * A1) .* fft(conjmatr .* conj(fz1))

    fz1 = 2 * abs.(z) .^ 2 * dz .+ z .^ 2 * conj(dz)
    champu3 = -1im * llambda * A2 .* exp.(-1im * t .* A1) .* (-A1 .^ 2) .* fft(conjmatr .* fz1)
    champv3 = -1im * llambda * A2 .* exp.(-1im * t .* A1) .* (-A1 .^ 2) .* fft(conjmatr .* conj(fz1))

    champu = champu1 .+ 2 * champu2 .+ champu3
    champv = champv1 .+ 2 * champv2 .+ champv3

    champu, champv

end


using Flux
using CuArrays

export NeuralNet

# This code is originally written by Avik Pal
struct ResidualBlock
    conv_layers
    norm_layers
    shortcut
end

Flux.@treelike ResidualBlock

# ResidualBlock Function allows us to define a Residual Block having any number of Convolution and Batch Normalization Layers
function ResidualBlock(filters, kernels::Array{Tuple{Int,Int}},
                       pads::Array{Tuple{Int,Int}},
                       strides::Array{Tuple{Int,Int}},
                       shortcut = identity)
    local conv_layers = []
    local norm_layers = []
    for i in 2:length(filters)
        @inbounds push!(conv_layers, Conv(kernels[i-1], filters[i-1]=>filters[i], pad = pads[i-1], stride = strides[i-1]) |> gpu)
        push!(norm_layers, BatchNorm(filters[i])|>gpu)
    end
    ResidualBlock(Tuple(conv_layers),Tuple(norm_layers),shortcut)
end

# Function converts the Array of scalar kernel, pad and stride values to tuples
function ResidualBlock(filters, kernels::Array{Int}, pads::Array{Int}, strides::Array{Int}, shortcut = identity)
    ResidualBlock(filters, [(i,i) for i in kernels], [(i,i) for i in pads], [(i,i) for i in strides], shortcut)
end

function (block::ResidualBlock)(input)
    local value = copy.(input)
    for i in 1:length(block.conv_layers)-1
        @inbounds value = relu.((block.norm_layers[i])((block.conv_layers[i])(value)))
    end
    @inbounds relu.(((block.norm_layers[end])((block.conv_layers[end])(value))) + block.shortcut(input))
end

struct NeuralNet
    base_net::Chain
    value::Chain
    policy::Chain
end

Flux.@treelike NeuralNet

function NeuralNet(N::Int = 9, tower_height::Int = 20, features::Int=256)

    res_block() = ResidualBlock([features,features,features], [3,3], [1,1], [1,1])
    # 9 residual blocks
    tower = [res_block() for i = 1:tower_height]

    # feature extraction
    base_net = Chain(Conv((3,3), 9=>features, pad=(1,1)),
                     BatchNorm(features, leakyrelu),
                     tower...) |> gpu

    # value head
    value = Chain(Conv((1,1), features=>1),
                  BatchNorm(1, relu), x->reshape(x, N^2, :),
                  Dense(N^2, features, leakyrelu),
                  Dense(features, 1, tanh)) |> gpu

    policy = Chain(Conv((1,1), features=>2),
                   BatchNorm(2, relu), x->reshape(x, 2N^2, size(x, 4)),
                   Dense(2N^2, 9), softmax) |> gpu

    NeuralNet(base_net, value, policy)
end

function Base.show(io::IO, nn::NeuralNet)
    compact = get(io, :compact, false)
    if !compact
        print(io, "NeuralNet()")
    end
end

function (nn::NeuralNet)(input::AbstractArray{Float32,4})
    features = nn.base_net(input)
    p = nn.policy(features)
    v = nn.value(features)
    return p, v
end


include("../src/LarSurf.jl")
using Plasm
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
# using LarSurf

block_size = [5, 5, 5]
margin_size = 0

## Artifical data
data3d = ones(Int16, 3, 3, 3)
data3d[2, 2, 2] = 10
# data3d

voxelsize_mm = [0.5, 1.0, 2.]
threshold=0

verts, trifaces = LarSurf.import_data3d(data3d, voxelsize_mm, 1)

println(verts)

println(trifaces)

println("======= To LAR ========")

V, EV, FE = LarSurf.to_lar(verts, trifaces)

println(V)
println(EV)
println(FE)

println("V: ", typeof(V), size(V))
println("EV: ", typeof(EV))
print("FV: ", typeof(FE))

@show EV;

# convert to list of lists
# trifaces_list = [trifaces[k, :] for k=1:size(trifaces,1)]
# this is not general
trifaces_list = [trifaces[k, 3:-1:1] for k=1:size(trifaces,1)]
Vt = permutedims(V, [2,1])

# Plasm.view(Lar.cuboid([1,1,1]))

Plasm.view(Vt, trifaces_list)
Plasm.viewexploded(Vt, trifaces_list)(2,2,2)


import NbodyGradient: get_orbital_elements

function Base.isapprox(a::Elements,b::Elements;tol=1e-10)
    fields = setdiff(fieldnames(Elements),[:a,:e,:ϖ])
    for i in fields
        af = getfield(a,i)
        bf = getfield(b,i)
        if abs(af - bf) > tol
            return false
        end
    end
    return true
end

@testset "Cartesian to Elements" begin

    # Get known Elements
    fname = "elements.txt"
    t0 = 7257.93115525-7300.0
    H = [4,1,1,1]
    ic = ElementsIC(t0,H,fname)

    a = Elements(ic.elements[1,:]...)
    b = Elements(ic.elements[2,:]...)
    c = Elements(ic.elements[3,:]...)
    d = Elements(ic.elements[4,:]...)
    system = [a,b,c,d]

    # Convert to Cartesian coordinates
    s = State(ic)

    # Convert back to elements
    elems = get_orbital_elements(s,ic)

    for i in eachindex(system)
        @test isapprox(elems[1],system[1])
    end
end


import Base: +, *, zero, cos
using Distances: sqeuclidean, SqEuclidean, Euclidean
using Base.Broadcast: broadcast_shape
using LinearAlgebra: isposdef, checksquare

abstract type Kernel end

abstract type BaseKernel <: Kernel end

# API exports
export Kernel, kernel, elementwise, pairwise, ew, pw, stretch, scale

# Kernel exports
export EQ, Exp, PerEQ, Matern12, Matern32, Matern52, RQ, Cosine, Linear, Poly, GammaExp,
    Wiener, WienerVelocity, Precomputed

export scale, stretch

#
# Base Kernels
#

"""
    ZeroKernel <: BaseKernel

A rank 0 `Kernel` that always returns zero.
"""
struct ZeroKernel{T<:Real} <: BaseKernel end
ZeroKernel() = ZeroKernel{Float64}()
zero(::Kernel) = ZeroKernel()

# Binary methods.
ew(k::ZeroKernel{T}, x::AV, x′::AV) where {T} = zeros(T, broadcast_shape(size(x), size(x′)))
pw(k::ZeroKernel{T}, x::AV, x′::AV) where {T} = zeros(T, length(x), length(x′))

# Unary methods.
ew(k::ZeroKernel{T}, x::AV) where {T} = zeros(T, length(x))
pw(k::ZeroKernel{T}, x::AV) where {T} = zeros(T, length(x), length(x))

"""
    OneKernel{T<:Real} <: BaseKernel

A rank 1 constant `BaseKernel`. Useful for consistency when creating composite Kernels,
but (almost certainly) shouldn't be used as a base `BaseKernel`.
"""
struct OneKernel{T<:Real} <: BaseKernel end
OneKernel() = OneKernel{Float64}()

# Binary methods.
ew(k::OneKernel{T}, x::AV, x′::AV) where {T} = ones(T, broadcast_shape(size(x), size(x′)))
pw(k::OneKernel{T}, x::AV, x′::AV) where {T} = ones(T, length(x), length(x′))

# Unary methods.
ew(k::OneKernel{T}, x::AV) where {T} = ones(T, length(x))
pw(k::OneKernel{T}, x::AV) where {T} = ones(T, length(x), length(x))

"""
    ConstKernel{T} <: BaseKernel

A rank 1 kernel that returns the same value `c` everywhere.
"""
struct ConstKernel{T} <: BaseKernel
    c::T
end
ConstKernel(c::Real) = ConstKernel(typeof(c)[c])
get_iparam(c::ConstKernel) = c.c

# Binary methods.
ew(k::ConstKernel, x::AV, x′::AV) = fill(k.c[1], broadcast_shape(size(x), size(x′))...)
pw(k::ConstKernel, x::AV, x′::AV) = fill(k.c[1], length(x), length(x′))

# Unary methods.
ew(k::ConstKernel, x::AV) = fill(k.c[1], length(x))
pw(k::ConstKernel, x::AV) = pw(k, x, x)

@doc raw"""
    EQ() <: BaseKernel

The standardised Exponentiated Quadratic kernel. a.k.a. the Radial Basis Function (RBF), or
Squared Exponential kernel.

``k(x, x^\prime) = \exp( -\frac{1}{2} || x - x^\prime||_2^2 )``

For length scales etc see [`stretch`](@ref), for variance see [`*`](@ref).
"""
struct EQ <: BaseKernel end

# Binary methods.
ew(::EQ, x::AV, x′::AV) = exp.(.-ew(SqEuclidean(), x, x′) ./ 2)
pw(::EQ, x::AV, x′::AV) = exp.(.-pw(SqEuclidean(), x, x′) ./ 2)

# Unary methods.
ew(::EQ, x::AV) = exp.(.-ew(SqEuclidean(), x) ./ 2)
pw(::EQ, x::AV) = exp.(.-pw(SqEuclidean(), x) ./ 2)

@doc raw"""
    PerEQ <: BaseKernel

The usual periodic kernel derived by mapping the input domain onto the unit circle.

`` k(x, x^\prime) = \exp (-2 (\sin (\pi | x - x^\prime |) / l)^2)``

For length scales etc see [`stretch`](@ref), for variance see [`*`](@ref).
"""
struct PerEQ{LT<:Parameter{<:Real}} <: BaseKernel
    l::LT
end
PerEQ(l::Real) = PerEQ(Parameter(l, Val(:pos)))

_pereq(d, l) = exp(-2.0*sin(π*d)^2 / l^2)

# Binary methods.
ew(k::PerEQ, x::AV, x′::AV) = _pereq.(ew(Euclidean(), x, x′), value(k.l))
pw(k::PerEQ, x::AV, x′::AV) = _pereq.(pw(Euclidean(), x, x′), value(k.l))

# Unary methods.
ew(k::PerEQ, x::AV) = _pereq.(ew(Euclidean(), x), value(k.l))
pw(k::PerEQ, x::AV) = _pereq.(pw(Euclidean(), x), value(k.l))

@doc raw"""
    Matern12 <: BaseKernel

The standardised Matern-1/2 / Exponential kernel:

`` k(x, x^\prime) = \exp(-||x - x^\prime||_2)``

For length scales etc see [`stretch`](@ref), for variance see [`*`](@ref).
"""
struct Matern12 <: BaseKernel end

# Binary methods
ew(k::Matern12, x::AV, x′::AV) = exp.(.-ew(Euclidean(), x, x′))
pw(k::Matern12, x::AV, x′::AV) = exp.(.-pw(Euclidean(), x, x′))

# Unary methods
ew(::Matern12, x::AV) = exp.(.-ew(Euclidean(), x))
pw(::Matern12, x::AV) = exp.(.-pw(Euclidean(), x))

"""
    Exp <: Kernel

The standardised Exponential kernel. Equivalent to [`Matern12`](@ref).

For length scales etc see [`stretch`](@ref), for variance see [`*`](@ref).
"""
const Exp = Matern12

"""
    Matern32 <: BaseKernel

The standardised Matern kernel with ν = 3 / 2.

For length scales etc see [`stretch`](@ref), for variance see [`*`](@ref).
"""
struct Matern32 <: BaseKernel end

function _matern32(d::Real)
    d = sqrt(3) * d
    return (1 + d) * exp(-d)
end

# Binary methods
ew(k::Matern32, x::AV, x′::AV) = _matern32.(ew(Euclidean(), x, x′))
pw(k::Matern32, x::AV, x′::AV) = _matern32.(pw(Euclidean(), x, x′))

# Unary methods
ew(k::Matern32, x::AV) = _matern32.(ew(Euclidean(), x))
pw(k::Matern32, x::AV) = _matern32.(pw(Euclidean(), x))

"""
    Matern52 <: BaseKernel

The standardised Matern kernel with ν = 5 / 2.

For length scales etc see [`stretch`](@ref), for variance see [`*`](@ref).
"""
struct Matern52 <: BaseKernel end

function _Matern52(d::Real)
    λ = sqrt(5) * d
    return (1 + λ + λ^2 / 3) * exp(-λ)
end

_Matern52(d::AbstractArray{<:Real}) = _Matern52.(d)

@adjoint function _Matern52(d::AbstractArray{<:Real})
    λ = sqrt(5) .* d
    b = exp.(-λ)
    return (1 .+ λ .+ λ.^2 ./ 3) .* b, Δ->(.-Δ .* sqrt(5) .* b .* λ .* (1 .+ λ) ./ 3,)
end

# Binary methods
ew(k::Matern52, x::AV, x′::AV) = _Matern52(ew(Euclidean(), x, x′))
pw(k::Matern52, x::AV, x′::AV) = _Matern52(pw(Euclidean(), x, x′))

# Unary methods
ew(k::Matern52, x::AV) = _Matern52(ew(Euclidean(), x))
pw(k::Matern52, x::AV) = _Matern52(pw(Euclidean(), x))

"""
    RQ <: BaseKernel

The standardised Rational Quadratic, with kurtosis `α`.

For length scales etc see [`stretch`](@ref), for variance see [`*`](@ref).
"""
struct RQ{Tα<:Parameter{<:Real}} <: BaseKernel
    α::Tα
end
RQ(α::Real) = RQ(Parameter(α, Val(:pos)))

_rq(d, α) = (1 + d / (2α))^(-α)

# I redefine adjoint function for `_rq`, since on my computer, Zygote's gradient
# of `_rq` returns a number of Complex type.
Zygote.@adjoint function _rq(d::dT, α::αT) where {dT<:Real, αT<:Real}
    y = _rq(d, α)
    return y, function (ȳ)
    	T = promote_type(dT, αT)
    	x = 1 + d / (2α)
    	-0.5*ȳ*y/x, ȳ*y*(d / (x*(2α)) - log(x+eps(T)))
    	end
end

# Binary methods.
ew(k::RQ, x::AV, x′::AV) = _rq.(ew(SqEuclidean(), x, x′), value(k.α))
pw(k::RQ, x::AV, x′::AV) = _rq.(pw(SqEuclidean(), x, x′), value(k.α))

# Unary methods.
ew(k::RQ, x::AV) = _rq.(ew(SqEuclidean(), x), value(k.α))
pw(k::RQ, x::AV) = _rq.(pw(SqEuclidean(), x), value(k.α))

"""
    Cosine <: BaseKernel

Cosine BaseKernel with period parameter `p`.
"""
struct Cosine{Tp<:Parameter{<:Real}} <: BaseKernel
    p::Tp
end
Cosine(p::Real) = Cosine(Parameter(p, Val(:pos)))

# Binary methods.
function ew(k::Cosine, x::AV{<:Real}, x′::AV{<:Real})
    return cos.(π .* ew(Euclidean(), x, x′) ./ value(k.p))
end
function pw(k::Cosine, x::AV{<:Real}, x′::AV{<:Real})
    return cos.(π .* pw(Euclidean(), x, x′) ./ value(k.p))
end

# Unary methods.
ew(k::Cosine, x::AV{<:Real}) = 1 .+ ew(Euclidean(), x)
pw(k::Cosine, x::AV{<:Real}) = cos.(π .* pw(Euclidean(), x) ./ value(k.p))

"""
    Linear{T<:Real} <: BaseKernel

The standardised linear kernel / dot-product kernel.
"""
struct Linear <: BaseKernel end

# Binary methods
ew(k::Linear, x::AV{<:Real}, x′::AV{<:Real}) = x .* x′
pw(k::Linear, x::AV{<:Real}, x′::AV{<:Real}) = x .* x′'
ew(k::Linear, x::ColVecs, x′::ColVecs) = reshape(sum(x.X .* x′.X; dims=1), :)
pw(k::Linear, x::ColVecs, x′::ColVecs) = x.X' * x′.X

# Unary methods
ew(k::Linear, x::AV{<:Real}) = x.^2
pw(k::Linear, x::AV{<:Real}) = x .* x'
ew(k::Linear, x::ColVecs) = reshape(sum(abs2.(x.X); dims=1), :)
pw(k::Linear, x::ColVecs) = x.X' * x.X

"""
    Poly{Tσ<:Real} <: BaseKernel

Inhomogeneous Polynomial kernel. `Poly(p, σ²)` creates a `Poly{p}` with variance σ²,
defined as
```julia
k(xl, xr) = (dot(xl, xr) + σ²)^p
```
"""
struct Poly{p, Tσ²<:Parameter{<:Real}} <: BaseKernel
    σ²::Tσ²
end
Poly(p::Int, σ²::T) where {T} = Poly{p, Parameter{T}}(Parameter(σ², Val(:pos)))

_poly(k, σ², p) = (σ² + k)^p
Zygote.@adjoint function _poly(k, σ², p)
    y = _poly(k, σ², p)
    return y, function(Δ)
        d = Δ * p * y / (σ² + k)
        return (d, d, nothing)
    end
end

# Binary methods
ew(k::Poly{p}, x::AV, x′::AV) where {p} = _poly.(ew(Linear(), x, x′), value(k.σ²), p)
pw(k::Poly{p}, x::AV, x′::AV) where {p} = _poly.(pw(Linear(), x, x′), value(k.σ²), p)

# Unary methods
ew(k::Poly{p}, x::AV) where {p} = _poly.(ew(Linear(), x), value(k.σ²), p)
pw(k::Poly{p}, x::AV) where {p} = _poly.(pw(Linear(), x), value(k.σ²), p)

"""
    GammaExp <: BaseKernel

The γ-Exponential kernel, 0 < γ ⩽ 2, is given by `k(xl, xr) = exp(-||xl - xr||^γ)`.
"""
struct GammaExp{Tγ<:Parameter{<:Real}} <: BaseKernel
    γ::Tγ
end
GammaExp(γ::T) where {T} = GammaExp(Parameter(γ, inv(Logit(zero(T), T(2.0)))))

# Binary methods
ew(k::GammaExp, x::AV, x′::AV) = exp.(.-ew(Euclidean(), x, x′).^(value(k.γ)))
pw(k::GammaExp, x::AV, x′::AV) = exp.(.-pw(Euclidean(), x, x′).^(value(k.γ)))

# Unary methods
ew(k::GammaExp, x::AV) = exp.(.-ew(Euclidean(), x).^(value(k.γ)))
pw(k::GammaExp, x::AV) = exp.(.-pw(Euclidean(), x).^(value(k.γ)))

"""
    Wiener <: BaseKernel

The standardised stationary Wiener-process kernel.
"""
struct Wiener <: BaseKernel end

_wiener(x::Real, x′::Real) = min(x, x′)

# Binary methods
ew(k::Wiener, x::AV{<:Real}, x′::AV{<:Real}) = _wiener.(x, x′)
pw(k::Wiener, x::AV{<:Real}, x′::AV{<:Real}) = _wiener.(x, x′')

# Unary methods
ew(k::Wiener, x::AV{<:Real}) = x
pw(k::Wiener, x::AV{<:Real}) = pw(k, x, x)

"""
    WienerVelocity <: BaseKernel

The standardised WienerVelocity kernel.
"""
struct WienerVelocity <: BaseKernel end

_wiener_vel(x::Real, x′::Real) = min(x, x′)^3 / 3 + abs(x - x′) * min(x, x′)^2 / 2

# Binary methods
ew(k::WienerVelocity, x::AV{<:Real}, x′::AV{<:Real}) = _wiener_vel.(x, x′)
pw(k::WienerVelocity, x::AV{<:Real}, x′::AV{<:Real}) = _wiener_vel.(x, x′')

# Unary methods
ew(k::WienerVelocity, x::AV{<:Real}) = ew(k, x, x)
pw(k::WienerVelocity, x::AV{<:Real}) = pw(k, x, x)

"""
    Noise{T<:Real} <: BaseKernel

The standardised aleatoric white-noise kernel. Isn't really a kernel, but never mind...
"""
struct Noise{T<:Real} <: BaseKernel end
Noise() = Noise{Int}()

# Binary methods.
ew(k::Noise{T}, x::AV, x′::AV) where {T} = zeros(T, broadcast_shape(size(x), size(x′))...)
pw(k::Noise{T}, x::AV, x′::AV) where {T} = zeros(T, length(x), length(x′))

# Unary methods.
ew(k::Noise{T}, x::AV) where {T} = ones(T, length(x))
pw(k::Noise{T}, x::AV) where {T} = diagm(0=>ones(T, length(x)))

"""
    Precomputed{T<:Real} <: Kernel

Using the values of a precomputed Gram matrix as a kernel.

Optionally checks if the Gram matrix is positive definite by setting
`checkpd=true`
"""
struct Precomputed{M<:AbstractMatrix{<:Real}} <: Kernel
    K::M
    function Precomputed(K::AbstractMatrix{<:Real}; checkpd=false)
        checksquare(K)
        checkpd && @assert isposdef(K) "M is not positive definite"
        return new{typeof(K)}(K)
    end
end

# Binary methods.
ew(k::Precomputed, x::AV{<:Integer}, x′::AV{<:Integer}) = [k.K[p, q] for (p, q) in zip(x, x′)]
pw(k::Precomputed, x::AV{<:Integer}, x′::AV{<:Integer}) = k.K[x, x′]

# Unary methods.
ew(k::Precomputed, x::AV{<:Integer}) = diag(k.K)[x]
pw(k::Precomputed, x::AV{<:Integer}) = k.K[x,x]

precomputed(K::AbstractMatrix) = Precomputed(K)
export precomputed

#
# Composite Kernels
#

"""
    Sum{Tkl<:Kernel, Tkr<:Kernel} <: Kernel

Represents the sum of two kernels `kl` and `kr` s.t. `k(x, x′) = kl(x, x′) + kr(x, x′)`.
"""
struct Sum{Tkl<:Kernel, Tkr<:Kernel} <: Kernel
    kl::Tkl
    kr::Tkr
end
child(s::Sum) = (s.kl, s.kr)
"""
    +(kl::Kernel, kr::Kernel)

Construct the kernel whose value is given by the sum of those of `kl` and `kr`.

```jldoctest
julia> kl, kr = EQ(), Matern32();

julia> x = randn(11);

julia> pw(kl + kr, x) == pw(kl, x) + pw(kr, x)
true
```
"""
+(kl::Kernel, kr::Kernel) = Sum(kl, kr)

# Binary methods
ew(k::Sum, x::AV, x′::AV) = ew(k.kl, x, x′) + ew(k.kr, x, x′)
pw(k::Sum, x::AV, x′::AV) = pw(k.kl, x, x′) + pw(k.kr, x, x′)

# Unary methods
ew(k::Sum, x::AV) = ew(k.kl, x) + ew(k.kr, x)
pw(k::Sum, x::AV) = pw(k.kl, x) + pw(k.kr, x)

"""
    Product{Tkl<:Kernel, Tkr<:Kernel} <: Kernel

Represents the product of two kernels `kl` and `kr` s.t. `k(x, x′) = kl(x, x′) kr(x, x′)`.
"""
struct Product{Tkl<:Kernel, Tkr<:Kernel} <: Kernel
    kl::Tkl
    kr::Tkr
end
child(p::Product) = (p.kl, p.kr)
"""
    *(kl::Kernel, kr::Kernel)

Construct the kernel whose value is given by the product of those of `kl` and `kr`.

```jldoctest
julia> kl, kr = EQ(), Matern32();

julia> x = randn(11);

julia> pw(kl * kr, x) == pw(kl, x) .* pw(kr, x)
true
```
"""
*(kl::Kernel, kr::Kernel) = Product(kl, kr)

# Binary methods
ew(k::Product, x::AV, x′::AV) = ew(k.kl, x, x′) .* ew(k.kr, x, x′)
pw(k::Product, x::AV, x′::AV) = pw(k.kl, x, x′) .* pw(k.kr, x, x′)

# Unary methods
ew(k::Product, x::AV) = ew(k.kl, x) .* ew(k.kr, x)
pw(k::Product, x::AV) = pw(k.kl, x) .* pw(k.kr, x)

"""
Scaled{Tσ²<:AV{<:Real}, Tk<:Kernel} <: Kernel

Scale the variance of `Kernel` `k` by `σ²` s.t. `(σ² * k)(x, x′) = σ² * k(x, x′)`.
"""
struct Scaled{Tσ²<:Parameter{<:Real}, Tk<:Kernel} <: Kernel
    σ²::Tσ²
    k::Tk
end
scale(k::Kernel, σ²::Real) = Scaled(Parameter(σ², Val(:pos)), k)

"""
    *(σ²::Real, k::Kernel)
    *(k::Kernel, σ²::Real)

The right way to choose the variance of a kernel. Specifically, construct a kernel that
scales the output of `k` by `σ²`:
```jldoctest
julia> k = EQ();

julia> x = randn(11);

julia> pw(0.5 * k, x) == 0.5 .* Stheno.pw(k, x)
true
```
"""
*(σ²::Real, k::Kernel) = scale(k, σ²)
*(k::Kernel, σ²) = σ² * k

# Binary methods.
ew(k::Scaled, x::AV, x′::AV) = value(k.σ²) .* ew(k.k, x, x′)
pw(k::Scaled, x::AV, x′::AV) = value(k.σ²) .* pw(k.k, x, x′)

# Unary methods.
ew(k::Scaled, x::AV) = value(k.σ²) .* ew(k.k, x)
pw(k::Scaled, x::AV) = value(k.σ²) .* pw(k.k, x)

"""
    Stretched{Tk<:Kernel} <: Kernel

Apply a length scale to a kernel. Specifically, `k(x, x′) = k(x/a, x′/a)`.
"""
struct Stretched{Ta<:Parameter{<:Real}, Tk<:Kernel} <: Kernel
    a::Ta
    k::Tk
end
"""
    stretch(k::Kernel, a::Union{Real, AbstractVecOrMat{<:Real})

The canonical way to give a notion of length-scale to a kernel, `a::Real`, to construct
an automatic relevance determination (ARD) kernel, `a::AbstractVector{<:Real}`, or a
discriminative factor-analysis kernel, `a::AbstractMatrix{<:Real}`.

# Length Scale

For `a::Real`, return the kernel given by:
```julia
stretch(k, a)(x, y) = k(a .* x, a .* y)
```
where `x` and `y` are both either `Real`s or `AbstractVector{<:Real}`s. e.g.
```jldoctest
xs = range(0.0, 10.0; length=2)
ys = range(0.5, 10.5; length=3)
k = stretch(EQ(), 0.5)
K = pairwise(k, xs, ys)

# output
2×3 Array{Float64,2}:
 0.969233    0.0227942  1.03485e-6
 1.26071e-5  0.0795595  0.969233
```
if `x` and `y` are `Real`s.

# Automatic Relevance Determination (ARD)

For `a::AbstractVector{<:Real}`, return the automatic relevance determination (ARD) kernel
given by:
```julia
stretch(k, a)(x, y) = k(a .* x, a .* y)
```
where `x` and `y` must be `AbstractVector{<:Real}`s. e.g.

```jldoctest
rng = MersenneTwister(123456)
xs = ColVecs(randn(rng, 2, 2)) # efficient vector-of-column-vectors
ys = ColVecs(randn(rng, 2, 3)) # efficient vector-of-column-vectors
length_scales = [0.1 0.5]
k = stretch(EQ(), 1 ./ length_scales)
K = pairwise(k, xs, ys)

# output
2×3 Array{Float64,2}:
 1.00858e-8   0.000170346  0.0439152
 2.79756e-15  2.0696e-74   9.91334e-31
```

# Discriminative Factor Analysis

For `a::AbstractMatrix{<:Real}`, return the kernel given by:
```julia
stretch(k, a)(x, y) = k(a * x, a * y)
```
where  `x` and `y` must be `AbstractVector{<:Real}`s. e.g.

```jldoctest
rng = MersenneTwister(123456)

# Efficiently represented vectors of column-vectors
xs = ColVecs(randn(rng, 5, 2))
ys = ColVecs(randn(rng, 5, 3))

# Project from 5 dimensions down onto 2
A = rand(rng, 2, 5)

# Construct kernel and compute covariance matrix between `xs` and `ys`.
k = stretch(EQ(), A)
K = pairwise(k, xs, ys)

# output
2×3 Array{Float64,2}:
 0.000289928  0.00519575  0.000127567
 0.0507073    0.495861    0.432123
```
"""
stretch(k::Kernel, a::Union{Real, AVM{<:Real}}) = Stretched(Parameter(a, Val(:pos)), k)

# Binary methods (scalar `a`, scalar-valued input)
function ew(k::Stretched{<:Parameter{<:Real, 0}}, x::AV{<:Real}, x′::AV{<:Real})
    return ew(k.k, value(k.a) .* x, value(k.a) .* x′)
end
function pw(k::Stretched{<:Parameter{<:Real, 0}}, x::AV{<:Real}, x′::AV{<:Real})
    return pw(k.k, value(k.a) .* x, value(k.a) .* x′)
end

# Unary methods (scalar `a`, scalar-valued input)
ew(k::Stretched{<:Parameter{<:Real, 0}}, x::AV{<:Real}) = ew(k.k, value(k.a) .* x)
pw(k::Stretched{<:Parameter{<:Real, 0}}, x::AV{<:Real}) = pw(k.k, value(k.a) .* x)

# Binary methods (scalar/vector `a`, vector-valued input)
function ew(k::Stretched{<:Union{Parameter{<:Real, 0}, Parameter{<:Real, 1}}}, x::ColVecs, x′::ColVecs)
    return ew(k.k, ColVecs(value(k.a) .* x.X), ColVecs(value(k.a) .* x′.X))
end
function pw(k::Stretched{<:Union{Parameter{<:Real, 0}, Parameter{<:Real, 1}}}, x::ColVecs, x′::ColVecs)
    return pw(k.k, ColVecs(value(k.a) .* x.X), ColVecs(value(k.a) .* x′.X))
end

# Unary methods (vector `a`, vector-valued input)
ew(k::Stretched{<:Union{Parameter{<:Real, 0}, Parameter{<:Real, 1}}}, x::ColVecs) = ew(k.k, ColVecs(value(k.a) .* x.X))
pw(k::Stretched{<:Union{Parameter{<:Real, 0}, Parameter{<:Real, 1}}}, x::ColVecs) = pw(k.k, ColVecs(value(k.a) .* x.X))

# Binary methods (matrix `a`, vector-valued input)
function ew(k::Stretched{<:Parameter{<:Real, 2}}, x::ColVecs, x′::ColVecs)
    return ew(k.k, ColVecs(value(k.a) * x.X), ColVecs(value(k.a) * x′.X))
end
function pw(k::Stretched{<:Parameter{<:Real, 2}}, x::ColVecs, x′::ColVecs)
    return pw(k.k, ColVecs(value(k.a) * x.X), ColVecs(value(k.a) * x′.X))
end

# Unary methods (matrix `a`, vector-valued input)
ew(k::Stretched{<:Parameter{<:Real, 2}}, x::ColVecs) = ew(k.k, ColVecs(value(k.a) * x.X))
pw(k::Stretched{<:Parameter{<:Real, 2}}, x::ColVecs) = pw(k.k, ColVecs(value(k.a) * x.X))

"""
    kernel(k::Kernel;  l::Real=nothing, s::Real=nothing)

Convenience functionality to provide a kernel with a length scale `l`, and to scale the
variance of `k` by `s`. Simply applies the [`stretch`](@ref stretch) and [`*`](@ref *)
functions.

```jldoctest
julia> k1 = kernel(EQ(); l=1.1, s=0.9);

julia> k2 = 0.9 * stretch(EQ(), 1 / 1.1);

julia> x = randn(11);

julia> pw(k1, x) == pw(k2, x)
true
```
"""
function kernel(k::Kernel; l::Union{Real, Nothing}=nothing, s::Union{Real, Nothing}=nothing)
    if l !== nothing
        k = stretch(k, 1 / l)
    end
    if s !== nothing
        k = s * k
    end
    return k
end


## Loading packages.
using MicroHH
using PyPlot
using MicroHH.StencilBuilder
using LoopVectorization

## Loading settings.
include("drycbl_settings.jl")
settings[1]["timeloop"]["start_time"] = settings[1]["timeloop"]["end_time"]

## Initialize the model.
n_domains = 1
m = Model("drycbl", n_domains, settings, Float32)
load_model!(m)

## Compute the scalar dissipation
function calc_scalar_dissipation!(
    slngrad, s,
    dxi, dyi, dzi, dzhi,
    is, ie, js, je, ks, ke)

    @fast3d begin
        @fd (slngrad, s) slngrad = log(gradx(gradx(s))^2 + grady(grady(s))^2 + gradz(gradz(s))^2)
    end
end

f = m.fields[1]; g = m.grid[1]
slngrad = zeros(size(f.s))
calc_scalar_dissipation!(
    slngrad, f.s,
    g.dxi, g.dyi, g.dzi, g.dzhi,
    g.is, g.ie, g.js, g.je, g.ks, g.ke)

x = @view g.x[g.is:g.ie]
y = @view g.y[g.js:g.je]
z = @view g.z[g.ks:g.ke]
slngrad_xy = @view slngrad[g.is:g.ie, g.js:g.je, 6]
slngrad_xz = @view slngrad[g.is:g.ie, g.js, g.ks:g.ke]

smin = -25; smax = -11

figure()
#pcolormesh(x, y, slngrad_xy', vmin=smin, vmax=smax, shading="nearest", cmap=plt.cm.magma)
pcolormesh(x, y, slngrad_xy', shading="nearest", cmap=plt.cm.magma)
xlabel("x (m)")
ylabel("y (m)")
colorbar()
tight_layout()
display(gcf())

figure()
#pcolormesh(x, z, slngrad_xz', vmin=smin, vmax=smax, shading="nearest", cmap=plt.cm.magma)
pcolormesh(x, z, slngrad_xz', shading="nearest", cmap=plt.cm.magma)
xlabel("x (m)")
ylabel("z (m)")
colorbar()
tight_layout()
display(gcf())

show()


# This is an attempt to include the required for n-body simulations fields into structures
# Indeed, tools of ConcreteAbstractions.jl seem to be more suitable for the fields inhereting
DiffEqBase.@def position_velocity_mass begin
    r::SVector{3,cType}
    v::SVector{3,cType}
    m::mType
end

abstract type Body
end

struct MassBody{cType<:Real, mType<:Real} <: Body
    @position_velocity_mass
end

struct ChargedParticle{cType<:Real, mType<:Real, qType<:Real} <: Body
    @position_velocity_mass
    q::qType
end

struct MagneticParticle{cType<:Real, mType<:Real, mmType<:Real} <: Body
    @position_velocity_mass
    mm::SVector{3,mmType}
end

struct WaterMolecule <: Body
    O::MassBody
    H1::MassBody
    H2::MassBody
end

function generate_bodies_in_cell_nodes(n::Integer, m::Real, v_dev::Real, L::Real; T=Float64)
    rng = MersenneTwister(n)
    velocities = v_dev * randn(rng, T, (3,n))
    bodies = MassBody[]
    count = 1
    dL = T(L / (ceil(n^(1/3))))
    for x = dL/2:dL:L, y = dL/2:dL:L, z = dL/2:dL:L
        if count > n
            break
        end
        r = SVector(x, y, z)
        v = SVector{3,T}(velocities[:,count])
        body = MassBody{T,T}(r, v, m)
        push!(bodies, body)
        count += 1
    end
    return bodies
end


export sampleTangent
export samplePoint
"""
    $SIGNATURES

Return a random sample as a tangent vector from a belief represented by coordinates on a manifold at point p.

Notes

"""
function sampleTangent end

# Sampling MKD
function sampleTangent(M::AbstractGroupManifold, x::ManifoldKernelDensity, p=mean(x))
    # get legacy matrix of coordinates and selected labels
    coords, lbls = sample(x.belief,1)
    X = hat(x.manifold, p, coords)
    return X
end

function sampleTangent(x::ManifoldKernelDensity, p=mean(x))
  return sampleTangent(x.manifold, x, p)
end

# Sampling Distributions
function sampleTangent(M::AbstractManifold, z::Distribution, p, basis::AbstractBasis)
  return get_vector(M, p, rand(z), basis)
end

function sampleTangent(M::AbstractGroupManifold, z::Distribution, p=identity_element(M))
  return hat(M, p, rand(z,1)[:]) #TODO find something better than (z,1)[:]
end

"""
    $SIGNATURES

Return a random sample point on a manifold from a belief represented by coordinates at point p.

Notes

"""
function samplePoint(M::AbstractManifold, sbelief, p, basis::AbstractBasis, retraction_method::AbstractRetractionMethod=ExponentialRetraction())
  X = sampleTangent(M, sbelief, p, basis)
  return retract(M, p, X, retraction_method)
end
function samplePoint(M::AbstractGroupManifold, sbelief, p=identity_element(M), retraction_method::AbstractRetractionMethod=ExponentialRetraction())
  X = sampleTangent(M, sbelief, p)
  return retract(M, p, X, retraction_method)
end

function samplePoint(M::AbstractGroupManifold, sbelief::ManifoldKernelDensity, p=identity_element(M, mean(sbelief)), retraction_method::AbstractRetractionMethod=ExponentialRetraction())
  X = sampleTangent(M, sbelief, p)
  return retract(M, p, X, retraction_method)
end

function samplePoint(x::ManifoldKernelDensity, p=mean(x))
  return samplePoint(x.manifold, x, p)
end

# FIXME: rather use manifolds
function samplePoint(distr::SamplableBelief)
  Base.depwarn("samplePoint(distr::SamplableBelief) should be replaced by samplePoint(M<:AbstractManifold, distr::SamplableBelief, ...)", :samplePoint)
  rand(distr,1)
end

## default getSample
# getSample(cf::CalcFactor{<:AbstractPrior}, N::Int=1) = ([samplePoint(getManifold(cf.factor), cf.factor.Z, ) for _=1:N], )
# getSample(cf::CalcFactor{<:AbstractRelative}, N::Int=1) = ([sampleTangent(getManifold(cf.factor), cf.factor.Z) for _=1:N], )


for p in ("Knet","ArgParse","Compat","GZip")
    Pkg.installed(p) == nothing && Pkg.add(p)
end

using Knet,ArgParse,Compat,GZip

z(P,i,x) = P[1][i]*x + P[2][i]
sig_der(x) = sigm(x)*(1-sigm(x))

#Analytical Solution of the Eq
diff_sol(x) = exp(-(x^2)/2)/(1+x+(x)^3) + x^2

#The phi trial solution
phi(P,x) = 1 + x*(predict(P,x))

#The Differential Equation function for Example 1
f(P,x) = x^3 + 2*x + (x^2)*((1+3*(x^2))/(1+x+(x^3))) - (phi(P,x))*(x + ((1+3*(x^2))/(1+x+x^3)))

 x_tr = linspace(0.0,1,10)
 x_tst = linspace(0.1,1,10)

function predict(P,x)
  w1, b1, w2 = P
  h = sigm(w1 * x .+ b1)
  return w2 * h
end

function loss(w,x,ygold)
    sumabs2(ygold - predict(w,x)) / size(ygold,2)
end

lossgradient = grad(loss)

function train(w, dtrn; lr=.09, epochs=500)
    for epoch=1:epochs
        for (x,y) in dtrn
            g = lossgradient(w, x, y)
            #println(" pred: ",ypred[1]," orig: ",ygold," error: "ypred-ygold)
            for i in 1:length(w)
                # w[i] -= lr * g[i]
                axpy!(-lr, g[i], w[i])
            end
        end
    end
    return w
end

function accuracy(w, dtst, pred=predict)
    ncorrect = ninstance = nerror = 0
    for (x, ygold) in dtst
        ypred = pred(w, x)[1]
        println(" X: ",x," Prediction: ",ypred[1]," True: ",ygold)
        nerror += (ygold - ypred[1])
        ninstance += 1
    end
    return (1-(nerror/ninstance), nerror/ninstance)
end

function init_params(;ftype=Float32,atype=KnetArray)
    #P = Vector{Vector{Float32}}(4)
    P = Array{Any}(3)
    P[1] = randn(Float32,10,1)
    P[2] = zeros(Float32,10,1)
    P[3] = randn(Float32,1,10)
    #P[4] = zeros(Float32,1,1)
    P = map(x -> convert(atype, x), P)

    return P
end

function batch(x, y, batchsize; atype=Array{Float32})
    data = Any[]
    x_ = convert(atype,x)
    y_ = convert(atype,y)
    for i=1:10
        push!(data, (x_[i],y_[i]))
    end
    return data
end

function main(args="")
    s = ArgParseSettings()
    s.description="mnist.jl (c) Deniz Yuret, 2016. Multi-layer perceptron model on the MNIST handwritten digit recognition problem from http://yann.lecun.com/exdb/mnist."
    s.exc_handler=ArgParse.debug_handler
    @add_arg_table s begin
        ("--seed"; arg_type=Int; default=-1; help="random number seed: use a nonnegative int for repeatable results")
        ("--batchsize"; arg_type=Int; default=1; help="minibatch size")
        ("--epochs"; arg_type=Int; default=500; help="number of epochs for training")
        ("--hidden"; nargs='*'; arg_type=Int; help="sizes of hidden layers, e.g. --hidden 128 64 for a net with two hidden layers")
        ("--lr"; arg_type=Float64; default=0.09; help="learning rate")
        ("--winit"; arg_type=Float64; default=0.1; help="w initialized with winit*randn()")
        ("--fast"; action=:store_true; help="skip loss printing for faster run")
        ("--atype"; default=(gpu()>=0 ? "KnetArray{Float32}" : "Array{Float32}"); help="array type: Array for cpu, KnetArray for gpu")
        ("--gcheck"; arg_type=Int; default=0; help="check N random gradients per parameter")
        # These are to experiment with sparse arrays
        # ("--xtype"; help="input array type: defaults to atype")
        # ("--ytype"; help="output array type: defaults to atype")
    end
    isa(args, AbstractString) && (args=split(args))
    o = parse_args(args, s; as_symbols=true)
    # if !o[:fast]
    #     println(s.description)
    #     println("opts=",[(k,v) for (k,v) in o]...)
    # end
    o[:seed] > 0 && srand(o[:seed])
    atype = eval(parse(o[:atype]))
    w = init_params()
    #w = weights(o[:hidden]...; atype=atype, winit=o[:winit])
    #if !isdefined(MNIST,:xtrn); loaddata(); end

    global dtrn = batch(x_tr, map(diff_sol,x_tr), o[:batchsize])
    global dtst = batch(x_tst, map(diff_sol,x_tst), o[:batchsize])
    report(epoch)=println((:epoch,epoch,:trn,accuracy(w,dtrn),:tst,accuracy(w,dtst)))
    if o[:fast]
        (train(w, dtrn; lr=o[:lr], epochs=o[:epochs]); gpu()>=0 && Knet.cudaDeviceSynchronize())
    else
        report(0)
        @time for epoch=1:o[:epochs]
            train(w, dtrn; lr=o[:lr], epochs=500)
            report(epoch)
            if o[:gcheck] > 0
                gradcheck(loss, w, first(dtrn)...; gcheck=o[:gcheck], verbose=true)
            end
        end
    end
    return w
end

# This allows both non-interactive (shell command) and interactive calls like:
# $ julia mnist.jl --epochs 10
# julia> MNIST.main("--epochs 10")
if VERSION >= v"0.5.0-dev+7720"
    PROGRAM_FILE == "mnist.jl" && main(ARGS)
else
    !isinteractive() && !isdefined(Core.Main,:load_only) && main(ARGS)
end

main()


const to = TimerOutput()

function run( nstep; npm = 100 )

    dt = 0.01
    alpha = 0.1
    kx   = 0.5
    dimx = 2*pi/kx
    dimy = 1
    nx   = 128  # nombre de pts suivant x
    ny   = 16   # nombre de pts suivant y
    mesh = TwoDGrid( dimx, nx, dimy, ny)
    dx, dy = mesh.dx, mesh.dy

    nbpart = npm*nx*ny
    println( " nbpart = $nbpart ")

    particles = ParticleGroup{2,2}(nbpart, charge = 1.0, mass = 1.0, n_weights = 1)

    sampler = LandauDamping(alpha, kx)

    sample!(particles, mesh, sampler)

    fdtd = FDTD(mesh)
    for i=1:nx, j=1:ny+1
        fdtd.ex[i,j] = alpha/kx * sin(kx*(mesh.x[i]+mesh.x[i+1])/2)
    end
    time = 0
    energy = Float64[compute_energy(fdtd, mesh)]
    t = Float64[time]

    kernel = CloudInCell()

    reset_timer!()
    @showprogress 1 for istep in 1:nstep

       if istep > 1
           @timeit to "fdtd" faraday!( fdtd, mesh, 0.5dt )
       end
       update_fields!(mesh, fdtd)
       #@timeit to "interpolation" interpolation!(particles, mesh)
       @timeit to "pushv" push_v!( particles, kernel, mesh, dt )
       @timeit to "pushx" push_x!( particles, mesh, 0.5dt)
       @timeit to "deposition" compute_current!( mesh, kernel, particles)
       @timeit to "pushx" push_x!( particles, mesh, 0.5dt)
       @timeit to "fdtd" faraday!(fdtd, mesh, 0.5dt)
       @timeit to "fdtd" ampere_maxwell!(fdtd, mesh, dt)
       time = time + dt
       push!(t, time)
       push!(energy, compute_energy(fdtd, mesh))

    end

    t, energy

end


#### Frank wolfe for matrix completion
using FirstOrderOptimization
import Base: axpy!, scale!

## Generate data

# make it deterministic
srand(10)

m, n, r, rtrue = 100, 100, 2, 2
Xtrue = randn(m, rtrue) * randn(rtrue, n)
tau = nucnorm(Xtrue)

pobs = .5
obs = sprand(m, n, pobs).>0
iobs = find(obs)
A = IndexingOperator(m, n, iobs)
b = A*Xtrue

nonsketched = true

params = FrankWolfeParams(maxiters = 200, reltol = 1e-1, stepsize = BacktrackingStepSize(1,.5,.3)) # DecreasingStepSize(2,1))

srand(10)

objective_sketched(z) = .5*sum((z - b).^2)
grad_objective_sketched(z) = LowRankOperator(A, z-b, transpose = Symbol[:T, :N]) # concretely, should be A'*(z-b)
function min_lin_st_nucnorm_sketched(G, tau)
	u,s,v = svds(Array(G), nsv=1) # for this case, g is a sparse matrix so representing it is O(m)
	return LowRankOperator(-tau*u, v')
end
const_nucnorm(z) = tau # we'll always saturate the constraint, don't bother computing it
# I can't think of a pretty way to compute
# <g, tilde_x - x> = <A'(z-b), u tau v' - x> = <z-b, A(u tau v') - A(x)> = <z-b, A(u tau v') - z>
# dot_g_w_tx_minus_x(g, tilde_x, x) = dot(g.factors[2], g.factors[1]*tilde_x - x)
dot_g_w_tx_minus_x(g, Delta, z) = dot(vec(g.factors[2]), g.factors[1]*Delta - z)
# we end up computing A*Delta twice with this scheme
update_var!(z, Delta, a) = (scale!(z, 1-a); axpy!(a, A*Delta, z))

z = zeros(length(b))
X_sketched = frank_wolfe_sketched(
	z,
	objective_sketched, grad_objective_sketched,
	const_nucnorm,
	tau,
	min_lin_st_nucnorm_sketched,
	AsymmetricSketch(m,n,r),
	params,
#	LB = 0,
	verbose=true
	)

@show vecnorm(Array(X_sketched) - Xtrue) / vecnorm(Xtrue)


using AdFem
using Test
using PyPlot
using PyCall
using LinearAlgebra

np = pyimport("numpy")
m = 40
n = 20
h = 0.1
bdnode = Int64[]
bdval1 = Float64[]
bdval2 = Float64[]
X = (0:m)*h
Y = (0:n)*h
X, Y = np.meshgrid(X,Y)
for i = 1:m+1
    for j = 1:n+1
        push!(bdnode, i+(j-1)*(n+1))
        x = (i-1)*h
        y = (j-1)*h
        push!(bdval1, x^2+y^2)
        push!(bdval2, x^2-y^2)
    end
end
bdnode = [bdnode;bdnode.+(m+1)*(n+1)]
bdval = [bdval1;bdval2]

Xv = []
Yv = []
for j = 1:n
    for i = 1:m
        push!(Xv, (i-1/2)*h)
        push!(Yv, (j-1/2)*h)
    end
end

include("Core.jl")


export Bernoulli

import Distributions: Categorical, probs

vague(::Type{ <: Categorical }, dims::Int) = Categorical(ones(dims) ./ dims)

prod_analytical_rule(::Type{ <: Categorical }, ::Type{ <: Categorical }) = ProdAnalyticalRuleAvailable()

function prod(::ProdAnalytical, left::Categorical, right::Categorical)
    # Multiplication of 2 categorical PMFs: p(z) = p(x) * p(y)
    mvec = clamp.(probvec(left) .* probvec(right), tiny, huge)
    norm = sum(mvec)
    return Categorical(mvec ./ norm)
end

probvec(dist::Categorical) = probs(dist)


"""Usage:
```
For Unconstrained Optimsation:
  f = X->((x,y)=X;x^3 + 2y + 5)
  A = IntervalBox(2..4, 2..3)
  (global_min, minimisers, info) = ibc_minimise(f, A)
  (global_max, maximisers, info) = ibc_maximise(f, A)

For Constrained Optimisation:
  f = X->((x,y)=X;-(x-4)^2-(y-4)^2)
  A = IntervalBox(-4..4, -4..4)

  vars = ModelingToolkit.@variables x y
  C1 = Constraint(vars, x+y, -Inf..4)
  C2 = Constraint(vars, x+3y, -Inf..9)

  (global_min, minimisers, info) = ibc_minimise(f, A, [C1, C2])
  (global_max, maximisers, info) = ibc_maximise(f, A, [C1, C2])

ibc_minimise/ibc_maximise find the global minimum/maximum value of the function in given search space by using Interval Bound & Contract(IBC) algorithm
```
"""
function ibc_minimise(f::Function , X::IntervalBox{N,T}; ibc_chnl = RemoteChannel(()->Channel{Tuple{IntervalBox{N,T}, T}}(0)),
               diffevol_chnl = RemoteChannel(()->Channel{Tuple{SVector{N, T}, T, Union{Nothing, IntervalBox{N, T}}}}(0)), debug = false, structure = SortedVector, tol=1e-6, ibc_ind = true) where{N, T}

    vars = [Variable(Symbol("x",i))() for i in 1:length(X)]
    g(x...) = f(x)
    C = BasicContractor(vars, g)

    working = structure([(X, inf(f(X)))], x->x[2]) # list of boxes with corresponding lower bound, arranged according to selected structure :

    minimizers = IntervalBox{N,T}[]
    global_min = ∞  # upper bound

    info = Information(0, 0, 0)
    num_bisections = 0

    while !isempty(working)

        info.iterations= info.iterations + 1

        if !ibc_ind
            if isready(ibc_chnl)
                from_diff = take!(ibc_chnl)     # Receiving best individual from ibc_minimise
                if debug
                    println("Box recevied from DifferentialEvolution: ", from_diff)
                end
                global_min = min(from_diff[2], global_min)
                info.de_to_ibc = info.de_to_ibc + 1
            end
       end

        (X, X_min) = popfirst!(working)

        gra = gradient(f, X.v)
        (lbb, fcb, cb) = bauman_form(X, f, gra)      # ----- second order form

        if global_min < lbb
            continue
        else
            X_min = max(X_min, lbb)
            if fcb < global_min
                global_min = fcb
                x_best = cb
            end
        end                                        # ------

        if debug
            println("New search-space : ", X)
        end

        A = -∞..global_min
        X = invokelatest(C, A, X)                        # Contracting the box by constraint f(X) < globla_min
        X_min = inf(f(X))

        if debug
            println("Contracted search_space: ", X)
        end

        if X_min > global_min    # X_min == inf(f(X))
            continue
        end

        # find candidate for upper bound of global minimum by just evaluating a point in the interval:
        m = sup(f(Interval.(mid.(X))))   # evaluate at midpoint of current interval

        if m < global_min
            global_min = m
            x_best = SVector(mid(X))
            if !ibc_ind
                if debug
                    println("Box send to DifferentialEvolution: ", x_best )
                end
                if info.iterations % 200 == 0
                    put!(diffevol_chnl, (x_best, global_min, X))  # sending best individual to diffevol
                else
                    put!(diffevol_chnl, (x_best, global_min, nothing))
                end
                info.ibc_to_de = info.ibc_to_de + 1
            end
        end

        filter_elements!(working , (X, global_min) )   # Remove all boxes whose lower bound is greater than the current one:

        if diam(X) < tol
            push!(minimizers, X)
        else
            X1, X2 = bisect(X)
            push!( working, (X1, inf(f(X1))) )
            push!( working, (X2, inf(f(X2))) )
            num_bisections += 1
        end
    end

    if debug
        println("IBC search is ended")
    end

    if !ibc_ind
        if isready(diffevol_chnl)
            take!(diffevol_chnl)
            put!(diffevol_chnl,(SVector(mid(X)), Inf, nothing) )
        else
            put!(diffevol_chnl,(SVector(mid(X)), Inf, nothing) )
        end
        if debug
            println("DifferentialEvolution is also terminated")
        end
        if isready(ibc_chnl)
            take!(ibc_chnl)
        end
    end

    lower_bound = minimum(inf.(f.(minimizers)))

    return Interval(lower_bound,global_min), minimizers, info

end

function ibc_maximise(f::Function , X::IntervalBox{N,T}; ibc_chnl = RemoteChannel(()->Channel{Tuple{IntervalBox{N,T}, T}}(0)),
               diffevol_chnl = RemoteChannel(()->Channel{Tuple{SVector{N, T}, T, Union{Nothing, IntervalBox{N, T}}}}(0)), debug = false, structure = SortedVector, tol=1e-6, ibc_ind = true) where{N, T}
    bound, minimizer, info = ibc_minimise(x -> -f(x), X, ibc_chnl = ibc_chnl, diffevol_chnl = diffevol_chnl, debug = debug, structure = structure, tol = tol, ibc_ind = true)
    return -bound, minimizer, info
end


#=
Contains info specific to finite volume such as
- cell enters
- face centers
- interpolation weights
-

Also has functions useful for FV such as
- cell to node / node to cell interpolation
-
=#

struct FVInfo
    cellCenters::Array{Float64}     # Coordinates of cell centers
    faceCenters::Array{Float64}     # Coordinates of face centers

    cell2node::Array{Array{Int64,1}}            # Neighboring cell indices for each node
    cell2nodeWeight::Array{Array{Float64,1}}    # Cell to node interpolation weights
end

etypetonv = [2, 3, 4, 4, 8, 6, 5, 2, 3, 4, 4, 8, 6, 5, 1, 4, 8, 6, 5]; # number of vertices for each type
etypetodim= [1, 2, 2, 3, 3, 3, 3, 1, 2, 2, 3, 3, 3, 3, 1, 2, 2, 3, 3]; # dimension of each type
etypetonf = [2, 3, 4, 4, 6, 5, 5, 2, 3, 4, 4, 6, 5, 5, 1, 4, 6, 5, 5]; # number of faces for element types
etypetoftype=[0,1, 1, 2, 3, 3, 3, 0, 1, 1, 2, 3, 3, 3, 0, 0, 0, 0, 0]; # type of faces for this element type

# Build the FVInfo from a given grid
function build_FV_info(grid)
    nel = size(grid.loc2glb, 2);
    nface = size(grid.face2glb, 3);
    nnode = size(grid.allnodes, 2);
    dim = size(grid.allnodes, 1);

    cell_centers = zeros(dim, nel);
    face_centers = zeros(dim, nface);
    cell2node = Array{Array{Int64,1},1}(undef,nnode);
    cell2nodeWeight = Array{Array{Float64,1},1}(undef,nnode);

    # Find cell centers
    for ei=1:nel
        center = zeros(dim);
        nvtx = size(grid.loc2glb, 1);
        for ni=1:nvtx
            center .+= grid.allnodes[:, grid.loc2glb[ni, ei]];
        end
        cell_centers[:, ei] = center ./ nvtx;
    end

    # Find face centers
    for fi=1:nface
        center = zeros(dim);
        nvtx = size(grid.face2glb, 1);
        for ni=1:nvtx
            center .+= grid.allnodes[:, grid.face2glb[ni, 1, fi]];
        end
        face_centers[:, fi] = center ./ nvtx;
    end

    # Build cell2node map
    # Check all elements to find ones touching this node.
    nvtx = size(grid.loc2glb, 1);
    for ni=1:nnode
        cell2node[ni] = zeros(Int,0);
        for ei=1:nel
            for nj=1:nvtx
                if grid.loc2glb[nj, ei] == ni
                    # This cell touches this node
                    push!(cell2node[ni], ei);
                    break;
                end
            end
        end
    end

    # Compute cell2node interpolation weights
    for ni=1:nnode
        cell2nodeWeight[ni] = zeros(length(cell2node[ni]));
        sumweight = 0;
        for ei=1:length(cell2node[ni])
            dist = norm(grid.allnodes[:,ni] - cell_centers[:, cell2node[ni][ei]]);
            sumweight += 1 / dist;
            cell2nodeWeight[ni][ei] = 1 / dist;
        end
        for ei=1:length(cell2node[ni])
            cell2nodeWeight[ni][ei] /= sumweight;
        end
    end

    return FVInfo(cell_centers, face_centers, cell2node, cell2nodeWeight);
end

# Interpolate nodal values from neighboring cells.
# Returns an array of nodal values.
function FV_cell_to_node(cell_values, node_values; dofs = 1)
    Nnodes = size(grid_data.allnodes,2);
    if node_values === nothing
        node_values = zeros(Nnodes * dofs);
    end

    for ni=1:Nnodes
        if dofs > 1
            st = (ni-1)*dofs + 1;
            en = ni*dofs;
            node_values[st:en] .= 0;
            ncell = length(fv_info.cell2nodeWeight[ni]);
            for ci=1:ncell
                cind = fv_info.cell2node[ni][ci];
                cst = (cind-1)*dofs + 1;
                cen = cind*dofs;
                node_values[st:en] .+= fv_info.cell2nodeWeight[ni][ci] .* cell_values[cst:cen];
            end

        else
            node_values[ni] = 0;
            ncell = length(fv_info.cell2nodeWeight[ni]);
            for ci=1:ncell
                node_values[ni] += fv_info.cell2nodeWeight[ni][ci] * cell_values[fv_info.cell2node[ni][ci]];
            end
        end
    end

    return node_values;
end

# Find cell averages from nodal values
# Returns an array of cell values.
function FV_node_to_cell(node_values, cell_values = nothing)
    Ncells = size(grid_data.loc2glb,2);
    if cell_values === nothing
        cell_values = zeros(Ncells);
    end

    for ci=1:Ncells
        cell_values[ci] = 0;
        nnode = length(grid_data.loc2glb[:,ci]);
        for ni=1:nnode
            cell_values[ci] += node_values[grid_data.loc2glb[ni, ci]];
        end
        cell_values[ci] /= nnode;
    end

    return cell_values;
end


using Test
using Parameters
using BenchmarkTools

#
import PartialWaveFunctions: CG_doublearg, clebschgordan_doublearg
ClGd_pwf(two_j1,two_m1,two_j2,two_m2,two_j,two_m) = clebschgordan_doublearg(two_j1,two_m1,two_j2,two_m2,two_j,two_m)

#
using GSL
function ClGd_gsl(two_j1,two_m1,two_j2,two_m2,two_j,two_m)
    factor = sqrt(two_j+1)*(mod(two_j1-two_j2+two_m,4)==2 ? -1 : +1)
    three_j = sf_coupling_3j(two_j1,two_j2,two_j,two_m1,two_m2,-two_m)
    return factor*three_j;
end

#
using SymPy
import PyCall
PyCall.pyimport_conda("sympy.physics.wigner","sympy")
import_from(sympy.physics.wigner)
#
ClGd_sympy(two_j1,two_m1,two_j2,two_m2,two_j,two_m) =
    convert(Float64, clebsch_gordan(Sym(two_j1)/2, Sym(two_j2)/2, Sym(two_j)/2, Sym(two_m1)/2, Sym(two_m2)/2, Sym(two_m)/2))

#
import HalfIntegers: half
using WignerSymbols
function ClGd_WS(two_j1,two_m1,two_j2,two_m2,two_j,two_m)
    ((abs(two_m1) > two_j1) || (abs(two_m2) > two_j2) || (abs(two_m ) > two_j )) && return 0.0
    return clebschgordan(half(two_j1),half(two_m1),half(two_j2),half(two_m2),half(two_j),half(two_m))
end

#                                      _|
#  _|  _|_|    _|_|_|  _|_|_|      _|_|_|
#  _|_|      _|    _|  _|    _|  _|    _|
#  _|        _|    _|  _|    _|  _|    _|
#  _|          _|_|_|  _|    _|    _|_|_|

function rand_clebsch(;two_j_max::Int=15)
    two_j1 = rand(0:two_j_max); two_m1 = rand(-two_j1:2:two_j1)
    two_j2 = rand(0:two_j_max); two_m2 = rand(-two_j2:2:two_j2)
    two_j = rand(abs(two_j2-two_j1):2:(two_j1+two_j2))
    two_m = two_m1+two_m2
    abs(two_m) > two_j && return rand_clebsch(; two_j_max=two_j_max)
    (two_j1=two_j1, two_j2=two_j2, two_j=two_j,
     two_m1=two_m1, two_m2=two_m2, two_m=two_m)
end

# fraction of vanishing clebshces ≈ 0.6%
sum(let
    @unpack two_j1, two_j2, two_j, two_m1, two_m2, two_m = rand_clebsch()
    ClGd_pwf(two_j1,two_m1,two_j2,two_m2,two_j,two_m) ≈ 0
end for _ in 1:10_000)

#                                                                    _|
#    _|_|_|    _|_|    _|_|_|  _|_|    _|_|_|      _|_|_|  _|  _|_|        _|_|_|    _|_|    _|_|_|
#  _|        _|    _|  _|    _|    _|  _|    _|  _|    _|  _|_|      _|  _|_|      _|    _|  _|    _|
#  _|        _|    _|  _|    _|    _|  _|    _|  _|    _|  _|        _|      _|_|  _|    _|  _|    _|
#    _|_|_|    _|_|    _|    _|    _|  _|_|_|      _|_|_|  _|        _|  _|_|_|      _|_|    _|    _|
#                                      _|
#                                      _|

for _ in 1:1_000
    @unpack two_j1, two_j2, two_j,
            two_m1, two_m2, two_m = rand_clebsch()
    #
    v1 = ClGd_pwf(two_j1,two_m1,two_j2,two_m2,two_j,two_m)
    v2 = ClGd_gsl(two_j1,two_m1,two_j2,two_m2,two_j,two_m)
    v3 = ClGd_sympy(two_j1,two_m1,two_j2,two_m2,two_j,two_m)
    v4 = ClGd_WS(two_j1,two_m1,two_j2,two_m2,two_j,two_m)
    #
    (abs(v1-v2) > 1e-5) && error("gsl is different: $(v1) $(v2)")
    (abs(v1-v3) > 1e-5) && error("sympy is different")
    (abs(v1-v4) > 1e-5) && error("WS is different")
end

#    _|      _|                  _|
#  _|_|_|_|      _|_|_|  _|_|        _|_|_|      _|_|_|
#    _|      _|  _|    _|    _|  _|  _|    _|  _|    _|
#    _|      _|  _|    _|    _|  _|  _|    _|  _|    _|
#      _|_|  _|  _|    _|    _|  _|  _|    _|    _|_|_|
#                                                    _|
#                                                _|_|

@btime for _ in 1:1_000
    @unpack two_j1, two_j2, two_j, two_m1, two_m2, two_m = rand_clebsch()
    ClGd_pwf(two_j1,two_m1,two_j2,two_m2,two_j,two_m)
end # 227.200 μs

@btime for _ in 1:1_000
    @unpack two_j1, two_j2, two_j, two_m1, two_m2, two_m = rand_clebsch()
    ClGd_gsl(two_j1,two_m1,two_j2,two_m2,two_j,two_m)
end # 921.8 μs

@btime for _ in 1:1_000
    @unpack two_j1, two_j2, two_j, two_m1, two_m2, two_m = rand_clebsch()
    ClGd_WS(two_j1,two_m1,two_j2,two_m2,two_j,two_m)
end # 7.483 ms

@btime for _ in 1:1_000
    @unpack two_j1, two_j2, two_j, two_m1, two_m2, two_m = rand_clebsch()
    ClGd_sympy(two_j1,two_m1,two_j2,two_m2,two_j,two_m)
end # 1.61 s

# Summary
# -------------
# function      | `PartialWaveFunctions.jl` | `GSL.jl` | `WignerSymbols.jl` | `SymPy.jl` |
# -------------
# random Clebsh-Gordan coef.(j<15) | 227.2 μs | 921.8 μs | 7.483 ms | 1.61 s |
# -------------


using FromFile
@from "Core.jl" import Options, Dataset, Node, copyNode
@from "Utils.jl" import getTime
@from "LossFunctions.jl" import scoreFunc

# Define a member of population by equation, score, and age
mutable struct PopMember{T<:Real}
    tree::Node
    score::T
    birth::Int
end

"""
    PopMember(t::Node, score::T)

Create a population member with a birth date at the current time.

# Arguments

- `t::Node`: The tree for the population member.
- `score::T`: The loss to assign this member.
"""
function PopMember(t::Node, score::T) where {T<:Real}
    PopMember{T}(t, score, getTime())
end

"""
    PopMember(dataset::Dataset{T}, baseline::T,
              t::Node, options::Options)

Create a population member with a birth date at the current time.
Automatically compute the score for this tree.

# Arguments

- `dataset::Dataset{T}`: The dataset to evaluate the tree on.
- `baseline::T`: The baseline loss.
- `t::Node`: The tree for the population member.
- `options::Options`: What options to use.
"""
function PopMember(dataset::Dataset{T},
                   baseline::T, t::Node,
                   options::Options) where {T<:Real}
    PopMember(t, scoreFunc(dataset, baseline, t, options))
end

function copyPopMember(p::PopMember{T}) where {T<:Real}
    tree = copyNode(p.tree)
    score = p.score
    birth = p.birth
    return PopMember{T}(tree, score, birth)
end


using Test, bad

import Distributions

p = 0:0.1:1
for n in 0:15, x = 0:n, pp in p
    @test Distributions.pdf.(Distributions.Binomial(n, pp), x) ≈ pmf(x, n, pp)
    @test Distributions.cdf.(Distributions.Binomial(n, pp), x) ≈ cdf(x, n, pp)
end

p = Beta(5, 7)

# check that marginal distribution (of x1) sums to 1
for xpartial in 0:10, npartial in xpartial:10
    pmff = pmf.(0:10, 10, p; partial = (xpartial, npartial) )
    @test sum(pmff) ≈ 1.0
    @test cumsum(pmff) ≈ cdf.(0:10, 10, p; partial = (xpartial, npartial) )
end

function simon(r1::Int, n1::Int, r::Int, n::Int)
    nn2 = zeros(Int, n1 + 1)
    nn2[(r1 + 2):end] .= n - n1
    cc = repeat([Inf], n1 + 1)
    cc[(r1 + 2):end] .= r
    cc2 = cc .- collect(0:n1)
    cc2[cc2 .< 0]     .= -Inf
    cc2[cc2 .>= nn2]  .= Inf
    nn2[cc2 .== -Inf] .= 0
    Design(nn2, cc2)
end

design = simon(4, 18, 10, 33)
XX     = sample_space(design)
x1, x2 = XX[:,1], XX[:,2]

# check that joint distribution sums to one for all possible partial x1 observations
for x1partial in 0:n1(design)
    pmf_x1x2 = pmf.(x2, n2.(design, x1), p) .*
        pmf.(x1, n1(design), p; partial = (x1partial, n1(design)) )
    @test 1.0 ≈ sum( pmf_x1x2 )
end

# check that conditional distribution (of x2 given x2) sums to 1
for x1 in 0:n1(design)
    pmff = pmf.(0:n2(design, x1), n2(design, x1), p)
    @test sum(pmff) ≈ 1.0
    @test all( cumsum(pmff) .≈ cdf.(0:n2(design, x1), n2(design, x1), p))
end

@test true


#=
  Copyright 2013 Artur Mkrtchyan

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
=#

function calculate_check_digit(bban, country_code)::String
    combined = bban * country_code * "00"
    check_digits = 98 - mod_97(combined)
    lpad(check_digits, 2, "0")
end

# [0..9] are matched to values:            [0..9]
# [a..z] and [A..Z] are matched to values: [10..35]
const zero_to_Z = vcat(['0':'9';], ['A':'Z';])
const a_to_z = ['a':'z';]
const c_to_numeric = merge(
    Dict(zero_to_Z[i] => i - 1 for i = 1:length(zero_to_Z)),
    Dict(a_to_z[i] => i + 9 for i = 1:length(a_to_z)),
)

function mod_97(str)
    total = 0
    for c in str
        c_numeric = c_to_numeric[c]
        total = (c_numeric > 9 ? total * 100 : total * 10) + c_numeric
        if total > 999999999
            total = mod(total, 97)
        end
    end
    mod(total, 97)
end

# according to @btime, about of order of magnitude more expensive
# then mod_97
function mod_97_bigint(str)
    numeric_str = map(c -> alphabet_to_numeric[c] , collect(str)) |> join
    numeric_bi = parse(BigInt, numeric_str)
    mod(numeric_bi, 97)
end


using Lora

println("    Testing GenericModel constructors...")

println("      Initialization without input arguments...")

m = GenericModel()

θ = BasicContUnvParameter(:θ, 1, signature=:low)
x = Data(:x, 2)
λ = Hyperparameter(:λ, 3)

add_vertex!(m, θ)
add_vertex!(m, x)
add_vertex!(m, λ)

add_edge!(m, x, θ)
add_edge!(m, λ, θ)

println("      Initialization via matrix-based specification of dependencies...")

m = GenericModel([θ, x, λ], [x θ; λ θ])

# println("    Testing plotting of GenericModel...")

# g = convert(GenericGraph, m)
# plot(g)


using Distributions

include("./src/smote_exs.jl")
include("./src/ub_smote.jl")

m = 150
X = rand(m, 10)
y = ones(m)
X = hcat(X, y)

smote_exs(X, 11)

w1 = [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0]

cases_needed(w1)
pct_needed(w1)

X = randn(150, 10);
y = vcat(zeros(140), ones(10));
X2, y2 = smote(X, y, 5);

length(y2)
countmap(y2)

X = randn(150, 10);
y = vcat(zeros(130), ones(20));
X2, y2 = smote(X, y, 5)

length(y2)
countmap(y2)


@testset "Betweenness" begin
    # self loops
    s2 = SimpleDiGraph(3)
    add_edge!(s2, 1, 2); add_edge!(s2, 2, 3); add_edge!(s2, 3, 3)
    s1 = SimpleGraph(s2)
    g3 = PathGraph(5)

    gint = loadgraph(joinpath(testdir, "testdata", "graph-50-500.jgz"), "graph-50-500")

    c = vec(readdlm(joinpath(testdir, "testdata", "graph-50-500-bc.txt"), ','))
    for g in testdigraphs(gint)
        z  = @inferred(betweenness_centrality(g))
        @test map(Float32, z)  == map(Float32, c)

        y  = @inferred(betweenness_centrality(g, endpoints=true, normalize=false))
        @test round.(y[1:3], digits=4) ==
            round.([122.10760591498584, 159.0072453120582, 176.39547945994505], digits=4)

        x  = @inferred(betweenness_centrality(g, 3))
        x2  = @inferred(betweenness_centrality(g, collect(1:20)))

        @test length(x) == 50
        @test length(x2) == 50
    end

    @test @inferred(betweenness_centrality(s1)) == [0, 1, 0]
    @test @inferred(betweenness_centrality(s2)) == [0, 0.5, 0]

    g = SimpleGraph(2)
    add_edge!(g, 1, 2)
    z  = @inferred(betweenness_centrality(g; normalize=true))
    @test z[1] == z[2] == 0.0
    z2  = @inferred(betweenness_centrality(g, vertices(g)))
    z3  = @inferred(betweenness_centrality(g, [vertices(g);]))
    @test z == z2 == z3

    z  = @inferred(betweenness_centrality(g3; normalize=false))
    @test z[1] == z[5] == 0.0

    # Weighted Graph tests
    g = Graph(6)
    add_edge!(g, 1, 2)
    add_edge!(g, 2, 3)
    add_edge!(g, 3, 4)
    add_edge!(g, 2, 5)
    add_edge!(g, 5, 6)
    add_edge!(g, 5, 4)

    distmx = [ 0.0  2.0  0.0  0.0  0.0  0.0;
               2.0  0.0  4.2  0.0  1.2  0.0;
               0.0  4.2  0.0  5.5  0.0  0.0;
               0.0  0.0  5.5  0.0  0.9  0.0;
               0.0  1.2  0.0  0.9  0.0  0.6;
               0.0  0.0  0.0  0.0  0.6  0.0;]

    @test isapprox(betweenness_centrality(g, vertices(g), distmx; normalize=false), [0.0,6.0,0.0,0.0,6.0,0.0])
    @test isapprox(betweenness_centrality(g, vertices(g), distmx; normalize=false, endpoints=true), [5.0,11.0,5.0,5.0,11.0,5.0])
    @test isapprox(betweenness_centrality(g, vertices(g), distmx; normalize=true), [0.0, 0.6000000000000001, 0.0, 0.0, 0.6000000000000001, 0.0])
    @test isapprox(betweenness_centrality(g, vertices(g), distmx; normalize=true, endpoints=true), [0.5,1.1,0.5,0.5,1.1,0.5])

    adjmx2 = [0 1 0; 1 0 1; 1 1 0] # digraph
    a2 = SimpleDiGraph(adjmx2)
    for g in testdigraphs(a2)
        distmx2 = [Inf 2.0 Inf; 3.2 Inf 4.2; 5.5 6.1 Inf]
        c2 = [0.24390243902439027,0.27027027027027023,0.1724137931034483]
        @test isapprox(betweenness_centrality(g, vertices(g), distmx2; normalize=false), [0.0,1.0,0.0])
        @test isapprox(betweenness_centrality(g, vertices(g), distmx2; normalize=false, endpoints=true), [4.0,5.0,4.0])
        @test isapprox(betweenness_centrality(g, vertices(g), distmx2; normalize=true), [0.0,0.5,0.0])
        @test isapprox(betweenness_centrality(g, vertices(g), distmx2; normalize=true, endpoints=true), [2.0,2.5,2.0])
    end
end


function RBMs.initialize!(rbm::ConvRBM, data::AbstractArray; ϵ::Real = 1e-6, λ::Real = 0.1)
    RBMs.initialize!(visible(rbm), data; ϵ)
    RBMs.initialize!(hidden(rbm))
    RBMs.initialize_w!(rbm, data; ϵ, λ)
    RBMs.zerosum!(rbm)
    return rbm
end

function RBMs.initialize!(rbm::ConvRBM; ϵ::Real = 1e-6, λ::Real = 0.1)
    RBMs.initialize!(visible(rbm))
    RBMs.initialize!(hidden(rbm))
    RBMs.initialize_w!(rbm; ϵ, λ)
    RBMs.zerosum!(rbm)
    return rbm
end

function RBMs.initialize_w!(rbm::ConvRBM, data::AbstractArray; ϵ::Real = 1e-6, λ::Real = 0.1)
    d = dot(data, data) / prod(vsizes(rbm, data).batch_size)
    randn!(weights(rbm))
    J = prod(kernel_size(rbm))
    K = prod(output_size(rbm, data))
    weights(rbm) .*= sqrt(K/J) * λ / √(d + ϵ)
    return rbm # does not impose zerosum
end

function RBMs.initialize_w!(rbm::ConvRBM; ϵ::Real = 1e-6, λ::Real = 0.1)
    @assert 0 < ϵ < 1/2
    μ = RBMs.transfer_mean(visible(rbm))
    ν = RBMs.transfer_var(visible(rbm))
    d = sum(ν + μ.^2)
    randn!(weights(rbm))
    J = prod(kernel_size(rbm))
    K = prod(output_size(rbm, data))
    weights(rbm) .*= sqrt(K/J) * λ / √(d + ϵ)
    return rbm
end


# samples Bernoully trials based on the simple randomn number generator

function BernoulliTrial(n, seed)

    rng = MersenneTwister(seed)
    A = [0, 1]

    return rand(rng, A, n)

end

# samples a coin toss, either win or lose

function coinTrial(n, seed)

    rng = MersenneTwister(seed)
    A = [-1, 1]

    return rand(rng, A, n)

end


"""
    get_constrained_lowdimset(cpa::CartesianProductArray{N, S}, P::AbstractPolyhedron{N}) where {N<:Real, S<:LazySet{N}}

Preprocess step for intersection between cartesian product array and H-polyhedron.
Returns low-dimensional HPolytope in constrained dimensions of original cpa,
constrained variables and variables in corresponding blocks, original block structure
of low-dimensional set and list of constrained blocks.

### Input

- `cpa` -- Cartesian product array of convex sets
- `P` -- polyhedron

### Output

A tuple of low-dimensional set, list of constrained dimensions, original block
structure of low-dimensional set and corresponding blocks indices.
"""
function get_constrained_lowdimset(cpa::CartesianProductArray{N, S}, P::AbstractPolyhedron{N}) where {N<:Real, S<:LazySet{N}}

    if isbounded(P)
        blocks, non_empty_length = block_to_dimension_indices(cpa)
    else
        constrained_vars = constrained_dimensions(P)
        blocks, non_empty_length = block_to_dimension_indices(cpa, constrained_vars)
    end

    array = Vector{S}()
    sizehint!(array, non_empty_length)
    vars = Vector{Int}()
    block_structure = Vector{UnitRange{Int}}()
    sizehint!(block_structure, non_empty_length)

    last_var = 1
    for i in 1:length(blocks)
        start_index, end_index = blocks[i]
        block_end = last_var + end_index - start_index
        if start_index != -1
            push!(array, cpa.array[i])
            append!(vars, start_index : end_index)
            push!(block_structure, last_var : block_end)
            last_var = block_end + 1
        end
    end

    cpa_low_dim = HPolytope(constraints_list(CartesianProductArray(array)));

    return cpa_low_dim, vars, block_structure, blocks
end


# function simple_multiply(x, y)
#    product = 0
#    for i in 1:x
#        product += y
#    end
#    product
# end

# function simpledivide(dividend, divisor)
#     quotient = 0
#     remainder = dividend

#     while remainder - divisor >= 0
#         remainder -= divisor
#         quotient += 1
#     end
#     quotient, remainder
# end

addend = 0
accumulator = 0
counter = 0
shifter = 0

function clear()
    global addend = 0
    global accumulator = 0
    global counter = 0
    global shifter = 0
    nothing
end

function setaddend!(x)
   global addend = x
end

function leftshift!()
   global shifter += 1
end

function rightshift!()
   global shifter -= 1
end

function crank!(n)
    for _ in 1:n
        global accumulator += addend * 10^shifter
        global counter += 10^shifter
    end

    accumulator, counter
end

function uncrank!(n)
    for _ in 1:n
        global accumulator -= addend * 10^shifter
        global counter -= 10^shifter
    end

    accumulator, counter
end