# Info object return after a SLOPE solve.
type SLOPESolveInfo
  runtime       # Runtime
  Aprods        # Number of products with A
  ATprods       # Number of products with A^T
  objPrimal     # Primal objective
  objDual  # Dual objective (possibly for infeasible dual point)
  infeas        # Dual infeasibility
  status        # Status: 1 = optimal, 2 = iterations
end

# Constants for exit status
const STATUS_RUNNING   = 0
const STATUS_OPTIMAL   = 1
const STATUS_ITERATIONS = 2
const STATUS_MSG = ["Optimal", "Iteration limit reached"]

# The default choice for lambdas is Benjamini-Hochberg sequence where q is the FDR:
using Distributions
NormalDistr = Normal()
bh_lambdafunc(i, p, q = 0.10) = quantile(NormalDistr, 1 - i*q/(2*p))

# Lets implement the stack-based FastProx SL1 alg of Bogdan et al to see if that helps...
# This is Alg 4 in the Bogdan2014 paper.
function fast_prox_sl1(yproxvals, lambdas)
  p = length(yproxvals)
  tuples = Any[]
  t = 0

  # Find optimal group levels
  for k in 1:p
    t += 1
    i, j, s = k, k, (yproxvals[k] - lambdas[k])
    w = abs(s)
    push!(tuples, (i, j, s, w))
    while (t > 1) && (tuples[t-1][4] .<= w)
      itm1, jtm1, stm1, wtm1 = tuples[t-1]
      it, jt, st, wt = tuples[t]
      wnew = (jtm1 - itm1 + 1)/(jt - itm1 + 1) * stm1 .+ (jt - it + 1)/(jt - itm1 + 1) * st
      tuples[t-1] = (itm1, jt, stm1 .+ st, abs(wnew))
      pop!(tuples) # Delete last entry
      t -= 1
    end
  end

  # Set entries in x for each block
  x = zeros(p)
  for l in 1:t
    for k in tuples[l][1]:tuples[l][2]
      x[k] = tuples[l][4]
    end
  end
  x
end

# The prox problem we solve assumes y are sorted and positive. We keep the permutation vector
# around so we can revert later.
sort_and_normalize(y) = begin
  absy = abs(y)
  return sortperm(absy, rev = true), sign(y), absy
end

# Go back from a sortperm to orig order.
function revert_sortperm(y, yperm)
  p = length(y)
  yr = zeros(p)
  for i in 1:p
    yi = yperm[i]
    yr[yi] = y[i]
  end
  yr
end

solve_prox_problem_with_fastproxsl1(yproxvals, lambdas) = begin
  yperm, sngy, absy = sort_and_normalize(yproxvals)
  ynormalized = absy[yperm]

  xvals = fast_prox_sl1(ynormalized, lambdas)

  # Reorder and put back the signs...
  sngy .* revert_sortperm(xvals, yperm)
end

# Test values for params
#iterations = 10000      # Maximum number of iterations
#q = 0.05
#verbosity = 1         # 0 = nothing, 1 = major, 2 = every
#optimIter = 1         # Iterations between optimality-condition checks
#gradIter = 20         # Iterations between full gradient computations
#tolInfeas = 1e-6       # Maximum allowed dual infeasibility
#tolRelGap = 1e-6       # Stopping criterion for relative primal-dual gap
#xInit = zeros(size(A,2))  # Initial value of x
#lambdas = nothing

# solve_SLOPE  Sorted L1 parameter estimation solver
#
# x, info = solve_SLOPE(A, b, lambda; options) solves the Slope problem
#
#     Minimize 1/2*||Ax-b||_2^2 + sum_i (lambda_i * |x|_[i])
#
# where |x|_[i] denotes the i-th largest entry in |x|. The entries in
# lambda must be nonnegative and in non-increasing order. When lambda is a
# scalar, the above formulation is reduced to the Lasso:
#
#     Minimize 1/2*||Ax-b||_2^2 + lambda * ||x||_1.
#
# Original MatLab version is:
# Copyright 2013, M. Bogdan, E. van den Berg, W. Su, and E.J. Candes
# Julia version by Robert Feldt, robert.feldt@gmail.com
# was based on Adlas.m in SLOPE Toolbox version 1.0 while the FastProxSL1 algorithm
# was implemented from scratch in Julia by Robert based on the Bogdan et al 2014 paper
# on SLOPE. We also excluded the Lasso from the implementation since there are already
# LASSO implementations for Julia.
#
#   The SLOPE Toolbox is free software: you can redistribute it
#   and/or  modify it under the terms of the GNU General Public License
#   as published by the Free Software Foundation, either version 3 of
#   the License, or (at your option) any later version.
#
#   The SLOPE Toolbox is distributed in the hope that it will
#   be useful, but WITHOUT ANY WARRANTY; without even the implied
#   warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#   See the GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with the SLOPE Toolbox. If not, see
#   <http://www.gnu.org/licenses/>.
#
function solve_SLOPE(A, b, lambdas = nothing;
  iterations = 10000,      # Maximum number of iterations
  q = 0.05,                 # Expected FDR, False Discovery Rate
  verbosity = 1,         # 0 = nothing, 1 = major, 2 = every
  optimIter = 1,         # Iterations between optimality-condition checks
  gradIter = 20,         # Iterations between full gradient computations
  tolInfeas = 1e-6,       # Maximum allowed dual infeasibility
  tolRelGap = 1e-6,       # Stopping criterion for relative primal-dual gap
  xInit = zeros(size(A,2))  # Initial value of x
  )

# Start timer
t0 = tic()

# Get problem dimension
n = size(A, 2)

# Create lambda sequence unless already done.
if lambdas == nothing
  lambdas = map(k -> bh_lambdafunc(k, n, q), 1:n)
end

# Ensure that lambda is non-increasing, nonnegative and has at least one non-zero entry.
if ((length(lambdas) > 1) && any(lambdas[2:end] .> lambdas[1:end-1]))
  throw("Lambdas must be non-increasing.")
end
if (lambdas[end] < 0)
  throw("Lambdas must be nonnegative")
elseif (lambdas[1] == 0)
  throw("Lambdas must have at least one nonnegative entry.")
end

# Get initial lower bound on the Lipschitz constant
x = randn(n,1)
x = x / norm(x,2)
x = A'*(A*x)
L = norm(x,2)

# Initialize parameters and iterates
t     = 1
eta    = 2
lambdas  = lambdas[:]
b     = b[:]
x     = xInit
y     = x
Ax    = A*x
fPrev  = Inf
iter   = 0

info = SLOPESolveInfo(
  0.0, # runtime is set after finishing
  2,   # Aprods
  1,   # ATprods
  0.0, # objPrimal is set later
  0.0, # info.objDual   is set later
  0.0, # Dual infeasibility
  STATUS_RUNNING
)

solve_prox = solve_prox_problem_with_fastproxsl1

if (verbosity > 0)
  @printf("%5s  %9s  %9s  %9s  %9s", "Iter", "||r||_2", "Gap", "Infeas.", "Rel. gap")
end

