#=
Copyright (c) 2015, Intel Corporation
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
- Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
- Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
=#

using DocOpt

function compute(src, dst, N)
  local z, y, x
  (x,y,z) = size(src)
  for n = 1:N
    for k = 1:z
      for j = 1:y
        for i = 1:x
          if (i == 1 || i == x || j == 1 || j == y || k == 1 || k == z)
            dst[i,j,k] = src[i,j,k]
          else
            dst[i,j,k] = (src[i,j,k-1] + src[i,j,k+1] + src[i,j-1,k] + src[i,j+1,k] + src[i-1,j,k] + src[i+1,j,k]) / 6
          end
        end
      end
    end
    tmp=src
    src=dst
    dst=tmp
  end
  return src
end

function initialize(x, y, z)
  src = Array(Float32, x, y, z)
  rand!(src)
  dst = copy(src)
  return src, dst
end

function checksum(src)
  local z, y, x
  (x, y, z) = size(src)
  sum(src) / (x * y * z)
end

function main()
    doc = """laplace-3d.jl

Laplace 6-point 3D stencil.

Usage:
  laplace-3d.jl -h | --help
  laplace-3d.jl [--size=<size>] [--iterations=<iterations>]

Options:
  -h --help                  Show this screen.
  --size=<size>              Specify a 3d array size (<size> x <size> x <size>) [default: 300].
  --iterations=<iterations>  Specify a number of iterations [default: 100].
"""
    arguments = docopt(doc)

    size = parse(Int, arguments["--size"])
    iterations = parse(Int, arguments["--iterations"])

    srand(0)

    function laplace_3d(iterations)
        local src, dst
        println("Run laplace-3d with size ", size, "x", size, "x", size, " for ", iterations, " iterations.")
        (src, dst) = initialize(size, size, size)
        tic()
        compute(src, dst, 0)
        println("SELFPRIMED ", toq())

        tic()
        dst = compute(src, dst, iterations)
        println("SELFTIMED ", toq())
        println("checksum: ", checksum(dst))
    end

    laplace_3d(iterations)

end

main()


using Pkg
Pkg.activate(pwd())

# # More complex networks
# ## Preparing data

using MLDatasets
using MLDatasets: MNIST

T = Float32
X_train, y_train = MNIST.traindata(T)
X_test, y_test = MNIST.testdata(T)

# ### Exercise:
# Plot the first 15 images of the digit 0 from the training set.
#
# **Hints:**:
# - The `ImageInspector` package written earlier provides the function
# `imageplot(X_train, inds; nrows=3)`, where `inds` are the desired indices.
# - To find the correct indices, use the function `findall`.
#
# ---
# ### Solution:

# ---
#
# ### Exercise:
# Write function `reshape_data`, which reshapes `X_train` and `X_test` into the correct
# size required by Flux.
#
# **Hint**: The function should work only on inputs with the correct size. This can be
# achieved by specifying the correct input type `X::AbstractArray{<:Real, 3}`.
#
# ---
# ### Solution:

# ---

using Flux
using Flux: onehotbatch, onecold

function load_data(dataset; T=Float32, onehot=false, classes=0:9)
    X_train, y_train = dataset.traindata(T)
    X_test, y_test = dataset.testdata(T)

    X_train = reshape_data(X_train)
    X_test = reshape_data(X_test)

    if onehot
        y_train = onehotbatch(y_train, classes)
        y_test = onehotbatch(y_test, classes)
    end

    return X_train, y_train, X_test, y_test
end

#+

X_train, y_train, X_test, y_test = load_data(MLDatasets.MNIST; T=T, onehot=true)

# ### Exercise:
# Try to load the CIFAR10 dataset via the `load_data` function and fix the error in one
# line of code.
#
# **Hint**: Use ` dataset = MLDatasets.CIFAR10`.
#
# ---
# ### Solution:

# ---
#
# ## Training and storing the network
# ### Exercise:
# Use the help of the function `DataLoader` to split the dataset into minibatches.
#
# **Hint**: It needs to be imported from Flux via `using Flux.Data: DataLoader`.
#
# ---
# ### Solution:

# ---

using Base.Iterators: partition
using Random

batches = map(partition(randperm(size(y_train, 2)), batchsize)) do inds
    return (X_train[:, :, :, inds], y_train[:, inds])
end

#+

batches2 = [(X_train[:, :, :, inds], y_train[:, inds]) for inds in partition(randperm(size(y_train, 2)), batchsize)]

#+

using Random
using BSON

Random.seed!(666)
m = Chain(
    Conv((2,2), 1=>16, relu),
    MaxPool((2,2)),
    Conv((2,2), 16=>8, relu),
    MaxPool((2,2)),
    flatten,
    Dense(288, size(y_train,1)),
    softmax,
)

#+

using Flux: crossentropy

L(X, y) = crossentropy(m(X), y)

#+

function train_model!(m, L, X, y;
        opt = Descent(0.1),
        batchsize = 128,
        n_epochs = 10,
        file_name = "")

    batches = DataLoader((X, y); batchsize, shuffle = true)

    for _ in 1:n_epochs
        Flux.train!(L, params(m), batches, opt)
    end

    !isempty(file_name) && BSON.bson(file_name, m=m)

    return
end

# ### Exercise:
# Train the model for one epoch and save it to `MNIST_simple.bson`. Print the accuracy on
# the testing set.
#
# ---
# ### Solution:

# ---
#
# ### Exercise:
# Write a function `train_or_load!(file_name, m, args...; ???)` checking whether the file
# `file_name` exists.
# - If it exists, it loads it and then copies its parameters into `m` using the function
# `Flux.loadparams!`.
# - If it does not exist, it trains it using `train_model!`.
# In both cases, the model `m` should be modified inside the `train_or_load!` function. Pay
# special attention to the optional arguments `???`.
#
# Use this function to load the model from `data/mnist.bson` and evaluate the performance
# at the testing set.
#
# ---
# ### Solution:


const FacElemQ = Union{FacElem{fmpq, FlintRationalField}, FacElem{fmpz, FlintIntegerRing}}

@doc Markdown.doc"""
    evaluate(x::FacElem{fmpq}) -> fmpq
    evaluate(x::FacElem{fmpz}) -> fmpz

Expands or evaluates the factored element, i.e. actually computes the
the element.
Works by first obtaining a simplified version of the power product
into coprime base elements.
"""
function evaluate(x::FacElem{fmpq})
  return evaluate_naive(simplify(x))
end

function evaluate(x::FacElem{fmpz})
  return evaluate_naive(simplify(x))
end
@doc Markdown.doc"""
    simplify(x::FacElem{fmpq}) -> FacElem{fmpq}
    simplify(x::FacElem{fmpz}) -> FacElem{fmpz}

Simplfies the factored element, i.e. arranges for the base to be coprime.
"""
function simplify(x::FacElem{fmpq})
  y = deepcopy(x)
  simplify!(y)
  return y
end

function simplify(x::FacElem{fmpz})
  y = deepcopy(x)
  simplify!(y)
  return y
end

function simplify!(x::FacElem{fmpq})
  if length(x.fac) <= 1
    return nothing
  end
  cp = vcat([denominator(y) for (y, v) in x if !iszero(v)], [numerator(y) for (y, v) in x if !iszero(v)])
  ev = Dict{fmpq, fmpz}()
  if isempty(cp)
    ev[fmpq(1)] = 0
    x.fac = ev
    return nothing
  end
  cp = coprime_base(cp)
  for p = cp
    if p == 1 || p == -1
      continue
    end
    v = fmpz(0)
    for (b, vb) in x
      if !iszero(vb)
        v += valuation(b, abs(p))*vb
      end
    end
    if v != 0
      ev[fmpq(abs(p))] = v
    end
  end
  f = b -> b < 0 && isodd(x.fac[b]) ? -1 : 1
  s = prod((f(v) for v in base(x)))
  if s == -1
    ev[fmpq(-1)] = 1
  else
    if length(ev) == 0
      ev[fmpq(1)] = 0
    end
  end
  x.fac = ev
  return nothing
end

function simplify!(x::FacElem{fmpz})
  if length(x.fac) == 0
    x.fac[fmpz(1)] = 0
    return
  end
  if length(x.fac) <= 1
    k,v = first(x.fac)
    if isone(k)
      x.fac[k] = 0
    elseif k == -1
      if isodd(v)
        x.fac[k] = 1
      else
        delete!(x.fac, k)
        x.fac[fmpz(1)] = 0
      end
    end
    return
  end
  cp = coprime_base(collect(base(x)))
  ev = Dict{fmpz, fmpz}()
  for p = cp
    if p == 1 || p == -1
      continue
    end
    v = fmpz(0)
    for (b, vb) in x
      v += valuation(b, abs(p))*vb
    end
    if v < 0
      throw(DomainError(v, "Negative valuation in simplify!"))
    end
    if v != 0
      ev[abs(p)] = v
    end
  end
  f = b -> b < 0 && isodd(x.fac[b]) ? -1 : 1
  s = prod(f(v) for v in base(x))
  if s == -1
    ev[-1] = 1
  else
    if length(ev) == 0
      ev[fmpz(1)] = 0
    end
  end
  x.fac = ev
  nothing
end

@doc Markdown.doc"""
    isone(x::FacElem{fmpq}) -> Bool
    isone(x::FacElem{fmpz}) -> Bool
Tests if $x$ represents $1$ without an evaluation.
"""
function isone(x::FacElem{fmpq})
  y = simplify(x)
  return all(iszero, values(y.fac)) || all(isone, keys(y.fac))
end

function isone(x::FacElem{fmpz})
  y = simplify(x)
  return all(iszero, values(y.fac)) || all(isone, keys(y.fac))
end

@doc Markdown.doc"""
    factor_coprime(x::FacElem{fmpz}) -> Fac{fmpz}
Computed a partial factorisation of $x$, ie. writes $x$ as a product
of pariwise coprime integers.
"""
function factor_coprime(x::FacElem{fmpz})
  x = deepcopy(x)
  simplify!(x)
  d = Dict(abs(p) => Int(v) for (p,v) = x.fac)
  if haskey(d, fmpz(-1))
    delete!(d, fmpz(-1))
    return Fac(fmpz(-1), d)
  else
    return Fac(fmpz(1), d)
  end
end

function abs(A::FacElemQ)
  B = empty(A.fac)
  for (k,v) = A.fac
    ak = abs(k)
    add_to_key!(B, ak, v)
  end
  if length(B) == 0
    return FacElem(Dict(one(base_ring(A)) => fmpz(1)))
  end
  return FacElem(B)
end

function ==(A::T, B::T) where {T <: FacElemQ}
  C = A*B^-1
  simplify!(C)
  return isone(C)
end

function isone(A::FacElemQ)
  C = simplify(A)
  return all(iszero, values(C.fac)) || all(isone, keys(C.fac))
end


function calc_bias_potential(umbrella_center, data)
    spring_constant = 200.0 * (pi/180.0)^2
    return spring_constant * (data - umbrella_center)^2
end


#########
# model #
#########

# TODO put this into FunctionalCollections:
import FunctionalCollections
Base.IndexStyle(::Type{<:FunctionalCollections.PersistentVector}) = IndexLinear()

@gen function datum(x::Float64, (grad)(inlier_std::Float64), (grad)(outlier_std), (grad)(slope), (grad)(intercept))::Float64
    is_outlier = @trace(bernoulli(0.5), :z)
    std = is_outlier ? inlier_std : outlier_std
    y = @trace(normal(x * slope + intercept, std), :y)
    return y
end

data = Map(datum)

@gen function model(xs::Vector{Float64})
    n = length(xs)
    inlier_std = exp(@trace(normal(0, 2), :log_inlier_std))
    outlier_std = exp(@trace(normal(0, 2), :log_outlier_std))
    slope = @trace(normal(0, 2), :slope)
    intercept = @trace(normal(0, 2), :intercept)
    ys = @trace(data(xs, fill(inlier_std, n), fill(outlier_std, n),
               fill(slope, n), fill(intercept, n)), :data)
    return ys
end


using BEAST
using CompScienceMeshes

Γ = readmesh(Pkg.dir("BEAST", "test", "assets","flatplate.in"))
RT = raviartthomas(Γ)

X = RT
Y = RT

κ = 1.0
t = Maxwell3D.singlelayer(gamma=im*κ)

BEAST.quadrule(op::BEAST.MaxwellOperator3D, g::BEAST.RTRefSpace, f::BEAST.RTRefSpace, i, τ, j, σ, qd) = BEAST.qrib(op, g, f, i, τ, j, σ, qd)
Zbogaert = assemble(t,X,Y)

BEAST.quadrule(op::BEAST.MaxwellOperator3D, g::BEAST.RTRefSpace, f::BEAST.RTRefSpace, i, τ, j, σ, qd) = BEAST.q(op, g, f, i, τ, j, σ, qd)
Zss = assemble(t,X,Y)

println(abs(Zbogaert)./abs(Zss))


"""
File defining Trajectory, SpatialTrajectory and related functions
"""

"""
A struct to store all data needed to iterate. principal_vectors have to be orthonormal
"""
struct Trajectory
    data_set
    n :: Int
    k :: Int
    error :: Float64
    indices :: Array{Int,1}
    principal_vectors :: Array{Float64,2}
