mutable struct FTestResult{N}
    nobs::Int
    ssr::NTuple{N, Float64}
    dof::NTuple{N, Int}
    r2::NTuple{N, Float64}
    fstat::NTuple{N, Float64}
    pval::NTuple{N, Float64}
end

"""A helper function to determine if mod1 is nested in mod2"""
function issubmodel(mod1::LinPredModel, mod2::LinPredModel; atol::Real=0.0)
    mod1.rr.y != mod2.rr.y && return false # Response variables must be equal

    # Test that models are nested
    pred1 = mod1.pp.X
    npreds1 = size(pred1, 2)
    pred2 = mod2.pp.X
    npreds2 = size(pred2, 2)
    # If model 1 has more predictors, it can't possibly be a submodel
    npreds1 > npreds2 && return false
    # Test min norm pred2*B - pred1 ≈ 0
    rtol = Base.rtoldefault(typeof(pred1[1,1]))
    nresp = size(pred2, 1)
    return norm(view(qr(pred2).Q'pred1, npreds2 + 1:nresp, :)) <= max(atol, rtol*norm(pred1))
end

_diffn(t::NTuple{N, T}) where {N, T} = ntuple(i->t[i]-t[i+1], N-1)

_diff(t::NTuple{N, T}) where {N, T} = ntuple(i->t[i+1]-t[i], N-1)

"""
    ftest(mod::LinearModel...; atol::Real=0.0)

For each sequential pair of linear models in `mod...`, perform an F-test to determine if
the one model fits significantly better than the other. Models must have been fitted
on the same data, and be nested either in forward or backward direction.

A table is returned containing consumed degrees of freedom (DOF),
difference in DOF from the preceding model, sum of squared residuals (SSR), difference in
SSR from the preceding model, R², difference in R² from the preceding model, and F-statistic
and p-value for the comparison between the two models.

!!! note
    This function can be used to perform an ANOVA by testing the relative fit of two models
    to the data

Optional keyword argument `atol` controls the numerical tolerance when testing whether
the models are nested.

# Examples

Suppose we want to compare the effects of two or more treatments on some result. Because
this is an ANOVA, our null hypothesis is that `Result ~ 1` fits the data as well as
`Result ~ 1 + Treatment`.

```jldoctest
julia> dat = DataFrame(Result=[1.1, 1.2, 1, 2.2, 1.9, 2, 0.9, 1, 1, 2.2, 2, 2],
                       Treatment=[1, 1, 1, 2, 2, 2, 1, 1, 1, 2, 2, 2],
                       Other=categorical([1, 1, 2, 1, 2, 1, 3, 1, 1, 2, 2, 1]));

julia> nullmodel = lm(@formula(Result ~ 1), dat);

julia> model = lm(@formula(Result ~ 1 + Treatment), dat);

julia> bigmodel = lm(@formula(Result ~ 1 + Treatment + Other), dat);

julia> ftest(nullmodel.model, model.model)
F-test: 2 models fitted on 12 observations
─────────────────────────────────────────────────────────────────
     DOF  ΔDOF     SSR     ΔSSR       R²     ΔR²        F*  p(>F)
─────────────────────────────────────────────────────────────────
[1]    2        3.2292           -0.0000
[2]    3     1  0.1283  -3.1008   0.9603  0.9603  241.6234  <1e-7
─────────────────────────────────────────────────────────────────

julia> ftest(nullmodel.model, model.model, bigmodel.model)
F-test: 3 models fitted on 12 observations
──────────────────────────────────────────────────────────────────
     DOF  ΔDOF     SSR     ΔSSR       R²     ΔR²        F*   p(>F)
──────────────────────────────────────────────────────────────────
[1]    2        3.2292           -0.0000
[2]    3     1  0.1283  -3.1008   0.9603  0.9603  241.6234   <1e-7
[3]    5     2  0.1017  -0.0266   0.9685  0.0082    1.0456  0.3950
──────────────────────────────────────────────────────────────────
```
"""
function ftest(mods::LinearModel...; atol::Real=0.0)
    if length(mods) < 2
        Base.depwarn("Passing less than two models to ftest is deprecated", :ftest)
    end
    if !all(==(nobs(mods[1])), nobs.(mods))
        throw(ArgumentError("F test is only valid for models fitted on the same data, " *
                            "but number of observations differ"))
    end
    forward = length(mods) == 1 || dof(mods[1]) <= dof(mods[2])
    if forward
        for i in 2:length(mods)
            if dof(mods[i-1]) >= dof(mods[i]) || !issubmodel(mods[i-1], mods[i], atol=atol)
                throw(ArgumentError("F test is only valid for nested models"))
            end
        end
    else
        for i in 2:length(mods)
            if dof(mods[i]) >= dof(mods[i-1]) || !issubmodel(mods[i], mods[i-1], atol=atol)
                throw(ArgumentError("F test is only valid for nested models"))
            end
        end
    end

    SSR = deviance.(mods)

    df = dof.(mods)
    Δdf = _diff(df)
    dfr = Int.(dof_residual.(mods))
    MSR1 = _diffn(SSR) ./ Δdf
    MSR2 = (SSR ./ dfr)
    if forward
        MSR2 = MSR2[2:end]
        dfr_big = dfr[2:end]
    else
        MSR2 = MSR2[1:end-1]
        dfr_big = dfr[1:end-1]
    end

    fstat = (NaN, (MSR1 ./ MSR2)...)
    pval = (NaN, ccdf.(FDist.(abs.(Δdf), dfr_big), abs.(fstat[2:end]))...)
    return FTestResult(Int(nobs(mods[1])), SSR, df, r2.(mods), fstat, pval)
end

function show(io::IO, ftr::FTestResult{N}) where N
    Δdof = _diff(ftr.dof)
    Δssr = _diff(ftr.ssr)
    ΔR² = _diff(ftr.r2)

    nc = 9
    nr = N
    outrows = Matrix{String}(undef, nr+1, nc)

    outrows[1, :] = ["", "DOF", "ΔDOF", "SSR", "ΔSSR",
                     "R²", "ΔR²", "F*", "p(>F)"]

    outrows[2, :] = ["[1]", @sprintf("%.0d", ftr.dof[1]), " ",
                     @sprintf("%.4f", ftr.ssr[1]), " ",
                     @sprintf("%.4f", ftr.r2[1]), " ", " ", " "]

    for i in 2:nr
        outrows[i+1, :] = ["[$i]",
                           @sprintf("%.0d", ftr.dof[i]), @sprintf("%.0d", Δdof[i-1]),
                           @sprintf("%.4f", ftr.ssr[i]), @sprintf("%.4f", Δssr[i-1]),
                           @sprintf("%.4f", ftr.r2[i]), @sprintf("%.4f", ΔR²[i-1]),
                           @sprintf("%.4f", ftr.fstat[i]), string(PValue(ftr.pval[i])) ]
    end
    colwidths = length.(outrows)
    max_colwidths = [maximum(view(colwidths, :, i)) for i in 1:nc]
    totwidth = sum(max_colwidths) + 2*8

    println(io, "F-test: $N models fitted on $(ftr.nobs) observations")
    println(io, '─'^totwidth)

    for r in 1:nr+1
        for c in 1:nc
            cur_cell = outrows[r, c]
            cur_cell_len = length(cur_cell)

            padding = " "^(max_colwidths[c]-cur_cell_len)
            if c > 1
                padding = "  "*padding
            end

            print(io, padding)
            print(io, cur_cell)
        end
        print(io, "\n")
        r == 1 && println(io, '─'^totwidth)
    end
    print(io, '─'^totwidth)
end


# # Bump On Tail
#
#md # [`notebook`](@__NBVIEWER_ROOT_URL__notebooks/bump_on_tail.ipynb)
#

import Splittings: advection!, UniformMesh
import Splittings: compute_rho, compute_e
import Splittings: CubicSpline, @Strang

using Plots
using LaTeXStrings

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

function push_t!( f, mesh1, v, dt )

    advection!( f, mesh1, v, 0.5dt, CubicSpline(), 1)

end

function push_v!( f, mesh1, mesh2, nrj, dt )

    rho = compute_rho(mesh2, f)
    e   = compute_e(mesh1, rho)
    advection!( f, mesh2, e, dt, CubicSpline(), 2)
    push!(nrj, 0.5*log(sum(e.*e)*mesh1.step))

end

function vlasov_poisson(mesh1  :: UniformMesh,
                        mesh2  :: UniformMesh,
                        f      :: Array{Float64,2},
                        nstep  :: Int64,
                        dt     :: Float64)

    x = mesh1.points
    v = mesh2.points

    nrj = Float64[]
    for istep in 1:nstep

        @Strang(  push_t!( f, mesh1, v, dt ),
                  push_v!( f, mesh1, mesh2, nrj, dt ))

    end
    nrj

end

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

α = 0.03
kx  = 0.3
x1min, x1max = 0.0, 2π / kx
n1, n2 = 32, 64
x2min, x2max = -9., 9.
mesh1 = UniformMesh(x1min, x1max, n1)
mesh2 = UniformMesh(x2min, x2max, n2)
f = zeros(Float64,(mesh1.length,mesh2.length))
for (i,x) in enumerate(mesh1.points), (j,v) in enumerate(mesh2.points)
    f[i,j]  = (1.0+α*cos(kx*x)) / (10*sqrt(2π)) * (9*exp(-0.5*v^2)+2*exp(-2*(v-4.5)^2))
end

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

nstep = 500
t = range(0.0, stop=50.0, length=nstep)
dt = t[2]
@elapsed nrj = vlasov_poisson( mesh1, mesh2, f, nstep, dt)

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

plot(t, nrj, label=L"\frac{1}{2} \log(∫e²dx)")
savefig("bot-plot.png"); nothing # hide

using Test                       #src
@testset "Bump On Tail" begin    #src
@test length(nrj) > 0            #src
end                              #src

#md # ![](bot-plot.png)


@testset "643.maximum-average-subarray-i.jl" begin
    @test find_max_average([1, 12, -5, -6, 50, 3], 4) == 51 / 4
    @test find_max_average([13, 20, -1, -6, 28, 30, 10, -4, -10, 22, 26, -21, -1, -9, -9,
                            -27, -3, 21, -13, -4], 7) == 102 / 7
end


module SpecialMatrices

using Compat
using Compat.LinearAlgebra
using FFTW

import Compat.LinearAlgebra: Matrix, inv
if VERSION >= v"0.7.0"
    import Compat.LinearAlgebra: mul!
end

import Base: getindex, isassigned, size, *
import Base.\

include("cauchy.jl") #Cauchy matrix
include("companion.jl") #Companion matrix
include("frobenius.jl") #Frobenius matrix
include("hankel.jl") #Hankel matrix
include("hilbert.jl") #Hilbert matrix
include("kahan.jl") #Kahan matrix
include("riemann.jl") #Riemann matrix
include("strang.jl") #Strang matrix
include("toeplitz.jl") #Toeplitz matrix, Circulant matrix
include("vandermonde.jl") #Vandermonde matrix

end # module


"""
return an table of the amount of each element in every molecule
as well as what the number of the rows and columns stand for.
"""

function chemical_parse(str::String)
    elems = Dict{String, Integer}()

    eldt = [
        begin
            ## remove parentheses and brackets
            flag = true
            while flag
                flag = false
                m = match(r".*?\((.*?)\)([0-9]*)", elemdata)
                m != nothing && ((flag=true; k = m[2]=="" ? 1 : parse(Int, m[2]));
                    elemdata = replace(elemdata, r"(\(.*?\)[0-9]*)",  m[1]^k, 1))
                m = match(r".*?\[(.*?)\]([0-9]*)", elemdata)
                m != nothing && ((flag=true; k = m[2]=="" ? 1 : parse(Int, m[2]));
                    elemdata = replace(elemdata, r"(\[.*?\][0-9]*)",  m[1]^k, 1))
                m = match(r".*?\{(.*?)\}([0-9]*)", elemdata)
                m != nothing && ((flag=true; k = m[2]=="" ? 1 : parse(Int, m[2]));
                    elemdata = replace(elemdata, r"(\{.*?\}[0-9]*)",  m[1]^k, 1))
            end
            ## arrange elements' names

            len = length(elemdata)
            for i in 1:len
                '0' ≤ elemdata[i] ≤ '9' && continue
                if 'A' ≤ elemdata[i] ≤ 'Z'
                    name = i < len && 'a' ≤ elemdata[i+1] ≤ 'z' ?
                        elemdata[i:i+1] : elemdata[i:i]
                    haskey(elems, name) || (elems[name] = length(elems)+1)
                end
            end
            elemdata
        end
        for elemdata in split(str)]

    arr = zeros(Int, length(elems), length(eldt))
    ## arr[<element>, <molecule>]
    for (k, elemdata) in enumerate(eldt)
        len = length(elemdata)
        i = 1
        while i ≤ len
            @assert  'A' ≤ elemdata[i] ≤ 'Z'
            name =
                if i < len && 'a' ≤ elemdata[i+1] ≤ 'z'
                    i += 1
                    elemdata[i-1:i]
                else
                    elemdata[i:i]
                end
            i += 1
            amount_str = ""
            while i ≤ len && '0' ≤ elemdata[i] ≤ '9'
                amount_str *= elemdata[i:i]
                i += 1
            end
            arr[elems[name], k] = amount_str=="" ? 1 : parse(Int, amount_str)
        end
    end
    row = split(str)
    revelems = Dict(v=>k for (k, v) in elems)
    column = [revelems[i] for i in 1:length(revelems)]
    return arr, row, column
end


flatten{T}(x::Vector{T}) = x
flatten(X) = reshape(X, prod(size(X)))

# convert vector of data samples into a single data matrix, flatten samples if needed
datamatrix{T,N}(samples::Vector{Array{T,N}}) =
    hcat([flatten(sample) for sample in samples]...)

# convert single data matrix into a vector of samples
# samplevec{T}(datamat::Matrix{T}, sampledims::(Int...)) =
#     Matrix{T}[reshape(datamat[:, i], sampledims) for i=1:size(datamat, 2)]

function gs_orthonorm(M)
    O = zeros(eltype(M), size(M))
    ncol = size(M, 2)
    tol = 100*eps(Float64)
    k = 1
    for i=1:ncol
        v = M[:, i]   # column to orthonormalize
        # subtract projections over previous vectors
        for j=1:k-1
            v = v - dot(O[:,j], v) * O[:,j]
        end
        # only keep nonzero vectors
        n = norm(v)
        if n > tol
            O[:, k] = v / n
            k = k + 1
        end
    end
    return O[:, 1:k-1]
end

function StatsBase.describe{T,N}(A::Array{T,N})
    println("size   : $(size(A))")
    println("nzs    : $(countnz(A) / length(A))")
    println("min    : $(minimum(A))")
    println("mean   : $(mean(A))")
    println("median : $(median(A))")
    println("max    : $(maximum(A))")
    println("sum    : $(sum(A))")
end


using Distributed
@everywhere using Mamba

## Data
asthma = Dict{Symbol, Any}(
  :y =>
    [210 60 0 1  1
     88 641 0 4 13
     1    0 0 0  1],
  :M =>
    [272, 746, 2]
)

## Model Specification
model = Model(

  y = Stochastic(2,
    (M, q) ->
      MultivariateDistribution[
        Multinomial(M[i], vec(q[i, :]))
        for i in 1:length(M)
      ],
    false
  ),

  q = Stochastic(2,
    M ->
      MultivariateDistribution[
        Dirichlet(ones(5))
        for i in 1:length(M)
      ],
    true
  )

)

## Initial Values
inits = [
  Dict{Symbol, Any}(
    :y => asthma[:y],
    :q => vcat([rand(Dirichlet(ones(5)))' for i in 1:3]...)
  )
  for i in 1:3
]

## Sampling Scheme
scheme = [SliceSimplex(:q)]
setsamplers!(model, scheme)

## MCMC Simulations
sim = mcmc(model, asthma, inits, 10000, burnin=2500, thin=2, chains=3)
describe(sim)


using CartesianFDM
using LinearAlgebra
using SparseArrays

n = (33, 32)
top = (nonperiodic(), periodic())

ctx = cartesianfdmcontext(top, n)

# square domain (0, 1) x (0, 1)
h = spacing.(top, n)
x = coordinate.(top, n)

# cylinder centered at (0.5, 0.5) with radius 0.25
cylinder(x, y, z...) = 0.25 - (2x - 1) ^ 2 - (2y - 1) ^ 2

# capacities
using Vofinit

h₀ = Cdouble.((h..., sum(h) / 2))
xₑ = Cdouble[0, 0, 0, 0]

using Base.Iterators

# volume
V = map(zip(CartesianIndices(n), product(x...))) do (ind, el)
    x₀ = Cdouble.((el..., 0.))
    prod(h) * getcc(cylinder, x₀, h₀, xₑ)
end

V = map(mask(ctx, reshape(V, :))) do val
    iszero(val) ? eps(val) : val
end

# area
A₁ = map(zip(CartesianIndices(n), product(x...))) do (ind, el)
    x₀ = Cdouble.((el..., 0.))
    h[2] * getcc(x₀, h₀, xₑ) do y...
        cylinder(x₀[1], y[2])
    end
end

A₂ = map(zip(CartesianIndices(n), product(x...))) do (ind, el)
    x₀ = Cdouble.((el..., 0.))
    h[1] * getcc(x₀, h₀, xₑ) do y...
        cylinder(y[1], x₀[2])
    end
end

A = mask(ctx, [reshape(A₁, :), reshape(A₂, :)])

# potential flow
W = reshape(map(CartesianIndices(n)) do ind
                ind[1] ≤ 2 || ind[1] ≥ n[1] - 3 ? 1. : 0.
            end, :)
_, U = potentialflow(ctx, A, V, W)

using Plots

heatmap(reshape(U[1], n...))

#=
T = scalar(:T, n)
D = ones(prod(n))

# secondary
G = gradient(dir, ctx, A, V, T, D)
L = divergence(dir, ctx, A, G)

#
M = map(V) do el
    el ≤ 10eps(el) ? prod(h) : el
end
=#

filename = "heat.jl"

# if file already exists
#include(filename)
# otherwise
#(heat, heat!) = build_function(L ./ 2M, T)
#open(filename, "w") do file
#    write(file, string(heat))
#    write(file, "\n")
#    write(file, string(heat!))
#end

#=
using DifferentialEquations

T₀ = zeros(prod(n))
tspan = (0.0, 1.0)
prob = ODEProblem(heat, T₀, tspan)
sol = solve(prob, Tsit5(), reltol=1e-8, abstol=1e-8)

using Plots

heatmap(reshape(sol(0.1), n...))
=#

#=
potential = eval(first(build_function(Δ, Φ)))

using NLsolve

sol = nlsolve(zeros(prod(n))) do res, Φ
    res .= potential(Φ)
end
phi = getproperty(sol, :zero)
=#

#J = Symbolics.jacobian(Δ, Φ)
#J = eval(first(build_function(sparse(Symbolics.jacobian(Δ, Φ)))))()

#phi = rand(prod(n))
#res1 = eval(potential)(phi)
#res2 = J * phi + eval(potential)(zeros(prod(n)))

#@assert norm(res1 - res2) < 1e-14

# also in CYLINDER!
#=
vec = diag(J)
broadcast!(vec, vec) do x
    iszero(x) ? -8one(x) : zero(x)
end

J .+= vec

phi = (-J) \ eval(potential)(zeros(prod(n)))
=#

#=
@variables t p
(rhs, rhs!) = build_function(L, T, p, t)

using DifferentialEquations

tspan = (0.0, 100.0)
T₀ = zeros(prod(n))
prob = ODEProblem(eval(rhs!), T₀, tspan)
sol = solve(prob)

using Plots

heatmap(reshape(sol(0.1), n...))
=#

### Step 1 -- thermal flow
#
# 1. Potential flow ;
# 2. Add convection ;
# 3. Neumann B. C. on the outflow ;
# 4. Robin B. C.
#
### Step 2 -- chemically reacting flow
#
# 1. More scalars with Fick's law ;
# 2. Heterogeneous reactions (with Robin B. C.) ;
# 3. Homogeneous reactions ;
# 4. Fluid-solid interaction with quasi-steady assumption.
#
### Step 2bis -- coupling with Navier-Stokes solver
#
### Step 3 -- senstivity analysis


@symbol_func function phi_res(cur_reactor::AbstractReactor)
  cur_phi = cur_reactor.C_ejima

  cur_phi *= cur_reactor.R_0

  cur_phi *= ( 4e-7 * cur_reactor.pi )

  cur_phi *= f_IN(cur_reactor) * cur_reactor.I_P

  cur_phi *= 1e6

  cur_phi
end


using Elemental, TSVD, LinearAlgebra, Test
using Random
A = Elemental.DistMatrix(Float64)
Elemental.gaussian!(A, 500, 200)

# Warning! If seed is not set, each worker will end up with a different initial vector
# making the Lanczos processes inconsistent
Random.seed!(123)
vals_Elemental = tsvd(A, 5)[2]
localA = Array(A)
vals_LAPACK    = svdvals(localA)[1:5]
@test vals_Elemental ≈ vals_LAPACK


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

import .Libjsp
using Plots

function fx(t)
    sin(2*pi*5500*t) + 0.5*sin(2*pi*2300*t + 3*pi/4) + 0.8*cos(2*pi*8150*t + 2*pi/7)
end

fsample = 24000
sampling_num = 256

ts = 1/fsample

x_n = zeros(Float64,sampling_num)

for i_sampling=0:sampling_num-1
    x_n[i_sampling+1] = fx(ts*i_sampling)
end

trangle_window = Libjsp.window("hanning", sampling_num)

x_window = x_n .* trangle_window

X_n = Libjsp.dft(sampling_num, x_window, fsample)

x0 = range(0, stop=sampling_num-1)

y0 = x_n.^2
y0 = 10 * broadcast(log10, y0)

y1 = x_window

y2 = broadcast(abs, X_n)
y2 = 10 * broadcast(log10, y2)

p0 = plot(x0, y0, marker = 2, linecolor= :blue)
p1 = plot(x0, y1, marker = 2, linecolor= :green)
p2 = plot(x0, y2, marker = 3, linecolor= :red)

plot(p0, p1, p2, layout=(3,1))


using FirstOrderOptimization

srand(1)

## test asymmetric sketch
m = 100; n = 110; r = 5

X = zeros(m,n)
s = AsymmetricSketch(m,n,r)
for iter=1:r
	Xi = LowRankOperator(rand(m),rand(n)')
	X += Array(Xi)
	additive_update!(s,Xi)
end
@show vecnorm(Array(reconstruct(s)) - X) / vecnorm(X)

## test symmetric sketch
n = 100; r = 5

X = zeros(n,n)
s = SymmetricSketch(n,r)
for iter=1:r
	v = rand(n)
	Xi = LowRankOperator(v,v')
	X += v*v'
	additive_update!(s,Xi)
end
@show vecnorm(Array(reconstruct(s)) - X) / vecnorm(X)


# matplotlib.use("macosx")

@testset "eval_f_on_gauss_pts" begin
    f = (x,y)->x^2+y^2
    F = eval_f_on_gauss_pts(f, m, n, h)
    pts = AdFem.get_gauss_points(m, n, h)
    close("all")
    scatter3D(pts[:,1], pts[:,2], F, marker=".")
    mesh(X, Y, (@. f(X, Y)), color="orange", alpha=0.5)
    savefig("eval_f_on_gauss_pts.png")
end

@testset "compute_fvm_mechanics_term" begin
    u = zeros(2*(m+1)*(n+1))
    k = 0
    for j = 1:n+1
        for i = 1:m+1
            x = (i-1)*h
            y = (j-1)*h
            k += 1
            u[k] = x^2+y^2
            u[k + (m+1)*(n+1)] = x^2-y^2
        end
    end
    eps_v = compute_fvm_mechanics_term(u, m, n, h)
    close("all")
    scatter3D(Xv, Yv, eps_v, marker=".")
    scatter3D(Xv, Yv, 2h^2 * (@. Xv-Yv), marker="+")
    savefig("compute_fvm_mechanics_term.png")
end

@testset "compute_interaction_matrix" begin
    A = compute_interaction_matrix(m, n, h)
    u = zeros(2*(m+1)*(n+1))
    k = 0
    for j = 1:n+1
        for i = 1:m+1
            x = (i-1)*h
            y = (j-1)*h
            k += 1
            u[k] = x^2+y^2
            u[k + (m+1)*(n+1)] = x^2-y^2
        end
    end
    p = zeros(m*n)
    k = 0
    for j = 1:n
        for i = 1:m
            k += 1
            x = (i-1/2)*h
            y = (j-1/2)*h
            p[k] =  x+y
        end
    end
    val = p'*A*u
    @test val ≈ 64
end

@testset "compute_fvm_source_term" begin
    f = (x,y)->x^2+y^2
    F = eval_f_on_gauss_pts(f, m, n, h)
    S = compute_fvm_source_term(F, m, n, h)
    Fs = f.(Xv, Yv)*h^2
    close("all")
    scatter3D(Xv, Yv, S, marker = ".")
    scatter3D(Xv, Yv, Fs, marker = "+")
    savefig("compute_fvm_source_term.png")
end

@testset "compute_fvm_tpfa_matrix" begin
    M = compute_fvm_tpfa_matrix(m, n, h)
    p = Xv.^2 + Yv.^2
    V = M * p
    close("all")
    scatter3D(Xv, Yv, V, marker = ".")
    scatter3D(Xv, Yv, 4h^2 * ones(m*n), marker = "+")
    savefig("compute_fvm_tpfa_matrix.png")
end

@testset "static linear elasticity" begin
    bdedge = []
    for i = 1:m
        push!(bdedge, [i i+1])
    end
    bdedge = vcat(bdedge...)

    bdnode = Int64[]
    for j = 1:n+1
        push!(bdnode, (j-1)*(m+1)+1)
        push!(bdnode, (j-1)*(m+1)+m+1)
    end
    for i = 2:m
        push!(bdnode, n*(m+1)+i)
    end

    F1 = eval_f_on_gauss_pts((x,y)->3.0, m, n, h)
    F2 = eval_f_on_gauss_pts((x,y)->-1.0, m, n, h)
    F = compute_fem_source_term(F1, F2, m, n, h)

    t1 = eval_f_on_boundary_edge((x,y)->-x-y, bdedge, m, n, h)
    t2 = eval_f_on_boundary_edge((x,y)->2y, bdedge, m, n, h)
    T = compute_fem_traction_term([t1 t2], bdedge, m, n, h)

    D = diagm(0=>[1,1,0.5])
    K = compute_fem_stiffness_matrix(D, m, n, h)
    rhs = T - F
    bdval = [eval_f_on_boundary_node((x,y)->x^2+y^2, bdnode, m, n, h);
            eval_f_on_boundary_node((x,y)->x^2-y^2, bdnode, m, n, h)]
    rhs[[bdnode;bdnode .+ (m+1)*(n+1)]] = bdval
    K, Kbd = fem_impose_Dirichlet_boundary_condition(K, bdnode, m, n, h)
    u = K\(rhs-Kbd*bdval)
    X, Y, U, V = plot_u(u, m, n, h)

    figure(figsize=[10,4])
    subplot(121)
    pcolormesh(X, Y, (@. X^2+Y^2-U), alpha=0.6); xlabel("x"); ylabel("y"); title("Error for u")
    colorbar()
    subplot(122)
    pcolormesh(X, Y, (@. X^2-Y^2-V), alpha=0.6); xlabel("x"); ylabel("y"); title("Error for v")
    colorbar()

end

@testset "Poisson" begin
    m = 40
    n = 20
    h = 0.1
    # bdedge = []
    # for i = 1:m
    #     push!(bdedge, [i i+1])
    # end
    # bdedge = vcat(bdedge...)

    bdnode = Int64[]
    for j = 1:n+1
        push!(bdnode, (j-1)*(m+1)+1)
        push!(bdnode, (j-1)*(m+1)+m+1)
    end
    for i = 2:m
        push!(bdnode, n*(m+1)+i)
        push!(bdnode, i)
    end

    K_ = [2.0 1.0
        1.0 2.0]
    K = compute_fem_stiffness_matrix1(K_, m, n, h)

    K, Kbd = fem_impose_Dirichlet_boundary_condition1(K, bdnode, m, n, h)
    F = eval_f_on_gauss_pts((x,y)->-8, m, n, h)
    rhs = compute_fem_source_term1(F, m, n, h)
    bdval = eval_f_on_boundary_node( (x,y)->(x^2+y^2), bdnode, m, n, h)
    rhs[bdnode] = bdval
    U = K\(rhs-Kbd*bdval)

    Uexact = zeros(n+1,m+1)
    for j = 1:n+1
        for i = 1:m+1
            x = (i-1)*h; y = (j-1)*h
            Uexact[j, i] = (x^2+y^2)
        end
    end
    pcolormesh(reshape(U, m+1, n+1)'-Uexact); colorbar()
end

@testset "heat equation" begin
    m = 40
    n = 20
    h = 0.01
    bdedge = []
    for i = 1:m
        push!(bdedge, [i i+1])
    end
    bdedge = vcat(bdedge...)

    bdnode = Int64[]
    for j = 1:n+1
        push!(bdnode, (j-1)*(m+1)+1)
        push!(bdnode, (j-1)*(m+1)+m+1)
    end
    for i = 2:m
        push!(bdnode, n*(m+1)+i)
    end

    ρ = eval_f_on_gauss_pts((x,y)->1+x^2+y^2, m, n, h)
    M = compute_fem_mass_matrix1(ρ, m, n, h)

    K_ = [2.0 1.0
        1.0 2.0]
    K = compute_fem_stiffness_matrix1(K_, m, n, h)

    NT = 200
    Δt = 1/NT
    A = M/Δt+K
    A, Abd = fem_impose_Dirichlet_boundary_condition1(A, bdnode, m, n, h)
    U = zeros((m+1)*(n+1), NT+1)
    for i = 1:m+1
        for j = 1:n+1
            x = (i-1)*h; y = (j-1)*h
            U[(j-1)*(m+1)+i, 1] = x^2+y^2
        end
    end
    for i = 1:NT
        F = eval_f_on_gauss_pts((x,y)->(-(1+x^2+y^2)*(x^2+y^2)-8)*exp(-i*Δt), m, n, h)
        F = compute_fem_source_term1(F, m, n, h)

        T = eval_f_on_boundary_edge((x,y)->-(2*x+4*y)*exp(-i*Δt), bdedge, m, n, h)
        T = compute_fem_flux_term1(T, bdedge, m, n, h)

        rhs = F  + M*U[:,i]/Δt + T
        bdval = eval_f_on_boundary_node( (x,y)->(x^2+y^2)*exp(-i*Δt), bdnode, m, n, h)
        rhs[bdnode] = bdval
        U[:,i+1] = A\(
            rhs - Abd*bdval
        )
    end

    Uexact = zeros(n+1,m+1)
    for j = 1:n+1
        for i = 1:m+1
            x = (i-1)*h; y = (j-1)*h
            Uexact[j, i] = (x^2+y^2)*exp(-1)
        end
    end
    pcolormesh(reshape(U[:,end], m+1, n+1)'-Uexact); colorbar()
end

@testset "elastodynamics" begin
    # alpha method
    # table 9.1.1, page 493
    β = 1/4; γ = 1/2
    a = b = 0.1
    m = 40
    n = 20
    h = 0.01
    NT = 200
    Δt = 1/NT
    bdedge = []
    for j = 1:n
        push!(bdedge, [(j-1)*(m+1)+m+1 j*(m+1)+m+1])
    end
    bdedge = vcat(bdedge...)

    bdnode = Int64[]
    for j = 1:n+1
        push!(bdnode, (j-1)*(m+1)+1)
    end

    M = compute_fem_mass_matrix1(m, n, h)
    S = spzeros((m+1)*(n+1), (m+1)*(n+1))
    M = [M S;S M]

    D = diagm(0=>[1,1,0.5])
    K = compute_fem_stiffness_matrix(D, m, n, h)
    C = a*M + b*K # damping matrix

    L = M + γ*Δt*C + β*Δt^2*K
    L, Lbd = fem_impose_Dirichlet_boundary_condition(L, bdnode, m, n, h)

    a = zeros(2(m+1)*(n+1))
    v = zeros(2(m+1)*(n+1))
    d = zeros(2(m+1)*(n+1))
    U = zeros(2(m+1)*(n+1),NT+1)
    for i = 1:NT
        T = eval_f_on_boundary_edge((x,y)->0.01, bdedge, m, n, h)
        # T = eval_f_on_boundary_edge((x,y)->0.0, bdedge, m, n, h)
        T = [zeros(length(T)) -T]
        T = compute_fem_traction_term(T, bdedge, m, n, h)
        f1 = eval_f_on_gauss_pts((x,y)->0., m, n, h)
        f2 = eval_f_on_gauss_pts((x,y)->0., m, n, h)
        # f2 = eval_f_on_gauss_pts((x,y)->0.1, m, n, h)
        F = compute_fem_source_term(f1, f2, m, n, h)

        rhs = F+T

        td = d + Δt*v + Δt^2/2*(1-2β)*a
        tv = v + (1-γ)*Δt*a
        rhs = rhs - C*tv - K*td
        rhs[[bdnode; bdnode.+(m+1)*(n+1)]] .= 0.0

        a = L\rhs
        d = td + β*Δt^2*a
        v = tv + γ*Δt*a
        U[:,i+1] = d
    end

    x = []
    y = []
    for j= 1:n+1
        for i = 1:m+1
            push!(x, (i-1)*h)
            push!(y, (j-1)*h)
        end
    end
    for i = 1:10:NT+1
        close("all")
        scatter(x+U[1:(m+1)*(n+1), i], y+U[(m+1)*(n+1)+1:end, i])
        xlabel("x")
        ylabel("y")
        k = string(i)
        k = repeat("0", 3-length(k))*k
        title("t = $k")
        ylim(-0.05,0.25)
        gca().invert_yaxis()
        savefig("u$k.png")
    end
end

@testset "tpfa" begin

    n = 50
    m = 2n
    h = 1/n
    X = zeros((m+1)*(n+1))
    Y = zeros((m+1)*(n+1))
    for i = 1:m+1
        for j = 1:n+1
            k = i+(j-1)*(m+1)
            X[k] = (i-1)*h
            Y[k] = (j-1)*h
        end
    end
    bd = bcedge("all", m, n, h)
    bdval = Float64[]
    for i = 1:size(bd,1)
        x = X[bd[i,1]]; y = Y[bd[i,1]]
        p1 = x^2+y^2
        x = X[bd[i,2]]; y = Y[bd[i,2]]
        p2 = x^2+y^2
        push!(bdval, (p1+p2)/2)
    end

    A, rhs = compute_fvm_tpfa_matrix(ones(m*n), bd, bdval, m, n, h)
    f = compute_fvm_source_term(4*ones(4*m*n), m, n, h)
    f -= rhs
    sol = A\f
    figure(figsize=(4,10))
    subplot(211)
    pcolormesh(LinRange(h,2,m)|>collect, LinRange(h,1,n)|>collect, reshape(sol, m, n)',
        vmin = -1.0, vmax = 3.5)
    gca().invert_yaxis()
    axis("equal")
    colorbar()
    subplot(212)
    x0 = LinRange(h,2,m)|>collect
    y0 = LinRange(h,1,n)|>collect
    X, Y = np.meshgrid(x0, y0)
    Z = @. X^2+Y^2
    pcolormesh(x0, y0, Z,
        vmin = -1.0, vmax = 3.5)
    gca().invert_yaxis()
    axis("equal")
    colorbar()

    figure()
    pcolormesh(log.(abs.(reshape(sol, m, n)'-Z))/log(10))
    colorbar()
end


@testset "Xor" begin
@testset "Xor basic all solutions" begin
    m = Model(optimizer_with_attributes(
        CS.Optimizer,
        "all_solutions" => true,
        "logging" => [],
    ))
    @variable(m, 1 <= x <= 5, Int)
    @variable(m, 1 <= y <= 5, Int)
    @constraint(m, (x <= 2) ⊻ (y <= 2))
    optimize!(m)

    num_sols = MOI.get(m, MOI.ResultCount())
    @test num_sols == count((i <= 2) ⊻ (j <= 2) for i=1:5, j=1:5)
    for i in 1:num_sols
        xval = convert(Int, JuMP.value(x; result=i))
        yval = convert(Int, JuMP.value(y; result=i))
        @test (xval <= 2) ⊻ (yval <= 2)
    end
end
end


"""
    generateData(problemDescriptor::ProblemDescriptor)

Helper function designed to generate experiment data for each problem under consideration.
It returns a tuple `(C, c, d, Γ, X)` where
1. `C` is a vector of nonnegative first stage costs
2. `c` is a vector of a nonnegative nominal second stage costs
3. `d` is a vector of maximal deviations of the costs from their nominal values
4. `Γ` is a budget, or the amount of uncertainty, which can be allocated to the second stage costs
5. `X` is a set of feasible solutions represented as a list functions, each of which accepts a list of JuMP variables as an argument and returns a JuMP linear constraint
"""
generateData(c::ProblemDescriptor) = error("No generateData method defined for data generator type $(typeof(c))")

function generateData(g::KnapsackProblemDescriptor)
    C = rand(1:20, g.n)
    c = rand(1:20, g.n)
    d = rand(1:100, g.n)
    Γ = floor(Int, 0.1 * sum(d))
    w = rand(1:20, g.n)
    W = floor(Int, 0.3 * sum(w))
    X = getKnapsackConstraints(w, W)
    (C, c, d, Γ, X)
end

function generateData(g::AssignmentProblemDescriptor)
    C = rand(1:20, g.n, g.n)
    c = rand(1:20, g.n, g.n)
    d = rand(1:100, g.n, g.n)
    Γ = floor(Int, 0.1 * sum(d))
    X = getAssignmentConstraints(g.n)
    (C, c, d, Γ, X)
end


module GraphLasso
using LinearAlgebra, SparseArrays
using DataFrames, Lasso, Graphs

export graphlasso

function graphlasso(S::Matrix{Float64}, α::Float64; tol::Float64=1e-5,
                    maxit::Int=1000, penalize_diag::Bool=true, verbose::Bool=false)
    p = size(S,1)
    adj = convert(SparseMatrixCSC{Int64,Int64}, abs.(S) .> α)
    blocks = connected_components(sparse2adjacencylist2(max.(adj, sparse(I, p, p))))
    if(verbose)
        print(blocks)
    end
    W = zeros(p,p)
    Θ = zeros(p,p)
    for block = blocks
        W[block,block], Θ[block,block] = fit_block(S[block,block], α, tol, maxit, penalize_diag)
    end
    return W, Θ
end

function fit_block(S::Matrix{Float64}, α::Float64, tol::Float64,
                   maxit::Int, penalize_diag::Bool)
    p = size(S,1)
    W = copy(S)
    if penalize_diag
        W += α*I
    end
    if size(S,1) == 1
        return W, inv(W)
    end
    Θ = zeros(p,p)
    i = 0
    β = zeros(p-1,p)
    while i < maxit
        i += 1
        W_old = copy(W)
        for j = 1:p
            inds = collect(1:p)
            splice!(inds, j)
            W11 = W[inds,inds]
            sqrtW11 = sqrt(W11)
            β[:,j] = fit(LassoPath, sqrtW11, sqrtW11 \ S[inds,j], λ=[α/(p-1)],
                         standardize=false, intercept=false).coefs
            W[inds,j] = W11 * β[:,j]
            W[j,inds] = W[inds,j]
        end
        if norm(W - W_old) < tol
            break
        end
    end
    if i == maxit
        warn("Maximum number of iterations reached, graphlasso failed to converge")
    end
    for j in 1:p
        inds = collect(1:p)
        splice!(inds, j)
        Θ[j,j] = 1/(W[j,j] - dot(W[inds,j],β[:,j]))
        Θ[inds,j] = -Θ[j,j] * β[:,j]
        Θ[j,inds] = Θ[inds,j]
    end
    return W, Θ
end

# Modify code from unmaintained Graph.jl
function sparse2adjacencylist2(A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti<:Integer}
    colptr = A.colptr
    rowval = A.rowval
    n = size(A, 1)
    adjlist = Array{Array{Ti,1}}(undef, n)
    s = 0
    for j in 1:n
        adjj = Ti[]
        sizehint!(adjj, colptr[j+1] - colptr[j] - 1)
        for k in colptr[j]:(colptr[j+1] - 1)
            rvk = A.rowval[k]
            if rvk != j push!(adjj, rvk) end
        end
        s += length(adjj)
        adjlist[j] = adjj
    end
    GenericAdjacencyList{Ti, UnitRange{Ti}, Vector{Vector{Ti}}}(!ishermitian(A),
                                                             one(Ti):convert(Ti,n),
                                                             s, adjlist)
end

end


function get_Kσs(sp::LinearElasticityProblem{dim, TT}, dofs, cellvalues) where {dim, TT}
    E = YoungsModulus(sp)
    ν = PoissonRatio(sp)
    dh = sp.ch.dh
    n = ndofs_per_cell(dh)
    global_dofs = zeros(Int, n)
    Kσs = [zeros(TT, n, n) for i in 1:getncells(dh.grid)]
    Kσ_e = zeros(TT, n, n)
    ψ_e = zeros(TT, dim*n, dim*n)
    G = zeros(TT, dim*n, n)
    δ = eye(TT, dim)
    ϵ = zeros(TT, dim, dim)
    σ = zeros(TT, dim, dim)
    u = zeros(TT, dim, dim)
    n_basefuncs = getnbasefunctions(cellvalues)
    for (cellidx, cell) in enumerate(CellIterator(dh))
        Kσ_e .= 0
        reinit!(cellvalues, cell)
        celldofs!(global_dofs, dh, cellidx)
        for q_point in 1:getnquadpoints(cellvalues)
            dΩ = getdetJdV(cellvalues, q_point)
            for d in 1:dim
                ψ_e[(d-1)*dim+1:d*dim, (d-1)*dim+1:d*dim] .= 0
            end
            for a in 1:n_basefuncs
                ∇ϕ = shape_gradient(cellvalues, q_point, a)
                _u = @view dofs[(@view global_dofs[dim*(a-1) + (1:dim)])]
                @einsum u[i,j] = _u[i]*∇ϕ[j]
                @einsum ϵ[i,j] = 1/2*(u[i,j] + u[j,i])
                @einsum σ[i,j] = E*ν/(1-ν^2)*δ[i,j]*ϵ[k,k] + E*ν*(1+ν)*ϵ[i,j]
                for d in 1:dim
                    ψ_e[(d-1)*dim+1:d*dim, (d-1)*dim+1:d*dim] .+= σ
                    G[(dim*(d-1)+1):(dim*d), (a-1)*dim+d] .= ∇ϕ
                end
            end
            Kσ_e .+= G'*ψ_e*G*dΩ
        end
        Kσs[cellidx] .= Kσ_e
    end

    return Kσs
end

function buckling(problem::LinearElasticityProblem{dim, T}, ginfo, einfo) where {dim, T}
    dh = problem.ch.dh

    u = ginfo.K \ ginfo.f
    Kσs = get_Kσs(problem, u, einfo.cellvalues)
    Kσ = deepcopy(ginfo.K)

    if Kσ isa Symmetric
        Kσ.data.nzval .= 0
        assembler = JuAFEM.AssemblerSparsityPattern(Kσ.data, T[], Int[], Int[])
    else
        Kσ.nzval .= 0
        assembler = JuAFEM.AssemblerSparsityPattern(Kσ, T[], Int[], Int[])
    end

    global_dofs = zeros(Int, ndofs_per_cell(dh))
    Kσ_e = zeros(T, size(Kσs[1]))
    celliteratortype = CellIterator{typeof(dh).parameters...}
    _celliterator::celliteratortype = CellIterator(dh)
    TK = eltype(Kσs)
    for (i,cell) in enumerate(_celliterator)
        celldofs!(global_dofs, dh, i)
        if TK <: Symmetric
            JuAFEM.assemble!(assembler, global_dofs, Kσs[i].data)
        else
            JuAFEM.assemble!(assembler, global_dofs, Kσs[i])
        end
    end

    return ginfo.K, Kσ
end


#------------------------------
#   Solve for the equilibrium interest rate in the Hugget economy.
#
#   First, solve the HJB equation to get the A ("transition") matrix, then
#   solve the discretized KF equations. Then iterate over interest rates to
#   get the asset supply function. Update interest rate until convergence.
#
#   This code uses an implicit method for the solution: use
#   an 'upwind' scheme consisting of forwards and backwards difference
#   approximation to the derivative of the value function.
#
#   Assume: CRRA utility function u(c) = c^(1-s)/(1-s)
#   Note that r < rho
#
#   Written by Ben Moll et al (2015). Optimized for speed #   by SeHyoun Ahn. Translation into Julia by James Graham
#------------------------------

# Define types for workspace variables
type Huggett_model
  sigma::Float64            # risk aversion coefficient
  rho::Float64              # discount rate.
  z1::Float64               # Income in state 1
  z2::Float64               # Income in state 2
  z::Array{Float64}        # Vector for the two incomes
  la1::Float64              # State 1 switching intensity
  la2::Float64              # State 2 switching intensity
  la::Array{Float64}       # Vector for the two intensities
  I::Int64                  # Number of grids for assets
  Ir::Int64                   # Number of interest rate iterations
  amin::Float64             # Min assets = borrowing constraint
  amax::Float64
  rmin::Float64             # Min interest rate
  rmax::Float64             # Max interest rate
  agrid::Vector{Float64}    # Grid space for assets
  aa::Array{Float64}       # Matrix for asset grid space
  zz::Array{Float64}       # Matrix for income grid space
  da::Float64               # This is the implied grid step size
  maxit::Int64
  crit::Float64
  crit_S::Float64             # Tolerance for asset mkt clearing
  Delta::Float64
  Aswitch::SparseMatrixCSC{Float64,Int64}    # Transition matrix
  v0::Array{Float64}       # Initial guess for value function
  r0::Float64                 # Intial guess for interest rate
end

# Create model function
function Huggett_model(;sigma::Float64=1.5,
  rho::Float64=0.05,
  I::Int64=1000,
  Ir::Int64=40,
  amin::Float64=-0.15,
  amax::Float64=5.0,
  rmin::Float64=0.01,
  rmax::Float64=0.04,
  z1::Float64=0.1,
  z2::Float64=0.2,
  la1::Float64=1.2,
  la2::Float64=1.5,
  maxit::Int64=100,
  crit::Float64=1e-6,
  crit_S::Float64=1e-5,
  Delta::Float64=1000.0,
  r0::Float64=0.03)

  z       = hcat(z1,z2)
  la      = hcat(la1,la2)
  agrid   = collect(linspace(amin,amax,I))
  aa      = hcat(agrid, agrid)
  zz      = ones(I,1)*z
  da      = (amax-amin)/(I-1)
  v0 = hcat( (z[1] + r0.*agrid).^(1-sigma)/(1-sigma)/rho, (z[2] + r0.*agrid).^(1-sigma)/(1-sigma)/rho )

  Aswitch =  vcat( hcat(-speye(I)*la[1], speye(I)*la[1]), hcat(speye(I)*la[2],-speye(I)*la[2]) )

  Huggett_model(sigma, rho, z1, z2, z, la1, la2, la, I, Ir, amin, amax, rmin, rmax, agrid, aa, zz, da, maxit, crit, crit_S, Delta, Aswitch,  v0, r0)

end

# Creates an instance of the model
mod = Huggett_model()

# Define Utility funciton
function util(x; sigma::Float64=mod.sigma)
  u = x.^(1-sigma)/(1-sigma)
  return u
end

function utilprime(x; sigma::Float64=mod.sigma)
  uprime = x.^(-sigma)
  return uprime
end

# Define inverse of derivative of utility function
function utilprimeinv(x; sigma::Float64=mod.sigma)
  uprimeinv = x.^(-1/sigma)
  return uprimeinv
end

# Function to compute the KFE
function KFE(A::SparseMatrixCSC{Float64,Int64}; I::Int64=mod.I, da::Float64=mod.da)
  AT          = A'
  b           = zeros(2*I,1)
  i_fix       = 1    # fix one value so that matrix is not singular
  b[i_fix]    = 0.1
  row         = hcat(zeros(1,i_fix-1), 1, zeros(1,2*I-i_fix))
  AT[i_fix,:] = row
  #Solve linear system
  gg          = AT\b          # A'*gg = b --> gg = (A')^(-1)*b
  g_sum       = gg'*ones(2*I,1)*da
  gg          = gg./g_sum     # Normalize so sum(gg)*da) = 1
  g           = hcat(gg[1:I], gg[I+1:2*I])
  #check1      = g[:,1]'*ones(mod.I,1)*mod.da
  #check2      = g[:,2]'*ones(mod.I,1)*mod.da
  return g
end

# Initialization
dVf         = zeros(mod.I,2)      # Forward difference in each state
dVb         = zeros(mod.I,2)      # Backward difference in each state
c           = zeros(mod.I,2)      # Consumption in each state
A           = sparse(zeros(2*mod.I,2*mod.I))   # Transition matrix
r_r         = zeros(mod.Ir,1)     # Interest rate guesses
rmin_r      = zeros(mod.Ir,1)     # Lower bound for bisection
rmax_r      = zeros(mod.Ir,1)     # Upper bound for bisection
g_r         = zeros(mod.I,2,mod.Ir)   # Stationary distribution
adot        = zeros(mod.I,2,mod.Ir)   # Policy function
V_r         = zeros(mod.I,2,mod.Ir)   # Value function
V           = zeros(mod.I,2)        # Value function
S           = zeros(mod.Ir,1)     # Aggregate asset supply
r           = mod.r0
dist        = zeros(mod.maxit,1)
r_end       = 0                   # Last iteration value

for ir in 1:mod.Ir
  r_r[ir]     = r
  rmin_r[ir]  = mod.rmin
  rmax_r[ir]  = mod.rmax

  if ir>1
    mod.v0 = V_r[:,:,ir-1]
  end

  v = mod.v0

  for n = 1:mod.maxit
    V               = v
    # forward difference
    dVf[1:mod.I-1,:]    = (V[2:mod.I,:]-V[1:mod.I-1,:])/mod.da
    dVf[mod.I,:]        = utilprime(mod.z + r.*mod.amax)
    # backward difference
    dVb[1,:]            = utilprime(mod.z + r.*mod.amin)  #state constraint boundary condition
    dVb[2:mod.I,:]      = (V[2:mod.I,:]-V[1:mod.I-1,:])/mod.da

    # I_concave       = dVb .> dVf #indicator whether value function is concave (problems arise if this is not the case)

    #consumption and savings with forward difference
    cf              = utilprimeinv(dVf)
    ssf             = mod.zz + r.*mod.aa - cf
    #consumption and savings with backward difference
    cb              = utilprimeinv(dVb)
    ssb             = mod.zz + r.*mod.aa - cb
    #consumption and derivative of value function at steady state
    c0              = mod.zz + r.*mod.aa
    dV0             =utilprime(c0)

    # dV_upwind makes a choice of forward or backward differences based on the sign of the drift
    If              = ssf .> 0    # positive drift --> forward difference
    Ib              = ssb .< 0    # negative drift --> backward difference
    I0              = (1-If-Ib)   # at steady state
    #make sure backward difference is used at amax
    #Ib(I,:) = 1; If(I,:) = 0;
    #STATE CONSTRAINT at amin: USE BOUNDARY CONDITION UNLESS muf > 0:
    #already taken care of automatically

    dV_Upwind       = dVf.*If + dVb.*Ib + dV0.*I0    #important to include third term
    c               = utilprimeinv(dV_Upwind)
    u               = util(c)

    #CONSTRUCT MATRIX
    X               = - min(ssb,0)/mod.da
    Y               = - max(ssf,0)/mod.da + min(ssb,0)/mod.da
    Z               = max(ssf,0)/mod.da

    A1              = spdiagm(Y[:,1],0,mod.I,mod.I) + spdiagm(X[2:mod.I,1],-1,mod.I,mod.I) + spdiagm(Z[1:mod.I-1,1],1,mod.I,mod.I)
    A2              = spdiagm(Y[:,2],0,mod.I,mod.I) + spdiagm(X[2:mod.I,2],-1,mod.I,mod.I) + spdiagm(Z[1:mod.I-1,2],1,mod.I,mod.I)
    A               = vcat(hcat(A1, sparse(zeros(mod.I,mod.I))), hcat(sparse(zeros(mod.I,mod.I)), A2)) + mod.Aswitch
    B               = (mod.rho + 1/mod.Delta)*speye(2*mod.I) - A

    u_stacked       = vcat(u[:,1], u[:,2])
    V_stacked       = vcat(V[:,1], V[:,2])

    b               = u_stacked + V_stacked/mod.Delta
    V_stacked       = B\b             # Solve linear system of eqns

    V               = hcat(V_stacked[1:mod.I], V_stacked[mod.I+1:2*mod.I])

    Vchange         = V - v
    v               = V

    dist[n] = maximum(maximum(abs(Vchange)))
    if dist[n]<mod.crit
      @printf("-----------------------\n")
      @printf("Interest rate = %f\n", r)
      @printf("Value Function Converged, Iteration = %i\n", n)
      break
    end

  end

  ## Solve the Kolmogorov Forward (Fokker-Planck) equations
  g = KFE(A)

  # Pack up output
  g_r[:,:,ir] = g
  adot[:,:,ir] = mod.zz + r.*mod.aa - c
  V_r[:,:,ir] = V
  S[ir] = (g[:,1]'*mod.agrid*mod.da + g[:,2]'*mod.agrid*mod.da)[1,1]

  #UPDATE INTEREST RATE
  r_end = ir

  if S[ir]>mod.crit_S
    @printf("Excess Supply\n")
    mod.rmax = r
    r = 0.5*(r+mod.rmin)
  elseif S[ir]<-mod.crit_S
    @printf("Excess Demand\n")
    mod.rmin = r
    r = 0.5*(r+mod.rmax)
  elseif abs(S[ir])<mod.crit_S
    @printf("Equilibrium Found, Interest rate = %f/n", r)
    break
  end

end

## EXPORT DATA FOR GRAPHING IN PYTHON
using DataFrames
df = DataFrame(adot[:,:,r_end])
writetable("adot.csv", df)

df = DataFrame(c)
writetable("cons.csv", df)

df = DataFrame(g_r[:,:,r_end])
writetable("dist.csv", df)

## GRAPHS
using Gadfly

# Savings Policy function
myplot1 = plot(layer(x=mod.agrid,y=adot[:,1,r_end],Geom.line,Theme(line_width=1.5mm)), layer(x=mod.agrid,y=adot[:,2,r_end],Geom.line, Theme(line_width=1.5mm,default_color=colorant"red")),
xintercept=[mod.amin], yintercept=[0], Geom.hline(color=colorant"black", size=0.5mm), Geom.vline(color=colorant"black", size=0.5mm),
Guide.title("Savings policy function"),
Guide.xlabel("Wealth, a"),
Guide.ylabel("Savings, si(a)"),
Coord.cartesian(xmin=mod.amin, xmax=mod.amax), Theme(major_label_font_size=6mm,
minor_label_font_size=4mm))

using Cairo
using Fontconfig

# Distributions
plot(layer(x=mod.agrid,y=g_r[:,1,r_end],Geom.line,Theme(line_width=1.5mm)), layer(x=mod.agrid,y=g_r[:,2,r_end],Geom.line, Theme(line_width=1.5mm,default_color=colorant"red")),
xintercept=[mod.amin], Geom.vline(color=colorant"black", size=0.5mm), yintercept=[0], Geom.hline(color=colorant"black", size=0.5mm),
Guide.title("Stationary distributions"),
Guide.xlabel("Wealth, a"),
Guide.ylabel("Density, gi(a)"),
Coord.cartesian(xmin=mod.amin, xmax=1), Theme(major_label_font_size=6mm,
minor_label_font_size=4mm))


# SPDX-License-Identifier: X11
# 2020-09-13
# Dist Max (500pt)

function _main()
  N = parse(Int, readline())
  X, Y = parse.(Int, split(readline()))
  A = [(X + Y, X - Y)]
  for i ∈ 2:N
    X, Y = parse.(Int, split(readline()))
    push!(A, (X + Y, X - Y))
  end
  return max(maximum(x -> x[1], A) - minimum(x -> x[1], A),
             maximum(x -> x[2], A) - minimum(x -> x[2], A))
end

println(_main())


# abstract HyperRectangle{N, T}

struct HyperRectangle{T}
    mins::Vector{T}
    maxes::Vector{T}
end

# Computes a bounding box around a point cloud
function compute_bbox(data::Vector{V}) where {V <: AbstractVector}
    T = eltype(V)
    n_dim = length(V)
    maxes = Vector{T}(undef, n_dim)
    mins = Vector{T}(undef, n_dim)
    @inbounds for j in 1:length(V)
        dim_max = typemin(T)
        dim_min = typemax(T)
        for k in 1:length(data)
            dim_max = max(data[k][j], dim_max)
            dim_min = min(data[k][j], dim_min)
        end
        maxes[j] = dim_max
        mins[j] = dim_min
    end
    return HyperRectangle(mins, maxes)
end

############################################
# Rectangle - Point functions
############################################
@inline function get_min_dim(rec::HyperRectangle, point::AbstractVector, dim::Int)
    @inbounds d = abs2(max(0, max(rec.mins[dim] - point[dim], point[dim] - rec.maxes[dim])))
    d
end

@inline function get_max_dim(rec::HyperRectangle, point::AbstractVector, dim::Int)
    @inbounds d = abs2(max(rec.maxes[dim] - point[dim], point[dim] - rec.mins[dim]))
    d
end

# Max distance between rectangle and point
@inline function get_max_distance(rec::HyperRectangle, point::AbstractVector{T}) where {T}
    max_dist = zero(T)
    @inbounds @simd for dim in eachindex(point)
        max_dist += abs2(max(rec.maxes[dim] - point[dim], point[dim] - rec.mins[dim]))
    end
    return max_dist
end

# Min distance between rectangle and point
@inline function get_min_distance(rec::HyperRectangle, point::AbstractVector{T}) where {T}
    min_dist = zero(T)
    @inbounds @simd for dim in eachindex(point)
        min_dist += abs2(max(0, max(rec.mins[dim] - point[dim], point[dim] - rec.maxes[dim])))
    end
    return min_dist
end


# Interpolation
xs = 1:5
ys = 1:8
g = Float64[x^2 * sin(y) for x in xs, y in ys]

gitp_quad2d = interpolate(g, BSpline(Quadratic(Line(OnCell()))))
using Gadfly
display(plot(x=xs,y=ys,z=g,Geom.contour))
display(plot(x=1:.1:5, y=1:.1:8, z=[gitp_quad2d(x,y) for x in 1:.1:5, y in 1:.1:8], Geom.contour))

xs = [1; 2]
ys = [3; 9]
g = ys
myinterp = LinearInterpolation([1 3; 2 4], BSpline(Linear()))
myinterp(1.5,1.5)

function grid_shift(A, row_vals, )

myinterp = LinearInterpolation([1 3; 2 4], BSpline(Linear()))
myinterp(1.5,1.5)

function interpolate_cols(X, Y, X_grid::Array; extrapolation_bc = Throw())
    #linear interpolation of a column of Y on that column of X, evaluated at point x
    Y_interp = similar(X_grid)
    for col in eachindex(X_grid[1,:])
        li = LinearInterpolation(X[:,col], Y[:,col],extrapolation_bc=extrapolation_bc)
        Y_interp[:,col] = li.(X_grid[:,col])
    end
    return Y_interp
end

A = [1. 3.; 2. 4.; 3. 5.]
A_vals = [6.; 7.; 8.]
Y_vals = [3; 5]

function invert_rule(A,A_vals)
    A_interp = similar(A)
    for col in eachindex(A[1,:])
        li = LinearInterpolation(A[:,col], A_vals, extrapolation_bc = Line())
        # use extrapolate and do different for each side fixed value/line
        A_interp[:,col] = li.(A_vals)
    end
    return A_interp
end

A_inv = invert_rule(A, A_vals)
li = interpolate(invert_rule(A, A_vals), BSpline(Linear()))
li(1.5,1.5)

function grid_shift(A, A_vals_old, A_vals_new, Y_vals_old, Y_vals_new)
    A_grid = zeros( eltype(A_inv[1]), length(A[:,1]), length(Y_vals) )

    # Find relative coordinates of new basis
    li = LinearInterpolation(Y_vals_old, 1:length(Y_vals_old), extrapolation_bc=Line())
    Y_ind = li(Y_vals_new)
    li = LinearInterpolation(A_vals_old, 1:length(A_vals_old), extrapolation_bc=Line())
    A_ind = li(A_vals_new)

    # Bi-linear interpolation on new coordinates
    li = interpolate(A_inv, BSpline(Linear()) )  #extrapolate with boundary rules
    [ A_grid[i,j] = li(A_ind, Y_ind) for  (i, A_ind) in enumerate(A_ind), (j, Y_ind) in enumerate(Y_ind)]
    return A_grid
end

A_inv = [1. 3.; 2. 4.; 2.5 5.5]
A_vals = [6.; 7.; 8.]
Y_vals = [3; 5]
Y_vals_new = [3.1; 4.9]
A_shift = grid_shift(A, A_vals, A_vals, Y_vals, Y_vals_new)
Plots.heatmap(A_inv)
Plots.heatmap(A_shift)
plot()
Plotly.plot(A_inv, st=:surface, camera=(-30,30))
Plotly.plot(A_shift, st=:surface, camera=(-30,30))

#############
plot(A_vals, Y_vals, A_inv, st=:scatter, camera=(45,45))
li = interpolate(A_inv, BSpline(Linear()))
plot(li, st=:surface, camera=(-30,30))
li(1.5,1.5)

A_vals_new = [6.5; 7.; 8.5]

Y_vals_new = [3; 5]
grid_shift(A_inv, A_vals, Y_vals, A_vals_new, Y_vals_new)

A_vals_old = [3.5;4.0;5.0;7.0]
scatter(A_vals_old)
scatter!(A_vals_new)
A_vals_new = [3.6;4.0;6;7.0]
li = LinearInterpolation(A_vals_old, 1:length(A_vals_old), extrapolation_bc=Line())
li(A_vals_new)

# Find indices for new grid bases
li_A = LinearInterpolation(A_vals_old, A_vals_new, extrapolation_bc=Line())
@show A_inds = li_A(A_vals_old) .- minimum(A_vals_old) .+ 1
li_Y = LinearInterpolation(Y_vals_old, Y_vals_new, extrapolation_bc=Line())
@show Y_inds = li_Y(Y_vals_old) .- minimum(Y_vals_old) .+ 1

A = [1. 3.; 2. 4.; 3. 5.]
A_vals = [6.; 7.; 8.]
Y_vals = [3; 5]
Y_vals_new = [4]
A_shift = grid_shift(A, A_vals, A_vals, Y_vals, Y_vals_new)
plot(repeat(A_vals,1,length(Y_vals)),repeat(Y_vals', length(A_vals),1), A, st=:surface, camera=(-30,30))
Plotly.plot(A_shift, st=:surface, camera=(-30,30))

A |> size
repeat(A_vals,1,length(Y_vals)) |> size
repeat(Y_vals', length(A_vals),1) |> size
using Plots
pyplot()

plot(A_vals, Y_vals, vec(A), st=:surface, camera=(-30,30))

vec(A)
A

A = [3. 4; 5 6; 7 7]
A_vals = [6.; 7.; 8.]
A_vals_new = [3.1;5;6.9]
Y_vals = [3; 5]
Y_vals_new = [2.5; 4.5]
li = interpolate([1;2;3], A_vals, BSpline(Linear()))
li = extrapolate(interpolate(A[:,1], A_vals, Gridded(Linear())), extrapolation_bc = (0.,Line()))
minimum(A[:,1])
maximum(A[:,1])
li(A_vals_new)
A_inv1 = invert_rule(A,A_vals)
A_inv2 = invert_rule(A,A_vals_new)
A_shift = grid_shift(A, A_vals, A_vals_new, Y_vals, Y_vals)
plotly()
plot(repeat(A_vals,length(Y_vals)),vec(repeat(Y_vals', length(A_vals),1 )),vec(A), st=:scatter, camera=(-30,30))
plot!(repeat(A_vals_new,length(Y_vals)),vec(repeat(Y_vals', length(A_vals_new),1 )),vec(A_shift), st=:scatter, camera=(-30,30))

plotly()
plot(repeat(A_vals,length(Y_vals)),vec(repeat(Y_vals', length(A_vals),1 )),vec(A_inv1), st=:scatter, camera=(-30,30))
plot!(repeat(A_vals_new,length(Y_vals)),vec(repeat(Y_vals', length(A_vals_new),1 )),vec(A_inv2), st=:scatter, camera=(-30,30))

plotly()
plot(repeat(A_vals,length(Y_vals)),vec(repeat(Y_vals', length(A_vals),1 )),vec(A), st=:scatter, camera=(-30,30))
plot!(repeat(A_vals,length(Y_vals)),vec(repeat(Y_vals', length(A_vals),1 )),vec(A_inv), st=:scatter, camera=(-30,30))

plot!(repeat(A_vals_new,length(Y_vals_new)),vec(repeat(Y_vals_new', length(A_vals_new),1 )),vec(A_shift), st=:scatter, camera=(-30,30))

using Plots, Interpolations
B = [x^2 for x in 1.:10]
scatter(1.:10,B, legend=:bottomright)
li = extrapolate(interpolate(B, BSpline(Linear())), 0.)
x2=0.:11.
B_interp = li(x2)
scatter(x2, B_interp)
B_interp = extrapolate(interpolate(B, BSpline(Linear())), 0.)(x2)


#####
##### An attempt to simulate a scenario similar to that reported by
##### Veron and Melville, JFM (2001)
#####

using Oceananigans
using Oceananigans.Units
using Printf
using SpecialFunctions
using GLMakie

#####
##### Domain
#####

@info "Building a grid..."

Nx = 128

grid = RegularRectilinearGrid(size = (2Nx, Nx),
                              halo = (3, 3),
                              y = (    0, 0.1),
                              z = (-0.05,   0),
                              topology = (Flat, Periodic, Bounded))

@show grid

#####
##### Stokes drift
#####

# Physical parameters
g = 9.81 # gravitational acceleration
T = 7.2e-5 # (kinematic) surface tension, m³ s⁻²

const k = 2π / 0.02
const ϵ = 0.3
const ω = sqrt(g * k + T * k^3)

@show a = ϵ / k
@inline uˢ(z, t) = ϵ^2 * ω / k * exp(2k * z)
@inline ∂z_uˢ(z, t) = 2 * ϵ^2 * ω * exp(2k * z)

surface_uˢ = ϵ^2 * ω / k
surface_∂z_uˢ = 2 * ϵ^2 * ω
@show surface_uˢ
@show 1 / surface_∂z_uˢ

#####
##### Basic state
#####

closure = IsotropicDiffusivity(ν=1.05e-6, κ=1.0e-6)

Re = 1e5
ν = closure.ν
t₀ = 12

hν = sqrt(ν * t₀)
const h = 1 / k
const U₀ = ν * Re / h # Re = U L / ν

@show Sh = surface_∂z_uˢ / (U₀ / h)
@show Re h hν U₀

@inline η(z, t) = - z / 2sqrt(ν * t)

@inline paolucci(x, y, z, t) = uˢ(z, t) +
    2 * √(ν * t / π) * (exp(-η(z, t)^2) - √π * η(z, t) * erfc(η(z, t)))

U = BackgroundField(paolucci)

model = NonhydrostaticModel(architecture = CPU(),
                            advection = CenteredSecondOrder(),
                            timestepper = :RungeKutta3,
                            grid = grid,
                            background_fields = (; u=U),
                            closure = closure,
                            stokes_drift = UniformStokesDrift(∂z_uˢ=∂z_uˢ),
                            coriolis = nothing,
                            tracers = nothing,
                            buoyancy = nothing)

@show Δt = 1e-2 * grid.Δy / surface_uˢ

u, v, w = model.velocities
ξ = ComputedField(∂z(v) - ∂y(w))
KE_op = @at (Center, Center, Center) 1/2 * (u^2 + v^2 + w^2)
KE = AveragedField(KE_op, dims=(1, 2, 3))

function progress(sim)
    compute!(KE)

    @info @sprintf("Iteration: %d, time: %.2e, perturbation KE: %.2e",
                   sim.model.clock.iteration,
                   sim.model.clock.time,
                   KE[1, 1, 1])

    return nothing
end

simulation = Simulation(model,
                        Δt = Δt,
                        #stop_time = 10/surface_∂z_uˢ,
                        stop_iteration = 100,
                        iteration_interval = 10,
                        progress = progress)

function grow_instability!(simulation, energy)
    ## Initialize
    simulation.model.clock.iteration = 0
    t₀ = simulation.model.clock.time = 20seconds
    compute!(energy)
    energy₀ = energy[1, 1, 1]

    ## Grow
    run!(simulation)

    ## Analyze
    compute!(energy)
    energy₁ = energy[1, 1, 1]
    Δτ = simulation.model.clock.time - t₀

    ## ½(u² + v²) ~ exp(2 σ Δτ)
    σ = growth_rate = log(energy₁ / energy₀) / 2Δτ

    return growth_rate
end

function rescale!(model, energy; target_kinetic_energy=1e-6)
    compute!(energy)
    rescale_factor = √(target_kinetic_energy / energy[1, 1, 1])
    [parent(q) .*= rescale_factor for q in model.velocities]

    ν = model.closure.ν
    U = maximum(abs, model.velocities.v)
    L = 1 / k
    Re = U * L / ν
    @info @sprintf("Reynolds number after rescaling: %.2e", Re)

    return nothing
end

convergence(σ) = length(σ) > 1 ? abs((σ[end] - σ[end-1]) / σ[end]) : 9.1e18 # pretty big (not Inf tho)

function estimate_growth_rate(simulation, energy, ξ;
                              target_kinetic_energy = 1e-6,
                              convergence_criterion = 1e-2,
                              maximum_iterations = 100)

    σ, power_method_data = [], []

    compute!(ξ)
    push!(power_method_data, (; ξ=collect(interior(ξ)[:, 1, :])))
    iterations = 0

    while convergence(σ) > convergence_criterion && iterations < maximum_iterations
        compute!(energy)

        @info @sprintf("About to start power method iteration %d; kinetic energy: %.2e", length(σ)+1, energy[1, 1, 1])

        push!(σ, grow_instability!(simulation, energy))
        compute!(energy)

        @info @sprintf("Power method iteration %d, kinetic energy: %.2e, σⁿ: %.2e, relative Δσ: %.2e",
                       length(σ), energy[1, 1, 1], σ[end], convergence(σ))

        compute!(ξ)
        rescale!(simulation.model, energy; target_kinetic_energy)
        push!(power_method_data, (; ξ=collect(interior(ξ)[:, 1, :])))

        iterations += 1
    end

    return σ, power_method_data
end

#####
##### Estimate growth rate
#####

Re′ = 1e-2 # perturbation Re
@show KE′ = (Re′ * ν / h)^2

noise(x, y, z) = randn()
set!(model, u=noise, v=noise, w=noise)
rescale!(model, KE, target_kinetic_energy=KE′)

growth_rates, power_method_data = estimate_growth_rate(simulation, KE, ξ, target_kinetic_energy=KE′)
@info "Power iterations converged! Estimated growth rate: $(growth_rates[end])"

simulation.model.clock.time = 0.0
simulation.model.clock.iteration = 0
simulation.stop_time = 50 / surface_∂z_uˢ

# Re = U h / ν => U² = (ν * Re / h)²
rescale!(model, KE, target_kinetic_energy=KE′)

simulation.output_writers[:fields] = JLD2OutputWriter(model, merge(model.velocities, (; ξ)),
                                                      schedule = IterationInterval(10),
                                                      prefix = "craik_leibovich_instability",
                                                      field_slicer = nothing,
                                                      force = true)

run!(simulation)

#####
##### Visualize
#####

filepath = "craik_leibovich_instability.jld2"

vt = FieldTimeSeries(filepath, "v")
wt = FieldTimeSeries(filepath, "w")
ξt = FieldTimeSeries(filepath, "ξ")

times = vt.times
Nt = length(times)

vn(n) = Array(interior(vt[n])[1, :, :])
wn(n) = Array(interior(wt[n])[1, :, :])
ξn(n) = Array(interior(ξt[n])[1, :, :])

n = Node(1)

max_v = + maximum(abs, vn(Nt))
max_w = + maximum(abs, wn(Nt))
max_ξ = + maximum(abs, ξn(Nt))

min_v = - max_v
min_w = - max_w
min_ξ = - max_ξ

v = @lift vn($n)
w = @lift wn($n)
ξ = @lift ξn($n)

fig = Figure(resolution = (1200, 1200))

ax = Axis(fig[1, 1], title="v")
hm = heatmap!(ax, v, colorrange=(min_v, max_v), colormap=:balance)
cb = Colorbar(fig[1, 2], hm)

ax = Axis(fig[2, 1], title="w")
hm = heatmap!(ax, w, colorrange=(min_w, max_w), colormap=:balance)
cb = Colorbar(fig[2, 2], hm)

ax = Axis(fig[3, 1], title="∂z(v) - ∂y(w)")
hm = heatmap!(ax, ξ, colorrange=(min_ξ, max_ξ), colormap=:balance)
cb = Colorbar(fig[3, 2], hm)

title_str = @lift "Langmuir instability at t = " * prettytime(times[$n])
ax_t = fig[0, :] = Label(fig, title_str)

record(fig, "craik_leibovich.mp4", 1:Nt, framerate=8) do i
    @info "Plotting iteration $i of $Nt..."
    n[] = i
end

display(fig)


#####
##### Functions for non-rotating models
#####

@inline x_f_cross_U(i, j, k, grid::AbstractGrid{T}, ::Nothing, U) where T = zero(T)
@inline y_f_cross_U(i, j, k, grid::AbstractGrid{T}, ::Nothing, U) where T = zero(T)
@inline z_f_cross_U(i, j, k, grid::AbstractGrid{T}, ::Nothing, U) where T = zero(T)

#####
##### The 'FPlane' approximation. This is equivalent to a model with a constant
##### rotation rate around its vertical axis.
#####

"""
    FPlane{T} <: AbstractRotation

A parameter object for constant rotation around a vertical axis.
"""
struct FPlane{T} <: AbstractRotation
    f :: T
end

"""
    FPlane([T=Float64;] f)

Returns a parameter object for constant rotation at the angular frequency
`2f`, and therefore with background vorticity `f`, around a vertical axis.

Also called `FPlane`, after the "f-plane" approximation for the local effect of
Earth's rotation in a planar coordinate system tangent to the Earth's surface.
"""
function FPlane(T=Float64; f)
    return FPlane{T}(f)
end

@inline fv(i, j, k, grid::RegularCartesianGrid{T}, f, v) where T =
    T(0.5) * f * (avgy_f2c(grid, v, i-1,  j, k) + avgy_f2c(grid, v, i, j, k))

@inline fu(i, j, k, grid::RegularCartesianGrid{T}, f, u) where T =
    T(0.5) * f * (avgx_f2c(grid, u, i,  j-1, k) + avgx_f2c(grid, u, i, j, k))

@inline x_f_cross_U(i, j, k, grid, rotation::FPlane, U) = -fv(i, j, k, grid, rotation.f, U.v)
@inline y_f_cross_U(i, j, k, grid, rotation::FPlane, U) =  fu(i, j, k, grid, rotation.f, U.u)
@inline z_f_cross_U(i, j, k, grid::AbstractGrid{T}, ::FPlane, U) where T = zero(T)


using Test
using Bridge
using Bridge: expint
using LinearAlgebra

e1taylor(x) = @evalpoly(x,
-0.5772156649015329,
 1.0,
-0.25,
 0.05555555555555555,
-0.010416666666666666,
 0.0016666666666666666,
-0.0002314814814814815,
 2.834467120181406e-5,
-3.1001984126984127e-6,
 3.0619243582206544e-7,
-2.7557319223985888e-8) - log(x)

@test norm(e1taylor(1.0) - 0.21938393439552029) < 1e-8
@test norm(e1taylor(0.6) - 0.4543795031894021) < 1e-11
@test norm(e1taylor(0.05) - 2.467898488509974369559902) < eps()

@test norm(expint(3.0) - 0.013048381094197037) < eps()
@test norm(expint(1.0) - 0.21938393439552029) == 0
@test norm(expint(0.6) - 0.4543795031894021) == 0
@test norm(expint(0.05) - 2.467898488509974369559902)  == 0


#!/usr/bin/env julia
""" Kullback-Leibler divergence functions and klUCB utilities, in naive Julia code.

- How to use it? From Julia, it's easy:

```
julia> using KullbackLeibler
julia> KullbackLeibler.klBern(0.5, 0.5)
0.0
julia> KullbackLeibler.klBern(0.1, 0.9)
1.757779...
julia> KullbackLeibler.klBern(0.9, 0.1)  # And this KL is symmetric
1.757779...
julia> KullbackLeibler.klBern(0.4, 0.5)
0.020135...
julia> KullbackLeibler.klBern(0.01, 0.99)
4.503217...
```

- You can also use it with the [`Distributions`](https://juliastats.github.io/Distributions.jl/) module:

```julia
julia> using Distributions
julia> Bern1 = Distributions.Bernoulli(0.33)
Distributions.Bernoulli{Float64}(p=0.33)
julia> Bern2 = Distributions.Bernoulli(0.42)
Distributions.Bernoulli{Float64}(p=0.42)
julia> ex_kl_1 = KL( Bern1, Bern2 )  # Calc KL divergence for Bernoulli R.V
0.017063...
julia> klBern(0.33, 0.42)  # same!
0.017063...
```

- Note: there is also a Python version, if you need. (Cf. https://github.com/Naereen/Kullback-Leibler-divergences-and-kl-UCB-indexes/)

- MIT Licensed, (C) 2018 Lilian Besson (Naereen), https://GitHub.com/Naereen/KullbackLeibler.jl

- Cf. https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence
- Reference: [Filippi, Cappé & Garivier - Allerton, 2011](https://arxiv.org/pdf/1004.5229.pdf), [Garivier & Cappé, 2011](https://arxiv.org/pdf/1102.2490.pdf), and [Kullback & Leibler, 1951](http://www.jstor.org/stable/2236703).
"""

__author__ = "Lilian Besson"
__version__ = "0.1"

# Name of module
module KullbackLeibler

# List of exported functions
export klBern, klBin, klPoisson, klExp, klGamma, klNegBin, klGauss, KL, klucb, klucbBern, klucbGauss, klucbPoisson, klucbExp, klucbGamma

eps = 1e-15  #: Threshold value: everything in [0, 1] is truncated to [eps, 1 - eps]

# --- Simple Kullback-Leibler divergence for known distributions

doc"""
    function klBern(x, y)

Kullback-Leibler divergence for Bernoulli distributions. https://en.wikipedia.org/wiki/Bernoulli_distribution#Kullback.E2.80.93Leibler_divergence

$$\mathrm{KL}(\mathcal{B}(x), \mathcal{B}(y)) = x \log(\frac{x}{y}) + (1-x) \log(\frac{1-x}{1-y}).$$

```julia
julia> klBern(0.5, 0.5)
0.0
julia> klBern(0.1, 0.9)
1.757779...
julia> klBern(0.9, 0.1)  # And this KL is symmetric
1.757779...
julia> klBern(0.4, 0.5)
0.020135...
julia> klBern(0.01, 0.99)
4.503217...
```

- Special values:

```julia
julia> klBern(0, 1)  # Should be +Inf, but 0 --> eps, 1 --> 1 - eps
34.539575...
```
"""
function klBern(x, y)
    x = min(max(x, eps), 1 - eps)
    y = min(max(y, eps), 1 - eps)
    return x * log(x / y) + (1 - x) * log((1 - x) / (1 - y))
end

doc"""
    function klBin(x, y, n)

Kullback-Leibler divergence for Binomial distributions. https://math.stackexchange.com/questions/320399/kullback-leibner-divergence-of-binomial-distributions

- It is simply the n times `klBern` on x and y.

$$\mathrm{KL}(\mathrm{Bin}(x, n), \mathrm{Bin}(y, n)) = n \times \left(x \log(\frac{x}{y}) + (1-x) \log(\frac{1-x}{1-y}) \right).$$

- **Warning**: The two distributions must have the same parameter n, and x, y are p, q in (0, 1).

```julia
julia> klBin(0.5, 0.5, 10)
0.0
julia> klBin(0.1, 0.9, 10)
17.57779...
julia> klBin(0.9, 0.1, 10)  # And this KL is symmetric
17.57779...
julia> klBin(0.4, 0.5, 10)
0.20135...
julia> klBin(0.01, 0.99, 10)
45.03217...
```

- Special values:

```julia
julia> klBin(0, 1, 10)  # Should be +Inf, but 0 --> eps, 1 --> 1 - eps
345.39575...
```
"""
function klBin(x, y, n)
    x = min(max(x, eps), 1 - eps)
    y = min(max(y, eps), 1 - eps)
    return n * (x * log(x / y) + (1 - x) * log((1 - x) / (1 - y)))
end

doc"""
    function klPoisson(x, y)

Kullback-Leibler divergence for Poison distributions. https://en.wikipedia.org/wiki/Poisson_distribution#Kullback.E2.80.93Leibler_divergence

$$\mathrm{KL}(\mathrm{Poisson}(x), \mathrm{Poisson}(y)) = y - x + x \times \log(\frac{x}{y}).$$

```julia
julia> klPoisson(3, 3)
0.0
julia> klPoisson(2, 1)
0.386294...
julia> klPoisson(1, 2)  # And this KL is non-symmetric
0.306852...
julia> klPoisson(3, 6)
0.920558...
julia> klPoisson(6, 8)
0.273907...
```

- Special values:

```julia
julia> klPoisson(1, 0)  # Should be +Inf, but 0 --> eps, 1 --> 1 - eps
33.538776...
julia> klPoisson(0, 0)
0.0
```
"""
function klPoisson(x, y)
    x = max(x, eps)
    y = max(y, eps)
    return y - x + x * log(x / y)
end

doc"""
    function klExp(x, y)

Kullback-Leibler divergence for exponential distributions. https://en.wikipedia.org/wiki/Exponential_distribution#Kullback.E2.80.93Leibler_divergence

$$\mathrm{KL}(\mathrm{Exp}(x), \mathrm{Exp}(y)) = \begin{cases}
\frac{x}{y} - 1 - \log(\frac{x}{y}) & \text{if} x > 0, y > 0\\
+\infty & \text{otherwise}
\end{cases}$$

```julia
julia> klExp(3, 3)
0.0
julia> klExp(3, 6)
0.193147...
julia> klExp(1, 2)  # Only the proportion between x and y is used
0.193147...
julia> klExp(2, 1)  # And this KL is non-symmetric
0.306852...
julia> klExp(4, 2)  # Only the proportion between x and y is used
0.306852...
julia> klExp(6, 8)
0.037682...
```

- x, y have to be positive:

```julia
julia> klExp(-3, 2)
Inf
julia> klExp(3, -2)
Inf
julia> klExp(-3, -2)
Inf
```
"""
function klExp(x, y)
    if (x <= 0) || (y <= 0)
        return Inf
    else
        x = max(x, eps)
        y = max(y, eps)
        return x / y - 1 - log(x / y)
    end
end

doc"""
    function klGamma(x, y, a=1)

Kullback-Leibler divergence for gamma distributions. https://en.wikipedia.org/wiki/Gamma_distribution#Kullback.E2.80.93Leibler_divergence

- It is simply the a times `klExp` on x and y.

.. math::

    \mathrm{KL}(\Gamma(x, a), \Gamma(y, a)) = \begin{cases}
    a \times \left( \frac{x}{y} - 1 - \log(\frac{x}{y}) \right) & \text{if} x > 0, y > 0\\
    +\infty & \text{otherwise}
    \end{cases}

- **Warning**: The two distributions must have the same parameter a.

```julia
julia> klGamma(3, 3)
0.0
julia> klGamma(3, 6)
0.193147...
julia> klGamma(1, 2)  # Only the proportion between x and y is used
0.193147...
julia> klGamma(2, 1)  # And this KL is non-symmetric
0.306852...
julia> klGamma(4, 2)  # Only the proportion between x and y is used
0.306852...
julia> klGamma(6, 8)
0.037682...
```

- x, y have to be positive:

```julia
julia> klGamma(-3, 2)
Inf
julia> klGamma(3, -2)
Inf
julia> klGamma(-3, -2)
Inf
```
"""
function klGamma(x, y, a=1)
    if (x <= 0) || (y <= 0)
        return Inf
    else
        x = max(x, eps)
        y = max(y, eps)
        return a * (x / y - 1 - log(x / y))
    end
end

doc"""
    function klNegBin(x, y, r=1)

Kullback-Leibler divergence for negative binomial distributions. https://en.wikipedia.org/wiki/Negative_binomial_distribution

$$\mathrm{KL}(\mathrm{NegBin}(x, r), \mathrm{NegBin}(y, r)) = r \times \log((r + x) / (r + y)) - x \times \log(y \times (r + x) / (x \times (r + y))).$$

- **Warning**: The two distributions must have the same parameter r.

```julia
julia> klNegBin(0.5, 0.5)
0.0
julia> klNegBin(0.1, 0.9)
-0.711611...
julia> klNegBin(0.9, 0.1)  # And this KL is non-symmetric
2.0321564...
julia> klNegBin(0.4, 0.5)
-0.130653...
julia> klNegBin(0.01, 0.99)
-0.717353...
```

- Special values:

```julia
julia> klBern(0, 1)  # Should be +Inf, but 0 --> eps, 1 --> 1 - eps
    34.539575...
```

- With other values for `r`:

```julia
julia> klNegBin(0.5, 0.5, r=2)
0.0
julia> klNegBin(0.1, 0.9, r=2)
-0.832991...
julia> klNegBin(0.1, 0.9, r=4)
-0.914890...
julia> klNegBin(0.9, 0.1, r=2)  # And this KL is non-symmetric
2.3325528...
julia> klNegBin(0.4, 0.5, r=2)
-0.154572...
julia> klNegBin(0.01, 0.99, r=2)
-0.836257...
```
"""
function klNegBin(x, y, r=1)
    x = max(x, eps)
    y = max(y, eps)
    return r * log((r + x) / (r + y)) - x * log(y * (r + x) / (x * (r + y)))
end

doc"""
    function klGauss(x, y, sig2x=0.25, sig2y=0.25)

Kullback-Leibler divergence for Gaussian distributions of means ``x`` and ``y`` and variances ``sig2x`` and ``sig2y``, $\nu_1 = \mathcal{N}(x, \sigma_x^2)$ and $\nu_2 = \mathcal{N}(y, \sigma_x^2)$:

$$\mathrm{KL}(\nu_1, \nu_2) = \frac{(x - y)^2}{2 \sigma_y^2} + \frac{1}{2}\left( \frac{\sigma_x^2}{\sigma_y^2} - 1 \log\left(\frac{\sigma_x^2}{\sigma_y^2}\right) \right).$$

See https://en.wikipedia.org/wiki/Normal_distribution#Other_properties

- By default, sig2y is assumed to be sig2x (same variance).

```julia
julia> klGauss(3, 3)
0.0
julia> klGauss(3, 6)
18.0
julia> klGauss(1, 2)
2.0
julia> klGauss(2, 1)  # And this KL is symmetric
2.0
julia> klGauss(4, 2)
8.0
julia> klGauss(6, 8)
8.0
```

- x, y can be negative:

```julia
julia> klGauss(-3, 2)
50.0
julia> klGauss(3, -2)
50.0
julia> klGauss(-3, -2)
2.0
julia> klGauss(3, 2)
2.0
```

- With other values for `sig2x`:

```julia
julia> klGauss(3, 3, sig2x=10)
0.0
julia> klGauss(3, 6, sig2x=10)
0.45
julia> klGauss(1, 2, sig2x=10)
0.05
julia> klGauss(2, 1, sig2x=10)  # And this KL is symmetric
0.05
julia> klGauss(4, 2, sig2x=10)
0.2
julia> klGauss(6, 8, sig2x=10)
0.2
```

- With different values for `sig2x` and `sig2y`:

```julia
julia> klGauss(0, 0, sig2x=0.25, sig2y=0.5)
-0.0284...
julia> klGauss(0, 0, sig2x=0.25, sig2y=1.0)
0.2243...
julia> klGauss(0, 0, sig2x=0.5, sig2y=0.25)  # not symmetric here!
1.1534...

julia> klGauss(0, 1, sig2x=0.25, sig2y=0.5)
0.9715...
julia> klGauss(0, 1, sig2x=0.25, sig2y=1.0)
0.7243...
julia> klGauss(0, 1, sig2x=0.5, sig2y=0.25)  # not symmetric here!
3.1534...

julia> klGauss(1, 0, sig2x=0.25, sig2y=0.5)
0.9715...
julia> klGauss(1, 0, sig2x=0.25, sig2y=1.0)
0.7243...
julia> klGauss(1, 0, sig2x=0.5, sig2y=0.25)  # not symmetric here!
3.1534...
```

- **Warning:** Using [`Policies.klUCB`](https://smpybandits.github.io/docs/Policies.klUCB.html) (and variants) with `klGauss` is equivalent to use [`Policies.UCB`](https://smpybandits.github.io/docs/Policies.UCB.html), so prefer the simpler version. (**this is only for the Python version**)
"""
function klGauss(x, y, sig2x=0.25, sig2y=0.25)
    if - eps < (sig2y - sig2x) < eps
        return (x - y) ^ 2 / (2.0 * sig2x)
    else
        return (x - y) ^ 2 / (2.0 * sig2y) + 0.5 * ((sig2x/sig2y)^2 - 1 - log(sig2x/sig2y))
    end
end

# --- generic KL divergence functions for Distributions
# See https://github.com/Naereen/KullbackLeibler.jl/issues/2

using Distributions

doc"""
    function KL( D1, D2 )

Common function to compute the Kullback-Leibler for some continuous and discrete 1D parametric distributions (of module Distributions).

See [doc for continuous distributions](https://juliastats.github.io/Distributions.jl/latest/univariate.html#Continuous-Distributions-1) and [doc for discrete distributions](https://juliastats.github.io/Distributions.jl/latest/univariate.html#Discrete-Distributions-1).

Currently supported: `Bernoulli`, `Binomial`, `Poisson`, `NegativeBinomial` (discrete) and `Exponential`, `Gamma`, `Gaussian` (continuous).

Examples:

- Bernoulli:
```julia
julia> Bern1 = Distributions.Bernoulli(0.33)
Distributions.Bernoulli{Float64}(p=0.33)
julia> Bern2 = Distributions.Bernoulli(0.42)
Distributions.Bernoulli{Float64}(p=0.42)
julia> ex_kl_1 = KL( Bern1, Bern2 )  # Calc KL divergence for Bernoulli R.V
0.017063...
julia> klBern(0.33, 0.42)  # same!
0.017063...
```

- Binomial:
```julia
julia> Bin1 = Distributions.Binomial(13, 0.33)
Distributions.Binomial{Float64}(n=13, p=0.33)
julia> Bin2 = Distributions.Binomial(13, 0.42)   # must have same parameter n
Distributions.Binomial{Float64}(n=13, p=0.42)
julia> ex_kl_2 = KL( Bin1, Bin2 )  # Calc KL divergence for Binomial R.V
0.221828...
```

- Poisson:
```julia
julia> Pois1 = Distributions.Poisson(0.33)
Distributions.Poisson{Float64}(λ=0.33)
julia> Pois2 = Distributions.Poisson(0.92)
Distributions.Poisson{Float64}(λ=0.92)
julia> ex_kl_3 = KL( Pois1, Pois2 )  # Calc KL divergence for Poisson R.V
0.251657
```

- Exponential:
```julia
julia> Exp1 = Distributions.Exponential( 0.33 )
Distributions.Exponential{Float64}(θ=0.33)
julia> Exp2 = Distributions.Exponential( 0.42 )
Distributions.Exponential{Float64}(θ=0.42)
julia> ex_kl_4 = KL( Exp1, Exp2 )  # Calc KL div for Exponential R.V
0.0268763
```

- Gamma:
```julia
julia> Gam1 = Distributions.Gamma(0.5, 0.33)
Distributions.Gamma{Float64}(α=0.5, θ=0.33)
julia> Gam2 = Distributions.Gamma(0.5, 0.42)     # must have same parameter α
Distributions.Gamma{Float64}(α=0.5, θ=0.42)
julia> ex_kl_5 = KL( Gam1, Gam2 )  # Calc KL divergence for Gamma R.V
0.0134381
```

- NegativeBinomial:
```julia
julia> NegBin1 = Distributions.NegativeBinomial(40, 0.33)
Distributions.NegativeBinomial{Float64}(r=40, p=0.33)
julia> NegBin2 = Distributions.NegativeBinomial(40, 0.99)  # must have same parameter r
Distributions.NegativeBinomial{Float64}(r=40, p=0.99)
julia> ex_kl_6 = KL( NegBin1, NegBin2 )       # Calc KL divergence for NegativeBinomial R.V
17.277824
```

- Gaussian:
```julia
julia> Gauss1 = Distributions.Gaussian(0.33, 1)
Distributions.Normal{Float64}(μ=0.33, σ=1.0)
julia> Gauss2 = Distributions.Gaussian(0.42, 1)      # can have same σ
Distributions.Normal{Float64}(μ=0.42, σ=1.0)
julia> ex_kl_7 = KL( Gauss1, Gauss2 )  # Calc KL divergence for Gaussian R.V
0.0040499...

julia> Gauss3 = Distributions.Gaussian(0.42, 10)     # but can also have a different σ
Distributions.Normal{Float64}(μ=0.33, σ=10.0)
julia> ex_kl_7 = KL( Gauss1, Gauss3 )  # Calc KL divergence for Gaussian R.V
0.656697...
```
"""
function KL(D1, D2)
    throw(ErrorException("Unsupported distributions D1 or D2. Only Bernoulli, Binomial, Poisson, Exponential, Gamma, NegativeBinomial, Gaussian are supported."))
end

KL( D1::Distributions.Bernoulli, D2::Distributions.Bernoulli ) = klBern( D1.p, D2.p )

function KL( D1::Distributions.Binomial, D2::Distributions.Binomial )
    assert(D1.n == D2.n)
    return klBin( D1.p, D2.p, D1.n )
end

KL( D1::Distributions.Poisson, D2::Distributions.Poisson ) = klPoisson( D1.λ, D2.λ )

KL( D1::Distributions.Exponential, D2::Distributions.Exponential ) = klExp( D1.θ, D2.θ )

function KL( D1::Distributions.Gamma, D2::Distributions.Gamma )
    assert(D1.α == D2.α)
    return klGamma( D1.θ, D2.θ, D1.α )
end

function KL( D1::Distributions.NegativeBinomial, D2::Distributions.NegativeBinomial )
    assert(D1.r == D2.r)
    return klGamma( D1.p, D2.p, D1.r )
end

KL( D1::Distributions.Gaussian, D2::Distributions.Gaussian ) = klGauss( D1.μ, D2.μ, D1.σ, D2.σ )

# --- klucb functions, for the kl-UCB policy

doc"""
    function klucb(x, d, kl, upperbound, lowerbound=-Inf, precision=1e-6, max_iterations=50)

The generic kl-UCB index computation.

- x: value of the cum reward,
- d: upper bound on the divergence,
- kl: the KL divergence to be used (`klBern`, `klGauss`, etc),
- upperbound, lowerbound=-Inf: the known bound of the values x,
- precision=1e-6: the threshold from where to stop the research,
- max_iterations: max number of iterations of the loop (safer to bound it to reduce time complexity).

- **Note**: It uses a **bisection search**, and one call to ``kl`` for each step of the bisection search.

For example, for `klucbBern`, the two steps are to first compute an upperbound (as precise as possible) and the compute the kl-UCB index:

```julia
julia> x, d = 0.9, 0.2   # mean x, exploration term d
julia> upperbound = min(1.0, klucbGauss(x, d, sig2x=0.25))  # variance 1/4 for [0,1] bounded distributions
julia> upperbound
1.0
julia> klucb(x, d, klBern, upperbound, lowerbound=0, precision=1e-3, max_iterations=10)
0.9941...
julia> klucb(x, d, klBern, upperbound, lowerbound=0, precision=1e-6, max_iterations=10)
0.994482...
julia> klucb(x, d, klBern, upperbound, lowerbound=0, precision=1e-3, max_iterations=50)
0.9941...
julia> klucb(x, d, klBern, upperbound, lowerbound=0, precision=1e-6, max_iterations=100)  # more and more precise!
0.994489...
```

- **Note**: See below for more examples for different KL divergence functions.
"""
function klucb(x, d, kl, upperbound, lowerbound=-Inf, precision=1e-6, max_iterations=50)
    value = max(x, lowerbound)
    u = upperbound
    _count_iteration = 0
    while (_count_iteration < max_iterations) && (u - value > precision)
        _count_iteration += 1
        m = (value + u) / 2.0
        if kl(x, m) > d
            u = m
        else
            value = m
        end
    end
    return (value + u) / 2.0
end

doc"""
    function klucbBern(x, d, precision=1e-6)

kl-UCB index computation for Bernoulli distributions, using `klucb`.

- Influence of x:

```julia
julia> klucbBern(0.1, 0.2)
0.378391...
julia> klucbBern(0.5, 0.2)
0.787088...
julia> klucbBern(0.9, 0.2)
0.994489...
```

- Influence of d:

```julia
julia> klucbBern(0.1, 0.4)
0.519475...
julia> klucbBern(0.1, 0.9)
0.734714...

julia> klucbBern(0.5, 0.4)
0.871035...
julia> klucbBern(0.5, 0.9)
0.956809...

julia> klucbBern(0.9, 0.4)
0.999285...
julia> klucbBern(0.9, 0.9)
0.999995...
```
"""
function klucbBern(x, d, precision=1e-6)
    upperbound = min(1.0, klucbGauss(x, d, sig2x=0.25))  # variance 1/4 for [0,1] bounded distributions
    # upperbound = min(1.0, klucbPoisson(x, d))  # also safe, and better ?
    return klucb(x, d, klBern, upperbound, precision)
end

doc"""
    function klucbGauss(x, d, sig2x=0.25, precision=0.0)

kl-UCB index computation for Gaussian distributions.

- **Note**: it does not require any search.

- **Warning**: it works only if the good variance constant is given.

- Influence of x:

```julia
julia> klucbGauss(0.1, 0.2)
0.416227...
julia> klucbGauss(0.5, 0.2)
0.816227...
julia> klucbGauss(0.9, 0.2)
1.216227...

- Influence of d:

```julia
julia> klucbGauss(0.1, 0.4)
0.547213...
julia> klucbGauss(0.1, 0.9)
0.770820...

julia> klucbGauss(0.5, 0.4)
0.947213...
julia> klucbGauss(0.5, 0.9)
1.170820...

julia> klucbGauss(0.9, 0.4)
1.347213...
julia> klucbGauss(0.9, 0.9)
1.570820...
```

- **Warning**: Using [`Policies.klUCB`](https://smpybandits.github.io/docs/Policies.klUCB.html) (and variants) with `klucbGauss` is equivalent to use [`Policies.UCB`](https://smpybandits.github.io/docs/Policies.UCB.html), so prefer the simpler version (**this is only for the Python version**).
"""
function klucbGauss(x, d, sig2x=0.25, precision=0.0)
    return x + sqrt(2 * sig2x * d)
end

doc"""
    function klucbPoisson(x, d, precision=1e-6)

kl-UCB index computation for Poisson distributions, using `klucb`.

- Influence of x:

```julia
julia> klucbPoisson(0.1, 0.2)
0.450523...
julia> klucbPoisson(0.5, 0.2)
1.089376...
julia> klucbPoisson(0.9, 0.2)
1.640112...
```

- Influence of d:

```julia
julia> klucbPoisson(0.1, 0.4)
0.693684...
julia> klucbPoisson(0.1, 0.9)
1.252796...

julia> klucbPoisson(0.5, 0.4)
1.422933...
julia> klucbPoisson(0.5, 0.9)
2.122985...

julia> klucbPoisson(0.9, 0.4)
2.033691...
julia> klucbPoisson(0.9, 0.9)
2.831573...
```
"""
function klucbPoisson(x, d, precision=1e-6)
    upperbound = x + d + sqrt(d * d + 2 * x * d)  # looks safe, to check: left (Gaussian) tail of Poisson dev
    return klucb(x, d, klPoisson, upperbound, precision)
end

doc"""
    function klucbExp(x, d, precision=1e-6)

kl-UCB index computation for exponential distributions, using `klucb`.

- Influence of x:

```julia
julia> klucbExp(0.1, 0.2)
0.202741...
julia> klucbExp(0.5, 0.2)
1.013706...
julia> klucbExp(0.9, 0.2)
1.824671...
```

- Influence of d:

```julia
julia> klucbExp(0.1, 0.4)
0.285792...
julia> klucbExp(0.1, 0.9)
0.559088...

julia> klucbExp(0.5, 0.4)
1.428962...
julia> klucbExp(0.5, 0.9)
2.795442...

julia> klucbExp(0.9, 0.4)
2.572132...
julia> klucbExp(0.9, 0.9)
5.031795...
```
"""
function klucbExp(x, d, precision=1e-6)
    if d < 0.77
        # XXX where does this value come from?
        upperbound = x / (1 + 2.0 / 3 * d - sqrt(4.0 / 9 * d * d + 2 * d))
        # safe, klexp(x,y) >= e^2/(2*(1-2e/3)) if x=y(1-e)
    else
        upperbound = x * exp(d + 1)
    end
    if d > 1.61
        # XXX where does this value come from?
        lowerbound = x * exp(d)
    else
        lowerbound = x / (1 + d - sqrt(d * d + 2 * d))
    end
    return klucb(x, d, klGamma, upperbound, lowerbound, precision)
end

doc"""
    function klucbGamma(x, d, precision=1e-6)

kl-UCB index computation for Gamma distributions, using `klucb`.

- Influence of x:

```julia
julia> klucbGamma(0.1, 0.2)
0.202...
julia> klucbGamma(0.5, 0.2)
1.013...
julia> klucbGamma(0.9, 0.2)
1.824...
```

- Influence of d:

```julia
julia> klucbGamma(0.1, 0.4)
0.285...
julia> klucbGamma(0.1, 0.9)
0.559...

julia> klucbGamma(0.5, 0.4)
1.428...
julia> klucbGamma(0.5, 0.9)
2.795...

julia> klucbGamma(0.9, 0.4)
2.572...
julia> klucbGamma(0.9, 0.9)
5.031...
```
"""# FIXME this one is wrong!
function klucbGamma(x, d, precision=1e-6)
    if d < 0.77
        # XXX where does this value come from?
        upperbound = x / (1 + 2.0 / 3 * d - sqrt(4.0 / 9 * d * d + 2 * d))
        # safe, klexp(x,y) >= e^2/(2*(1-2e/3)) if x=y(1-e)
    else
        upperbound = x * exp(d + 1)
    end
    if d > 1.61
        # XXX where does this value come from
        lowerbound = x * exp(d)
    else
        lowerbound = x / (1 + d - sqrt(d * d + 2 * d))
    end
    # FIXME specify the value for a !
    return klucb(x, d, klGamma, max(upperbound, 1e2), min(-1e2, lowerbound), precision)
end

end

# --- Debugging

if "test" in ARGS
    # Code for debugging purposes.
    println("FIXME write these tests!")
    println("\nDone for tests of 'kullback_leibler.jl' ...")
end


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

    # center/radius/high/low/generators
    c = N[0, 0]
    r = N[0, 0]
    s = Singleton(c)
    @test center(s) == c
    @test radius_hyperrectangle(s) == r
    @test high(s) == r
    @test low(s) == -r
    for i in 1:2
        @test radius_hyperrectangle(s, i) == r[i]
        @test high(s, i) == r[i]
        @test low(s, i) == -r[i]
    end
    @test collect(generators(s)) == Vector{Vector{N}}()
    @test genmat(s) == Matrix{N}(undef, 2, 0)
    @test ngens(s) == 0

    # 1D singleton
    s = Singleton(N[1])
    d = N[1]
    @test σ(d, s) == N[1]
    d = N[-1]
    @test σ(d, s) == N[1]
    d = N[0]
    @test σ(d, s) == N[1]

    # 2D singleton
    s = Singleton(N[1, 2])
    d = N[1, 0]
    @test σ(d, s) == N[1, 2]
    d = N[-1, 0.5]
    @test σ(d, s) == N[1, 2]
    d = N[0, 0]
    @test σ(d, s) == N[1, 2]

    # test that abstract vectors can be used as well
    S = Singleton(SingleEntryVector(1, 1000, 0.5))
    @test element(S) == SingleEntryVector(1, 1000, 0.5)

    # boundedness
    @test isbounded(s)

    # element function
    @test element(s) == s.element
    for i in 1:2
        @test element(s, i) == s.element[i]
    end

    # isempty
    @test !isempty(s)

    # isuniversal
    answer, w = isuniversal(s, true)
    @test !isuniversal(s) && !answer && w ∉ s

    # membership
    S = Singleton(N[1, 1])
    @test N[0.9, 1.1] ∉ S
    @test N[1, 1] ∈ S

    # check membership error message
    @test_throws ErrorException S.element ⊆ S

    # an_element function
    @test an_element(S) ∈ S

    # vertices_list
    @test vertices_list(S)[1] ∈ S

    # radius_hyperrectangle
    @test iszero(radius_hyperrectangle(S))

    # high and low
    @test high(S) == N[1, 1]
    @test low(S) == N[1, 1]

    # concrete linear map
    M = N[0 1; -1 0]
    @test element(linear_map(M, S)) == an_element(M * S)

    # translation
    @test translate(s, N[1, 2]) == Singleton(N[2, 4])

    # subset
    s1 = Singleton(N[0, 1])
    s2 = Singleton(N[0, 3])
    p1 = VPolygon([N[0,0], N[0, 2]])
    p2 = VPolygon([N[0,0], N[0, 2], N[2, 0]])
    b = BallInf(N[0, 1], N(1))
    @test s1 ⊆ p1 && ⊆(s1, p1, true)[1]
    subset, point = ⊆(s2, p2, true)
    @test s2 ⊈ p2 && !subset && point ∈ s2 && point ∉ p2
    @test s1 ⊆ s1 && ⊆(s1, s1, true)[1]
    subset, point = ⊆(s1, s2, true)
    @test s1 ⊈ s2 && !subset && point ∈ s1 && point ∉ s2
    subset, point = ⊆(s1, b, true)
    @test subset && s1 ⊆ b
    subset, point = ⊆(s2, b, true)
    @test s2 ⊈ b && !subset && point ∈ s2 && point ∉ b

    # intersection
    S1 = Singleton(N[1, 1])
    S2 = Singleton(N[0, 0])
    S3 = ZeroSet{N}(2)
    H = BallInf(N[1, 1], N(0.5))
    M = LinearMap(N[1 0; 0 1], H)
    @test intersection(S1, H) == intersection(H, S1) == S1
    @test intersection(S2, H) == intersection(H, S2) == EmptySet{N}(2)
    @test intersection(S3, H) == intersection(H, S3) == EmptySet{N}(2)
    @test is_intersection_empty(S1, S2) && is_intersection_empty(S1, S2, true)[1]
    intersection_empty, point = is_intersection_empty(S2, S3, true)
    @test !is_intersection_empty(S2, S3) && !intersection_empty &&
        point ∈ S2 && point ∈ S3
    for X in [H, M]
        intersection_empty, point = is_intersection_empty(S1, X, true)
        @test !is_intersection_empty(S1, X) && !intersection_empty &&
            point ∈ S1 && point ∈ X
        @test is_intersection_empty(S2, X) &&
            is_intersection_empty(S2, X, true)[1]
    end
end