# Main loop
while (true)

  # Compute the gradient at f(y)
  if (mod(iter, gradIter) == 0) # Includes first iterations
   r = A*y - b
  else
   r = (Ax + ((tPrev - 1) / t) * (Ax - AxPrev)) - b
  end
  g = A' * r
  f = dot(r, r) / 2

  # Increment iteration count
  iter += 1

  # Check optimality conditions
  if ((mod(iter, optimIter) == 0))
    # Compute 'dual', check infeasibility and gap
    gs    = sort(abs(g), rev = true)
    ys    = sort(abs(y), rev = true)
    info.infeas = max(maximum(cumsum(gs-lambdas)),0)

    # Compute primal and dual objective
    info.objPrimal = ( f + lambdas'*ys)[1]
    info.objDual   = (-f - r'*b       )[1]

    if (verbosity > 0)
      str = @sprintf("  %9.2e  %9.2e  %9.2e", info.objPrimal - info.objDual, info.infeas/lambdas[1],
        abs(info.objPrimal - info.objDual) / max(1,info.objPrimal))
    end

    # Check primal-dual gap
    if ((abs(info.objPrimal - info.objDual)/max(1,info.objPrimal) < tolRelGap) &&
        (info.infeas < tolInfeas * lambdas[1]))
      info.status = STATUS_OPTIMAL;
    end
  else
    str = ""
  end

  if ((verbosity == 2) || ((verbosity == 1) && (mod(iter,optimIter) == 0)))
    @printf("%5d  %9.2e%s\n", iter, f, str)
  end

  # Stopping criteria
  if (info.status == STATUS_RUNNING)
    if (iter >= iterations)
      info.status = STATUS_ITERATIONS
    end
  end

  if (info.status != 0)
    if (verbosity > 0)
      @printf("Exiting with status %d -- %s\n", info.status, STATUS_MSG[info.status])
    end
    break;
  end

  # Keep copies of previous values
  AxPrev = Ax
  xPrev  = x
  fPrev  = f
  tPrev  = t

  # Lipschitz search
  while (true)
    # Compute prox mapping
    x = solve_prox(y - (1/L)*g, lambdas/L)
    d = x - y

    Ax = A*x
    r  = Ax-b
    f  = dot(r, r)/2
    q  = fPrev + dot(d, g) + (L/2)*dot(d, d)

    info.Aprods += 1

    if (q >= f*(1-1e-12))
      break
    else
      L = L * eta
    end
  end

  # Update
  t = (1 + sqrt(1 + 4*t^2)) / 2
  y = x + ((tPrev - 1) / t) * (x - xPrev)
end

# Final touchup of info values
info.runtime = toc()
info.Aprods = info.Aprods + ceil(iter / gradIter)
info.ATprods  = info.ATprods + iter

# ...and we are finished!
return y, info

end # function solve_SLOPE

function rand_sparse_problem(p; numactive = iceil(log(p)),
  n = int(p/2), amplitude = 4.1, sigma = 1.0, shuffle = true)

  beta = vcat(amplitude*ones(numactive)+randn(numactive), zeros(p-numactive))
  indices = collect(1:p)
  if shuffle
    shuffle!(indices)
    beta = beta[indices]
  end
  X = randn(n, p)
  errors = sigma*randn(n)
  y = X * beta .+ errors
  return beta, X, y, indices, n, p, numactive, sigma, amplitude, errors
end

# Find the sharpest difference between the found values and use that as the estimate
# of the num active values.
function estimate_active(xhat)
  relsize = abs(xhat) / maximum(abs(xhat))
  relsizeperm = sortperm(relsize, rev=true)
  srelsize = relsize[relsizeperm]
  relfactors = srelsize[1:end-1] ./ srelsize[2:end]
  numactive = find(rf -> rf == maximum(relfactors), relfactors)[1]
  return relsizeperm[1:numactive]
end

# First estimate the number of active coefficients by a SLOPE solve then
# do a linreg on only those factors.
function sparse_regression_SLOPE(A, b; iterations = 250, options...)
  xhat, info = solve_SLOPE(A, b; iterations = iterations, options...)
  active = estimate_active(xhat)
  Asub = A[:, active]
  res = linreg(Asub, b)
  xhatsub = zeros(size(A, 2))
  xhatsub[active] = res[2:end]
  return xhatsub, res[2:end], active
end

if length(ARGS) >= 1
  numcoefs = int(ARGS[1])
else
  numcoefs = 1000
end
if length(ARGS) >= 2
  numact = int(ARGS[2])
else
  numact = 10
end
if length(ARGS) >= 3
  maxiterations = int(ARGS[3])
else
  maxiterations = 500
end

x, A, b, indices, n, p, numactive, sigma, amplitude, errors = rand_sparse_problem(numcoefs;
  shuffle = false, numactive = numact);

# Split in train and test set
idxs = shuffle(collect(1:n))
trainsize = ifloor(0.80*n)
trainidxs = idxs[1:trainsize]
testidxs = idxs[(trainsize+1):end]
Atrain = A[trainidxs, :]
btrain = b[trainidxs]
Atest = A[testidxs, :]
btest = b[testidxs]

xhat, xselected, selected = sparse_regression_SLOPE(Atrain, btrain; iterations = maxiterations)

println(x[1:numactive])
println(xhat[1:numactive])
println("Norm(x .- xhat) = ", norm(x .- xhat))
mape(y, yhat) = mean(100.0 * abs((yhat - y)/y))
@printf("Num active = %d, Num active selected = %d\n", numactive, length(selected))
@printf("Train MAPE = %.2f%%\n", mape(btrain, Atrain*xhat))
@printf("Test MAPE = %.2f%%\n", mape(btest, Atest*xhat))


using InvariantMeasures
using ValidatedNumerics

using Plots
using LaTeXStrings
using StatsPlots

pgfplotsx()

function onegrid(T, Btype, size)

    time_assembling = @elapsed begin
        D = Mod1Dynamic(T)
        # different backend, a tad slower
        # D = mod1_dynamic(x -> x->2*x+0.5*x*(1-x))
        B = Btype(size)
        Q = DiscretizedOperator(B, D)
    end

    time_norms = @elapsed norms = powernormbounds(B, D; Q=Q)
    time_eigen = @elapsed w = invariant_vector(B, Q)
    time_error = @elapsed error = distance_from_invariant(B, D, Q, w, norms)

    times = [time_error time_eigen time_norms time_assembling 0]

    return times, error, (B, D, norms, time_assembling+time_norms)
end

function twogrid(Btype, size, (B, D, norms, time_coarse))
    time_assembling_fine = @elapsed begin
        B_fine = Btype(size)
        Q_fine = DiscretizedOperator(B_fine, D)
    end

    time_norms_fine = @elapsed norms_fine = finepowernormbounds(B, B_fine, D, norms; Q_fine=Q_fine)
    time_eigen_fine = @elapsed w_fine = invariant_vector(B_fine, Q_fine)
    time_error_fine = @elapsed error_fine = distance_from_invariant(B_fine, D, Q_fine, w_fine, norms_fine)

    times_fine = [time_error_fine time_eigen_fine time_norms_fine time_assembling_fine time_coarse]
    return times_fine, error_fine
end

function time_convergence_plot(T, Btype, k_onegrid, k_twogrid)
    n_onegrid = 2 .^ k_onegrid

    times_onegrid = fill(NaN, length(n_onegrid), 5)
    errors_onegrid = fill(NaN, size(n_onegrid))
    onegrid(T, Btype, n_onegrid[1]) #warmup
    for i in 1:length(n_onegrid)
        times_onegrid[i,:], errors_onegrid[i], _ = onegrid(T, Btype, n_onegrid[i])
    end

    n_twogrid = 2 .^ k_twogrid
    _, _, coarse_data = onegrid(T, Btype, 1024)
    times_twogrid = fill(NaN, length(n_twogrid), 5)
    errors_twogrid = fill(NaN, size(n_twogrid))

    twogrid(Btype, n_twogrid[1], coarse_data) #warmup
    for i in 1:length(n_twogrid)
        times_twogrid[i,:], errors_twogrid[i] = twogrid(Btype, n_twogrid[i], coarse_data)
    end

    p1 = groupedbar(
        times_onegrid,
        bar_position = :stack,
        legend = false,
        label = ["err" "eigen" "norm" "matrix" "coarse"],
        title = "CPU time breakdown (s)",
        xticks = (1:length(n_onegrid), LaTeXString.(raw"2^{" .* string.(k_onegrid) .* raw"}")),
        link = :y,
    )

    p2 = groupedbar(
        times_twogrid,
        bar_position = :stack,
        legend = false,
        label = ["err" "eigen" "norm" "matrix" "coarse"],
        title = "CPU time breakdown (s)",
        xticks = (1:length(n_twogrid), LaTeXString.(raw"2^{" .* string.(k_twogrid) .* raw"}")),
        link = :y,
    )

    p3 = plot(
        n_onegrid,
        errors_onegrid,
        title = "Error",
        mark = :dot,
        yscale = :log10,
        xscale = :log10,
        xticks = (n_onegrid, LaTeXString.(raw"2^{" .* string.(k_onegrid) .* raw"}" )),
        label = "One-grid strategy",
        legend = :bottomleft,
        link = :y,
    )

    p4 = plot(
        n_twogrid,
        errors_twogrid,
        title = "Error",
        mark = :dot,
        yscale = :log10,
        xscale = :log10,
        color = :red,
        xticks = (n_twogrid, LaTeXString.(raw"2^{" .* string.(k_twogrid) .* raw"}")),
        label = "Two-grid strategy",
        legend = :bottomleft,
        link = :y,
    )

    p = plot(p1, p2, p3, p4)

end

p = time_convergence_plot(x->2x+0.5x*(1-x), Ulam, 8:11, 11:14)

# large-scale version:
# p = time_convergence_plot(x->2x+0.5x*(1-x), Ulam, 8:14, 13:26); savefig(p, "time_convergence_plot.tikz"); savefig(p, "time_convergence_plot.pdf");

# savefig(p, "time_convergence_plot.tikz")


### A Pluto.jl notebook ###
# v0.16.1

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 02329bbc-15a4-11ec-2522-df175bfb53e0
begin
    import Pkg
    # careful: this is _not_ a reproducible environment
    # activate the global environment
    Pkg.activate()

    using PlutoUI
	using FlightGNC
	using ComponentArrays, Transducers
	using Plots
end

# ╔═╡ 50becda1-a0e6-4e4c-98ca-a72992a5fc42
include("main_missiles.jl")

# ╔═╡ b9c304e7-9034-44d4-a4a6-a5d34d0baf58
begin
	include("sim_missiles.jl")

	# To avoid redefinition of a named (or an anonymous) function at each time tuning parameters and running simulation,
	# place the custom definition of a bias command function in a different cell from the cell containing sim_plot.

	custom_bias(s,x,t) = zeros(3)                           # option 1
	# function custom_bias(s,x,t)  return zeros(3) end      # option 2
	# custom_bias = (s,x,t) -> zeros(3)                     # option 3
end

# ╔═╡ b4aa1cc3-07de-4717-a39e-8a9c4b1c8512
md"""
# Missile Guidance Simulation
2D/3D Vector-Form Biased Proportional Navigation Guidance Law is implemented.
The bias command currently implemented considers target being stationary.

## Vector BPNG Equations
* Three-dimensional lead angle
$\sigma = \cos^{-1}\left(\hat{\mathbf{r}} \cdot \hat{\mathbf{v}}\right) \in \left[0,\pi\right]$

* Manoeuvre plane unit normal vector
$\hat{\mathbf{k}} = \frac{\mathbf{r} \times \mathbf{v}}{\left\|\mathbf{r} \times \mathbf{v}\right\|}$

* Predicted final velocity direction vector
$\begin{aligned}
	\hat{\mathbf{v}}_{f_{pred}} &= \hat{\mathbf{v}} \cos\left(\frac{N}{N-1}\sigma\right) - \left(\hat{\mathbf{k}}\times\hat{\mathbf{v}}\right)\sin\left(\frac{N}{N-1}\sigma\right)\\
	&= \hat{\mathbf{r}} \cos\left(\frac{1}{N-1}\sigma\right) - \left(\hat{\mathbf{k}}\times\hat{\mathbf{r}}\right)\sin\left(\frac{1}{N-1}\sigma\right)
\end{aligned}$

* Predicted final velocity direction error angle
$e = \cos^{-1}\left(\hat{\mathbf{v}}_{f_{pred}} \cdot \hat{\mathbf{v}}_{f_{d}}\right) \in \left[0,\pi\right]$

* Desired angular velocity of predicted final velocity direction vector
$\boldsymbol{\omega}_{f_{d}} = -K_{e}\left(r\right) \dot{r} e   \frac{\hat{\mathbf{v}}_{f_{pred}} \times \hat{\mathbf{v}}_{f_{d}}}{\left\|\hat{\mathbf{v}}_{f_{pred}} \times \hat{\mathbf{v}}_{f_{d}}\right\|} + K_{roll}\left(r\right) \hat{\mathbf{v}}_{f_{d}}$

* Bias angular velocity command
$\boldsymbol{\omega}_{bias} = K_{\sigma}\left(\sigma\right) \left[-\left(N-1\right) \left(\boldsymbol{\omega}_{f_{d}}\cdot \hat{\mathbf{k}}\right) \hat{\mathbf{k}}  + \sin\sigma \,\boldsymbol{\omega}_{f_{d}} \cdot \left\{\hat{\mathbf{r}} + \cot\left(\frac{1}{N-1}\sigma\right)\left(\hat{\mathbf{k}}\times\hat{\mathbf{r}}\right)\right\} \left(\hat{\mathbf{v}} \times \hat{\mathbf{k}}\right)\right]$

* LOS angular velocity
$\boldsymbol{\omega}_{r}^{r\perp} = \frac{\mathbf{v} \times \mathbf{r}}{\left\| \mathbf{r} \right\|^{2}} = -\frac{v\sin\sigma}{r}\hat{\mathbf{k}}$

* Lateral acceleration command
$\mathbf{a}_{BPNG}^{v\perp} = \left(N\boldsymbol{\omega}_{r}^{r\perp} + \boldsymbol{\omega}_{bias}\right) \times \mathbf{v}$

"""

# ╔═╡ 9fcf0c9b-d559-402f-883c-c53f61815c94
begin
	md"""
	 $\alpha$ = $(@bind α Slider(0.01:0.01:1; default=1, show_value=true))

	 $n$ = $(@bind n Slider(0 : 0.1 : 10; default=1, show_value=true))

	 $\gamma_{f_{d}}$ [deg] = $(@bind γ_f_d_deg Slider(-90:1:90; default=0, show_value=true))

	 $\chi_{f_{d}}$ [deg] = $(@bind χ_f_d_deg Slider(0:1:180; default=180, show_value=true))

	 $\sigma_{M_{\lim}}$ [deg] = $(@bind σ_M_lim_deg Slider(0:1:90; default=60, show_value=true))
	"""
end

# ╔═╡ b75636f4-b7f2-4356-9e2a-d90942c0071e
begin
	s_BPNG.Bias    = Bias_IACG_StationaryTarget
	s_BPNG.s_Bias  = ComponentArray(α = 1, δ = 0.01, n = n, r_ref = 10E3, k = 9, m = 10)
	s_BPNG.σ_M_lim = deg2rad(σ_M_lim_deg)
	γ_f_d = deg2rad(γ_f_d_deg)
	χ_f_d = deg2rad(χ_f_d_deg)
	s_BPNG.v̂_f_d = [cos(γ_f_d)*sin(χ_f_d); cos(γ_f_d)*cos(χ_f_d); sin(γ_f_d)]
	df, f = sim_plot(p_M_0, v_M_0, p_T_0, v_T_0, s_BPNG, A_M_max)
	f
end

# ╔═╡ bce435b0-c6d4-4ac7-93aa-727ed08c73e5
begin
	t = df.time
	r = df.sol |> Map(datum -> datum.r) |> collect
	plot(t,r, label = :false)
	annotate!((0.9,0.9), "min r = "*string(minimum(r)) )
end

# ╔═╡ Cell order:
# ╟─b4aa1cc3-07de-4717-a39e-8a9c4b1c8512
# ╟─02329bbc-15a4-11ec-2522-df175bfb53e0
# ╠═50becda1-a0e6-4e4c-98ca-a72992a5fc42
# ╠═b9c304e7-9034-44d4-a4a6-a5d34d0baf58
# ╟─9fcf0c9b-d559-402f-883c-c53f61815c94
# ╠═b75636f4-b7f2-4356-9e2a-d90942c0071e
# ╟─bce435b0-c6d4-4ac7-93aa-727ed08c73e5


using DifferentialDynamicProgramming
using Test, Statistics, LinearAlgebra

@info("Compile time is high for this package, this is expected and is not an error.")
# write your own tests here

include("test_readme.jl")
demo_linear()
demo_linear_kl(kl_step=100)
demo_pendcart()

demoQP()

# include(Pkg.dir("GuidedPolicySearch","examples","bb.jl"))


module MCOPlots

    # ----------------------------------------------------------------------------------
    #   Imports
    # ----------------------------------------------------------------------------------
    using MonteCarloObservable; const MCO = MonteCarloObservable;
    using StatsBase
    using PyPlot

    import PyPlot: plot

    # ----------------------------------------------------------------------------------
    #   Plots
    # ----------------------------------------------------------------------------------
    """
        plot(obs::Observable{T}[; errors=true, digits=3])

    Plot the observable's time series.
    """
    function plot(obs::Observable{T}; errors=true, digits=3) where T
        ts = timeseries(obs)
        Xmean = mean(obs)
        err = error(obs)
        Xstd = std(ts)

        fig, ax = subplots(1,1)
        ax[:plot](ts, ".-")
        ax[:set_ylabel](name(obs))
        ax[:set_xlabel]("Monte Carlo time \$ t \$")
        # ax[:set_yticks]([])
        ax[:axhline](Xmean, color="black", label="\$ $(round.(Xmean, digits=digits)) \$ (mean)", linewidth=2.0)

        if errors
            ax[:axhline](Xmean+err, color="r", label="\$ \\pm $(round.(err, digits=digits)) \$ (σ error)", linewidth=2.0)
            ax[:axhline](Xmean-err, color="r", linewidth=2.0)

            ax[:axhline](Xmean+2*err, color="r", alpha=.3, label="\$ \\pm $(round.(2*err, digits=digits)) \$ (2σ error)", linewidth=2.0)
            ax[:axhline](Xmean-2*err, color="r", alpha=.3, linewidth=2.0)

            ax[:axhline](Xmean+Xstd, color="g", label="\$ \\pm $(round.(Xstd, digits=digits)) \$ (std)", linewidth=2.0)
            ax[:axhline](Xmean-Xstd, color="g", linewidth=2.0)
        end

        ax[:legend](frameon=true, loc="best")

        tight_layout()
        nothing
    end

    ################################################
    """
        hist(obs::Observable{T}[; errors=true, digits=3])

    Plot a histogram of the observable's time series.
    """
    function hist(obs::Observable{T}; errors=true, digits=3, customtitle="") where T
        ts = timeseries(obs)
        Xmean = mean(obs)
        err = error(obs)
        Xstd = std(ts)

        fig, ax = subplots(1,1)
        ax[:hist](ts, 50, color="gray", alpha=.5, density=1)
        ax[:set_ylabel]("Frequency")
        ax[:set_xlabel](name(obs))
        ax[:set_yticks]([])
        ax[:axvline](Xmean, color="black", label="\$ $(round.(Xmean, digits=digits)) \$ (mean)", linewidth=2.0)

        if errors
            ax[:axvline](Xmean+err, color="r", label="\$ \\pm $(round.(err, digits=digits)) \$ (σ error)", linewidth=2.0)
            ax[:axvline](Xmean-err, color="r", linewidth=2.0)

            ax[:axvline](Xmean+2*err, color="r", alpha=.3, label="\$ \\pm $(round.(2*err, digits=digits)) \$ (2σ error)", linewidth=2.0)
            ax[:axvline](Xmean-2*err, color="r", alpha=.3, linewidth=2.0)

            ax[:axvline](Xmean+Xstd, color="g", label="\$ \\pm $(round.(Xstd, digits=digits)) \$ (std)", linewidth=2.0)
            ax[:axvline](Xmean-Xstd, color="g", linewidth=2.0)
        end

        ax[:legend](frameon=true, loc="best")

        if customtitle != ""
            title(customtitle)
        end

        tight_layout()
        nothing
    end

    ################################################
    """
        corrplot(obs::Observable{T})

    Plot the autocorrelation function of the observable.
    """
    function corrplot(obs::Observable{T}; showtau=false) where T
        ts = timeseries(obs)

        fig, ax = subplots(1,1)
        ax[:plot](autocor(ts), "-", color="k", linewidth=2.0)
        ax[:set_xlabel]("Monte Carlo time \$ t \$")
        ax[:set_ylabel]("Autocorrelation of $(name(obs))")

        if showtau
            t = tau(obs)
            tfinder = tau(obs, MCO.Rplateaufinder(obs)[2])
            tsum = sum(StatsBase.autocor(ts))
            ax[:axvline](x=t, color="red", label="tau ≈ $(round(t,2))")
            ax[:axvline](x=tfinder, color="orange", label="tau (finder) ≈ $(round(tfinder,2))")
            ax[:axvline](x=tsum, color="green", label="tau (sum) ≈ $(round(tsum,2))")
            # @show sum(StatsBase.autocor(ts))
            # @show Rplateaufinder(obs)[2]
            # @show tau(obs)
            ax[:axhline](y=exp(-1), color="gray", alpha=.2, label="1/e")
            ax[:legend]()
        end

        tight_layout()
        nothing
    end

    ################################################
    """
        binningplot(obs::Observable{T}[; min_nbins=32])

    Creates a plot of the binning error coefficient `R` as a function of bin size.

    The coefficient `R` should (up to statistical fluctuations) show a plateau for larger bin sizes,
    indicating that the bin averages have become independent.
    For correlated data one has `R>≈1` and `sqrt(R)` quantifies how much one would have underestimated
    the one-sigma errorbar.

    See [`binning_error`](@ref).
    """
    function binningplot(obs::Observable{T}; min_nbins=50) where T
        ts = timeseries(obs)

        bss, R, means = MCO.R_function(ts, min_nbins=min_nbins)

        figure()
        plot(bss, R, "m.-", label="R")
        plot(bss, means, ".-", label="R means")
        ylabel("error coefficient")
        legend()
        xlabel("bin size")
        tight_layout();
        nothing
    end

    ################################################
    """
        errorplot(obs::Observable{T})

    Creates a plot of the statistical standard error as a function of bin size.

    The standard error should (up to statistical fluctuations) show a plateau for larger bin sizes,
    indicating that the bin averages have become independent.

    See [`binning_error`](@ref), [`error`](@ref).
    """
    function errorplot(obs::Observable; customtitle::String="")
        ts = timeseries(obs)

        bss, R, means = MCO.R_function(ts, min_nbins=50)
        rawerrors = [MCO.binning_error_from_R(ts, r) for r in R]
        meanerrors = [MCO.binning_error_from_R(ts, r) for r in means]

        figure()
        plot(bss, rawerrors, "m.-", label="raw")
        plot(bss, meanerrors, ".-", label="means")
        ylabel("error")
        legend()
        xlabel("bin size")
        if customtitle != ""
            title(customtitle)
        end
        tight_layout();
        nothing
    end

    # ----------------------------------------------------------------------------------
    #   Exports
    # ----------------------------------------------------------------------------------
    export plot, hist, binningplot, errorplot, corrplot

end

    # # --------------------------------------
    # #           Playground
    # # --------------------------------------
    # # using Plots instead of PyPlot

    # function plot_histogram_Plots(obs::Observable{T}; errors=true, digits=3) where T
    #     ts = timeseries(obs)
    #     Xmean = mean(obs)
    #     err = std(obs)
    #     Xstd = std(ts)

    #     histogram(ts, framestyle=:box, grid=false, normed=true, label="", color="lightgrey")
    #     # ax[:hist](ts, 50, color="gray", alpha=.5, normed=1)
    #     ylabel!("Frequency")
    #     xlabel!(name(obs))
    #     yticks!(Float64[])
    #     vline!([Xmean], color="black", linewidth=2.0, label="\$ $(round.(Xmean, digits=digits)) \$ (mean)")

    #     if errors
    #         vline!([Xmean+err, Xmean-err], color="red", linewidth=2.0, label="\$ \\pm $(round.(err, digits=digits)) \$ (σ error)")
    #         vline!([Xmean+2*err, Xmean-2*err], color="red", linewidth=2.0, label="\$ \\pm $(round.(2*err, digits=digits)) \$ (2σ error)", alpha=.3)
    #         vline!([Xmean+Xstd, Xmean-Xstd], color="green", linewidth=2.0, label="\$ \\pm $(round.(Xstd, digits=digits)) \$ (std)", alpha=.3)
    #     end
    #     nothing
    # end


using SeaPearl

#

struct InputData
    A::Matrix{Int} #number of lines
end

"""
    model_latin(input_file; limit=1)

return the SeaPearl model for to the Latin completion problem using SeaPearl.AllDifferent
The input file can come from the Latin DataGen in SeaPearl.

# Arguments
- `InputData`: Matrix containting the uncompleted game
- 'limit' : Int, giving the number of solutions after which it will stop searching. if nothing given, it will looks for all the solutions
"""
function model_latin(input::InputData; limit=1)
    trailer = SeaPearl.Trailer()
    model = SeaPearl.CPModel(trailer)
    model.limit.numberOfSolutions = limit
    A = InputData.A
    N = size(A)[1]
    puzzle = Matrix{SeaPearl.AbstractIntVar}(undef, N,N)
    for i = 1:N
        for j in 1:N
            puzzle[i,j] = SeaPearl.IntVar(1, N, "puzzle_"*string(i)*","*string(j), trailer)
            SeaPearl.addVariable!(cpmodel, puzzle[i,j]; branchable=true)
            if A[i,j]>0 push!(model.constraints,SeaPearl.EqualConstant(puzzle[i,j], A[i,j], trailer)) end
        end
    end
    for i in 1:N
        push!(cpmodel.constraints, SeaPearl.AllDifferent(puzzle[i,:], cpmodel.trailer))
        push!(cpmodel.constraints, SeaPearl.AllDifferent(puzzle[:,i], cpmodel.trailer))
    end
    return model
end

"""
    solve_latin!(model::SeaPearl.CPModel; variableSelection=SeaPearl.MinDomainVariableSelection{false}(), valueSelection=SeaPearl.BasicHeuristic())

Solve the SeaPearl model for to the Latin Completion problem, using an existing model

# Arguments
- `model::SeaPearl.CPModel`: model (from model_latin)
- 'variableSelection': SeaPearl variable selection. By default: SeaPearl.MinDomainVariableSelection{false}().
- 'valueSelection': SeaPearl value selection. By default: =SeaPearl.BasicHeuristic().
"""
function solve_latin!(model::SeaPearl.CPModel; variableSelection=SeaPearl.MinDomainVariableSelection{false}(), valueSelection=SeaPearl.BasicHeuristic())
    status = @time SeaPearl.solve!(model; variableHeuristic=variableSelection, valueSelection=valueSelection)
    return model
end

"""
    print_latin(model::SeaPearl.CPModel)

Nice print the latin game when solved

# Arguments
- `model::SeaPearl.CPModel`: model (from model_latin)
"""
function print_latin(model)
    sol= model.statistics.solutions[1]
    n = oftype(1,sqrt(length(sol)))
    tableau = Matrix{Int}(undef,n,n)
    for i in 1:n, j in 1:n
        tableau[i, j]= sol["puzzle_"*string(i)*","*string(j)]
    end
    print(" ")
    for k in 1:5*n
        print("-")
    end
    println()
    for i in 1:n
        print("|")
        for j in 1:n
            printstyled(" "*string(tableau[i,j],pad=2)*" ", color=tableau[i,j])
            print("|")
        end
        println()
    end
    print(" ")
    for k in 1:5*n
        print("-")
    end
end


### A Pluto.jl notebook ###
# v0.12.11

using Markdown
using InteractiveUtils

# ╔═╡ b4949654-ff30-11ea-225c-2ffc47a35a31
using Pkg, DrWatson

# ╔═╡ b494da42-ff30-11ea-25d2-37857c757c9f
begin
	#@quickactivate "StatisticalRethinkingStan"
	using StanSample
	using StatisticalRethinking
	using PlutoUI
end

# ╔═╡ b0e1bc4a-ff2f-11ea-1dca-f78fefb58b8f
md"## Clip-06-16s.jl"

# ╔═╡ b495756a-ff30-11ea-378a-236707f95773
begin
	N = 100
	df = DataFrame(
		:h0 => rand(Normal(10,2 ), N),
		:treatment => vcat(zeros(Int, Int(N/2)), ones(Int, Int(N/2)))
	);
	df.fungus = [rand(Binomial(1, 0.5 - 0.4 * df[i, :treatment]), 1)[1] for i in 1:N]
	df.h1 = [df[i, :h0] + rand(Normal(5 - 3 * df[i, :fungus]), 1)[1] for i in 1:N]
end;

# ╔═╡ 221f1e8a-00f6-11eb-3504-a91f4659a8a6
with_terminal() do
	precis(df; io=String)
end

# ╔═╡ b4a2fe56-ff30-11ea-255b-bbe592530dac
stan6_7 = "
data {
  int <lower=1> N;
  vector[N] h0;
  vector[N] h1;
  vector[N] treatment;
  vector[N] fungus;
}
parameters{
  real a;
  real bt;
  real bf;
  real<lower=0> sigma;
}
model {
  vector[N] mu;
  vector[N] p;
  a ~ lognormal(0, 0.2);
  bt ~ normal(0, 0.5);
  bf ~ normal(0, 0.5);
  sigma ~ exponential(1);
  for ( i in 1:N ) {
    p[i] = a + bt*treatment[i] + bf*fungus[i];
    mu[i] = h0[i] * p[i];
  }
  h1 ~ normal(mu, sigma);
}
";

# ╔═╡ b4a44e3c-ff30-11ea-0868-93935a280415
begin
	m6_7s = SampleModel("m6.7s", stan6_7)
	m6_7_data = Dict(
	  :N => nrow(df),
	  :h0 => df[:, :h0],
	  :h1 => df[:, :h1],
	  :fungus => df[:, :fungus],
	  :treatment => df[:, :treatment]
	)
	rc6_7s = stan_sample(m6_7s; data=m6_7_data)
	success(rc6_7s) && (post6_7s_df = read_samples(m6_7s, :dataframe))
end;

# ╔═╡ b4a98408-ff30-11ea-2918-9b5d39692695
success(rc6_7s) && Particles(post6_7s_df)

# ╔═╡ b4b7702a-ff30-11ea-169a-159df3080e8b
success(rc6_7s) && (Text(precis(post6_7s_df; io=String)))

# ╔═╡ b4b8301e-ff30-11ea-2a3c-a11e3745616e
if success(rc6_7s)
	(part6_7s, fig6_7s) = plot_model_coef([m6_7s], [:a, :bt, :bf])
	fig6_7s
end

# ╔═╡ b4c852fa-ff30-11ea-18d0-47dfe321ba8a
md"## End of clip-06-16s.jl"

# ╔═╡ Cell order:
# ╟─b0e1bc4a-ff2f-11ea-1dca-f78fefb58b8f
# ╠═b4949654-ff30-11ea-225c-2ffc47a35a31
# ╠═b494da42-ff30-11ea-25d2-37857c757c9f
# ╠═b495756a-ff30-11ea-378a-236707f95773
# ╠═221f1e8a-00f6-11eb-3504-a91f4659a8a6
# ╠═b4a2fe56-ff30-11ea-255b-bbe592530dac
# ╠═b4a44e3c-ff30-11ea-0868-93935a280415
# ╠═b4a98408-ff30-11ea-2918-9b5d39692695
# ╠═b4b7702a-ff30-11ea-169a-159df3080e8b
# ╠═b4b8301e-ff30-11ea-2a3c-a11e3745616e
# ╟─b4c852fa-ff30-11ea-18d0-47dfe321ba8a


using NearlyNewton
using Test
using StaticArrays
#using Optim
#using LineSearches
using Printf
using LinearAlgebra: norm, I

include("noalloc.jl")

@testset "NearlyNewton.jl" begin

    function theta(x)
       if x[1] > 0
           return atan(x[2] / x[1]) / (2.0 * pi)
       else
           return (pi + atan(x[2] / x[1])) / (2.0 * pi)
       end
    end
    f(x) = 100.0 * ((x[3] - 10.0 * theta(x))^2 + (sqrt(x[1]^2 + x[2]^2) - 1.0)^2) + x[3]^2

    function f∇f!(∇f, x)
        if !(∇f==nothing)
            if ( x[1]^2 + x[2]^2 == 0 )
                dtdx1 = 0;
                dtdx2 = 0;
            else
                dtdx1 = - x[2] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) );
                dtdx2 =   x[1] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) );
            end
            ∇f[1] = -2000.0*(x[3]-10.0*theta(x))*dtdx1 +
                200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[1]/sqrt( x[1]^2+x[2]^2 );
            ∇f[2] = -2000.0*(x[3]-10.0*theta(x))*dtdx2 +
                200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[2]/sqrt( x[1]^2+x[2]^2 );
            ∇f[3] =  200.0*(x[3]-10.0*theta(x)) + 2.0*x[3];
        end

        fx = f(x)
        return ∇f==nothing ? fx : (fx, ∇f)
    end

    function f∇f(∇f, x)
        if !(∇f == nothing)
            gx = similar(x)
            return f∇f!(gx, x)
        else
            return f∇f!(∇f, x)
        end
    end
    function f∇fs(∇f, x)
        if !(∇f == nothing)
            if ( x[1]^2 + x[2]^2 == 0 )
                dtdx1 = 0;
                dtdx2 = 0;
            else
                dtdx1 = - x[2] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) )
                dtdx2 =   x[1] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) )
            end

            s1 = -2000.0*(x[3]-10.0*theta(x))*dtdx1 +
                200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[1]/sqrt( x[1]^2+x[2]^2 )
            s2 = -2000.0*(x[3]-10.0*theta(x))*dtdx2 +
                200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[2]/sqrt( x[1]^2+x[2]^2 )
            s3 = 200.0*(x[3]-10.0*theta(x)) + 2.0*x[3]
            ∇f = @SVector [s1, s2, s3]
            return f(x), ∇f
        else
            return f(x)
        end
    end

    x0 = [-1.0, 0.0, 0.0]
    xopt = [1.0, 0.0, 0.0]
    x0s = @SVector [-1.0, 0.0, 0.0]

    println("Starting  from: ", x0)
    println("Targeting     : ", xopt)

    shortname(::GradientDescent) = "GD  "
    shortname(::InverseApprox) = "(inverse)"
    shortname(::NearlyNewton.DirectApprox) = " (direct)"

    function printed_minimize(f∇f, x0, method, B0=nothing)
        res = minimize(f∇f, x0, method, B0, NearlyNewton.OptOptions())
        print("NN  $(shortname(method)) $(shortname(method.approx)): ")
        @printf("%2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    end
    function printed_minimize!(f∇f!, x0, method, B0=nothing)
        res = minimize!(f∇f!, x0, method, B0, NearlyNewton.OptOptions())
        print("NN! $(shortname(method)) $(shortname(method.approx)): ")
        @printf("%2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    end

    printed_minimize(f∇f, x0, GradientDescent(InverseApprox()))
    printed_minimize!(f∇f!, copy(x0), GradientDescent(InverseApprox()))
    res = minimize(f∇fs, x0s, GradientDescent(InverseApprox()))
    @printf("NN  GD(S) (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])

    printed_minimize(f∇f, x0, GradientDescent(DirectApprox()))
    printed_minimize!(f∇f!, copy(x0), GradientDescent(DirectApprox()))
    res = minimize(f∇fs, x0s, GradientDescent(DirectApprox()))
    @printf("NN  GD(S)  (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    println()

    res = minimize(f∇f, x0, BFGS(InverseApprox()), nothing, NearlyNewton.OptOptions())
    @test res[4] == 30
    @printf("NN  BFGS    (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(x0), (BFGS(InverseApprox()), BackTracking()), nothing, NearlyNewton.OptOptions())
    @test res[4] == 30
    @printf("NN! BFGS    (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇fs, x0s, BFGS(InverseApprox()))
    @test res[4] == 30
    @printf("NN  BFGS(S) (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])

    res = minimize(f∇f, x0, BFGS(DirectApprox()), nothing, NearlyNewton.OptOptions())
    @printf("NN  BFGS    (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(x0), BFGS(DirectApprox()), nothing, NearlyNewton.OptOptions())
    @printf("NN! BFGS    (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇fs, x0s, BFGS(DirectApprox()))
    @printf("NN  BFGS(S) (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    println()
    # res = optimize(f∇f!, copy(x0), Optim.BFGS(linesearch=BackTracking()))
    # @time optimize(f∇f!, copy(x0), Optim.BFGS(linesearch=BackTracking()))
    # @printf("OT! BFGS (inverse): %2.2e  %2.2e\n", norm(Optim.minimizer(res)-xopt,Inf), Optim.g_residual(res))

    res = minimize(f∇f, x0, DFP(InverseApprox()))
    @printf("NN  DFP    (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(x0), DFP(InverseApprox()))
    @printf("NN! DFP    (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇fs, x0s, DFP(InverseApprox()))
    @printf("NN  DFP(S) (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇f, x0, DFP(DirectApprox()))
    @printf("NN  DFP    (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(x0), DFP(DirectApprox()))
    @printf("NN! DFP    (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇fs, x0s, DFP(DirectApprox()))
    @printf("NN  DFP(S) (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    println()

    res = minimize(f∇f, x0, SR1(InverseApprox()))
    @printf("NN  SR1     (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(x0), SR1(InverseApprox()))
    @printf("NN! SR1     (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇fs, x0s, SR1(InverseApprox()))
    @printf("NN  SR1(S)  (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇f, x0, SR1(DirectApprox()))
    @printf("NN  SR1     (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(x0), SR1(DirectApprox()))
    @printf("NN! SR1     (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇fs, x0s, SR1(DirectApprox()))
    @printf("NN  SR1(S)  (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])

    xrand = rand(3)
    xrands = SVector{3}(xrand)
    println("\nFrom a random point: ", xrand)
    res = minimize(f∇f, xrand, GradientDescent(InverseApprox()))
    @printf("NN  GD   (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(xrand), GradientDescent(InverseApprox()))
    @printf("NN! GD   (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇f, xrand, GradientDescent(DirectApprox()))
    @printf("NN  GD    (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(xrand), GradientDescent(DirectApprox()))
    @printf("NN! GD    (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    println()
    res = minimize(f∇f, xrand, BFGS(InverseApprox()))
    @printf("NN  BFGS (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(xrand), (BFGS(InverseApprox()), BackTracking()))
    @printf("NN! BFGS (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇f, xrand, BFGS(DirectApprox()))
    @printf("NN  BFGS  (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(xrand), BFGS(DirectApprox()))
    @printf("NN! BFGS  (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    println()

    res = minimize(f∇f, xrand, DFP(InverseApprox()))
    @printf("NN  DFP  (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(xrand), DFP(InverseApprox()))
    @printf("NN! DFP  (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇f, xrand, DFP(DirectApprox()))
    @printf("NN  DFP   (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(xrand), DFP(DirectApprox()))
    @printf("NN! DFP   (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    println()

    res = minimize(f∇f, xrand, SR1(InverseApprox()))
    @printf("NN  SR1  (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(xrand), SR1(InverseApprox()))
    @printf("NN! SR1  (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇fs, xrand, SR1(InverseApprox()))
    @printf("NN  SR1(S)   (inverse): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇f, xrand, SR1(DirectApprox()))
    @printf("NN  SR1   (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize!(f∇f!, copy(xrand), SR1(DirectApprox()))
    @printf("NN! SR1   (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])
    res = minimize(f∇fs, xrand, SR1(DirectApprox()))
    @printf("NN  SR1(S)   (direct): %2.2e  %2.2e  %d\n", norm(res[1]-xopt,Inf), norm(res[3], Inf), res[4])

    println()

    function himmelblau!(∇f, x)
        if !(∇f == nothing)
            ∇f[1] = 4.0 * x[1]^3 + 4.0 * x[1] * x[2] -
                44.0 * x[1] + 2.0 * x[1] + 2.0 * x[2]^2 - 14.0
            ∇f[2] = 2.0 * x[1]^2 + 2.0 * x[2] - 22.0 +
                4.0 * x[1] * x[2] + 4.0 * x[2]^3 - 28.0 * x[2]
        end

        fx = (x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2
        return ∇f == nothing ? fx : (fx, ∇f)
    end

    function himmelblaus(∇f, x)
        f = (x[1]^2 + x[2] - 11)^2 + (x[1] + x[2]^2 - 7)^2
        if !(∇f == nothing)
            ∇f1 = 4.0 * x[1]^3 + 4.0 * x[1] * x[2] -
                44.0 * x[1] + 2.0 * x[1] + 2.0 * x[2]^2 - 14.0
            ∇f2 = 2.0 * x[1]^2 + 2.0 * x[2] - 22.0 +
                4.0 * x[1] * x[2] + 4.0 * x[2]^3 - 28.0 * x[2]
            return f, @SVector([∇f1, ∇f2])
        else
            return f
        end
    end

    function himmelblau(∇f, x)
        g = ∇f == nothing ? ∇f : similar(x)

        return himmelblau!(g, x)
    end

    println("\nHimmelblau function")
    x0 = [3.0, 1.0]
    x0s = @SVector([3.0, 1.0])
    res = minimize(himmelblau, x0, GradientDescent(InverseApprox()))
    @printf("NN  GD    (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!, copy(x0), GradientDescent(InverseApprox()))
    @printf("NN! GD    (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblaus, x0s, GradientDescent(InverseApprox()))
    @printf("NN  GD(S) (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    println()
    res = minimize(himmelblau, x0, GradientDescent(DirectApprox()))
    @printf("NN  GD    (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!, copy(x0), GradientDescent(DirectApprox()))
    @printf("NN! GD    (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblaus, x0s, GradientDescent(DirectApprox()))
    @printf("NN  GD(S) (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    println()

    res = minimize(himmelblau, x0, BFGS(InverseApprox()))
    @printf("NN  BFGS    (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!, copy(x0), BFGS(InverseApprox()))
    @printf("NN! BFGS    (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, x0s, BFGS(InverseApprox()))
    @printf("NN  BFGS(S) (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    println()
    res = minimize(himmelblau, x0, BFGS(DirectApprox()))
    @printf("NN  BFGS    (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!, copy(x0), BFGS(DirectApprox()))
    @printf("NN! BFGS    (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, x0s, BFGS(DirectApprox()))
    @printf("NN  BFGS(S) (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    println()

    res = minimize(himmelblau, x0, DFP(InverseApprox()))
    @printf("NN  DFP  (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!, copy(x0), DFP(InverseApprox()))
    @printf("NN! DFP  (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, x0s, DFP(InverseApprox()))
    @printf("NN  DFP(S)  (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, x0, DFP(DirectApprox()))
    println()
    @printf("NN  DFP   (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!, copy(x0), DFP(DirectApprox()))
    @printf("NN! DFP   (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, x0s, DFP(DirectApprox()))
    @printf("NN  DFP(S)   (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    println()

    res = minimize(himmelblau, x0, SR1(InverseApprox()))
    @printf("NN  SR1  (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!, copy(x0), SR1(InverseApprox()))
    @printf("NN! SR1  (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, x0s, SR1(InverseApprox()))
    @printf("NN  SR1(S)  (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, x0, SR1(DirectApprox()))
    @printf("NN  SR1   (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!, copy(x0), SR1(DirectApprox()))
    @printf("NN! SR1   (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, x0s, SR1(DirectApprox()))
    @printf("NN  SR1(S)   (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    println()
    xrand = rand(2)
    println("\nFrom a random point: ", xrand)

    res = minimize(himmelblau, xrand, GradientDescent(InverseApprox()))
    @printf("NN  GD   (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!, copy(xrand), GradientDescent(InverseApprox()))
    @printf("NN! GD   (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, xrand, GradientDescent(DirectApprox()))
    @printf("NN  GD    (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!, copy(xrand), GradientDescent(DirectApprox()))
    @printf("NN! GD    (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])

    res = minimize(himmelblau, xrand, BFGS(InverseApprox()))
    @printf("NN   BFGS (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!, copy(xrand), BFGS(InverseApprox()))
    @printf("NN!  BFGS (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, xrand, BFGS(DirectApprox()))
    @printf("NN(S) BFGS (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, xrand, BFGS(DirectApprox()))
    @printf("NN  BFGS  (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!, copy(xrand), BFGS(DirectApprox()))
    @printf("NN! BFGS  (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])

    res = minimize(himmelblau, xrand, DFP(InverseApprox()))
    @printf("NN  DFP  (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!,  copy(xrand), DFP(InverseApprox()))
    @printf("NN! DFP  (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, xrand, DFP(DirectApprox()))
    @printf("NN  DFP   (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!,  copy(xrand), DFP(DirectApprox()))
    @printf("NN! DFP   (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])

    res = minimize(himmelblau, xrand, SR1(InverseApprox()))
    @printf("NN  SR1  (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!,  copy(xrand), SR1(InverseApprox()))
    @printf("NN! SR1  (inverse): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize(himmelblau, xrand, SR1(DirectApprox()))
    @printf("NN  SR1   (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])
    res = minimize!(himmelblau!,  copy(xrand), SR1(DirectApprox()))
    @printf("NN! SR1   (direct): %2.2e  %d\n", norm(res[3], Inf), res[4])

end

include("geometrytypes.jl")
include("newton.jl")


typ(n) = Symbol(string("Series", n))
elt(n) = Symbol(string("c", n))
mem(s,n) = Expr(:., s, Expr(:quote, elt(n)))

function generate_type(n::Integer)
  @assert n > 0
  LastType = typ(n-1)
  Typ = typ(n)

  defn = :(immutable $Typ{T} <: AbstractSeries{T, $n} c0::T end)

  for i = 1:n
    push!(defn.args[3].args, :($(elt(i))::T))
  end

  eval(defn)

  if (n > 1)
    @eval restrict(x::$Typ) = $LastType($([mem(:x, i) for i = 0:n-1]...))
    @eval polyder(x::$Typ) = $LastType($([:($i*$(mem(:x, i))) for i = 1:n]...))
    @eval polyint{T}(x::$LastType{T}) = $Typ(zero(T),
      $([:($(mem(:x, i-1))/$i) for i = 1:n]...)
    )
  else
    @eval restrict(x::$Typ) = x.c0
    @eval polyder(x::$Typ) = x.c1
    @eval polyint(x::Real) = Series1(zero(typeof(x)), x)
  end

  @eval $Typ(x::Real) = $Typ(x, $([:(zero(x)) for i = 1:n]...))

  syms = [gensym() for i=0:n]

  @eval series($(syms...)) = $Typ(promote($(syms...))...)

  @eval convert{T<:Real}(::Type{$Typ{T}}, x::Real) = $Typ(convert(T, x))

  #Horner's method
  polyval_body = mem(:x, n)
  for m = n-1:-1:0
    polyval_body = :($(mem(:x, m)) + eps*$polyval_body)
  end
  @eval polyval(x::$Typ, eps::Real) = $polyval_body

  @eval constant(x::$Typ) = x.c0

  @eval +(x::$Typ, y::$Typ) = $Typ(
    $([:($(mem(:x, i)) + $(mem(:y, i))) for i = 0:n]...)
  )
  @eval +{T<:Real, S}(c::T, x::$Typ{S}) = $Typ{promote_type(T, S)}(
    c + $(mem(:x, 0)),
    $([mem(:x, i) for i = 1:n]...)
  )
  @eval +{T<:Real, S}(x::$Typ{S}, c::T) = $Typ{promote_type(T, S)}(
    $(mem(:x, 0)) + c,
    $([mem(:x, i) for i = 1:n]...)
  )

  @eval -(x::$Typ) = $Typ($([:(-$(mem(:x, i))) for i = 0:n]...))

  @eval -(x::$Typ, y::$Typ) = $Typ(
    $([:($(mem(:x, i)) - $(mem(:y, i))) for i = 0:n]...)
  )
  @eval -{T<:Real, S}(c::T, x::$Typ{S}) = $Typ{promote_type(T, S)}(
    c - $(mem(:x, 0)),
    $([:(-$(mem(:x, i))) for i = 1:n]...)
  )
  @eval -{T<:Real, S}(x::$Typ{S}, c::T) = $Typ{promote_type(T, S)}(
    $(mem(:x, 0)) - c,
    $([:($(mem(:x, i))) for i = 1:n]...)
  )

  times_term(i) = :(+($([:($(mem(:x, j))*$(mem(:y, i-j))) for j = 0:i]...)))
  @eval *(x::$Typ, y::$Typ) = $Typ($([times_term(i) for i = 0:n]...))

  # Disambiguate bool multiplication. Begrudgingly...
  # See https://github.com/JuliaLang/julia/issues/5611
  @eval *(x::Bool, z::$Typ) = ifelse(x,z,zero(z))
  @eval *(z::$Typ, x::Bool) = x * z

  @eval *(c::Real, x::$Typ) = $Typ($([:(c*$(mem(:x, i))) for i = 0:n]...))
  @eval *(x::$Typ, c::Real) = $Typ($([:($(mem(:x, i))*c) for i = 0:n]...))

  @eval /(x::$Typ, c::Real) = $Typ($([:($(mem(:x, i))/c) for i = 0:n]...))

  # TODO probably a better way than eval to get a type from a string
  return eval(Typ)
end


using UnPack, OrderedCollections, ForwardDiff, JLD2

mutable struct TextbookNK{T <: Real}
    β::T
    σ::T
    ψ::T
    η::T
    ϵ::T
    ϕ::T
    ρₐ::T
    σₐ::T
    ρᵢ::T
    σᵢ::T
    ϕ_π::T
    π̃_ss::T
end

function TextbookNK(; β::T = .99, σ::T = 2., ψ::T = 1., η::T = 1., ϵ::T = 4.45, ϕ::T = .7,
                    ρₐ::T = 0.9, σₐ::T = .004, ρᵢ::T = .7, σᵢ::T = .025 / 4.,
                    ϕ_π::T = 1.5, π̃_ss::T = 0.) where {T <: Real}
    return TextbookNK{T}(β, σ, ψ, η, ϵ, ϕ, ρₐ, σₐ, ρᵢ, σᵢ, ϕ_π, π̃_ss)
end

function textbook_nk(m::TextbookNK{T}) where {T <: Real}
    @unpack β, σ, ψ, η, ϵ, ϕ, ρₐ, σₐ, ρᵢ, σᵢ, ϕ_π, π̃_ss = m
    ĩ_ss = π̃_ss - log(β)

    # On notation: x̃ = log(1 + x); x′ = 𝔼ₜ[xₜ₊₁]
    S  = OrderedDict{Symbol, Int}(:a => 1, :ĩ₋₁ => 2, :v₋₁ => 3, :i_sh => 4) # State Variables
    J  = OrderedDict{Symbol, Int}(:c => 1, :π̃ => 2, :n => 3, :w => 4, :mc => 5, :v => 6,
                                  :x₁ => 7, :x₂ => 8, :ĩ => 9) # Jump variables
    E  = OrderedDict{Symbol, Int}(:euler => 1, :mrs => 2, :eq_mc => 3, :output => 4,
                                  :dispersion => 5, :phillips_curve => 6, :eq_x₁ => 7,
                                  :eq_x₂ => 8, :eq_mp => 9) # Equations
    SH = OrderedDict{Symbol, Int}(:εₐ => 1, :εᵢ => 2) # Exogenous shocks

    @unpack a, ĩ₋₁, v₋₁, i_sh = S
    @unpack c, π̃, n, w, mc, v, x₁, x₂, ĩ = J
    @unpack euler, mrs, eq_mc, output, dispersion, phillips_curve, eq_x₁, eq_x₂, eq_mp = E
    @unpack εₐ, εᵢ = SH

    Nz = length(S)
    Ny = length(J)
    Nε = length(SH)

    function μ(F, z, y)
        F_type  = eltype(F)
        F[a]    = ρₐ * z[a]
        F[ĩ₋₁]  = y[ĩ]
        F[v₋₁]  = y[v]
        F[i_sh] = zero(F_type)
    end

    function ξ(F, z, y)
        F_type = eltype(F)
        π̃_star = log(ϵ / (ϵ - 1.)) + y[π̃] + (y[x₁] - y[x₂])
        F[euler] = log(β) + σ * y[c] + y[ĩ]
        F[mrs] = log(ψ) + η * y[n] - (-σ * y[n] + y[w])
        F[eq_mc] = y[w] - (z[a] + y[mc])
        F[output] = y[c] - (z[a] + y[n] - y[v])
        F[dispersion] = y[v] - (ϵ * y[π̃] + log((1. - ϕ) * exp(π̃_star)^(-ϵ) + ϕ * exp(z[v₋₁])))
        F[phillips_curve] = (1. - ϵ) * y[π̃] - log((1. - ϕ) * exp(π̃_star)^(1 - ϵ) + ϕ)
        F[eq_x₁] = log(ϕ) + log(β) - log(exp(y[x₁]) - exp((1. - σ) * y[c] + y[mc]))
        F[eq_x₂] = log(ϕ) + log(β) - log(exp(y[x₂]) - exp((1. - σ) * y[c]))
        F[eq_mp] = y[ĩ] - ((1. - ρᵢ) * ĩ_ss + ρᵢ * z[ĩ₋₁]  + (1 - ρᵢ) * ϕ_π * (y[π̃] - π̃_ss) + z[i_sh])
    end

    # The cache is initialized as zeros so we only need to fill non-zero elements
    Λ = zeros(T, Nz, Ny)

    # The cache is initialized as zeros so we only need to fill non-zero elements
    function Σ(F, z)
        F[a, εₐ]    = σₐ
        F[i_sh, εᵢ] = σᵢ
    end

    function ccgf(F, α, z)
        # F .= .5 * diag(α * α') # slower but this is the underlying math
        sum!(F, α.^2) # faster implementation
        F .*= .5
    end

    Γ₅ = zeros(T, Ny, Nz)

    Γ₆ = zeros(T, Ny, Ny)
    Γ₆[euler, c] = -σ
    Γ₆[euler, π̃] = -one(T)
    Γ₆[eq_x₁, x₁] = one(T)
    Γ₆[eq_x₁, π̃] = ϵ
    Γ₆[eq_x₂, x₂] = one(T)
    Γ₆[eq_x₂, π̃] = (ϵ - 1.)

    Ψ = zeros(T, Ny, Nz)

    # Deterministic steady state as initial guess

    # z
    a0    = 0.
    ĩ₋₁0  = ĩ_ss
    v₋₁0  = 0.
    i_sh0 = 0.
    z     = [a0, ĩ₋₁0, v₋₁0, i_sh0]

    # y
    ĩ0  = ĩ_ss
    π̃0  = π̃_ss
    v0  = 0.
    mc0 = log((ϵ - 1.) / ϵ)
    x₁0 = log(exp(π̃0) * exp(mc0) / (1. - ϕ * β))
    x₂0 = log(1. / (1. - ϕ * β))
    n0  = (1 / (η + σ)) * log(1. / ψ * (exp(v0))^σ * exp(mc0))
    c0  = n0 - v0
    w0  = a0 + mc0
    y   = [c0, π̃0, n0, w0, mc0, v0, x₁0, x₂0, ĩ0]

    return RiskAdjustedLinearization(μ, Λ, Σ, ξ, Γ₅, Γ₆, ccgf, vec(z), vec(y), Ψ, Nε)
end

nk_cₜ(m, zₜ) = exp(m.y[1] + (m.Ψ * (zₜ - m.z))[1])

# Evaluates euler equation in log terms
function nk_logSDFxR(m, zₜ, εₜ₊₁, Cₜ; β::T = .99, σ::T = 2.) where {T <: Real}
    yₜ = m.y + m.Ψ * (zₜ - m.z)
    zₜ₊₁, yₜ₊₁ = simulate(m, εₜ₊₁, zₜ)

    return log(β) - σ * (yₜ₊₁[1] - log(Cₜ)) + yₜ[9] - yₜ₊₁[2]
end

# Calculate Euler equation via quadrature
std_norm_mean = zeros(2)
std_norm_sig  = ones(2)
nk_𝔼_quadrature(f::Function) = gausshermite_expectation(f, std_norm_mean, std_norm_sig, 10)


struct BlochMcConnellWorkspace{T<:Real,N}
    A::BlochMcConnellDynamicsMatrix{T,N}
    expmworkspace::MatrixExponentialWorkspace{T,N}

    BlochMcConnellWorkspace(T::Type{<:Real}, N) =
        new{T,N}(BlochMcConnellDynamicsMatrix{T}(N),
                 MatrixExponentialWorkspace{T}(N))
end

BlochMcConnellWorkspace(::Type{SpinMC{T,N}}) where {T,N} = BlochMcConnellWorkspace(T, N)
BlochMcConnellWorkspace(spin::SpinMC) = BlochMcConnellWorkspace(typeof(spin))

"""
    freeprecess(t, M0, T1, T2, Δf)

Simulate free-precession, i.e., relaxation and off-resonance precession. Returns
`(A, B)` such that `A * M + B` applies free-precession to the magnetization `M`.

For an in-place version, see [`freeprecess!`](@ref).

# Arguments
- `t::Real`: Duration of free-precession (ms)
- `M0::Real`: Equilibrium magnetization
- `T1::Real`: Spin-lattice recovery time constant (ms)
- `T2::Real`: Spin-spin recovery time constant (ms)
- `Δf::Real`: Off-resonance frequency (Hz)

# Examples
```jldoctest
julia> (A, B) = freeprecess(100, 1, 1000, 100, 3.75); A * Magnetization(1, 0, 0) + B
Magnetization vector with eltype Float64:
 Mx = -0.2601300475114444
 My = -0.2601300475114445
 Mz = 0.09516258196404048
```
"""
function freeprecess(t::Real, M0::Real, T1::Real, T2::Real, Δf::Real)

    A = FreePrecessionMatrix()
    B = Magnetization()
    freeprecess!(A, B, t, M0, T1, T2, Δf)
    return (A, B)

end

"""
    freeprecess!(A, B, t, M0, T1, T2, Δf)
    freeprecess!(A, B, spin, t, [nothing])
    freeprecess!(A, B, spinmc, t, [workspace])
    freeprecess!(A, B, spin, t, grad, [nothing])
    freeprecess!(A, B, spinmc, t, grad, [workspace])

Simulate free-precession, overwriting `A` and `B` (in-place version of
[`freeprecess`](@ref)).
"""
function freeprecess!(A, B, t, M0, T1, T2, Δf)

    E2 = exp(-t / T2)
    θ = 2π * Δf * t / 1000
    (s, c) = sincos(θ)

    A.E1 = exp(-t / T1)
    A.E2cosθ = E2 * c
    A.E2sinθ = E2 * s

    B.x = 0
    B.y = 0
    B.z = M0 * (1 - A.E1)

    return nothing

end

"""
    freeprecess(spin, t, [nothing])
    freeprecess(spinmc, t, [workspace])
    freeprecess(spin, t, grad, [nothing])
    freeprecess(spinmc, t, grad, [workspace])

Simulate free-precession for the given spin for time `t` ms, optionally in the
presence of a B0 gradient. Returns `(A, B)` such that `A * M + B` applies
free-precession to the magnetization `M`.

For `SpinMC` objects, `workspace isa BlochMcConnellWorkspace`. Pass in `nothing`
instead to use an approximate matrix exponential to solve the Bloch-McConnell
equation.

For an in-place version, see [`freeprecess!`](@ref).

# Examples
```jldoctest
julia> s = Spin(Magnetization(1, 0, 0), 1, 1000, 100, 3.75)
Spin{Float64}:
 M = Magnetization(1.0, 0.0, 0.0)
 M0 = 1.0
 T1 = 1000.0 ms
 T2 = 100.0 ms
 Δf = 3.75 Hz
 pos = Position(0.0, 0.0, 0.0) cm

julia> (A, B) = freeprecess(s, 100); A * s.M + B
Magnetization vector with eltype Float64:
 Mx = -0.2601300475114444
 My = -0.2601300475114445
 Mz = 0.09516258196404048

julia> s = Spin(Magnetization(1, 0, 0), 1, 1000, 100, 0, Position(0, 0, 3.75))
Spin{Float64}:
 M = Magnetization(1.0, 0.0, 0.0)
 M0 = 1.0
 T1 = 1000.0 ms
 T2 = 100.0 ms
 Δf = 0.0 Hz
 pos = Position(0.0, 0.0, 3.75) cm

julia> (A, B) = freeprecess(s, 100, Gradient(0, 0, 1/GAMBAR)); A * s.M + B
Magnetization vector with eltype Float64:
 Mx = -0.2601300475114444
 My = -0.2601300475114445
 Mz = 0.09516258196404048
```
"""
function freeprecess(spin::Spin, t, ::Nothing = nothing)

    A = FreePrecessionMatrix{eltype(spin)}()
    B = Magnetization{eltype(spin)}()
    freeprecess!(A, B, spin, t)
    return (A, B)

end

function freeprecess(spin::SpinMC{T,N}, t, workspace = BlochMcConnellWorkspace(spin)) where {T,N}

    A = BlochMcConnellMatrix{T}(N)
    B = MagnetizationMC{T}(N)
    freeprecess!(A, B, spin, t, workspace)
    return (A, B)

end

function freeprecess!(
    A::FreePrecessionMatrix,
    B::Magnetization,
    spin::Spin,
    t::Real,
    workspace::Nothing = nothing
)

    freeprecess!(A, B, t, spin.M0, spin.T1, spin.T2, spin.Δf)

end

function freeprecess!(
    A::BlochMcConnellMatrix,
    B::MagnetizationMC,
    spin::SpinMC,
    t::Real,
    workspace::Union{Nothing,<:BlochMcConnellWorkspace} = BlochMcConnellWorkspace(spin)
)

    expm!(A, workspace, spin, t)
    subtractmul!(B, I, A, spin.Meq)
    return nothing

end

function freeprecess(spin::Spin, t::Real, grad::Gradient, ::Nothing = nothing)

    A = FreePrecessionMatrix{eltype(spin)}()
    B = Magnetization{eltype(spin)}()
    freeprecess!(A, B, spin, t, grad)
    return (A, B)

end

function freeprecess(spin::SpinMC{T,N}, t::Real, grad::Gradient, workspace = BlochMcConnellWorkspace(spin)) where {T,N}

    A = BlochMcConnellMatrix{T}(N)
    B = MagnetizationMC{T}(N)
    freeprecess!(A, B, spin, t, grad, workspace)
    return (A, B)

end

function freeprecess!(
    A::FreePrecessionMatrix,
    B::Magnetization,
    spin::Spin,
    t::Real,
    grad::Gradient,
    workspace::Nothing = nothing
)

    gradfreq = gradient_frequency(grad, spin.pos) # Hz
    freeprecess!(A, B, t, spin.M0, spin.T1, spin.T2, spin.Δf + gradfreq)

end

function freeprecess!(
    A::BlochMcConnellMatrix,
    B::MagnetizationMC,
    spin::SpinMC,
    t::Real,
    grad::Gradient,
    workspace::Union{Nothing,<:BlochMcConnellWorkspace} = BlochMcConnellWorkspace(spin)
)

    gradfreq = gradient_frequency(grad, spin.pos) # Hz
    expm!(A, workspace, spin, t, gradfreq)
    subtractmul!(B, I, A, spin.Meq)
    return nothing

end

# Exact matrix exponential
function expm!(expAt, workspace, spin, t, gradfreq = 0)

    index = 0
    for i = 1:spin.N, j = 1:spin.N

        if i == j

            r_out = sum(spin.r[i][k] for k = 1:spin.N) # 1/ms

            A = workspace.A.A[i]
            A.R1 = -1 / spin.T1[i] - r_out # 1/ms
            A.R2 = -1 / spin.T2[i] - r_out # 1/ms
            A.Δω = 2π * (spin.Δf[i] + gradfreq) / 1000 # rad/ms

        else

            index += 1
            workspace.A.E[index].r = spin.r[i][j] # 1/ms

        end

    end

    mul!(workspace.A, t)
    expm!(expAt, workspace.A, workspace.expmworkspace)
    return nothing

end

# Approximate matrix exponential
# See page 2 of http://doi.org/10.1137/0714065
# (unnumbered equation with o(||E||) term)
# Referring to A and E from the paper:
# - A is block diagonal where each block along the diagonal is the traditional
#   Bloch matrix for the corresponding compartment (except the 1/T1 and 1/T2
#   terms also include exchange out of the compartment)
# - E is an "inverse" block diagonal matrix (i.e., the diagonal blocks are 0)
#   where the off-diagonal blocks are diagonal matrices with the exchange terms
# With this choice of A, e^At has an analytical solution. The code in this
# function follows the analytical solution of the aforementioned equation after
# manually taking e^At and doing the matrix-matrix products and integrating
# over s.
function expm!(expAt, ::Nothing, spin, t, gradfreq = 0)

    for j = 1:spin.N, i = 1:spin.N

        A = expAt.A[i][j]

        if i == j

            r_out = sum(spin.r[i][k] for k = 1:spin.N) # 1/ms
            E1 = exp(-t * (1 / spin.T1[i] + r_out))
            E2 = exp(-t * (1 / spin.T2[i] + r_out))
            θ = 2π * (spin.Δf[i] + gradfreq) * t / 1000 # rad
            (s, c) = sincos(θ)
            E2c = E2 * c
            E2s = E2 * s

            A.a11 = E2c
            A.a21 = -E2s
            A.a31 = 0
            A.a12 = E2s
            A.a22 = E2c
            A.a32 = 0
            A.a13 = 0
            A.a23 = 0
            A.a33 = E1

        else

            r_out_i = sum(spin.r[i][k] for k = 1:spin.N) # 1/ms
            r_out_j = sum(spin.r[j][k] for k = 1:spin.N) # 1/ms
            R1i = 1 / spin.T1[i] + r_out_i # 1/ms
            R1j = 1 / spin.T1[j] + r_out_j # 1/ms
            R2i = 1 / spin.T2[i] + r_out_i # 1/ms
            R2j = 1 / spin.T2[j] + r_out_j # 1/ms
            rji = spin.r[j][i]

            R2ji = R2j - R2i
            R2ji² = R2ji^2
            E2i = exp(-t * R2i)
            E2ji = exp(-t * R2ji)
            Δωji = 2π * (spin.Δf[j] - spin.Δf[i]) / 1000 # rad/ms
            Δωji² = Δωji^2
            θi = 2π * (spin.Δf[i] + gradfreq) * t / 1000 # rad
            θj = 2π * (spin.Δf[j] + gradfreq) * t / 1000 # rad
            (si, ci) = sincos(θi)
            (sj, cj) = sincos(θj)
            tmpc = rji * E2i * ((ci - E2ji * cj) / R2ji + Δωji * (E2ji * sj - si) / R2ji²) / (1 + Δωji² / R2ji²)
            tmps = rji * E2i * ((si - E2ji * sj) / R2ji + Δωji * (ci - E2ji * cj) / R2ji²) / (1 + Δωji² / R2ji²)

            A.a11 = tmpc
            A.a21 = -tmps
            A.a31 = 0
            A.a12 = tmps
            A.a22 = tmpc
            A.a32 = 0
            A.a13 = 0
            A.a23 = 0
            A.a33 = rji * E2i * (1 - E2ji) / R2ji

        end

    end

    return nothing

end

"""
    freeprecess!(spin, ...)

Apply free-precession to the given spin, overwriting the spin's magnetization
vector.
"""
function freeprecess!(spin::AbstractSpin, args...)

    (A, B) = freeprecess(spin, args...)
    BtoM = copy(spin.M)
    applydynamics!(spin, BtoM, A, B)

end

"""
    combine!(A, B, A1, B1, A2, B2)
    combine!(A, A1, A2)

Combine the matrices and vectors that describe the dynamics of a spin into one
matrix and one vector, overwriting `A` and `B`. The dynamics described by `A1`
and `B1` apply first, then those described by `A2` and `B2`. In other words,
`A = A2 * A1` and `B = A2 * B1 + B2`.

# Examples
```jldoctest
julia> s = Spin(1, 1000, 100, 3.75);

julia> A = BlochMatrix(); B = Magnetization();

julia> (A1, B1) = excite(s, InstantaneousRF(π/2));

julia> (A2, B2) = freeprecess(s, 100);

julia> combine!(A, B, A1, B1, A2, B2); A * s.M + B
Magnetization vector with eltype Float64:
 Mx = -0.2601300475114444
 My = -0.2601300475114445
 Mz = 0.09516258196404054
```
"""
function combine!(A, B, A1, B1, A2, B2)

    mul!(A, A2, A1)
    mul!(B, A2, B1)
    add!(B, B2)

end

function combine!(A, B, A1, B1, A2, ::Nothing)

    mul!(A, A2, A1)
    mul!(B, A2, B1)

end

function combine!(A, B, A1, ::Nothing, A2, B2)

    mul!(A, A2, A1)
    copyto!(B, B2)

end

function combine!(A, B, A1, B1, ::Nothing, ::Nothing)

    copyto!(A, A1)
    copyto!(B, B1)

end

function combine!(A, B, ::Nothing, ::Nothing, A2, B2)

    copyto!(A, A2)
    copyto!(B, B2)

end

combine!(A, A1, A2) = mul!(A, A2, A1)

"""
    applydynamics!(spin, BtoM, A, [B])

Apply dynamics to the given spin, overwriting the spin's magnetization vector.
`BtoM` is used to store intermediate results (and is thus overwritten).

# Examples
```jldoctest
julia> s = Spin(1, 1000, 100, 3.75); s.M
Magnetization vector with eltype Float64:
 Mx = 0.0
 My = 0.0
 Mz = 1.0

julia> BtoM = Magnetization();

julia> (A,) = excite(s, InstantaneousRF(π/2)); applydynamics!(s, BtoM, A); s.M
Magnetization vector with eltype Float64:
 Mx = 1.0
 My = 0.0
 Mz = 6.123233995736766e-17

julia> (A, B) = freeprecess(s, 100); applydynamics!(s, BtoM, A, B); s.M
Magnetization vector with eltype Float64:
 Mx = -0.2601300475114444
 My = -0.2601300475114445
 Mz = 0.09516258196404054
```
"""
function applydynamics!(spin::AbstractSpin, BtoM, A, B)

    copyto!(BtoM, B)
    muladd!(BtoM, A, spin.M)
    copyto!(spin.M, BtoM)
    return nothing

end

function applydynamics!(spin::AbstractSpin, BtoM, A, ::Nothing = nothing)

    mul!(BtoM, A, spin.M)
    copyto!(spin.M, BtoM)
    return nothing

end


struct DeivVecTag end

# J(f(x))*v
function auto_jacvec!(
    dy,
    f,
    x,
    v,
    cache1 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x))))),
    cache2 = similar(cache1),
)
    cache1 .= Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x)))))
    f(cache2, cache1)
    vecdy = _vec(dy)
    vecdy .= partials.(_vec(cache2), 1)
end

_vec(v) = vec(v)
_vec(v::AbstractVector) = v

function auto_jacvec(f, x, v)
    vv = reshape(v, axes(x))
    y = ForwardDiff.Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(vv)))
    vec(partials.(vec(f(y)), 1))
end

function num_jacvec!(
    dy,
    f,
    x,
    v,
    cache1 = similar(v),
    cache2 = similar(v);
    compute_f0 = true,
)
    vv = reshape(v, axes(x))
    compute_f0 && (f(cache1, x))
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    @. x += ϵ * vv
    f(cache2, x)
    @. x -= ϵ * vv
    vecdy = _vec(dy)
    veccache1 = _vec(cache1)
    veccache2 = _vec(cache2)
    @. vecdy = (veccache2 - veccache1) / ϵ
end

function num_jacvec(f, x, v, f0 = nothing)
    vv = reshape(v, axes(x))
    f0 === nothing ? _f0 = f(x) : _f0 = f0
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(minimum(x)))
    vec((f(x .+ ϵ .* vv) .- _f0) ./ ϵ)
end

function num_hesvec!(
    dy,
    f,
    x,
    v,
    cache1 = similar(v),
    cache2 = similar(v),
    cache3 = similar(v),
)
    cache = FiniteDiff.GradientCache(v[1], cache1, Val{:central})
    g = let f = f, cache = cache
        (dx, x) -> FiniteDiff.finite_difference_gradient!(dx, f, x, cache)
    end
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    @. x += ϵ * v
    g(cache2, x)
    @. x -= 2ϵ * v
    g(cache3, x)
    @. dy = (cache2 - cache3) / (2ϵ)
end

function num_hesvec(f, x, v)
    g = (x) -> FiniteDiff.finite_difference_gradient(f, x)
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    x += ϵ * v
    gxp = g(x)
    x -= 2ϵ * v
    gxm = g(x)
    (gxp - gxm) / (2ϵ)
end

function numauto_hesvec!(
    dy,
    f,
    x,
    v,
    cache = ForwardDiff.GradientConfig(f, v),
    cache1 = similar(v),
    cache2 = similar(v),
)
    g = let f = f, x = x, cache = cache
        g = (dx, x) -> ForwardDiff.gradient!(dx, f, x, cache)
    end
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    @. x += ϵ * v
    g(cache1, x)
    @. x -= 2ϵ * v
    g(cache2, x)
    @. dy = (cache1 - cache2) / (2ϵ)
end

function numauto_hesvec(f, x, v)
    g = (x) -> ForwardDiff.gradient(f, x)
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    x += ϵ * v
    gxp = g(x)
    x -= 2ϵ * v
    gxm = g(x)
    (gxp - gxm) / (2ϵ)
end

function autonum_hesvec!(
    dy,
    f,
    x,
    v,
    cache1 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x))))),
    cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x))))),
)
    cache = FiniteDiff.GradientCache(v[1], cache1, Val{:central})
    g = (dx, x) -> FiniteDiff.finite_difference_gradient!(dx, f, x, cache)
    cache1 .= Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x)))))
    g(cache2, cache1)
    dy .= partials.(cache2, 1)
end

function autonum_hesvec(f, x, v)
    g = (x) -> FiniteDiff.finite_difference_gradient(f, x)
    partials.(g(Dual{DeivVecTag}.(x, v)), 1)
end

function num_hesvecgrad!(dy, g, x, v, cache2 = similar(v), cache3 = similar(v))
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    @. x += ϵ * v
    g(cache2, x)
    @. x -= 2ϵ * v
    g(cache3, x)
    @. dy = (cache2 - cache3) / (2ϵ)
end

function num_hesvecgrad(g, x, v)
    T = eltype(x)
    # Should it be min? max? mean?
    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
    x += ϵ * v
    gxp = g(x)
    x -= 2ϵ * v
    gxm = g(x)
    (gxp - gxm) / (2ϵ)
end

function auto_hesvecgrad!(
    dy,
    g,
    x,
    v,
    cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x))))),
    cache3 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x))))),
)
    cache2 .= Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x)))))
    g(cache3, cache2)
    dy .= partials.(cache3, 1)
end

function auto_hesvecgrad(g, x, v)
    y = Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x)))))
    partials.(g(y), 1)
end

### Operator Forms

struct JacVec{F,T1,T2,xType}
    f::F
    cache1::T1
    cache2::T2
    x::xType
    autodiff::Bool
end

function JacVec(f, x::AbstractArray; autodiff = true)
    if autodiff
        cache1 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(x)))
        cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(x)))
    else
        cache1 = similar(x)
        cache2 = similar(x)
    end
    JacVec(f, cache1, cache2, x, autodiff)
end

Base.eltype(L::JacVec) = eltype(L.x)
Base.size(L::JacVec) = (length(L.cache1), length(L.cache1))
Base.size(L::JacVec, i::Int) = length(L.cache1)
Base.:*(L::JacVec, v::AbstractVector) =
    L.autodiff ? auto_jacvec(_x -> L.f(_x), L.x, v) :
    num_jacvec(_x -> L.f(_x), L.x, v)

function LinearAlgebra.mul!(dy::AbstractVector, L::JacVec, v::AbstractVector)
    if L.autodiff
        auto_jacvec!(dy, (_y, _x) -> L.f(_y, _x), L.x, v, L.cache1, L.cache2)
    else
        num_jacvec!(dy, (_y, _x) -> L.f(_y, _x), L.x, v, L.cache1, L.cache2)
    end
end

struct HesVec{F,T1,T2,xType}
    f::F
    cache1::T1
    cache2::T2
    cache3::T2
    x::xType
    autodiff::Bool
end

function HesVec(f, x::AbstractArray; autodiff = true)
    if autodiff
        cache1 = ForwardDiff.GradientConfig(f, x)
        cache2 = similar(x)
        cache3 = similar(x)
    else
        cache1 = similar(x)
        cache2 = similar(x)
        cache3 = similar(x)
    end
    HesVec(f, cache1, cache2, cache3, x, autodiff)
end

Base.size(L::HesVec) = (length(L.cache2), length(L.cache2))
Base.size(L::HesVec, i::Int) = length(L.cache2)
Base.:*(L::HesVec, v::AbstractVector) =
    L.autodiff ? numauto_hesvec(L.f, L.x, v) : num_hesvec(L.f, L.x, v)

function LinearAlgebra.mul!(dy::AbstractVector, L::HesVec, v::AbstractVector)
    if L.autodiff
        numauto_hesvec!(dy, L.f, L.x, v, L.cache1, L.cache2, L.cache3)
    else
        num_hesvec!(dy, L.f, L.x, v, L.cache1, L.cache2, L.cache3)
    end
end

struct HesVecGrad{G,T1,T2,uType}
    g::G
    cache1::T1
    cache2::T2
    x::uType
    autodiff::Bool
end

function HesVecGrad(g, x::AbstractArray; autodiff = false)
    if autodiff
        cache1 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(x)))
        cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(),eltype(x))),eltype(x),1}.(x, ForwardDiff.Partials.(tuple.(x)))
    else
        cache1 = similar(x)
        cache2 = similar(x)
    end
    HesVecGrad(g, cache1, cache2, x, autodiff)
end

Base.size(L::HesVecGrad) = (length(L.cache2), length(L.cache2))
Base.size(L::HesVecGrad, i::Int) = length(L.cache2)
Base.:*(L::HesVecGrad, v::AbstractVector) =
    L.autodiff ? auto_hesvecgrad(L.g, L.x, v) : num_hesvecgrad(L.g, L.x, v)

function LinearAlgebra.mul!(
    dy::AbstractVector,
    L::HesVecGrad,
    v::AbstractVector,
)
    if L.autodiff
        auto_hesvecgrad!(dy, L.g, L.x, v, L.cache1, L.cache2)
    else
        num_hesvecgrad!(dy, L.g, L.x, v, L.cache1, L.cache2)
    end
end


function flip(point::AbstractArray, prob::Float64)
    @assert zero(prob)<=prob<=one(prob) "The probability must be in range [0,1]"
    α = deepcopy(point)
    if prob == zero(prob)
        return α
    elseif prob == one(prob)
        return  flip.(α)
    end
    dist = Geometric(prob)
    index = rand(dist)+1
    while index <= length(α)
        α[index] = flip(α[index])
        index = index + rand(dist)+1
    end
    α
end

flip(bit::Bool) = ~bit

function staticBitMutation(prob::Float64)
    function _Mutation(instance, candidate, fitness)
        tmp = flip(candidate.point, prob)
        (point=tmp, fitness=fitness(instance, tmp))
    end
    _Mutation
end


using Test, LinearAlgebra, Random, DiffEqOperators
using DiffEqBase: is_constant
using DiffEqOperators: DiffEqScaledOperator, DiffEqOperatorCombination, DiffEqOperatorComposition

@testset "Operator Compostion" begin
  Random.seed!(0)
  A1 = rand(2,3)
  A2 = rand(3,2)
  b = rand()
  B = rand(2,2)
  L = DiffEqArrayOperator(A1) * DiffEqArrayOperator(A2) + DiffEqScalar(b) * DiffEqArrayOperator(B)

  # Structure
  @test isa(L, DiffEqOperatorCombination)
  L1, L2 = getops(L)
  @test isa(L1, DiffEqOperatorComposition)
  @test isa(L2, DiffEqScaledOperator)

  # Operations
  Lfull = Matrix(L)
  @test opnorm(L) ≈ opnorm(Lfull)
  @test size(L) == size(Lfull)
  @test L[1,2] ≈ Lfull[1,2]
  u = [1.0, 2.0]; du = zeros(2)
  @test L * u ≈ Lfull * u
  mul!(du, L, u); @test du ≈ Lfull * u
  Lf = factorize(L)
  ldiv!(du, Lf, u); @test Lfull * du ≈ u
  @test exp(L) ≈ exp(Lfull)
end

@testset "Mutable Composite Operators" begin
  A = zeros(2,2); fA = (A,u,p,t) -> fill!(A, t)
  B = zeros(2,2); fB = (B,u,p,t) -> fill!(B, 2t)
  L = DiffEqArrayOperator(A; update_func=fA) * DiffEqArrayOperator(B; update_func=fB)
  @test is_constant(L) == false
  u = [1.0, 2.0]
  t = 1.1
  @test L(u,nothing,t) ≈ fill(t, (2,2)) * (fill(2t, (2,2)) * u)
end


using BenchmarkTools, ResumableFunctions
using ResumableFunctions

const n = 93

const N = BigInt

function direct(a::N, b::N)
  b, a+b
end

function test_direct(n::Int)
  a, b = zero(N), one(N)
  for _ in 1:n-1
    a, b = direct(a, b)
  end
  a
end

@resumable function fibonacci_resumable(n::Int)
  a, b = zero(N), one(N)
  for _ in 1:n
    @yield a
    a, b = b, a + b
  end
end

@noinline function test_resumable(n::Int)
  a = 0
  for v in fibonacci_resumable(n)
    a = v
  end
  a
end

function fibonacci_channel(n::Int, ch::Channel)
  a, b = zero(N), one(N)
  for _ in 1:n
    put!(ch, a)
    a, b = b, a + b
  end
end

@noinline function test_channel(n::Int, csize::Int)
  fib_channel = Channel(c -> fibonacci_channel(n, c); ctype=N, csize=csize)
  a = 0
  for v in fib_channel
    a = v
  end
  a
end

struct Generator
  callee :: Task
  function Generator(f::Function, args...; kwargs...)
      ct = current_task()
      new(@task begin task_local_storage(:caller, ct); f(args...; kwargs...); yield(ct) end)
  end
end

function Base.iterate(gen::Generator, state=nothing)
  ret = yieldto(gen.callee)
  if ret === nothing return nothing end
  ret, nothing
end

function consume(gen::Generator, val=nothing)
  yieldto(gen.callee, val)
end

function produce(val=nothing)
  t = task_local_storage(:caller)
  yieldto(t, val)
end

function fibonacci_task(n::Int)
  a,b = zero(N), one(N)
  for _ in 1:n
      produce(a)
      a, b = b, a+b
  end
end

@noinline function test_task(n::Int)
  a = 0
  for v in Generator(fibonacci_task, n)
    a = v
  end
  a
end

function fibonacci_closure()
  a, b = zero(N), one(N)
  function()
    tmp = a
    a, b = b, a + b
    tmp
  end
end

@noinline function test_closure(n::Int)
  fib_closure = fibonacci_closure()
  a = 0
  for _ in 1:n
    a = fib_closure()
  end
  a
end

function fibonacci_closure_opt()
  a = Ref(zero(N))
  b = Ref(one(N))
  function()
    tmp = a[]
    a[], b[] = b[], a[] + b[]
    tmp
  end
end

@noinline function test_closure_opt(n::Int)
  fib_closure = fibonacci_closure_opt()
  a = 0
  for _ in 1:n
    a = fib_closure()
  end
  a
end

function fibonacci_closure_stm(n::Int)
  _state = Ref(0x00)
  a = Ref{N}()
  b = Ref{N}()
  _iterstate_1 = Ref{Int}()
  _iterator_1 = Ref{UnitRange{Int}}()

  function(_arg::Any=nothing)
    _state[] === 0x00 && @goto _STATE_0
    _state[] === 0x01 && @goto _STATE_1
    error("@resumable function has stopped!")
    @label _STATE_0
    _state[] = 0xff
    _arg isa Exception && throw(_arg)
    a[], b[] = zero(Int), one(Int)
    _iterator_1[] = 1:n
    _iteratornext_1 = iterate(_iterator_1[])
    while _iteratornext_1 !== nothing
      (_, _iterstate_1[]) = _iteratornext_1
      _state[] = 0x01
      return a[]
      @label _STATE_1
      _state[] = 0xff
      _arg isa Exception && throw(_arg)
      (a[], b[]) = (b[], a[] + b[])
      _iteratornext_1 = iterate(_iterator_1[], _iterstate_1[])
    end
  end
end

const FibClosure = typeof(fibonacci_closure_stm(n))

function Base.iterate(f::FibClosure, state=nothing)
  a = f()
  f._state[] === 0xff && return nothing
  a, nothing
end

@noinline function test_closure_stm(n::Int)
  fib_closure = fibonacci_closure_stm(n)
  a = 0
  for v in fibonacci_closure_stm(n)
    a = v
  end
  a
end

struct FibN
  n::Int
end

function Base.iterate(f::FibN, state::Tuple{N, N, Int}=(zero(N), one(N), 1))
  a, b, iters = state
  iters > f.n && return nothing
  a, (b, a + b, iters + 1)
end

@noinline function test_iteration_protocol(n::Int)
  a = 0
  for v in FibN(n)
    a = v
  end
  a
end

isinteractive() || begin
  println("Direct: ")
  @btime test_direct($n)
  @assert test_direct(n) == 7540113804746346429

  println("ResumableFunctions: ")
  @btime test_resumable($n)
  @assert test_resumable(n) == 7540113804746346429

  println("Channels csize=0: ")
  @btime test_channel($n, $0)
  @assert test_channel(n, 0) == 7540113804746346429

  println("Channels csize=1: ")
  @btime test_channel($n, $1)
  @assert test_channel(n, 1) == 7540113804746346429

  println("Channels csize=20: ")
  @btime test_channel($n, $20)
  @assert test_channel(n, 20) == 7540113804746346429

  println("Channels csize=100: ")
  @btime test_channel($n, $100)
  @assert test_channel(n, 100) == 7540113804746346429

  println("Task scheduling")
  @btime test_task($n)
  @assert test_task(n) == 7540113804746346429

  println("Closure: ")
  @btime test_closure($n)
  @assert test_closure(n) == 7540113804746346429

  println("Closure optimised: ")
  @btime test_closure_opt($n)
  @assert test_closure_opt(n) == 7540113804746346429

  println("Closure statemachine: ")
  @btime test_closure_stm($n)
  @assert test_closure_stm(n) == 7540113804746346429

  println("Iteration protocol: ")
  @btime test_iteration_protocol($n)
  @assert test_iteration_protocol(n) == 7540113804746346429
end


using Base.Test

include("grains.jl")

@testset "On squares" begin
    @testset "On square $s" for s = UInt64(1):64
        @test on_square(s) == 2^(s-1)
        @test total_after(s) == 2^s - 1
    end
end

@testset "Invalid values" begin
    @testset "Zero" begin
        @test_throws DomainError on_square(0)
        @test_throws DomainError total_after(0)
    end

    @testset "Negative" begin
        @test_throws DomainError on_square(-1)
        @test_throws DomainError total_after(-1)
    end

    @testset "Greater than 64" begin
        @test_throws DomainError on_square(65)
        @test_throws DomainError total_after(65)
    end
end


"""
    nonsym(nbOne::Ti, size::Ti, diag_l::Ti, diag_u::Ti, spectrum::AbstractVector{Tv}) where {Tv<:Complex, Ti<:Integer}

Generate a non symmetric matrix with default nilpotent matrix and default initialization of matrix.

The usage of `nonsym` is the same as the one of `nonherm`, please refers to that part for more detais.

Please keep in mind the constraints of given spectrum vector and initialization of matrix, click [here](getting_started.md) for more details.

"""
function nonsym(
    nbOne::Ti,
    size::Ti,
    diag_l::Ti,
    diag_u::Ti,
    spectrum::AbstractVector{Tv}
) where {Tv<:Complex, Ti<:Integer}

    Tv2 = eltype(real(spectrum[1]))
    Am = spzeros(Tv2, size, size)
    initMat!(Am, diag_l, diag_u, size)
    nilp = Nilp(nbOne, size)

    return nonsym(size, diag_l, diag_u, spectrum, Am, nilp)

end

"""
    nonsym(nbOne::Ti, size::Ti, diag_l::Ti, diag_u::Ti, spectrum::AbstractVector{Tv}, Am::SparseMatrixCSC{Tv2, Ti},) where {Tv<:Complex, Tv2<:Real, Ti<:Integer}

Generate a non symmetric matrix with default nilpotent matrix and user-provided initialization of matrix.

The usage of `nonsym` is the same as the one of `nonherm`, please refers to that part for more detais.

Please keep in mind the constraints of given spectrum vector and initialization of matrix, click [here](getting_started.md) for more details.

"""
function nonsym(
    nbOne::Ti,
    size::Ti,
    diag_l::Ti,
    diag_u::Ti,
    spectrum::AbstractVector{Tv},
    Am::SparseMatrixCSC{Tv2, Ti},
) where {Tv<:Complex, Tv2<:Real, Ti<:Integer}

    nilp = Nilp(nbOne, size)

    return nonsym(size, diag_l, diag_u, spectrum, Am, nilp)

end

"""
    nonsym(size::Ti, diag_l::Ti, diag_u::Ti, spectrum::AbstractVector{Tv}, Am::SparseMatrixCSC{Tv2, Ti}, nilp::Nilpotent{Ti}) where {Tv<:Complex, Tv2<:Real, Ti<:Integer}

Generate a non symmetric matrix with user-provided nilpotent matrix and initialization of matrix.

The usage of `nonsym` is the same as the one of `nonherm`, please refers to that part for more detais.

Please keep in mind the constraints of given spectrum vector and initialization of matrix, click [here](getting_started.md) for more details.

"""
function nonsym(
    size::Ti,
    diag_l::Ti,
    diag_u::Ti,
    spectrum::AbstractVector{Tv},
    Am::SparseMatrixCSC{Tv2, Ti},
    nilp::Nilpotent{Ti}
) where {Tv<:Complex, Tv2<:Real, Ti<:Integer}

    checkSpectrum(spectrum, size)

    for i = 1:size
        Am[i, i] = real(spectrum[i])
    end

    idx = 1
    while idx < size
        if imag(spectrum[idx]) == 0
            step = 1
        else
            step = 2
        end
        Am[idx, idx+1] = abs(imag(spectrum[idx]))
        Am[idx+1,idx] = -abs(imag(spectrum[idx]))
        idx += step
    end

    matAop = Am

    fact = 1

    for k = 1 : 2 * (nilp.degree-1)
        nilpMat = NilpMat(nilp)
        matAop = matAop * nilpMat - nilpMat * matAop
        fact /= factorial(k,k)
        Am += fact * matAop
    end

    return Am
end


@testset "compactify $Q" for Q in (Rational, FastRational)
    mid = Q(100, 111)
    tol = Q(1, 11)
    res = Q(5, 6)
    @test compactify(low=mid - tol, high=mid + tol) == res
    @test compactify(mid, tol) == res
    @test compactify(mid, eps()) == mid
    @test_throws ArgumentError compactify(mid, eps()/3)
    tm = typemax(Int64)>>12
    low = Q(tm, tm-1)
    high = Q(tm+1, tm-1)
    mid = Q(2tm+1, 2tm-2)
    tol = Q(1, 2tm-2)
    @test compactify(mid, tol) == Q(tm>>1+2, tm>>1+1) # rounding error effect?
    @test compactify(low=low, high=high) == Q(tm>>1+1, tm>>1) # better result
    #                                                                          >>>> !!FIXME!! throws MethodError
    # @test_throws ArgumentError compactify(mid, Q(0))
    @test compactify(low=high, high=high) == high
end


#=
  to_num in Julia + ConstraintSolver.jl

  to_num(List, Base, Num) converts a list of integers to a number for
  a base Base. It is bidirectional but it is really recommended that
  the length of List is fixed.

  See examples below.

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#
using ConstraintSolver, JuMP
using Cbc, GLPK
const CS = ConstraintSolver
include("constraints_utils.jl")

#
# to_num and increasing are defined in constraints_utils.jl
#
function to_num_test(n=5,base=10)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> true,
                                                            # "all_optimal_solutions"=>true,
                                                            "logging"=>[],

                                                            # "traverse_strategy"=>:BFS,
                                                            "traverse_strategy"=>:DFS, # <-
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf, # <-

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>16,

                                                            # "lp_optimizer" => cbc_optimizer, # 45.8s
                                                            # "lp_optimizer" => glpk_optimizer, # 21.5s
                                        ))

    @variable(model, 0 <= x[1:n] <= base-1, Int)
    @variable(model, 0 <= v <= 10^n-1, Int)

    to_num(model, x, base, v)
    increasing(model, x)
    @constraint(model, x in CS.AllDifferent())

    # Solve the problem
    println("solve")
    optimize!(model)

    status = JuMP.termination_status(model)
    println("status:$status")
    if status == MOI.OPTIMAL
        xx = convert.(Integer,JuMP.value.(x))
        vv = convert.(Integer,JuMP.value.(v))

        println("x:$xx")
        println("v:$vv")

        num_sols = MOI.get(model, MOI.ResultCount())
        println("\nnum_sols:$num_sols\n")

        for sol in 1:num_sols
            println("solution #$sol")
            xx = convert.(Integer,JuMP.value.(x,result=sol))
            vv = convert.(Integer,JuMP.value.(v,result=sol))
            println("$xx  <=> $vv\n")

        end
    end
end

to_num_test()


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:3] <= 1)

                     @variable(m, -8.232182974914569 <= q <= 6.070970235703473)

                     add_NL_constraint(m, :(swish(-0.17705543456567296 + 0.5634748900178774*swish(-0.9628319588335761 + 0.5826659544327617*$(x[1]) + -0.6099922924008303*$(x[2]) + 0.9101904633653426*$(x[3])) + 0.6106029608527352*swish(0.06629760484384972 + -0.7340864295918403*$(x[1]) + 0.2309908240081593*$(x[2]) + 0.7867070707759196*$(x[3])) + 0.5384388086098779*swish(-0.07370832879656497 + -0.32073040453400337*$(x[1]) + -0.07649911788086428*$(x[2]) + 0.03499506906851524*$(x[3]))) + swish(0.08053767290922664 + 0.911477397877237*swish(-0.9628319588335761 + 0.5826659544327617*$(x[1]) + -0.6099922924008303*$(x[2]) + 0.9101904633653426*$(x[3])) + 0.7134526118957916*swish(0.06629760484384972 + -0.7340864295918403*$(x[1]) + 0.2309908240081593*$(x[2]) + 0.7867070707759196*$(x[3])) + -0.34097520379571256*swish(-0.07370832879656497 + -0.32073040453400337*$(x[1]) + -0.07649911788086428*$(x[2]) + 0.03499506906851524*$(x[3]))) + swish(0.5992621446321471 + 0.15304126986204158*swish(-0.9628319588335761 + 0.5826659544327617*$(x[1]) + -0.6099922924008303*$(x[2]) + 0.9101904633653426*$(x[3])) + -0.37134199071135*swish(0.06629760484384972 + -0.7340864295918403*$(x[1]) + 0.2309908240081593*$(x[2]) + 0.7867070707759196*$(x[3])) + 0.8747239793986465*swish(-0.07370832879656497 + -0.32073040453400337*$(x[1]) + -0.07649911788086428*$(x[2]) + 0.03499506906851524*$(x[3]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


module ShaleDrillingLikelihood_Production_Test

using ShaleDrillingLikelihood
using Test
using StatsFuns
# using Calculus
using FiniteDiff: finite_difference_gradient
using Optim
using Random
using LinearAlgebra
using BenchmarkTools
using ForwardDiff

using ShaleDrillingLikelihood: _num_x,
    idx_produce_ψ, idx_produce_β, idx_produce_σ2η, idx_produce_σ2u,
    theta_produce, theta_produce_ψ, theta_produce_β, theta_produce_σ2η, theta_produce_σ2u,
    simloglik_produce!, grad_simloglik_produce!, loglik_produce_scalars,
    DataProduce, ObservationGroupProduce, ObservationProduce,
    _x, _y, _xsum, _nu, _i, update_nu!, update_xpnu!, _qm, _psi2,
    SimulationDraws, psi2_wtd_sum_and_sumsq, Observation,
    _nparm, update!

fdgradient(f, x0) = finite_difference_gradient(f, x0, Val{:central})

println("Starting production likelihood tests")

@testset "Production simulated likelihood" begin

    Random.seed!(1234)

    k = 3
    beta = rand(k)
    sigmas = (0.5, 0.3, 0.4,)
    theta = vcat(sigmas[1], beta, sigmas[2:3]...)
    num_i = 10_000
    M = 500

    data = DataProduce(num_i, 10, 10:20, theta)
    allsim = SimulationDraws(M,num_i)

    @test length(theta) == _nparm(data)
    @test _num_x(data) == k

    @test     theta_produce_ψ(  data,theta)  == theta[1]
    @test all(theta_produce_β(  data,theta) .== theta[1 .+ (1:k)])
    @test     theta_produce_σ2η(data,theta)  == theta[end-1]
    @test     theta_produce_σ2u(data,theta)  == theta[end]

    ff(x)          = grad_simloglik_produce!(zeros(length(x)), data, x, allsim, false)
    ffgg!(grad, x) = grad_simloglik_produce!(grad,             data, x, allsim, true)

    tmpgrad = similar(theta)
    ff(theta)

    fd = fdgradient(ff, theta)
    # ad = ForwardDiff.gradient(ff, theta)
    fill!(tmpgrad, 0)
    ffgg!(tmpgrad, theta)
    # @test ad ≈ tmpgrad
    @test norm(fd .- tmpgrad, Inf) < 2e-3
    # @test norm(ad .- tmpgrad, Inf) < 2e-5
    @show fd .- tmpgrad

    od = OnceDifferentiable(ff, ffgg!, ffgg!, theta)
    res = optimize(od, theta, BFGS(), Optim.Options(time_limit = 40.0))
    @test norm(res.minimizer .- theta, Inf) < 0.02
end

println("Beginning speed tests")

M = 500
num_i = 2000
sigmas = (0.5, 0.3, 0.4,)
beta = randn(3)
theta = vcat(sigmas[1], beta, sigmas[2:3]...)
data = DataProduce(num_i, 10, 10:20, theta)
sim = SimulationDraws(M, num_i)
simi = view(sim, 1)
_qm(sim) .= softmax(randn(M))

# @show @benchmark psi2_wtd_sum_and_sumsq(simi)
psi2_wtd_sum_and_sumsq(simi)

obs = Observation(data,1)
simloglik_produce!(obs, theta, simi)

end # module


@testset "interpolations" begin

for interpolation in (Lagrange{1, RefCube, 1}(),
                      Lagrange{1, RefCube, 2}(),
                      Lagrange{2, RefCube, 1}(),
                      Lagrange{2, RefCube, 2}(),
                      Lagrange{2, RefTetrahedron, 1}(),
                      Lagrange{2, RefTetrahedron, 2}(),
                      Lagrange{3, RefCube, 1}(),
                      Serendipity{2, RefCube, 2}(),
                      Lagrange{3, RefTetrahedron, 1}(),
                      Lagrange{3, RefTetrahedron, 2}())

    # Test of utility functions
    ndim = JuAFEM.getdim(interpolation)
    r_shape = JuAFEM.getrefshape(interpolation)
    func_order = JuAFEM.getorder(interpolation)
    @test typeof(interpolation) <: Interpolation{ndim,r_shape,func_order}
    @test typeof(JuAFEM.getlowerdim(interpolation)) <: Interpolation{ndim-1}
    @test typeof(JuAFEM.getlowerorder(interpolation)) <: Interpolation{ndim,r_shape,func_order-1}

    n_basefuncs = getnbasefunctions(interpolation)
    x = rand(Tensor{1, ndim})
    f = (x) -> JuAFEM.value(interpolation, Tensor{1, ndim}(x))
    @test vec(ForwardDiff.jacobian(f, Array(x))') ≈
           reinterpret(Float64, JuAFEM.derivative(interpolation, x))
    @test sum(JuAFEM.value(interpolation, x)) ≈ 1.0

    coords = JuAFEM.reference_coordinates(interpolation)
    for node in 1:n_basefuncs
        N_node = JuAFEM.value(interpolation, coords[node])
        for k in 1:node
            if k == node
                @test N_node[k] ≈ 1.0
            else
                @test N_node[k] ≈ 0.0
            end
        end
    end
end

end


#================================================

Basic typedefs, function / macro defintions, etc. for the module

=================================================#

using Base.Threads

#===========================
Constants
===========================#

const DEFAULT_P = 0.5
const DEFAULT_MAX_HEIGHT = 64
const DEFAULT_NODE_CAPACITY = 256

const FLAG_IS_LEFT_SENTINEL = 0x80
const FLAG_IS_RIGHT_SENTINEL = 0x40
const IS_SENTINEL = FLAG_IS_LEFT_SENTINEL | FLAG_IS_RIGHT_SENTINEL

const NODE_NOT_FOUND = -1

#===========================
Abstract typedefs
===========================#

abstract type AbstractNode{T,M} end
abstract type AbstractSkipList{T,M} end

abstract type LeftSentinel{T,M} end
abstract type RightSentinel{T,M} end

#===========================
Helper functions
===========================#

"""
    random_height(p, args...)

Samples a number from a geometric distribution with parameter ``p`` and uses it
for the height of a new node in a Skiplist.

# Arguments

# Examples
"""
function random_height(p, args...; max_height = DEFAULT_MAX_HEIGHT)
    # This function uses the fact that the c.d.f. of a geometric distribution
    # is 1 - (1 - p)^k. To generate the height for a new node in a skip list,
    # we want it to be distributed as a geometric RV plus one.
    #
    # To perform this sampling, we randomly sample X ∈ [0,1], and find the
    # smallest value of k for which cdf(k) > X. We observe that
    #
    #           1 - (1 - p)^k   ≥    X                          =>
    #           (1 - p)^k       ≤    1 - X                      =>
    #           k log(1 - p)    ≤    log(1 - X)                 =>
    #           k               ≥    log(1 - X) / log(1 - p)
    #
    # (The inequality is flipped in the last step since log(1 - p) is necessarily
    # negative.) We can simplify this further by observing that Y = 1 - X has
    # the same distribution as X (i.e., Uniform([0,1])). As a result, to sample
    # a new random number, all we need to do is find the smallest integer k
    # satisfying k ≥ log(Y) / log(1 - p) for some Y ~ Uniform([0,1]), which
    # implies that
    #
    #           k = ⌈log(Y) / log(1 - p)⌉
    #

    p_scaler = 1 / log(1 - p)
    Y = rand(args...)
    @.(ceil(Int64, log(Y) * p_scaler) |> x -> min(max_height, x))
end

function _check_mode(M)
    if M != :List && M != :Set
        "SkipList mode $M is not recognized. Valid options are :List and :Set" |>
        ErrorException |>
        throw
    end
end