end

    function show(io :: IO, traj :: Trajectory)
        return print(io, "error: $(traj.error)\nindices: $(traj.indices)")
    end

    "start with indices"
    function Trajectory(data_set, n, k, indices :: AbstractArray{I}) where I <: Integer
        @assert length(indices) == n
        error, svd = truncated_svd(hcat(data_set[indices]...), k)
        Trajectory(data_set, n, k, error, sort(indices), svd.U)
    end

    "pick the points closest to the span of vectors"
    function Trajectory(data_set, n, k, vectors :: Array{T, 2}) where T <: Number
        vectors = Matrix(qr(vectors).Q) # to orthonormalize principal_vectors
        indices, S = find_n_closest_pts(data_set, vectors, n)
        error, svd = truncated_svd(hcat(S...), k)
        Trajectory(data_set, n, k, error, sort(indices), svd.U)
    end

    "pick k points from the dataset at random and pick the rest n-k points to be the closest to the span of the k random points"
    function Trajectory(data_set, n, k)
        principal_vectors = hcat(sample(data_set, k)...)
        Trajectory(data_set, n, k, principal_vectors)
    end

    "perform one step: find n closest points and compute principal vectors. Return a new Trajectory object"
    function L2_step(trajectory)
        @unpack data_set, n, k, principal_vectors = trajectory
        indices, S = find_n_closest_pts(data_set, principal_vectors, n)
        error, svd = truncated_svd(hcat(S...), k)
        Trajectory(data_set, n, k, error, sort(indices), svd.U)
    end

    "perform one step: find n closest points and compute principal vectors. Return a new Trajectory object"
    function L2_step_local(trajectory)
        @unpack data_set, n, k, principal_vectors, indices = trajectory
        worst_index = maximum(i -> (distance_squared(data_set[i], principal_vectors), i), indices)[2]
        indices_minus_worst = setdiff(indices, worst_index)
        indices_to_search = setdiff(1:length(data_set), indices_minus_worst)
        new_ind = minimum(i -> (distance_squared(data_set[i], principal_vectors), i), indices_to_search)[2]
        indices = union(indices_minus_worst, new_ind)
        S = data_set[indices]
        error, svd = truncated_svd(hcat(S...), k)
        d1 = distance_squared( data_set[worst_index], principal_vectors)
        d2 = distance_squared( data_set[new_ind], principal_vectors)
        Trajectory(data_set, n, k, error, sort(indices), svd.U)
    end

"""
A struct to store all data needed to iterate, including a graph. principal_vectors have to be orthonormal
"""
struct SpatialTrajectory
    graph
    data_set
    n :: Int
    k :: Int
    λ :: Float64
    error :: Float64
    indices :: Array{Int,1}
    principal_vectors :: Array{Float64,2}
end

    function show(io :: IO, traj :: SpatialTrajectory)
        return print(io, "error: $(traj.error)\nindices: $(traj.indices)\nprincipal vectors: $(traj.principal_vectors)")
    end


using ImageFiltering, ImageCore, ComputationalResources
using LinearAlgebra
using Test

@testset "TriggsSdika" begin
    @testset "1d" begin
        l = 1000
        x = [1, 2, 3, 5, 10, 20, l>>1, l-19, l-9, l-4, l-2, l-1, l]
        σs = [3.1, 5, 10.0, 20, 50.0, 100.0]
        # For plotting:
        # aσ = Array{Any}(undef, length(x))
        # for i = 1:length(x)
        #     aσ[i] = Array{Float64}(undef, l, 2*length(σs))
        # end
        for (n,σ) in enumerate(σs)
            kernel = KernelFactors.IIRGaussian(σ)
            A = [kernel.a'; I zeros(2)]
            B = [kernel.b'; I zeros(2)]
            I1 = zeros(3,3); I1[1,1] = 1
            @test kernel.M ≈ (I1+B*kernel.M*A)  # Triggs & Sdika, Eq. 8
            for (i,c) in enumerate(x)
                a = zeros(l)
                a[c] = 1
                af = exp.(-((1:l).-c).^2/(2*σ^2))/(σ*√(2π))
                @inferred(imfilter!(a, a, (kernel,), Fill(0)))
                @test norm(a-af) < 0.1*norm(af)
                # aσ[i][:,2n-1:2n] = [af a]
            end
        end

        kernel = KernelFactors.IIRGaussian(2.0)
        @test_throws DimensionMismatch imfilter([1.0, 2.0], (kernel,))
    end

    @testset "commutivity" begin
        img = 1:8
        k1 = KernelFactors.IIRGaussian(2)
        k2 = centered(ones(3)/3)
        for border in (Pad(:replicate), Fill(0))
            @test imfilter(img, (k1, k2), border) ≈ imfilter(img, (k2, k1), border)
        end
    end

    @testset "images" begin
        imgf = zeros(5, 7); imgf[3,4] = 1
        imgg = fill(Gray{Float32}(0), 5, 7); imgg[3,4] = 1
        imgc = fill(RGB{Float64}(0,0,0), 5, 7); imgc[3,4] = RGB(1,0,0)
        σ = 5
        x = -2:2
        y = (-3:3)'
        kernel = KernelFactors.IIRGaussian((σ, σ))
        for img in (imgf, imgg, imgc)
            imgcmp = img[3,4] .* exp.(-(x.^2 .+ y.^2)/(2*σ^2))/(σ^2*2*pi)
            border = Fill(zero(eltype(img)))
            img0 = copy(img)
            imgfilt = @inferred(imfilter(img, kernel, border))
            @test sum(abs2, imgcmp - imgfilt) < 0.2^2*sum(abs2, imgcmp)
            @test img == img0
            @test imgfilt != img
        end
        ret = imfilter(imgf, KernelFactors.IIRGaussian((σ,0)), "replicate")
        @test ret != imgf
        ret = imfilter(imgf, KernelFactors.IIRGaussian((0,σ)), "replicate")
        @test ret != imgf
        out = similar(ret)
        imfilter!(CPU1(Algorithm.IIR()), out, imgf, KernelFactors.IIRGaussian(σ), 2, "replicate")
        @test out == ret

        # for code coverage
        # TODO: make into an actual test
        kerng = KernelFactors.IIRGaussian(σ)
        kern = (kerng,kerng)
        imfilter(imgf, kern, NA())

        # When the image has NaNs
        imgfnan = copy(imgf)
        imgfnum = copy(imgf)
        imgfnan[1,1] = NaN
        imgfnum[1,1] = 0
        imgfden = fill(1, axes(imgfnum))
        imgfden[1,1] = 0
        retnum, retden = imfilter(imgfnum, kernel, Fill(0.0)), imfilter(imgfden, kernel, Fill(0.0))
        ret = imfilter(imgfnan, kernel, NA())
        ret[1,1] = retnum[1,1] = 0
        @test ret ≈ retnum./retden
    end

    @testset "OffsetArrays" begin
        A = reshape(collect(1:100*100),100,100)
        kern = ntuple(d->KernelFactors.IIRGaussian(5.0), 2)
        B = imfilter(A, kern, NA())
        for i = -5:5
            for j = -5:5
                C = OffsetArray(A,-1,-1)
                D = imfilter(C, kern, NA())
                @test collect(B) == collect(D)
            end
        end
    end
end

nothing


"""
    support(cvxpoly, dir)

Compute the point of contact between a convex object and its
supporting hyperplane defined by the given normal direction as
a StaticArray.
"""
function support(vertices::SMatrix, dir::AbstractVector)
    @inbounds vertices[:, argmax(vertices'*dir)]
end

# TODO support for ordered vertices / simplices in a convex hull
#  struct CvxHull{A<:SMatrix}
#      pts::A
#  end
#  function support(vertices::CvxHull{SMatrix{2, M}}, dir::SVector{2, T}) where {M, T}
#      li = 1; ri = M
#      left = vertices[:, 1]'*dir
#      right = vertices[:, M]'*dir
#      cnt = 10
#      while ri - li > 2
#          mi = div(li + ri, 2)
#          mid = vertices[:, mi]'*dir
#          if left < right
#              left = mid
#              li = mi
#          else
#              right = mid
#              ri = mi
#          end
#          cnt -= 1
#      end
#      if left < right
#          return vertices[:, ri]
#      else
#          return vertices[:, li]
#      end
#  end


function define_arithmetic(machine)
    define(machine, "+",        gen_op_nn_n(+))
    define(machine, "-",        gen_op_nn_n(-))
    define(machine, "*",        gen_op_nn_n(*))
    define(machine, "/",        gen_op_nn_n(/))
    define(machine, "<",        gen_op_nn_n(less))
    define(machine, ">",        gen_op_nn_n(more))
    define(machine, "=",        gen_op_nn_n(equal))
end

less(a, b) = a < b ? -1 : 0
more(a, b) = less(b, a)
equal(a, b) = a == b ? -1 : 0

function gen_op_nn_n(operation)
    return (machine, parent, myidx) -> begin
        arg2 = pop!(machine.stack)
        arg1 = pop!(machine.stack)
        push!(machine.stack, operation(arg1, arg2))
        return 1
    end
end


using MatrixNetworks

type Biconnected_components_output
    bcc_edge_list::Vector{Int64}  #list of bcc edges
    components::Vector{Int64} #Indices of bcc corresponding to the edges in List
    number::Int64 # Total number of Biconnected components
    A::MatrixNetwork #MatrixNetwork
end

I = [1,2,2,5,5,1,5,3,3,6,6,5,4,5]
J = [2,1,5,2,1,5,3,5,6,3,5,6,5,4]
V = [1,1,1,1,1,1,1,1,1,1,1,1,1,1]

A = sparse(I,J,V)

function BiCC_util(A::MatrixNetwork, vertex::Int64, bcc_edges::Vector{Int64}, low::Vector{Int64}, dt::Vector{Int64}, pred::Vector{Int64}, components::Vector{Int64}, cs::Vector{Int64}, cn::Int64, t::Int64, css::Int64, bcc_edges_count::Int64)
    t = t+1
    dt[vertex] = t
    low[vertex] = t
    children = 0
    for i=A.rp[vertex]:A.rp[vertex+1]-1
        adj = A.ci[i];
        if dt[adj] == 0
            children += 1
            pred[adj] = vertex
            css = css+1
            cs[2*css-1] = vertex
            cs[2*css] = adj
            (cn,t,css,bcc_edges_count) = BiCC_util(A, adj, bcc_edges, low, dt, pred, components, cs, cn, t, css, bcc_edges_count)
            if low[vertex] > low[adj]
                low[vertex] = low[adj]
            end
            if ((dt[vertex] == 1 && children > 1) || (dt[vertex] > 1 && low[adj] >= dt[vertex]))
                while css > 0
                    bcc_edges_count += 1
                    bcc_edges[2*bcc_edges_count-1] = cs[2*css-1]
                    bcc_edges[2*bcc_edges_count] = cs[2*css]
                    components[bcc_edges_count] = cn
                    if (cs[2*css-1] == vertex && cs[2*css] == adj)
                        css -= 1
                        break
                    end
                    css -= 1
                end
                cn += 1
            end
        elseif (adj != pred[vertex] && dt[adj] < low[vertex])
            low[vertex] = dt[adj]
            css = css+1
            cs[2*css-1] = vertex
            cs[2*css] = adj
        end
    end
    return (cn,t,css,bcc_edges_count)
end

function BiCC(A::MatrixNetwork)
    bcc_edges=zeros(Int64,2*length(A.ci))
    bcc_edges_count=0
    cn=1
    low=zeros(Int64,A.n)
    dt=zeros(Int64,A.n)
    pred=zeros(Int64,A.n)
    components=-1 * ones(Int64,length(A.ci))
    t=0
    cs=zeros(Int64,2*length(A.ci))
    css=0
    for i=1:A.n
        if dt[i] == 0
            (cn,t,css,bcc_edges_count) = BiCC_util(A, i, bcc_edges, low, dt, pred, components, cs, cn, t, css, bcc_edges_count)
        end
    end
    while css > 0
        bcc_edges_count += 1
        bcc_edges[2*bcc_edges_count-1] = cs[2*css-1]
        bcc_edges[2*bcc_edges_count] = cs[2*css]
	components[bcc_edges_count] = cn
        css -= 1
    end
    print(bcc_edges)
    print(components)
end

function dfs_test()
    print(A);
B = MatrixNetwork(A);
print(B)
   print(BiCC(B))
end

dfs_test();


@testset "5_rotation" begin
  @testset "rotation_ckks" begin
    @testset "EncryptionParameters" begin
      @test_nowarn EncryptionParameters(SchemeType.ckks)
    end
    enc_parms = EncryptionParameters(SchemeType.ckks)

    @testset "polynomial modulus degree" begin
      @test_nowarn set_poly_modulus_degree!(enc_parms, 8192)
      @test poly_modulus_degree(enc_parms) == 8192
    end

    @testset "coefficient modulus" begin
      @test_nowarn coeff_modulus_create(8192, [40, 40, 40, 40, 40])
      @test_nowarn set_coeff_modulus!(enc_parms, coeff_modulus_create(8192, [40, 40, 40, 40, 40]))
    end

    @testset "SEALContext" begin
      @test_nowarn SEALContext(enc_parms)
    end
    context = SEALContext(enc_parms)

    @testset "KeyGenerator" begin
      @test_nowarn KeyGenerator(context)
    end
    keygen = KeyGenerator(context)

    @testset "PublicKey" begin
      @test_nowarn PublicKey()
    end
    public_key_ = PublicKey()

    @testset "create_public_key" begin
      @test_nowarn create_public_key!(public_key_, keygen)
    end

    @testset "SecretKey" begin
      @test_nowarn secret_key(keygen)
    end
    secret_key_ = secret_key(keygen)

    @testset "RelinKeys" begin
      @test_nowarn RelinKeys()
    end
    relin_keys_ = RelinKeys()

    @testset "create_relin_keys" begin
      @test_nowarn create_relin_keys!(relin_keys_, keygen)
    end

    @testset "GaloisKeys" begin
      @test_nowarn GaloisKeys()
    end
    galois_keys_ = GaloisKeys()

    @testset "create_galois_keys" begin
      @test_nowarn create_galois_keys!(galois_keys_, keygen)
    end

    @testset "Encryptor" begin
      @test_nowarn Encryptor(context, public_key_)
    end
    encryptor = Encryptor(context, public_key_)

    @testset "Evaluator" begin
      @test_nowarn Evaluator(context)
    end
    evaluator = Evaluator(context)

    @testset "Decryptor" begin
      @test_nowarn Decryptor(context, secret_key_)
    end
    decryptor = Decryptor(context, secret_key_)

    @testset "CKKSEncoder" begin
      @test_nowarn CKKSEncoder(context)
    end
    encoder = CKKSEncoder(context)

    slot_count_ = 4096
    @testset "slot_count" begin
      @test slot_count(encoder) == slot_count_
    end

    input = collect(range(0.0, 1.0, length=slot_count_))
    plain = Plaintext()
    encrypted = Ciphertext()
    initial_scale = 2.0^50
    @testset "encode and encrypt" begin
      @test_nowarn encode!(plain, input, initial_scale, encoder)
      @test_nowarn encrypt!(encrypted, plain, encryptor)
    end

    rotated = Ciphertext()
    @testset "rotate 2 steps left" begin
      @test_nowarn rotate_vector!(rotated, encrypted, 2, galois_keys_, evaluator)
    end

    result = similar(input)
    @testset "decrypt and decode" begin
      @test_nowarn decrypt!(plain, rotated, decryptor)
      @test_nowarn decode!(result, plain, encoder)
    end

    @testset "compare results" begin
      true_result = circshift(input, -2)
      @test isapprox(result, true_result)
    end
  end
end


#=
  False coin in Turing.jl

  From Peter Winkler
  Celebration of Mind (Gathering For Gardner):
  "CoM Apr 2021 - Drawing from Urns - Peter Winkler"
  https://www.youtube.com/watch?v=A0GxvYPTKDk&list=PL5_D5QQxvWKVtX4Vklvhsf5KTYZBLDSwO&index=3

  @3:43
  """
  You have two quarters.
  Both are sitting on a table heads up.
  One of them is an ordinary quarter, but the
  other has a head on the other side as well.

  You get two flips to help you guess which is which.

  Should you flip each coin once or one coin twice.
  """

  According to this model it's better to flip the same coin twice
  than once for each coin. But that's under the condition that
  all the flips shows head. If either flip shows tail then it's
  settled that it's the ordinary coin.

  """
  Distributions of variable coinA (num:0)
  1.00000 =>    8066  (0.806600)
  2.00000 =>    1934  (0.193400)

  Distributions of variable sideA (num:0)
  1.00000 =>   10000  (1.000000)

  Distributions of variable coinB (num:0)
  2.00000 =>    8066  (0.806600)
  1.00000 =>    1934  (0.193400)

  Distributions of variable sideB (num:0)
  1.00000 =>   10000  (1.000000)

  Distributions of variable coinSameIsFake (num:0)
  1.00000 =>    8066  (0.806600)
  0.00000 =>    1934  (0.193400)

  Distributions of variable coinBothIsFake (num:0)
  0.00000 =>    8066  (0.806600)
  1.00000 =>    1934  (0.193400)

  """

  (See false_coin.jl for an easier variant of the problem.)

=#

using Turing
include("jl_utils.jl")

@model function false_coin2()
    #  The two coins
    coin1 = 1 #  fake coin: head,head
    coin2 = 2 #  proper coin: head,tail
    coins = [coin1,coin2];

    #  Sides of the coins
    head = 1
    tail = 2

    function flip_coin(cond)
        return cond ? Categorical([1.0,0.0]) :  Categorical([0.5,0.5])
    end

    #  The setup: both coins are head (no real flipping)
    coinA ~ Categorical([0.5,0.5])
    sideA ~ flip_coin(coinA == coin1)
    true ~ Dirac(sideA == head)

    coinB ~ Dirac(coinA == coin1 ? coin2 : coin1)
    sideB ~ flip_coin(coinB == coin1);
    true ~ Dirac(sideB == head)

    #  condition(coinA == coin1 || coinB == coin1)

    #  Flip same coin
    coinSame_1 ~ Dirac(coinA)
    coinSame_1_flip ~ flip_coin(coinSame_1 == coin1);
    coinSame_2 ~ Dirac(coinSame_1)
    coinSame_2_flip ~ flip_coin(coinSame_2 == coin1);
    true ~ Dirac(coinSame_1_flip == head)
    true ~ Dirac(coinSame_2_flip == head)

    #  Flip both coins
    coinBoth_1 ~ Dirac(coinA)
    coinBoth_1_flip ~ flip_coin(coinBoth_1 == coin1);
    coinBoth_2 ~ Dirac(coinB)
    coinBoth_2_flip ~ flip_coin(coinBoth_2 == coin1);
    true ~ Dirac(coinBoth_1_flip == head)
    true ~ Dirac(coinBoth_2_flip == head)

    coinSameIsFake ~ Dirac(coinA ==coin1)
    coinBothIsFake ~ Dirac(coinB ==coin1)
end

model = false_coin2()
# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, IS(), 10_000)
# chns = sample(model, SMC(), 10_000)

display(chns)
show_var_dist_pct(chns,:coinA)
show_var_dist_pct(chns,:sideA)
show_var_dist_pct(chns,:coinB)
show_var_dist_pct(chns,:sideB)

show_var_dist_pct(chns,:coinSameIsFake)
show_var_dist_pct(chns,:coinBothIsFake)


module Threshold
using ..Util, ..WaveletTypes, ..Transforms
export
    # threshold
    threshold!,
    threshold,
    HardTH,
    SoftTH,
    SemiSoftTH,
    SteinTH,
    BiggestTH,
    PosTH,
    NegTH,
    # denoising
    DNFT,
    VisuShrink,
    denoise,
    noisest,
    # basis functions
    matchingpursuit,
    bestbasistree,
    # entropy
    coefentropy,
    Entropy,
    ShannonEntropy,
    LogEnergyEntropy

# THRESHOLD TYPES AND FUNCTIONS

abstract THType
immutable HardTH     <: THType end
immutable SoftTH     <: THType end
immutable SemiSoftTH <: THType end
immutable SteinTH    <: THType end
immutable BiggestTH  <: THType end
immutable PosTH      <: THType end
immutable NegTH      <: THType end

const DEF_TH = HardTH()

# biggest m-term approximation (best m-term approximation for orthogonal transforms)
# result is m-sparse
function threshold!{T<:Number}(x::AbstractArray{T}, TH::BiggestTH, m::Int)
    @assert m >= 0
    n = length(x)
    m > n && (m = n)
    ind = sortperm(x, alg=QuickSort, by=abs)
    @inbounds begin
        for i = 1:n-m
            x[ind[i]] = 0
        end
    end
    return x
end

# hard
function threshold!{T<:Number}(x::AbstractArray{T}, TH::HardTH, t::Real)
    @assert t >= 0
    @inbounds begin
        for i = 1:length(x)
            if abs(x[i]) <= t
                x[i] = 0
            end
        end
    end
    return x
end

# soft
function threshold!{T<:Number}(x::AbstractArray{T}, TH::SoftTH, t::Real)
    @assert t >= 0
    @inbounds begin
        for i = 1:length(x)
            sh = abs(x[i]) - t
            if sh < 0
                x[i] = 0
            else
                x[i] = sign(x[i])*sh
            end
        end
    end
    return x
end

# semisoft
function threshold!{T<:Number}(x::AbstractArray{T}, TH::SemiSoftTH, t::Real)
    @assert t >= 0
    @inbounds begin
        for i = 1:length(x)
            if x[i] <= 2*t
                sh = abs(x[i]) - t
                if sh < 0
                    x[i] = 0
                elseif sh - t < 0
                    x[i] = sign(x[i])*sh*2
                end
            end
        end
    end
    return x
end

# stein
function threshold!{T<:Number}(x::AbstractArray{T}, TH::SteinTH, t::Real)
    @assert t >= 0
    @inbounds begin
        for i = 1:length(x)
            sh = 1 - t*t/(x[i]*x[i])
            if sh < 0
                x[i] = 0
            else
                x[i] = x[i]*sh
            end
        end
    end
    return x
end

# shrink negative elements to 0
function threshold!{T<:Number}(x::AbstractArray{T}, TH::NegTH)
    @inbounds begin
        for i = 1:length(x)
            if x[i] < 0
                x[i] = 0
            end
        end
    end
    return x
end

# shrink positive elements to 0
function threshold!{T<:Number}(x::AbstractArray{T}, TH::PosTH)
    @inbounds begin
        for i = 1:length(x)
            if x[i] > 0
                x[i] = 0
            end
        end
    end
    return x
end

# the non inplace functions
function threshold{T<:Number}(x::AbstractArray{T}, TH::THType, t::Real)
    y = Array(T, size(x))
    return threshold!(copy!(y,x), TH, t)
end
function threshold{T<:Number}(x::AbstractArray{T}, TH::THType)
    y = Array(T, size(x))
    return threshold!(copy!(y,x), TH)
end

# DENOISING

abstract DNFT

type VisuShrink <: DNFT
    th::THType      # threshold type
    t::Float64      # threshold for noise level sigma=1, use sigma*t in application
end
# define type for signal length n
function VisuShrink(n::Int)
    return VisuShrink(DEF_TH, sqrt(2*log(n)))
end

const DEF_WAVELET = waveletfilter("sym5")    # default wavelet type

# denoise signal x by thresholding in wavelet space
function denoise{T<:DiscreteWavelet,S<:DNFT}(x::AbstractArray;
                                    wt::Union(T,Nothing)=DEF_WAVELET,
                                    level::Int=max(nscales(size(x,1))-6,1),
                                    dnt::S=VisuShrink(size(x,1)),
                                    sigma::Real=noisest(x, wt=wt),
                                    TI::Bool=false,
                                    nspin::Union(Int,Tuple)=tuple([8 for i=1:length(size(x))]...) )
    @assert iscube(x)

    if TI
        wt == nothing && error("TI not supported with wt=nothing")
        y = zeros(eltype(x), size(x))
        xt = Array(eltype(x), size(x))
        L = level2tl(size(x,1), level)
        pns = prod(nspin)

        if ndims(x) == 1
            z = Array(eltype(x), size(x))
            for i = 1:pns
                shift = nspin2circ(nspin, i)[1]
                circshift!(z, x, shift)

                dwt!(xt, z, wt, L, true)
                threshold!(xt, dnt.th, sigma*dnt.t)
                dwt!(z, xt, wt, L, false)

                circshift!(xt, z, -shift)
                arrayadd!(y, xt)
            end
        else # ndims > 1
            for i = 1:pns
                shift = nspin2circ(nspin, i)
                z = circshift(x, shift)

                dwt!(xt, z, wt, L, true)
                threshold!(xt, dnt.th, sigma*dnt.t)
                dwt!(z, xt, wt, L, false)

                z = circshift(z, -shift)
                arrayadd!(y, z)
            end
        end
        scale!(y, 1/pns)
    else # !TI
        if wt == nothing
            y = copy(x)
            threshold!(y, dnt.th, sigma*dnt.t)
        else
            L = level2tl(size(x,1),level)
            y = dwt(x, wt, L)
            threshold!(y, dnt.th, sigma*dnt.t)
            dwt!(y, wt, L, false)
        end
    end

    return y
end
# add z to y
function arrayadd!(y::AbstractArray, z::AbstractArray)
    @assert length(y) == length(z)
    for i = 1:length(y)
        @inbounds y[i] += z[i]
    end
    return y
end

# estimate the std. dev. of the signal noise, assuming Gaussian distribution
function noisest{T<:DiscreteWavelet}(x::AbstractArray; wt::Union(T,Nothing)=DEF_WAVELET)
    if wt == nothing
        y = x
    else
        y = dwt(x, wt, 1)
    end
    ind = detailrange(maxlevel(size(y,1)))
    dr = y[ind]
    return mad!(dr)/0.6745
end
# Median absolute deviation
function mad!(y::AbstractArray)
    m = median!(y)
    for i in 1:length(y)
        y[i] = abs(y[i]-m)
    end
    return median!(y, checknan=false)
end
function mad(x::AbstractArray)
    y = copy(x)
    mad!(y)
end

# convert index i to a circshift array starting at 0 shift
nspin2circ(nspin::Int, i::Int) = nspin2circ((nspin,), i)
function nspin2circ(nspin::Tuple, i::Int)
    c1 = ind2sub(nspin,i)
    c = Array(Int, length(c1))
    for k = 1:length(c1)
        c[k] = c1[k]-1
    end
    return c
end

# BASIS FUNCTIONS

# Matching Pursuit
# see: Mallat (2009) p.642 "A wavelet tour of signal processing"
# find sparse vector y such that ||x - f(y)|| < tol approximately
# f is the operation of a M by N (M<N) dictionary/matrix
# ft is a function defining the transpose of f
function matchingpursuit(x::AbstractVector, f::Function, ft::Function, tol::Real, nmax::Int=-1, oop::Bool=false, N::Int=0)
    @assert nmax >= -1
    @assert tol > 0
    r = x
    n = 1

    if !oop
        y = zeros(eltype(x), length(ft(x)))
    else # out of place functions f and ft
        y = zeros(eltype(x), N)
        tmp = Array(eltype(x), N)
        ftr = Array(eltype(x), N)
        aphi = Array(eltype(x), length(x))
    end
    spat = zeros(eltype(x), length(y))  # sparse for atom computation
    nmax == -1 && (nmax = length(y))

    while vecnorm(r) > tol && n <= nmax
        # find largest inner product
        !oop && (ftr = ft(r))
        oop  && ft(ftr, r, tmp)
        i = findmaxabs(ftr)

        # project on i-th atom
        spat[i] = ftr[i]
        !oop && (aphi = f(spat))
        oop  && f(aphi, spat, tmp)
        spat[i] = 0

        # update residual, r = r - aphi
        broadcast!(-, r, r, aphi)

        y[i] += ftr[i]
        n += 1
    end
    return y
end
function findmaxabs(x::AbstractVector)
    m = abs(x[1])
    k = 1
    @inbounds for i = 1:length(x)
        if abs(x[i]) > m
            k = i
            m = abs(x[i])
        end
    end
    return k
end

# ENTROPY MEASURES

abstract Entropy
immutable ShannonEntropy <: Entropy end  #Coifman-Wickerhauser
immutable LogEnergyEntropy <: Entropy end

# Entropy measures: Additive with coefentropy(0) = 0
# all coefs assumed to be on [-1,1] after normalization with nrm
# given x and y, where x has "more concentrated energy" than y
# then coefentropy(x, et, norm) <= coefentropy(y, et, norm) should be satisfied.

function coefentropy{T<:FloatingPoint}(x::T, et::ShannonEntropy, nrm::T)
    s = (x/nrm)^2
    if s == 0.0
        return -zero(T)
    else
        return -s*log(s)
    end
end
function coefentropy{T<:FloatingPoint}(x::T, et::LogEnergyEntropy, nrm::T)
    s = (x/nrm)^2
    if s == 0.0
        return -zero(T)
    else
        return -log(s)
    end
end
function coefentropy{T<:FloatingPoint}(x::AbstractArray{T}, et::Entropy, nrm::T=vecnorm(x))
    @assert nrm >= 0
    sum = zero(T)
    nrm == sum && return sum
    for i in 1:length(x)
        @inbounds sum += coefentropy(x[i], et, nrm)
    end
    return sum
end

function bestbasistree{T<:FloatingPoint}(y::AbstractVector{T}, wt::DiscreteWavelet, L::Integer=nscales(y), et::Entropy=ShannonEntropy())
    bestbasistree(y, wt, maketree(length(y), L, :full), et)
end
function bestbasistree{T<:FloatingPoint}(y::AbstractVector{T}, wt::DiscreteWavelet, tree::BitVector, et::Entropy=ShannonEntropy())

    n = length(y)
    J = nscales(n)
    @assert isdyadic(y)
    @assert isvalidtree(y, tree)
    tree[1] || return besttree      # do nothing

    besttree = copy(tree)
    x = copy(y)
    tmp = Array(T, size(y))
    entr = Array(T, length(tree)+length(y))
    nrm = vecnorm(y)

    L = J
    k = 1
    while L > 0
        ix = 1
        Lfw = J-L
        nj = detailn(tl2level(n, Lfw))

        dtmp = Transforms.unsafe_vectorslice(tmp, 1, nj)
        while ix <= n
            dx = Transforms.unsafe_vectorslice(x, ix, nj)

            entr[k] = coefentropy(dx, et, nrm)

            dwt!(dtmp, dx, wt, 1, true)
            copy!(dx, dtmp)

            ix += nj
            k += 1
        end
        L -= 1
    end

    for i = 1:length(y)
        entr[k] = coefentropy(y[i], et, nrm)
        k += 1
    end

    # make the best tree
    for i=1:length(tree)
        if (i>1 && !besttree[i>>1]) || !tree[i]  # parent is 0 or input tree-node is 0
            besttree[i] = false
        else
            if entr[i] <= bestsubtree(entr, i)
                besttree[i] = false
            else
                besttree[i] = true
            end
        end
    end

    @assert isvalidtree(y, besttree)
    return besttree
end

# the entropy of best subtree
function bestsubtree(entr::Array, i::Int)
    n = length(entr)

    if i<<1+1 > n  # bottom of tree
        return entr[i]
    else
        if (i<<1+1)<<1+1 > n  # above bottom of tree
            sum = entr[i<<1]
            sum += entr[i<<1+1]
        else
            sum = bestsubtree(entr, i<<1)
            sum += bestsubtree(entr, i<<1+1)
        end

        return min(entr[i], sum)
    end
end

end


#=
Parallel direct estimation of the probability for the Schlögl chemical reaction network
with X(0)= x₀ to reach X(T) ∈ A, the target set.

This reaction network corresponds to:
A + 2S → 3S
3S → A + 2S
B → S
S → B

Parameters were chosen so that the system is bistable, as in:
https://doi.org/10.1063/1.5017955
=#
using Distributed
# nw = 4; # number of workers
# addprocs(nw);

using Statistics
using HypothesisTests
using Printf

@everywhere include("schlogl_setup.jl")

Random.seed!(100)
@everywhere n_trials = 10^6;
n_in_set = @sync @distributed (+) for j in 1:n_trials
    X₀ = copy(x₀);
    f(integrate_schlogl(X₀, Tmax));
end
CI = confint(BinomialTest(n_in_set, n_trials));
@printf("P Estimate = %g\n",n_in_set/n_trials);
@printf("95%% CI = (%g,%g)\n", CI[1], CI[2])


module LaplaceApproximations

using Zygote
using LinearAlgebra

export
    laplace_approx,
    GPML

include("GPML.jl")

"""
    mnll, u = laplace_approx(nll, u0; grtol = sqrt(eps(eltype(u0))))

NOT ROBUST! Uses standard Newton updates with no checking or scaling.

For negative log-likelihood function `nll`, calculate the approximate marginal
likelihood after integrating out the parameters in `u` using the Laplace
approximation. Returns the marginal negative log-likelihood and the values of
`u` at the minimum. The function `nll` should accept only the argument `u`.
Newton steps are taken until the maximum gradient component is less than `grtol`
or `maxit`, whichever comes first.
"""
function laplace_approx(nll::Function, u0;
                        grtol = sqrt(eps(eltype(u0))),
                        maxit = 100)
    u = copy(u0)
    gr = Zygote.gradient(nll, u)[1]
    he = Zygote.hessian(nll, u)
    nit = 0
    while(true)
        nit += 1
        u .-= he \ gr
        gr .= Zygote.gradient(nll, u)[1]
        he .= Zygote.hessian(nll, u)
        maximum(gr) < grtol && break
        if nit ≥ maxit
            @warn "Gradient tolerance not reached in $nit Newton steps" gr = gr
            break
        end
    end
    mnll = -length(u) / 2 * log(2π) + logdet(he) / 2 + nll(u)
    return (mnll, u)
end

end


using GCMAES, Plots

xs = Float32.(collect(-6:0.1:6))
lo = fill(1f0, size(xs))
hi = fill(2f0, size(xs))
trans = BoxLinQuadTransform(lo, hi)
plot(xs, GCMAES.transform(trans, xs))
plot(xs, GCMAES.inverse(trans, xs))


@testset "chainrules" begin
    x = randn()
    z = x + exp(randn())
    y = z + exp(randn())
    test_frule(DistributionsAD.uniformlogpdf, x, y, z)
    test_rrule(DistributionsAD.uniformlogpdf, x, y, z)

    # StatsFuns: https://github.com/JuliaStats/StatsFuns.jl/pull/106
    x = exp(randn())
    y = exp(randn())
    z = logistic(randn())
    test_frule(StatsFuns.betalogpdf, x, y, z)
    test_rrule(StatsFuns.betalogpdf, x, y, z)

    x = exp(randn())
    y = exp(randn())
    z = exp(randn())
    test_frule(StatsFuns.gammalogpdf, x, y, z)
    test_rrule(StatsFuns.gammalogpdf, x, y, z)

    x = exp(randn())
    y = exp(randn())
    test_frule(StatsFuns.chisqlogpdf, x, y)
    test_rrule(StatsFuns.chisqlogpdf, x, y)

    x = exp(randn())
    y = exp(randn())
    z = exp(randn())
    test_frule(StatsFuns.fdistlogpdf, x, y, z)
    test_rrule(StatsFuns.fdistlogpdf, x, y, z)

    x = exp(randn())
    y = randn()
    test_frule(StatsFuns.tdistlogpdf, x, y)
    test_rrule(StatsFuns.tdistlogpdf, x, y)

    # TODO: Re-enable if https://github.com/JuliaMath/SpecialFunctions.jl/pull/325 is fixed
    # use `BigFloat` to avoid Rmath implementation in finite differencing check
    # (returns `NaN` for non-integer values)
    #n = rand(1:100)
    #x = BigFloat(n)
    #y = big(logistic(randn()))
    #z = BigFloat(rand(1:n))
    #test_frule(StatsFuns.binomlogpdf, x, y, z)
    #test_rrule(StatsFuns.binomlogpdf, x, y, z)

    x = big(exp(randn()))
    y = BigFloat(rand(1:100))
    test_frule(StatsFuns.poislogpdf, x, y)
    test_rrule(StatsFuns.poislogpdf, x, y)

    _, pb = rrule(StatsFuns.poislogpdf, 0.0, 0.0)
    _, x̄1, _ = pb(1)
    @test x̄1 == -1
    _, pb = rrule(StatsFuns.poislogpdf, 0.0, 1.0)
    _, x̄1, _ = pb(1)
    @test x̄1 == Inf

    # PoissonBinomial
    test_rrule(Distributions.poissonbinomial_pdf_fft, rand(50))
    if isdefined(Distributions, :poissonbinomial_pdf)
        test_rrule(Distributions.poissonbinomial_pdf, rand(50))
    end
end


"""
    slidingwindow(L::LT, f::FT, d::IT) where {LT <: SimpleSDMLayer, FT <: Function, IT <: Number}

This function will replace the value at any cell by applying the function `f` to
the array of cells values that are within a distance `d` (in kilometers) from
the focal cell. This is, for example, useful to use an average to smooth out the
layers. The distance is estimated using the haversine distance, assuming that
the radius of the Earth is 6371.0 km. This means that the size of the window
will vary a little bit across latitudes, but this is far better than using a
number of cells, which would have dramatic consequences near the poles.

It *always* returns a `SimpleSDMResponse`, and the cells containing `nothing`
will also not contain a value in the output. This is *different* from the
behavior of `coarsen`, which tends to expand the area of the layer in which we
have data.

This function is currently relatively slow. Performance improvements will arrive
at some point.
"""
function slidingwindow(layer::LT, f::FT, d::IT; threaded::Bool=Threads.nthreads()>1) where {LT <: SimpleSDMLayer, FT <: Function, IT <: Number}
    # We infer the return type from a call to the function on the first three elements
    return_type = typeof(f(collect(layer)[1:min(3, length(layer))]))

    # New layer using typed similar
    N = similar(layer, return_type)

    # Vector of all positions with a value
    filled_positions = keys(layer)

    # We then filter in the occupied positions
    if threaded
        Threads.@threads for pos in filled_positions
            N[pos] = f(_sliding_values(layer, pos, d))
        end
    else
        for pos in filled_positions
           N[pos] = f(_sliding_values(layer, pos, d))
        end
    end

    # And we return the object
    return N
end

function _sliding_values(layer, pt, d; R=6371.0)
    # Bounding box (approx.) for the sliding window of length d at the given point
    max_lat = min(layer.top, pt[2]+(180.0*d)/(π*R))
    min_lat = max(layer.bottom, pt[2]-(180.0*d)/(π*R))
    max_lon = min(layer.right, pt[1]+(360.0*d)/(π*R))
    min_lon = max(layer.left, pt[1]-(360.0*d)/(π*R))

    # Extracted layer for the sliding window
    _tmp = clip(layer; left=min_lon, right=max_lon, top=max_lat, bottom=min_lat)

    # Filter the correct positions
    filled_positions = keys(_tmp)
    neighbors = filter(p -> Distances.haversine((pt[1], pt[2]), (p[1], p[2]), R) < d, filled_positions)

    # Return
    return _tmp[neighbors]
end


using DynamicGrids

# Write the Game of life
#
# Use the dsctiption at:
# https://en.wikipedia.org/wiki/Conway%27s_Game_of_Life

# It could be something like this
life = Neighbors(Moore(1)) do data, hood, active, I
    if active
        sum(hood) in (2, 3)
    else
        sum(hood) == 3
    end
end

init = rand(Bool, 1000, 1000)
repl_output = REPLOutput(init; tspan=1:100, fps=25)
sim!(repl_output, life; boundary=Wrap())

# How fast is this thing?
#
# Benchmark it!
using BenchmarkTools
result_output = ResultOutput(init; tspan=1:100)
@btime sim!(result_output, life; boundary=Wrap())

# But actually, this laptop has 8 cores:
Threads.nthreads()
sim!(repl_output, life; boundary=Wrap(), proc=ThreadedCPU())
@btime sim!(result_output, life; boundary=Wrap(), proc=ThreadedCPU())

# And a GPU!
# Loading CUDAKernels gives us CUDA GPU capability
using CUDAKernels
sim!(repl_output, life; boundary=Wrap(), proc=CuGPU())
@btime sim!(result_output, life; boundary=Wrap(), proc=CuGPU())

# Forest Fire

# Use the numnered rules in
# https://en.wikipedia.org/wiki/Forest-fire_model

# To get you started, the grid will have 3 states:
const EMPTY = 0
const BURNING = 1
const ALIVE = 2

# forest fire with parameters for sliders
forest_fire = let f = Param(0.001; bounds=(0.0, 1.0)),
                  p = Param(0.01; bounds=(0.0, 1.0))
    Neighbors() do _, hood, state, I
        if state == ALIVE
            if BURNING in hood
                BURNING
            else
                rand() < f ? BURNING : ALIVE
            end
        elseif state == BURNING
            EMPTY
        else # EMPTY
            rand() < p ? ALIVE : EMPTY
        end
    end
end

init = fill(ALIVE, 150, 300)
repl_output = REPLOutput(init; tspan=1:200)
ruleset = Ruleset(forest_fire; boundary=Wrap())
sim!(repl_output, forest_fire)

using DynamicGridsInteract, ColorSchemes, Colors
# Live output
output = ElectronOutput(init;
    ruleset=ruleset,
    tspan=1:2000,
    scheme=ColorSchemes.brg,
    minval=EMPTY,
    maxval=ALIVE,
    zerocolor=RGB(0),
    throttle=0.5,
)


using AmbientForcing
using NetworkDynamics, Distributions
using Graphs
using OrdinaryDiffEq

# This is an example from NetworkDynamics
# https://github.com/PIK-ICoN/NetworkDynamics.jl/blob/main/examples/kuramoto_plasticity.jl
# This only works with NetworkDynamics v0.5.0 or newer

# Further examples can be found in the test folder

const N_plastic = 10 # number of nodes
k = 4  # average degree
g = barabasi_albert(N_plastic, k)

#=
  Berner, Rico, Eckehard Schöll, and Serhiy Yanchuk.
  "Multiclusters in Networks of Adaptively Coupled Phase Oscillators."
  SIAM Journal on Applied Dynamical Systems 18.4 (2019): 2227-2266.
=#

# The coupling function is modeled by a differential algebraic equation with mass matrix
# 0 * de[1] = e[2] * sin(v_s[1] - v_d[1] + α) / N - e[1] is equivalent to
# e[1] = e[2] * sin(v_s[1] - v_d[1] + α) / N
function kuramoto_plastic_edge!(de, e, v_s, v_d, p, t)
    de[1] =  e[2] * sin(v_s[1] - v_d[1] + α) / N_plastic - e[1]
    de[2] = - ϵ * (sin(v_s[1] - v_d[1] + β) + e[2])

    nothing
end

function kuramoto_plastic_vertex!(dv, v, edges, p, t)
    dv .= 0
    for e in edges
        dv .-= e[1]
    end
end

# Global parameters need to be const for type stability
const ϵ = 0.1
const α = .2π
const β = -.95π

# NetworkDynamics Setup
plasticvertex = ODEVertex(f = kuramoto_plastic_vertex!, dim =1)
mass_matrix_plasticedge = zeros(2,2)
mass_matrix_plasticedge[2,2] = 1. # First variables is set to 0

plasticedge = ODEEdge(f = kuramoto_plastic_edge!, dim=2, sym=[:e, :de], coupling=:undirected,mass_matrix = mass_matrix_plasticedge);
kuramoto_plastic! = network_dynamics(plasticvertex, plasticedge, g)

# Ambient Forcing starts here

# Using a random initial condition x0 violates the constraints!
# The constraints are fulfilled when g(x) ≈ 0.
x0_plastic = rand(106)
g_nd = constraint_equations(kuramoto_plastic!)
sum(g_nd(x0_plastic))

# Using zeros as the initial conditions for the ambient forcing algo
x0_plastic = zeros(106)
g_nd = constraint_equations(kuramoto_plastic!)
g_nd(x0_plastic)

# Perturbing all variables at once
Frand = random_force(kuramoto_plastic!, [0.0, 1.0], Uniform)
afoprob = ambient_forcing_problem(kuramoto_plastic!, x0_plastic, 2.0, Frand)
z_new = ambient_forcing(afoprob, x0_plastic, 2.0, Frand)

# As we can see the constraints are not violated!
sum(g_nd(z_new))

# Perturbing only the variables e_22 and de_22
idx = idx_exclusive(kuramoto_plastic!, ["e_22", "de_22"])
Frand = random_force(kuramoto_plastic!, [0.0, 1.0], Uniform, idx)
z_new = ambient_forcing(afoprob, x0_plastic, 2.0, Frand)

# Still the constraints are fulfilled!
sum(g_nd(z_new))


using Test
using LinearAlgebra
using UnPack
using MOONS.Grids
using MOONS.Fields

#####
##### CellCenter <-> CellFace
#####

function interp_simple!(cent::CellCenter, face::CellFace{1}, grid::Grid)
    for k in over_points(grid.c[3], All(), Center1D())
        for j in over_points(grid.c[2], All(), Center1D())
            for i in over_points(grid.c[1], All(), Center1D())
                D = 1//2
                U = 1//2
                cent[i, j, k] = D*face[i, j, k] + U*face[i+1, j, k]
            end
        end
    end
end

function interp_simple!(face::CellFace{1}, cent::CellCenter, grid::Grid)
    face .= 0 # boundaries are not populated.
    for k in over_points(grid.c[3], All(), Center1D())
        for j in over_points(grid.c[2], All(), Center1D())
            for i in over_points(grid.c[1], Boundary(), Vertex1D())
                D = grid.c[1].interp.diag[i-1]
                U = grid.c[1].interp.upper[i-1]
                face[i, j, k] = D*cent[i-1, j, k] + U*cent[i, j, k]
            end
        end
    end
end

@testset "Interpolations - Simple implementation Center <-> Face" begin
    FT = Float64
    n, a, b, β = 10, FT(0.0), FT(1.0), FT(1.01)
    grid = Cube(a, b, n; warpfun=Roberts_both, args=(β,))

    # Face -> Center
    face = CellFace(grid, 1)
    cent_correct = CellCenter(grid)
    cent = CellCenter(grid)
    @inbounds for local_grid in GridField(grid,face)
        @unpack ijk,x,y,z = local_grid
        face[ijk] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
    end
    interp_simple!(cent_correct, face, grid)
    interp!(cent, face, grid)
    @test all(cent .≈ cent_correct)

    # Center -> Face
    cent = CellCenter(grid)
    face_correct = CellFace(grid, 1)
    face = CellFace(grid, 1)
    @inbounds for local_grid in GridField(grid,cent)
        @unpack ijk,x,y,z = local_grid
        cent[ijk] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
    end
    interp_simple!(face_correct, cent, grid)
    interp!(face, cent, grid)
    @test all(face .≈ face_correct)
end

#####
##### CellCorner <-> CellEdge
#####

function interp_simple!(corn::CellCorner, edge::CellEdge{1}, grid::Grid)
    corn .= 0 # boundaries are not populated.
    for k in over_points(grid.c[3], All(), Vertex1D())
        for j in over_points(grid.c[2], All(), Vertex1D())
            for i in over_points(grid.c[1], Boundary(), Vertex1D())
                D = grid.c[1].interp.diag[i-1]
                U = grid.c[1].interp.upper[i-1]
                corn[i, j, k] = D*edge[i-1, j, k] + U*edge[i, j, k]
            end
        end
    end
end

function interp_simple!(edge::CellEdge{1}, corn::CellCorner, grid::Grid)
    for k in over_points(grid.c[3], All(), Vertex1D())
        for j in over_points(grid.c[2], All(), Vertex1D())
            for i in over_points(grid.c[1], All(), Center1D())
                D = 1//2
                U = 1//2
                edge[i, j, k] = D*corn[i, j, k] + U*corn[i+1, j, k]
            end
        end
    end
end

@testset "Interpolations - Simple implementation Corner <-> Edge" begin
    FT = Float64
    n, a, b, β = 10, FT(0.0), FT(1.0), FT(1.01)
    grid = Cube(a, b, n; warpfun=Roberts_both, args=(β,))

    # Edge -> Corner
    edge = CellEdge(grid, 1)
    corn_correct = CellCorner(grid)
    corn = CellCorner(grid)
    @inbounds for local_grid in GridField(grid,edge)
        @unpack ijk,x,y,z = local_grid
        edge[ijk] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
    end
    interp_simple!(corn_correct, edge, grid)
    interp!(corn, edge, grid)
    @test all(corn .≈ corn_correct)

    # corner -> edge
    corn = CellCorner(grid)
    edge_correct = CellEdge(grid, 1)
    edge = CellEdge(grid, 1)
    @inbounds for local_grid in GridField(grid,corn)
        @unpack ijk,x,y,z = local_grid
        corn[ijk] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
    end
    interp_simple!(edge_correct, corn, grid)
    interp!(edge, corn, grid)
    @test all(edge .≈ edge_correct)
end

#####
##### CellFace <-> CellEdge
#####

function interp_simple!(face::CellFace{1}, edge::CellEdge{2}, grid::Grid)
    for k in over_points(grid.c[3], All(), Center1D())
        for j in over_points(grid.c[2], All(), Center1D())
            for i in over_points(grid.c[1], All(), Vertex1D())
                D = 1//2
                U = 1//2
                face[i, j, k] = D*edge[i, j, k] + U*edge[i, j, k+1]
            end
        end
    end
end

function interp_simple!(edge::CellEdge{2}, face::CellFace{1}, grid::Grid)
    edge .= 0 # boundaries are not populated.
    for k in over_points(grid.c[3], Boundary(), Vertex1D())
        for j in over_points(grid.c[2], All(), Center1D())
            for i in over_points(grid.c[1], All(), Vertex1D())
                D = grid.c[3].interp.diag[k-1]
                U = grid.c[3].interp.upper[k-1]
                edge[i, j, k] = D*face[i, j, k-1] + U*face[i, j, k]
            end
        end
    end
end

@testset "Interpolations - Simple implementation Face <-> Edge" begin
    FT = Float64
    n, a, b, β = 10, FT(0.0), FT(1.0), FT(1.01)
    grid = Cube(a, b, n; warpfun=Roberts_both, args=(β,))

    # Edge -> Face
    edge = CellEdge(grid, 2)
    face_correct = CellFace(grid, 1)
    face = CellFace(grid, 1)
    @inbounds for local_grid in GridField(grid,edge)
        @unpack ijk,x,y,z = local_grid
        edge[ijk] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
    end
    interp_simple!(face_correct, edge, grid)
    interp!(face, edge, grid)
    @test all(face .≈ face_correct)

    # Face -> edge
    face = CellFace(grid, 1)
    edge_correct = CellEdge(grid, 2)
    edge = CellEdge(grid, 2)
    @inbounds for local_grid in GridField(grid,face)
        @unpack ijk,x,y,z = local_grid
        face[ijk] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
    end
    interp_simple!(edge_correct, face, grid)
    interp!(edge, face, grid)
    @test all(edge .≈ edge_correct)
end

#####
##### nested / non-base interpolations
#####

@testset "Interpolations - nested / non-base interps" begin
    FT = Float64
    n, a, b, β = 10, FT(0.0), FT(1.0), FT(1.01)
    grid = Cube(a, b, n; warpfun=Roberts_both, args=(β,))

    edge = (CellEdge(grid, 1), CellEdge(grid, 2), CellEdge(grid, 3))
    face = (CellFace(grid, 1), CellFace(grid, 2), CellFace(grid, 3))
    cent = CellCenter(grid)
    corn = CellCorner(grid)
    @inbounds for local_grid in GridField(grid, cent)
        @unpack ijk,x,y,z = local_grid
        cent[ijk] = sin(3*π*x)* sin(3*π*y)* sin(3*π*z)
    end

    # Just making sure these are callable:
    interp!(face[1], edge[1], grid)
    interp!(edge[1], face[1], grid)

    interp!(cent, edge[1], grid)
    interp!(edge[1], cent, grid)

    interp!(corn, face[1], grid)
    interp!(face[1], corn, grid)

    interp!(cent, corn, grid)
    interp!(corn, cent, grid)
end


import Oceananigans.TimeSteppers: calculate_pressure_correction!, pressure_correct_velocities!

calculate_pressure_correction!(::HydrostaticFreeSurfaceModel, Δt) = nothing

#####
##### Barotropic pressure correction for models with a free surface
#####

const HFSM = HydrostaticFreeSurfaceModel
const ExplicitFreeSurfaceHFSM = HFSM{<:Any, <:Any, <:Any, <:ExplicitFreeSurface}
const ImplicitFreeSurfaceHFSM = HFSM{<:Any, <:Any, <:Any, <:ImplicitFreeSurface}

pressure_correct_velocities!(model::ExplicitFreeSurfaceHFSM, Δt; kwargs...) = nothing

#####
##### Barotropic pressure correction for models with a free surface
#####

function pressure_correct_velocities!(model::ImplicitFreeSurfaceHFSM, Δt;
                                      dependencies = device_event(model.architecture))

    event = launch!(model.architecture, model.grid, :xyz,
                    _barotropic_pressure_correction,
                    model.velocities,
                    model.grid,
                    Δt,
                    model.free_surface.gravitational_acceleration,
                    model.free_surface.η,
                    dependencies = dependencies)

    wait(device(model.architecture), event)

    return nothing
end

@kernel function _barotropic_pressure_correction(U, grid, Δt, g, η)
    i, j, k = @index(Global, NTuple)

    @inbounds begin
        U.u[i, j, k] -= g * Δt * ∂xᶠᶜᵃ(i, j, k, grid, η)
        U.v[i, j, k] -= g * Δt * ∂yᶜᶠᵃ(i, j, k, grid, η)
    end
end


function linear_compartment_model(graph, sinks)
    """
    Input:
        - graph - graph of the model network represented via adjacency lists
        - sinks - the indices of nodes having a sink
    Output: the corresponding ODE object where each parameter a_ij is replaced
            with a_ij + b_ij * x_0, where x_0 is a constant input encoded as a constant output
    """
    n = length(graph)
    x_vars_names = ["x$i" for i in 0:n]
    edges_vars_names = Array{String, 1}()
    for i in 1:n
        for j in graph[i]
            push!(edges_vars_names, "a_$(i)_$(j)")
            push!(edges_vars_names, "b_$(i)_$(j)")
        end
    end
    for s in sinks
        push!(edges_vars_names, "a_$(s)_0")
        push!(edges_vars_names, "b_$(s)_0")
    end
    R, vars = PolynomialRing(QQ, vcat(x_vars_names, edges_vars_names))
    x_vars = vars[2:(n + 1)]
    x0 = vars[1]
    equations = Dict{fmpq_mpoly, Union{fmpq_mpoly, Generic.Frac{fmpq_mpoly}}}(x => R(0) for x in x_vars)
    equations[x0] = R(0)
    for i in 1:n
        for j in graph[i]
            rate = str_to_var("a_$(i)_$(j)", R) + str_to_var("b_$(i)_$(j)", R) * x0

            if i != j
                equations[x_vars[j]] += x_vars[i] * rate
                equations[x_vars[i]] -= x_vars[i] * rate
            else
                equations[x_vars[i]] -= x_vars[i] * rate
            end
        end
        if i in sinks
            rate = str_to_var("a_$(i)_0", R) + str_to_var("b_$(i)_0", R) * x0
            equations[x_vars[i]] += -x_vars[i] * rate
        end
    end
    return ODE{fmpq_mpoly}(equations)
end

#------------------------------------------------------------------------------

function bicycle(n)
    """
    Generates a bidirected cycle of length n
    """
    graph = []
    for i in 1:n
        prev = (i == 1) ? n : (i - 1)
        next = (i == n) ? 1 : i + 1
        push!(graph, [prev, next])
    end
    return graph
end

#------------------------------------------------------------------------------

function cycle(n)
    """
    Single directed cycle
    """
    graph = [[(i == n) ? 1 : (i + 1)] for i in 1:n]
    return graph
end

#------------------------------------------------------------------------------

function catenary(n)
    """
    Bidirected chain from 1 to n
    """
    graph = [[] for i in 1:n]
    for i in 1:n
        if i != 1
            push!(graph[i], i - 1)
        end
        if i != n
            push!(graph[i], i + 1)
        end
    end
    return graph
end

#------------------------------------------------------------------------------

function mammilary(n)
    """
    Bidirected 'star' with center at 1 and rays to 2, ..., n
    """
    graph = []
    push!(graph, [i for i in 2:n])
    for i in 2:n
        push!(graph, [1])
    end
    return graph
end

#------------------------------------------------------------------------------


using EventSimulation
using Distributions
using Printf
using Random

# Objectives of the example:
# * show how SimQueue object can be used
# * handling interruptions of actions

mutable struct Customer
    name::String
    arrival_time::Float64
    renege_a::Action{Float64} # this is initially left #undef

    function Customer(name, arrival_time)
        x = new()
        x.name = name
        x.arrival_time = arrival_time
        return x
    end
end

mutable struct BankState <: AbstractState
    all_cust::Int # how many customers we want to process
    n::Int # how many customers were processed
    reneged_count::Int # number of customers that reneged
    arrival::Exponential{Float64}
    patience::Uniform{Float64}
    cashing::Exponential{Float64}
    counter::SimQueue{Customer}
    report::Bool
end

function source(s::Scheduler)
    s.state.n += 1
    c = Customer(string("Customer", s.state.n), s.now)
    s.state.report && @printf("%6.4f %s: Here I am\n", s.now, c.name)
    provide!(s, s.state.counter, c)
    # we remember the action to remove it from event_queue if needed
    c.renege_a = register!(s, x -> renege(x, c), rand(s.state.patience))
    if s.state.n < s.state.all_cust
        register!(s, source, rand(s.state.arrival))
    end
end

function renege(s, c)
    s.state.report && @printf("%6.4f %s: RENEGED after %6.4f\n",
                              s.now, c.name, s.now-c.arrival_time)
    s.state.reneged_count += 1
    # customer reneged and leaves the queue
    withdraw!(s.state.counter, c) || error("should be succesffull")
end

function serve(s, c)
    s.state.report && @printf("%6.4f %s: waited %6.4f\n",
                              s.now, c.name, s.now-c.arrival_time)
    # customer got to service so we interrupt renege action
    interrupt!(s, c.renege_a) || error("should be succesfull")
    register!(s, x -> finish(x, c), rand(s.state.cashing))
end

function finish(s, c)
    s.state.report && @printf("%6.4f %s: Finished\n", s.now, c.name)
    request!(s, s.state.counter, serve)
end

function run(ccount, report)
    bs = BankState(ccount, 0, 0, Exponential(10.0),
                   Uniform(1.0, 3.0), Exponential(12.0),
                   SimQueue{Customer}(), report)
    report && println("Bank renege")
    s = Scheduler(bs)
    request!(s, bs.counter, serve)
    source(s)
    go!(s, Inf)
    println("Renege probability: ", bs.reneged_count / bs.all_cust)
end

Random.seed!(1)
@time run(10, true)
@time run(1_000_000, false)


@testset "distribution_free.jl" begin
    @testset "bounds" begin
        # default bounds for base kernels
        CalibrationTests.uniformbound(ExponentialKernel()) == 1
        CalibrationTests.uniformbound(SqExponentialKernel()) == 1
        CalibrationTests.uniformbound(WhiteKernel()) == 1

        # default bounds for kernels with input transformations
        CalibrationTests.uniformbound(SqExponentialKernel() ∘ ScaleTransform(rand())) == 1
        CalibrationTests.uniformbound(ExponentialKernel() ∘ ScaleTransform(rand(10))) == 1

        # default bounds for scaled kernels
        CalibrationTests.uniformbound(42 * ExponentialKernel()) == 42

        # default bounds for tensor product kernels
        kernel = (3.2 * SqExponentialKernel()) ⊗ (2.7 * WhiteKernel())
        CalibrationTests.uniformbound(kernel) == 3.2 * 2.7

        # default bounds for kernel terms
        CalibrationTests.uniformbound(BlockUnbiasedSKCE(kernel)) == 2 * 3.2 * 2.7
    end

    @testset "estimator and estimates" begin
        kernel = (ExponentialKernel() ∘ ScaleTransform(0.1)) ⊗ WhiteKernel()

        for skce in (BiasedSKCE(kernel), UnbiasedSKCE(kernel), BlockUnbiasedSKCE(kernel))
            for nclasses in (2, 10, 100), nsamples in (10, 50, 100)
                # sample predictions and targets
                dist = Dirichlet(nclasses, 1)
                predictions = [rand(dist) for _ in 1:nsamples]
                targets_consistent = [
                    rand(Categorical(prediction)) for prediction in predictions
                ]
                targets_onlyone = ones(Int, length(predictions))

                # for both sets of targets
                for targets in (targets_consistent, targets_onlyone)
                    test = DistributionFreeSKCETest(skce, predictions, targets)

                    @test test.estimator == skce
                    @test test.n == nsamples
                    @test test.estimate ≈ skce(predictions, targets)
                    @test test.bound == CalibrationTests.uniformbound(skce)
                end
            end
        end
    end

    @testset "consistency" begin
        kernel = (ExponentialKernel() ∘ ScaleTransform(0.1)) ⊗ WhiteKernel()
        αs = 0.05:0.1:0.95
        nsamples = 100

        pvalues_consistent = Vector{Float64}(undef, 100)

        for skce in (BiasedSKCE(kernel), UnbiasedSKCE(kernel), BlockUnbiasedSKCE(kernel))
            for nclasses in (2, 10)
                dist = Dirichlet(nclasses, 1)
                predictions = [Vector{Float64}(undef, nclasses) for _ in 1:nsamples]
                targets_consistent = Vector{Int}(undef, nsamples)

                for i in eachindex(pvalues_consistent)
                    # sample predictions and targets
                    for j in 1:nsamples
                        rand!(dist, predictions[j])
                        targets_consistent[j] = rand(Categorical(predictions[j]))
                    end

                    # define test
                    test_consistent = DistributionFreeSKCETest(
                        skce, predictions, targets_consistent
                    )

                    # estimate pvalue
                    pvalues_consistent[i] = pvalue(test_consistent)
                end

                # compute empirical test errors
                errors = ecdf(pvalues_consistent).(αs)
                @test all(((α, p),) -> p < α, zip(αs, errors))
            end
        end
    end
end


@testset "Checking converging data (Real)" begin
    # block of maximally correlated values:
    N_corr = 16

    # number of blocks
    N_blocks = 131_072 # 2^17

    uncorrelated = rand(N_blocks) # rand(rng, N_blocks)
    av = mean(uncorrelated)
    _var = var(uncorrelated)
    stderr = sqrt(_var/N_blocks)

    correlated = [x for x in uncorrelated for _ in 1:N_corr]
    @test BinningAnalysis.unbinned_tau(correlated) ≈ 7.5 rtol = 0.05

    @testset "FullBinner" begin
        B = FullBinner(correlated)

        # Strongly/Maximally correlated

        # Note:
        # With xs = [sample, sample] we have
        # 1/N sum(xs) = 1/N (sum(xs[1:N/2] + sum(xs[N/2+1:N])) = 2/N sum(xs[1:N/2])
        # thus both the mean and variance should not vary when repeating data
        # this doesn't seem to be exactly true for Statistics.var(vcat(sample, sample))
        # this isn't exactly true for BinningAnalysis due to truncation at odd binning levels

        # Expectations:
        # - mean matches
        # - var matches up to binsize = N_corr for power binsizes
        # - varN = var / N
        # - std_error = sqrt(var / N) grows with sqrt(N)
        # - autocorrelation time should be 0.5 * (binsize-1) (i.e. maximal)

        for binsize in 2 .^ (0:3)
            target = 0.5binsize - 0.5

            @test mean(B, binsize) ≈ av
            @test var(B, binsize) ≈ _var rtol=0.01
            @test varN(B, binsize) * N_blocks ≈ _var * binsize / N_corr rtol=0.01
            @test std_error(B, binsize) * sqrt(N_corr / binsize) ≈ stderr rtol=0.01
            @test tau(B, binsize) ≈ target rtol=0.01
        end

        # Not correlated

        # Note:
        # If we expand the variance of pair binned data
        # \frac{1}{N} \sum_{i=1}^N (\frac{x_{2i-1} + x_{2i}}{2} - \mu)^2
        # and pull out jalf the unbinned variance
        # \frac{1}{2} \frac{1}{2N} \sum_{i=1}^{2N} (x_i - \mu)^2
        # we are left with an additional term
        # \frac{1}{2} \frac{1}{N} \sum_{i=1}^N (x_{2i-1} x_{2i} - 2\mu (x_{2i-1} x_{2i}) + \mu^2)
        # if we assume no correlation we can simplify the first term as
        # \langle x_{2i-1} x_{2i} \rangle = \langle x_{2i-1} \rangle \langle x_{2i} \rangle = \mu^2
        # and the term drops, leaving us with
        # var(x_{pair binned}) ≈ \frac{1}{2} var(x)

        # At and after binning out correlations
        # - mean matches
        # - var ≈ unbinned var / 2
        # - varN = var / N with N = N_blocks * N_corr / binsize
        # - std_error = sqrt(var / N) should be constant as the change in N is
        #                             cancelled by the change in var
        # - tau should be roughly constant

        for binsize in 2 .^ (4:10)
            f = N_corr / binsize

            @test mean(B, binsize) ≈ av
            @test var(B, binsize) ≈ f * _var rtol=0.075
            @test varN(B, binsize) * N_blocks ≈ _var rtol=0.075
            @test std_error(B, binsize) ≈ stderr rtol=0.075
            @test tau(B, binsize) ≈ 0.5(N_corr - 1) rtol = 0.075
        end
    end

    @testset "LogBinner" begin
        B = LogBinner(correlated)

        # Strongly/Maximally correlated
        for lvl in 1:4
            binsize = 2^(lvl-1)
            target = 0.5 * (binsize - 1)

            @test mean(B, lvl) ≈ av
            @test var(B, lvl) ≈ _var rtol=0.01
            @test varN(B, lvl) * N_blocks ≈ _var * binsize / N_corr rtol=0.01
            @test std_error(B, lvl) * sqrt(N_corr / binsize) ≈ stderr rtol=0.01
            @test tau(B, lvl) ≈ target rtol=0.01
        end

        # Not correlated
        for lvl in 5:11
            binsize = 2.0 ^ (lvl-1)
            f = N_corr / binsize

            @test mean(B, lvl) ≈ av
            @test var(B, lvl) ≈ f * _var rtol=0.075
            @test varN(B, lvl) * N_blocks ≈ _var rtol=0.075
            @test std_error(B, lvl) ≈ stderr rtol=0.075
            @test tau(B, lvl) ≈ 0.5(N_corr - 1) rtol = 0.075
        end
    end

    @testset "ErrorPropagator" begin
        B = ErrorPropagator(correlated)

        # Strongly/Maximally correlated
        for lvl in 1:4
            binsize = 2^(lvl-1)
            target = 0.5 * (binsize - 1)

            @test mean(B, 1, lvl) ≈ av
            @test var(B, 1, lvl) ≈ _var rtol=0.01
            @test varN(B, 1, lvl) * N_blocks ≈ _var * binsize / N_corr rtol=0.01
            @test std_error(B, 1, lvl) * sqrt(N_corr / binsize) ≈ stderr rtol=0.01
            @test tau(B, 1, lvl) ≈ target rtol=0.01
        end

        # Not correlated
        for lvl in 5:11
            binsize = 2.0 ^ (lvl-1)
            f = N_corr / binsize

            @test mean(B, 1, lvl) ≈ av
            @test var(B, 1, lvl) ≈ f * _var rtol=0.075
            @test varN(B, 1, lvl) * N_blocks ≈ _var rtol=0.075
            @test std_error(B, 1, lvl) ≈ stderr rtol=0.075
            @test tau(B, 1, lvl) ≈ 0.5(N_corr - 1) rtol = 0.075
        end
    end
end

@testset "Ising Energies" begin
    # This data comes from an AFM Ising model on a triangular lattice, L = 8, T = 2.0
    Es = open("test.data", "r") do f
        Float64[read(f, Float64) for _ in 1:2^14]
    end

    av = mean(Es)
    τ = BinningAnalysis.unbinned_tau(Es)
    stderr = sqrt(var(Es) * (1 + 2τ) / length(Es))

    @testset "FullBinner" begin
        B = FullBinner(Es)
        N = BinningAnalysis._reliable_level(B)

        @test mean(B) ≈ av
        @test tau(B) ≈ τ rtol=0.5
        @test std_error(B) ≈ stderr rtol=0.3
    end

    @testset "LogBinner" begin
        B = LogBinner(Es)

        @test mean(B) ≈ av
        @test tau(B) ≈ τ rtol=0.5
        @test std_error(B) ≈ stderr rtol=0.3
    end

    @testset "ErrorPropagator" begin
        B = ErrorPropagator(Es)

        @test mean(B, 1) ≈ av
        @test tau(B, 1) ≈ τ rtol=0.5
        @test std_error(B, 1) ≈ stderr rtol=0.3
    end
end


DP = SymbolicUtils.DynamicPolynomials
Bijections = SymbolicUtils.Bijections

# extracting underlying polynomial and coefficient type from Polyforms
underlyingpoly(x::Number)    = x
underlyingpoly(pf::PolyForm) = pf.p
coefftype(x::Number)    = typeof(x)
coefftype(pf::PolyForm) = DP.coefficienttype(underlyingpoly(pf))

#=
Converts an array of symbolic polynomials
into an array of DynamicPolynomials.Polynomials
=#
function symbol_to_poly(sympolys::AbstractArray)
    @assert !isempty(sympolys) "Empty input."

    # standardize input
    stdsympolys = map(unwrap, sympolys)
    sort!(stdsympolys, lt=(<ₑ))

    pvar2sym  = Bijections.Bijection{Any,Any}()
    sym2term  = Dict{Sym,Any}()
    polyforms = map(f -> PolyForm(f, pvar2sym, sym2term), stdsympolys)

    # Discover common coefficient type
    commontype = mapreduce(coefftype, promote_type, polyforms, init=Int)
    @assert commontype <: Union{Integer,Rational} "Only integer and rational coefficients are supported as input."

    # Convert all to DP.Polynomial, so that coefficients are of same type,
    # and constants are treated as polynomials
    # We also need this because Groebner does not support abstract types as input
    polynoms = Vector{DP.Polynomial{true,commontype}}(undef, length(sympolys))
    for (i, pf) in enumerate(polyforms)
        polynoms[i] = underlyingpoly(pf)
    end

    polynoms, pvar2sym, sym2term
end

#=
Converts an array of AbstractPolynomialLike`s into an array of
symbolic expressions mapping variables w.r.t pvar2sym
=#
function poly_to_symbol(polys, pvar2sym, sym2term, ::Type{T}) where {T}
    map(f -> PolyForm{T}(f, pvar2sym, sym2term), polys)
end

"""

groebner_basis(polynomials; reduced=true)

Computes a Groebner basis of the ideal generated by the given `polynomials`.

If `reduced` is set, the basis is reduced and *unique*.

"""
function groebner_basis(polynomials; reduced=true)
    polynoms, pvar2sym, sym2term = symbol_to_poly(polynomials)

    basis = groebner(polynoms, reduced=reduced)

    # polynomials is nonemtpy
    T = symtype(first(polynomials))
    poly_to_symbol(basis, pvar2sym, sym2term, T)
end


# Quadratic form of 2 variables, used for representing the transition costs ℓ
struct QuadraticForm{T}
    P::SMatrix{2,2,T,4}
    q::SVector{2,T}
    r::T
end
#TODO Remove when StaticArrays properly handle scalars
function QuadraticForm(P::SMatrix{2,2,T,4}, q::AbstractVecOrMat{T}, r::T) where T
    return QuadraticForm(P, SVector{2,T}(q), r)
end

(qf::QuadraticForm)(y::Number, yp::Number) = [y,yp]'*qf.P*[y,yp] + qf.q'*[y, yp] + qf.r
(qf::QuadraticForm)(x::AbstractVector) = x'*qf.P*x + qf.q'*x + qf.r

+(qf1::QuadraticForm, qf2::QuadraticForm) = QuadraticForm(qf1.P+qf2.P, qf1.q+qf2.q, qf1.r+qf2.r)
==(qf1::QuadraticForm, qf2::QuadraticForm) = (qf1.P==qf2.P) && (qf1.q==qf2.q) && (qf1.r==qf2.r)


using CartesianGeneticProgramming
using Cambrian
import Cambrian.mutate
using StatsBase
using Base.Iterators: repeated
include("../graphing/graph_utils.jl")
"""
A simple example  for calculating parity
"""

 #bitArray parity
 parity(x) = isodd(sum(x))

 function create_testcases(sz)
    train_cases = []
    for i in 0:(2^sz-1)
      bin = digits(i, base=2, pad=sz)
      push!(train_cases, (bin, parity(bin)))
      @show (bin, parity(bin))
    end
    return train_cases
 end

 NUMCLASSES = 10
 WIDTH = 8 #60

 train = create_testcases(WIDTH)
 trainX = [ x[1] for x in train ]
 trainX = hcat(trainX...)
 trainY = [ x[2] for x in train ]

X, Y = trainX, trainY

function evaluate(ind::CGPInd, X::AbstractArray, Y::AbstractArray)
    accuracy = 0.0
    for i in 1:size(X, 2)
        out = process(ind, float(X[:, i]))
        if out[1] == Int(Y[i])
            accuracy += 1
        end
    end
    [accuracy / size(X, 1)]
end

function error_count(e, X::AbstractArray, Y::AbstractArray)
   failed = 0
   for i in 1:size(X,2)
      y_hat = process(e.population[5], float(X[:,i]))
      if y_hat[1] != Int(Y[i])
           failed += 1
      end
   end
   failed
end

cfg = get_config("cfg/parity.yaml")
fit(i::CGPInd) = evaluate(i, X, Y)
mutate(i::CGPInd) = goldman_mutate(cfg, i)
e = CGPEvolution(cfg, fit)
println("run!")
run!(e)

@show e.population[end].nodes
errors = error_count(e,X,Y )
println("$errors errors out of $(size(X,2)) entries")

#generate dot file
dot_array = walk_nodes(e.population[5])
dot_file  = "parity_graph.dot"
open(dot_file, "w") do f
   for i in dot_array
      println(f, i)
   end
end
# to view from commandline:
# $ dot -Tpng parity_graph.dot -o parity_graph.png
# $ display parity_graph.png


Qx,x = FlintQQ["x"]
K,a = NumberField(x^6+108)
OK = ring_of_integers(K)
lp = prime_decomposition(OK, 5)
P = lp[1][1]
L = Localization(OK, P)

@testset "OrderLocalization" begin

  @testset "Constructor" begin

    @test parent_type(OrdLocElem{nf_elem}) == OrdLoc{nf_elem}

    lp = prime_decomposition(OK, 3)
    P = lp[1][1]
    L = Localization(OK, P)

    @test elem_type(L) == OrdLocElem{nf_elem}
    @test nf(L) == K
    @test nf(L()) == K

end

#=
  @testset "Canonicalisation" begin
  for i in 1:300
    a = rand(L)
    c = canonical_unit(a)
    @test is_unit(c)
  end
end

    @testset "Valuation" begin

      @test valuation(L(5)) == 1
      @test valuation(L(5//5)) == 0
      @test valuation(L(5//12)) == 1
      @test valuation(L(500)) == 3

      for i in 1:300
        a = rand(L)
        while a == L()
          a = rand(L)
        end
        @test valuation(a) >= 0
      end

    end
    =#

    @testset "Parent object call overloading" begin

      @test parent(L(3)) == L
      @test L(6//3) == L(2)
      @test L(K(5)) == L(5)
      @test L(OK(5)) == L(5)
      @test_throws ErrorException L(24//90)
    end

    @testset "GCDX" begin
      for i in 1:550
        a = rand(L)
        b = rand(L)
        (g,u,v) = gcdx(a,b)
        @test g == u*a + v*b
      end
    end

    @testset "GCD" begin
      for i in 1:550
        a = rand(L)
        b = rand(L)
        g = gcd(a,b)
        @test divides(a,g)[1] && divides(b,g)[1]
      end
    end

    @testset "Exact division" begin

      @test divides(L(15),L(3)) == (true, L(5))
      @test divides(L(15),L(30)) == (true, L(1//2))
      @test divides(L(15),L(25)) == (false, L(0))
      @test_throws ErrorException divexact(L(8),L(10))
      for i in 1:300
        a = rand(L)
        b = rand(L)
        d = divides(a,b)
        if d[1]
          @test a == d[2] * b
        end
      end
    end

    @testset "Inversion" begin

      @test inv(L(8)) == L(1//8)
      @test inv(L(23)) == L(1//23)
      @test_throws ErrorException inv(L(40))
      @test_throws DivideError inv(L())
    end

    @testset "Binary operators" begin

      @test L(19) + L(1) == L(20)
      @test L(19) - L(1) == L(18)
      @test L(3) * L(4//2) == L(6)

      @test L(18//2) + L(2//1) == L(11)
      @test L(18//3) - L(1) == L(5)
      @test L(32) * L(4) == L(128)
    end

    @testset "Basic manipulation" begin

      @test iszero(L()) == true
      @test iszero(L(0)) == true
      @test isone(L(2//3)) == false
      @test is_unit(L(24)) == true

      @test iszero(L(4//3)) == false
      @test isone(L(13//13)) == true
      @test is_unit(L(50//2)) == false
    end
end


include("../src/PrivateMultiplicativeWeights.jl")

using PrivateMultiplicativeWeights
using Distributions
using Plots
#plotlyjs()
#plotly()
#pyplot()
gr()

d_continuous = Truncated(MixtureModel(Normal, [(.275, .075), (.65, .1)], [0.65, 0.35]), 0, 1)
xs = linspace(d_continuous.lower, d_continuous.upper, 300)
ys = pdf.(d_continuous, xs)
plot(xs, ys);

num_samples = 1000
domain_size = 100
samples = rand(d_continuous, num_samples)

function histogram_from_samples(samples, num_bins)
    @assert 0 <= minimum(samples)
    @assert maximum(samples) <= 1
    h = zeros(num_bins)
    new_samples = samples*num_bins
    for x in round.(Int64, new_samples)
        h[x] += 1
    end
    collect(1:num_bins)/num_bins, h
end

xs, ys = histogram_from_samples(samples, domain_size)
bar(xs, ys, width=1/domain_size);
#gui()

mw = mwem(RangeQueries(domain_size), Histogram(ys/sum(ys), num_samples))
bar(xs, mw.synthetic.weights*num_samples, color="red", alpha=0.3, width=1/domain_size);
#gui()

mw = mwem(RangeQueries(domain_size), Histogram(ys/sum(ys), num_samples), MWParameters(epsilon=100, iterations=80))
bar(xs, mw.synthetic.weights*num_samples, color="red", alpha=0.3, width=1/domain_size);
#gui()

a = bar(xs, ys)#, alpha=1/domain_size))
bar!(xs, mw.synthetic.weights*num_samples, color="red", alpha=0.3, width=1/domain_size)

display(a)


# `backward` functions return gates storing derivatives

backward(g::XGate, ::AbstractMatrix) = g
backward(g::YGate, ::AbstractMatrix) = g
backward(g::ZGate, ::AbstractMatrix) = g

backward(g::HadamardGate, ::AbstractMatrix) = g

backward(g::SGate, ::AbstractMatrix) = g
backward(g::TGate, ::AbstractMatrix) = g

backward(g::SdagGate, ::AbstractMatrix) = g
backward(g::TdagGate, ::AbstractMatrix) = g

backward(g::MatrixGate, ::AbstractMatrix) = g

function backward(g::RxGate, Δ::AbstractMatrix)
    c = cos(g.θ[] / 2)
    s = sin(g.θ[] / 2)
    # using conjugated derivative matrix; factor 2 cancels 1/2 from θ/2
    RxGate(sum(real([-s im * c; im * c -s] .* Δ)))
end

function backward(g::RyGate, Δ::AbstractMatrix)
    c = cos(g.θ[] / 2)
    s = sin(g.θ[] / 2)
    # using conjugated derivative matrix; factor 2 cancels 1/2 from θ/2
    RyGate(sum(real([-s -c; c -s] .* Δ)))
end

function backward(g::RzGate, Δ::AbstractMatrix)
    # using conjugated derivative matrix; factor 2 cancels 1/2 from θ/2
    eθ = exp(im * g.θ[] / 2)
    RzGate(real(im * eθ * Δ[1, 1] - im * conj(eθ) * Δ[2, 2]))
end

function backward(g::RotationGate, Δ::AbstractMatrix)
    θ = norm(g.nθ)
    if θ == 0
        σ = Matrix{ComplexF64}[matrix(X), matrix(Y), matrix(Z)]
        RotationGate([real(sum(conj(-im * σ[i]) .* Δ)) for i in 1:3])
    else
        n = g.nθ / θ
        c = cos(θ / 2)
        s = sin(θ / 2)
        dRθ = -s * I - im * c * pauli_vector(n...)
        dn = (I - reshape(kron(n, n), 3, 3)) / θ
        # using conjugated derivative matrix
        RotationGate([2 * real(sum(conj(0.5 * n[i] * dRθ - im * s * pauli_vector(dn[:, i]...)) .* Δ)) for i in 1:3])
    end
end

function backward(g::PhaseShiftGate, Δ::AbstractMatrix)
    # using conjugated derivative matrix
    PhaseShiftGate(2 * real(-im * Base.exp(-im * g.ϕ[1]) * Δ[2, 2]))
end

backward(g::SwapGate, ::AbstractMatrix) = g

function backward(g::EntanglementXXGate, Δ::AbstractMatrix)
    c = cos(g.θ[] / 2)
    s = sin(g.θ[] / 2)
    # using conjugated derivative matrix; factor 2 cancels 1/2 from θ/2
    EntanglementXXGate(sum(real([-s 0 0 im * c; 0 -s im * c 0; 0 im * c -s 0; im * c 0 0 -s] .* Δ)))
end

function backward(g::EntanglementYYGate, Δ::AbstractMatrix)
    c = cos(g.θ[] / 2)
    s = sin(g.θ[] / 2)
    # using conjugated derivative matrix; factor 2 cancels 1/2 from θ/2
    EntanglementYYGate(sum(real([-s 0 0 -im * c; 0 -s im * c 0; 0 im * c -s 0; -im * c 0 0 -s] .* Δ)))
end

function backward(g::EntanglementZZGate, Δ::AbstractMatrix)
    eθ = exp(im * g.θ[] / 2)
    # using conjugated derivative matrix; factor 2 cancels 1/2 from θ/2
    EntanglementZZGate(sum(real(im * eθ * Δ[1, 1] - im * conj(eθ) * Δ[2, 2] - im * conj(eθ) * Δ[3, 3] + im * eθ * Δ[4, 4])))
end

function backward(g::ControlledGate{G}, Δ::AbstractMatrix) where {G}
    # Note: target qubits correspond to fastest varying indices
    ControlledGate(backward(g.U, Δ[end-2^target_wires(g)+1:end, end-2^target_wires(g)+1:end]), control_wires(g))
end

function backward(cg::CircuitGate{M,G}, ψ::AbstractVector, Δ::AbstractVector, N::Int) where {M,G}
    ρ = rdm(N, cg.iwire, Δ, ψ)
    CircuitGate{M,G}(cg.iwire, backward(cg.gate, ρ))
end

function backward(m::Moment, ψ::AbstractVector, Δ::AbstractVector, N::Int)
    gates = AbstractCircuitGate[]
    for cg in reverse(m.gates)
        Udag = Base.adjoint(cg)
        ψ = apply(ψ, Udag)
        # backward step of quantum state
        pushfirst!(gates, backward(cg, ψ, Δ, N))
        Δ = apply(Δ, Udag)
    end
    return Moment(gates), ψ, Δ
end

function backward(moments::Vector{Moment}, ψ::AbstractVector, Δ::AbstractVector, N::Int)
    dmoments = Moment[]
    for moment in reverse(moments)
        # backward step of quantum state
        (m, ψ, Δ) = backward(moment, ψ, Δ, N)
        pushfirst!(dmoments, m)
    end
    return dmoments, Δ
end

"""
    gradients(c::Circuit{N}, ψ::AbstractVector, Δ::AbstractVector{<:Real}) where {N}

Perform a backward pass to compute gradients of a (fictitious) cost function with
respect to the circuit parameters of `c` and input wavefunction `ψ`. `Δ` contains
the cost function gradients with respect to the circuit outputs (measurement operator averages).
The gradients with respect to the circuit parameters are returned in a duplicate circuit;
the overall return value is the tuple (dc::Circuit{N}, dψ::AbstractVector).
"""
function gradients(c::Circuit{N}, ψ::AbstractVector, Δ::AbstractVector{<:Real}) where {N}
    # length of circuit output vector must match gradient vector
    @assert length(Δ) == length(c.meas)
    # forward pass through unitary gates
    ψ = apply(ψ, c.moments)
    # gradient (conjugated Wirtinger derivatives) of cost function with respect to ψ
    ψbar = sum([Δ[i] * (sparse_matrix(c.meas[i]) * ψ) for i in 1:length(Δ)])
    # backward pass through unitary gates
    dcgc, ψbar = backward(c.moments, ψ, ψbar, N)
    # TODO: efficiently represent Kronecker product without explicitly storing matrix entries
    dmeas = MeasurementOperator.([Δ[i] * reshape(kron(conj(ψ), ψ), length(ψ), length(ψ)) for i in 1:length(Δ)], (Tuple(1:N),))
    return Circuit{N}(dcgc, dmeas), ψbar
end


@testset "Finitely generated abelian groups" begin
  include("GrpAb/GrpAbFinGen.jl")
  include("GrpAb/Elem.jl")
  include("GrpAb/SubgroupEnum.jl")
  include("GrpAb/Map.jl")
  include("GrpAb/Lattice.jl")
end


abstract type Weight end

#=
    Sum over the bins that are within `threshold`% of the minimum.
    A large value means the well is shallow.
=#
struct ThresholdSum <: Weight
    threshold::Float64
    function ThresholdSum(threshold::Real)
        0 < threshold && return new(1.0 + threshold)
        throw(ArgumentError("Threshold must be less than 1"))
    end
end

#=
  Sum over the bins that are within the rectangle inscribing the
  the bins which are `threshold`% of the minimum.
  A large value means the well is shallow. Punished asymmetric wells
  more than ThresholdSum
=#
struct ThresholdArea <: Weight
    threshold::Float64
    function ThresholdArea(threshold::Real)
        0 < threshold && return new(1.0 + threshold)
        throw(ArgumentError("Threshold must be less than 1"))
    end
end

struct ThresholdSumShift <: Weight
    threshold::Float64
    shift::Float64
    function ThresholdSumShift(threshold::Real)
        0 < threshold && return new(1.0 + threshold, 1.0)
        throw(ArgumentError("Threshold must be less than 1"))
    end
end

#=
   Sum over the bins within the `radius` distance of the
   minimum, minus the minimum. A high score means the well is sharp.
=#
struct WellSum <: Weight
    radius::Int
end

function (w::ThresholdSum)(x)
    min = minimum(x)
    sum(x .≤ w.threshold*min)
end

function (w::ThresholdArea)(x)
    min = minimum(x)
    thresholded = x .≤ w.threshold*min
    maxx = maxy = -1
    miny = minx = 1e5
    for x in 1:size(thresholded, 1), y in 1:size(thresholded, 2)
        if thresholded[x, y]
            if x < minx
                minx = x
            end
            if x > maxx
                maxx = x
            end
            if y < miny
                miny = y
            end
            if y > maxy
                maxy = y
            end
        end
    end
    sidex = maxx - minx
    sidey = maxy - miny
    (sidex > 0 ? sidex : 1) * (sidey > 0 ? sidey : 1)
end

function (w::ThresholdSumShift)(x)
    min = minimum(x) + w.shift
    sum(x .≤ w.threshold*min)
end

function (w::WellSum)(x)
    truth = argmin(x)
    r = w.radius
    xstart, ystart = truth.I .- r
    xstop, ystop = truth.I .+ r
    xstart = xstart > 0 ? xstart : 1
    ystart = ystart > 0 ? ystart : 1
    xstop = xstop < size(x, 1) ? xstop : size(x, 1)
    ystop = ystop < size(x, 2) ? ystop : size(x, 2)
    sum(@view x[xstart:xstop, ystart:ystop]) - minimum(x)
end


@inline CI(a...) = CartesianIndex(a...)
"""
    δ(i,N::Int)
    δ(i,I::CartesianIndex{N}) where {N}

Return a CartesianIndex of dimension `N` which is one at index `i` and zero elsewhere.
"""
@inline δ(i,N::Int) = CI(ntuple(j -> j==i ? 1 : 0, N))
@inline δ(i,I::CartesianIndex{N}) where {N} = δ(i,N)

"""
    inside(dims)
    inside(a) = inside(size(a))

Return CartesianIndices range excluding the ghost-cells on the boundaries of
a _scalar_ array `a` with `dims=size(a)`.
"""
@inline inside(dims::NTuple{N}) where {N} = CartesianIndices(ntuple(i-> 2:dims[i]-1,N))
@inline inside(a::AbstractArray) = inside(size(a))

"""
    inside_u(dims,j)

Return CartesianIndices range excluding the ghost-cells on the boundaries of
a _vector_ array on face `j` with size `dims`.
"""
function inside_u(dims::NTuple{N},j) where {N}
    CartesianIndices(ntuple( i-> i==j ? (3:dims[i]-1) : (2:dims[i]), N))
end
splitn(n) = Base.front(n),n[end]
size_u(u) = splitn(size(u))

"""
    L₂(a)

L₂ norm of array `a` excluding ghosts.
"""
L₂(a) = sum(@inbounds(abs2(a[I])) for I ∈ inside(a))

"""
    @inside <expr>

Simple macro to automate efficient loops over cells excluding ghosts. For example

    @inside p[I] = sum(I.I)

becomes

    @inbounds @simd for I ∈ inside(p)
        p[I] = sum(I.I)
    end
"""
macro inside(ex)
    a,I = Meta.parse.(split(string(ex.args[1]),union("[",",","]")))
    return quote
        WaterLily.@loop $ex over $I ∈ inside($a)
    end |> esc
end

macro loop(args...)
    ex,_,itr = args
    op,I,R = itr.args
    @assert op ∈ (:(∈),:(in))

    if Threads.nthreads()!=1    #multithread mode
        return quote
            @inbounds Threads.@threads for $I ∈ $R
                $ex
            end
        end |> esc
    else
        return quote            #single thread mode
            @inbounds @simd for $I ∈ $R
                $ex
            end
        end |> esc
    end
end

function median(a,b,c)
    if a>b
        b>=c && return b
        a>c && return c
    else
        b<=c && return b
        a<c && return c
    end
    return a
end

using StaticArrays
"""
    loc(i,I)

Location in space of the cell at CartesianIndex `I` at face `i`.
Using `i=0` returns the cell center s.t. `loc = I`.
"""
@inline loc(i,I) = SVector(I.I .- 0.5 .* δ(i,I).I)

"""
    apply!(f, c)

Apply a vector function `f(i,x)` to the faces of a uniform staggered array `c`.
"""
function apply!(f,c)
    N,n = size_u(c)
    for i ∈ 1:n
        @loop c[I,i] = f(i,loc(i,I)) over I ∈ CartesianIndices(N)
    end
end

"""
    slice(dims,i,j,low=1)

Return `CartesianIndices` range slicing through an array of size `dims` in
dimension `j` at index `i`. `low` optionally sets the lower extent of the range
in the other dimensions.
"""
function slice(dims::NTuple{N},i,j,low=1) where N
    CartesianIndices(ntuple( k-> k==j ? (i:i) : (low:dims[k]), N))
end

"""
    BC!(a,A,f=1)

Apply boundary conditions to the ghost cells of a _vector_ field. A Dirichlet
condition `a[I,i]=f*A[i]` is applied to the vector component _normal_ to the domain
boundary. For example `aₓ(x)=f*Aₓ ∀ x ∈ minmax(X)`. A zero Nuemann condition
is applied to the tangential components.
"""
function BC!(a,A,f=1)
    N,n = size_u(a)
    for j ∈ 1:n, i ∈ 1:n
        if i==j # Normal direction, Dirichlet
            for s ∈ (1,2,N[j])
                @loop a[I,i] = f*A[i] over I ∈ slice(N,s,j)
            end
        else    # Tangential directions, Neumann
            @loop a[I,i] = a[I+δ(j,I),i] over I ∈ slice(N,1,j)
            @loop a[I,i] = a[I-δ(j,I),i] over I ∈ slice(N,N[j],j)
        end
    end
end

"""
    BC!(a)

Apply zero Nuemann boundary conditions to the ghost cells of a _scalar_ field.
"""
function BC!(a)
    N = size(a)
    for j ∈ 1:length(N)
        @loop a[I] = a[I+δ(j,I)] over I ∈ slice(N,1,j)
        @loop a[I] = a[I-δ(j,I)] over I ∈ slice(N,N[j],j)
    end
end


"""
    NLevelBasis(N)

Basis for a system consisting of N states.
"""
struct NLevelBasis{T} <: Basis
    shape::Vector{T}
    N::T
    function NLevelBasis(N::T) where T<:Integer
        if N < 1
            throw(DimensionMismatch())
        end
        new{T}([N], N)
    end
end

==(b1::NLevelBasis, b2::NLevelBasis) = b1.N == b2.N

"""
    transition(b::NLevelBasis, to::Integer, from::Integer)

Transition operator ``|\\mathrm{to}⟩⟨\\mathrm{from}|``.
"""
function transition(b::NLevelBasis, to::Integer, from::Integer)
    if to < 1 || b.N < to
        throw(BoundsError("'to' index has to be between 1 and b.N"))
    end
    if from < 1 || b.N < from
        throw(BoundsError("'from' index has to be between 1 and b.N"))
    end
    op = SparseOperator(b)
    op.data[to, from] = 1.
    op
end

"""
    nlevelstate(b::NLevelBasis, n::Integer)

State where the system is completely in the n-th level.
"""
function nlevelstate(b::NLevelBasis, n::Integer)
    if n < 1 || b.N < n
        throw(BoundsError("n has to be between 1 and b.N"))
    end
    basisstate(b, n)
end