using JuMP, GLPK, Cbc

function graph_coloring(; fast = true, benchmark = false)
    m = Model(optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0))

    @variable(m, max_color, Int)
    @variable(m, c[1:49, 1:4], Bin)

    @objective(m, Min, max_color)

    for s in 1:49
        @constraint(m, sum(c[s, 1:4]) == 1)
    end

    if fast
        for s in 1:49
            @constraint(m, sum(c[s, i] * i for i in 1:4) <= max_color)
        end
    else
        for s in 1:49, i in 1:4
            @constraint(m, max_color >= i * c[s, i])
        end
    end

    for i in 1:4
        @constraint(m, c[1, i] + c[11, i] <= 1)
        @constraint(m, c[1, i] + c[8, i] <= 1)
        @constraint(m, c[11, i] + c[8, i] <= 1)
        @constraint(m, c[11, i] + c[22, i] <= 1)
        @constraint(m, c[11, i] + c[24, i] <= 1)
        @constraint(m, c[24, i] + c[22, i] <= 1)
        @constraint(m, c[24, i] + c[36, i] <= 1)
        @constraint(m, c[22, i] + c[8, i] <= 1)
        @constraint(m, c[22, i] + c[23, i] <= 1)
        @constraint(m, c[22, i] + c[36, i] <= 1)
        @constraint(m, c[8, i] + c[2, i] <= 1)
        @constraint(m, c[8, i] + c[6, i] <= 1)
        @constraint(m, c[8, i] + c[23, i] <= 1)
        @constraint(m, c[23, i] + c[6, i] <= 1)
        @constraint(m, c[23, i] + c[31, i] <= 1)
        @constraint(m, c[23, i] + c[41, i] <= 1)
        @constraint(m, c[23, i] + c[36, i] <= 1)
        @constraint(m, c[36, i] + c[31, i] <= 1)
        @constraint(m, c[36, i] + c[41, i] <= 1)
        @constraint(m, c[2, i] + c[4, i] <= 1)
        @constraint(m, c[2, i] + c[5, i] <= 1)
        @constraint(m, c[2, i] + c[6, i] <= 1)
        @constraint(m, c[6, i] + c[5, i] <= 1)
        @constraint(m, c[6, i] + c[15, i] <= 1)
        @constraint(m, c[6, i] + c[31, i] <= 1)
        @constraint(m, c[31, i] + c[15, i] <= 1)
        @constraint(m, c[31, i] + c[33, i] <= 1)
        @constraint(m, c[31, i] + c[37, i] <= 1)
        @constraint(m, c[31, i] + c[41, i] <= 1)
        @constraint(m, c[41, i] + c[37, i] <= 1)
        @constraint(m, c[41, i] + c[40, i] <= 1)
        @constraint(m, c[4, i] + c[10, i] <= 1)
        @constraint(m, c[4, i] + c[5, i] <= 1)
        @constraint(m, c[5, i] + c[10, i] <= 1)
        @constraint(m, c[5, i] + c[13, i] <= 1)
        @constraint(m, c[5, i] + c[15, i] <= 1)
        @constraint(m, c[15, i] + c[13, i] <= 1)
        @constraint(m, c[15, i] + c[35, i] <= 1)
        @constraint(m, c[15, i] + c[33, i] <= 1)
        @constraint(m, c[33, i] + c[35, i] <= 1)
        @constraint(m, c[33, i] + c[37, i] <= 1)
        @constraint(m, c[37, i] + c[46, i] <= 1)
        @constraint(m, c[37, i] + c[40, i] <= 1)
        @constraint(m, c[40, i] + c[46, i] <= 1)
        @constraint(m, c[40, i] + c[47, i] <= 1)
        @constraint(m, c[10, i] + c[7, i] <= 1)
        @constraint(m, c[10, i] + c[13, i] <= 1)
        @constraint(m, c[13, i] + c[7, i] <= 1)
        @constraint(m, c[13, i] + c[26, i] <= 1)
        @constraint(m, c[13, i] + c[35, i] <= 1)
        @constraint(m, c[35, i] + c[26, i] <= 1)
        @constraint(m, c[35, i] + c[32, i] <= 1)
        @constraint(m, c[35, i] + c[39, i] <= 1)
        @constraint(m, c[35, i] + c[46, i] <= 1)
        @constraint(m, c[46, i] + c[39, i] <= 1)
        @constraint(m, c[46, i] + c[43, i] <= 1)
        @constraint(m, c[46, i] + c[47, i] <= 1)
        @constraint(m, c[47, i] + c[43, i] <= 1)
        @constraint(m, c[7, i] + c[26, i] <= 1)
        @constraint(m, c[7, i] + c[49, i] <= 1)
        @constraint(m, c[26, i] + c[21, i] <= 1)
        @constraint(m, c[26, i] + c[32, i] <= 1)
        @constraint(m, c[32, i] + c[21, i] <= 1)
        @constraint(m, c[32, i] + c[25, i] <= 1)
        @constraint(m, c[32, i] + c[29, i] <= 1)
        @constraint(m, c[32, i] + c[34, i] <= 1)
        @constraint(m, c[32, i] + c[39, i] <= 1)
        @constraint(m, c[39, i] + c[34, i] <= 1)
        @constraint(m, c[39, i] + c[38, i] <= 1)
        @constraint(m, c[39, i] + c[44, i] <= 1)
        @constraint(m, c[39, i] + c[42, i] <= 1)
        @constraint(m, c[39, i] + c[43, i] <= 1)
        @constraint(m, c[43, i] + c[42, i] <= 1)
        @constraint(m, c[49, i] + c[21, i] <= 1)
        @constraint(m, c[49, i] + c[25, i] <= 1)
        @constraint(m, c[21, i] + c[25, i] <= 1)
        @constraint(m, c[42, i] + c[44, i] <= 1)
        @constraint(m, c[42, i] + c[48, i] <= 1)
        @constraint(m, c[25, i] + c[17, i] <= 1)
        @constraint(m, c[25, i] + c[29, i] <= 1)
        @constraint(m, c[3, i] + c[12, i] <= 1)
        @constraint(m, c[12, i] + c[9, i] <= 1)
        @constraint(m, c[12, i] + c[14, i] <= 1)
        @constraint(m, c[9, i] + c[14, i] <= 1)
        @constraint(m, c[9, i] + c[16, i] <= 1)
        @constraint(m, c[14, i] + c[16, i] <= 1)
        @constraint(m, c[14, i] + c[19, i] <= 1)
        @constraint(m, c[14, i] + c[18, i] <= 1)
        @constraint(m, c[18, i] + c[19, i] <= 1)
        @constraint(m, c[16, i] + c[18, i] <= 1)
        @constraint(m, c[16, i] + c[17, i] <= 1)
        @constraint(m, c[16, i] + c[20, i] <= 1)
        @constraint(m, c[17, i] + c[20, i] <= 1)
        @constraint(m, c[17, i] + c[28, i] <= 1)
        @constraint(m, c[17, i] + c[30, i] <= 1)
        @constraint(m, c[17, i] + c[29, i] <= 1)
        @constraint(m, c[20, i] + c[28, i] <= 1)
        @constraint(m, c[30, i] + c[27, i] <= 1)
        @constraint(m, c[30, i] + c[29, i] <= 1)
        @constraint(m, c[30, i] + c[34, i] <= 1)
        @constraint(m, c[27, i] + c[34, i] <= 1)
        @constraint(m, c[29, i] + c[34, i] <= 1)
        @constraint(m, c[34, i] + c[38, i] <= 1)
        @constraint(m, c[38, i] + c[45, i] <= 1)
        @constraint(m, c[45, i] + c[44, i] <= 1)
        @constraint(m, c[44, i] + c[48, i] <= 1)
    end

    optimize!(m)
    if !benchmark
        println(JuMP.objective_value(m))
        println(JuMP.termination_status(m))
        _, color_washington = findmax(JuMP.value.(c[1, :]))
        println("color_washington: ", color_washington)
    end
end


using ApproxFun, Plots
import ApproxFun: factor, component

L = 1.0
T = 0.1
d = Interval(0.,T) × Interval(0.,L)
Dx = Derivative(d,[0,1])
Dt = Derivative(d,[1,0])

α = 1.0

B = [ldirichlet() ⊗ I;
     I ⊗ ldirichlet();
     I ⊗ rdirichlet();
     (I ⊗ rneumann()) - α*(I ⊗ lneumann());]

q0 = Fun(x->x^2*(L-x)^2 , factor(d,2))
f0 = zeros(component(∂(d),1))
g0 = zeros(component(∂(d),3))

u = \([B;Dt + Dx^3] , [q0;f0;g0;0;0] ; tolerance=1E-5)
surface(u)


"""
   mₑ : Electron rest mass (in eV)
"""
const mₑ = ustrip((ElectronMass * SpeedOfLightInVacuum^2) |> u"eV")

"""
    λₑ(E::Float64)

Wavelength of an electron in cm. (non-relativistic)

  * E in eV
"""
function λₑ(E::Float64)
    return 1.0e2 * ustrip(PlanckConstant / sqrt(2.0 * ElectronMass * ElementaryCharge * E * u"eV")) # cm
end

"""
    kₑ(E::Float64)

Electron wavenumber (inverse wavelength) in rad⋅cm⁻¹.
"""
kₑ(E::Float64) = 2π / λₑ(E)
# = 0.01*sqrt(2.0*9.109e-31*1.602e-19*E)/1.05457e-34


using Revise # reduce need for recompile
using Plots
using LinearAlgebra

push!(LOAD_PATH, "./src") # user defined modules
using CommonUtils
using Basis1D
using Basis2DTri
using UniformTriMesh

"Approximation parameters"
N   = 3 # The order of approximation
K1D = 16
CFL = .5
T   = 2.0 # endtime

"Mesh related variables"
VX,VY,EToV = uniform_tri_mesh(K1D,K1D)
FToF = connect_mesh(EToV,tri_face_vertices())
Nfaces,K = size(FToF)

"Construct matrices on reference elements"
r, s = nodes_2D(N)
V = vandermonde_2D(N, r, s)
Vr, Vs = grad_vandermonde_2D(N, r, s)
M = inv(V*V')
Dr = Vr/V
Ds = Vs/V

"Nodes on faces, and face node coordinate"
r1D, w1D = gauss_quad(0,0,N)
Nfp = length(r1D)
e = ones(Nfp,1)
z = zeros(Nfp,1)
rf = [r1D; -r1D; -e];
sf = [-e; r1D; -r1D];
wf = vec(repeat(w1D,3,1));
nrJ = [z; e; -e]
nsJ = [-e; e; z]
Vf = vandermonde_2D(N,rf,sf)/V
LIFT = M\(transpose(Vf)*diagm(wf)) # lift matrix

"Construct global coordinates"
r1,s1 = nodes_2D(1)
V1 = vandermonde_2D(1,r,s)/vandermonde_2D(1,r1,s1)
x = V1*VX[transpose(EToV)]
y = V1*VY[transpose(EToV)]

"Connectivity maps"
xf,yf = (x->Vf*x).((x,y))
mapM,mapP,mapB = build_node_maps((xf,yf),FToF)
mapM = reshape(mapM,Nfp*Nfaces,K)
mapP = reshape(mapP,Nfp*Nfaces,K)

"Make periodic"
LX,LY = (x->maximum(x)-minimum(x)).((VX,VY)) # find lengths of domain
mapPB = build_periodic_boundary_maps(xf,yf,LX,LY,Nfaces*K,mapM,mapP,mapB)
mapP[mapB] = mapPB

"Geometric factors and surface normals"
rxJ, sxJ, ryJ, syJ, J = geometric_factors(x, y, Dr, Ds)
nxJ = (Vf*rxJ).*nrJ + (Vf*sxJ).*nsJ;
nyJ = (Vf*ryJ).*nrJ + (Vf*syJ).*nsJ;
sJ = @. sqrt(nxJ^2 + nyJ^2)

"initial conditions"
u = @. exp(-25*(x^2+y^2))

"Time integration"
rk4a,rk4b,rk4c = rk45_coeffs()
CN = (N+1)*(N+2)/2  # estimated trace constant
dt = CFL * 2 / (CN*K1D)
Nsteps = convert(Int,ceil(T/dt))
dt = T/Nsteps

"pack arguments into tuples"
ops = (Dr,Ds,LIFT,Vf)
vgeo = (rxJ,sxJ,ryJ,syJ,J)
fgeo = (nxJ,nyJ,sJ)
nodemaps = (mapP,mapB)

function rhs(u,ops,vgeo,fgeo,nodemaps)
    # unpack args
    Dr,Ds,LIFT,Vf = ops
    rxJ,sxJ,ryJ,syJ,J = vgeo
    nxJ,nyJ,sJ = fgeo
    (mapP,mapB) = nodemaps

    uM = Vf*u # can replace with nodal extraction
    uP = uM[mapP]
    du = uP-uM

    ur = Dr*u
    us = Ds*u
    ux = @. rxJ*ur + sxJ*us

    tau = .5
    rhsu = ux + LIFT*(@. .5*du*nxJ - tau*du*abs(nxJ))

    return -rhsu./J
end

resu = zeros(size(x))
for i = 1:Nsteps
    global u, resu # for scoping - these variables are updated

    for INTRK = 1:5
        rhsu = rhs(u,ops,vgeo,fgeo,nodemaps)
        @. resu = rk4a[INTRK]*resu + dt*rhsu
        @. u += rk4b[INTRK]*resu
    end

    if i%10==0 || i==Nsteps
        println("Number of time steps $i out of $Nsteps")
    end
end

"plotting nodes"

gr(aspect_ratio=1,legend=false,
   markerstrokewidth=0,markersize=2)

rp, sp = equi_nodes_2D(10)
Vp = vandermonde_2D(N,rp,sp)/V
vv = Vp*u
scatter(Vp*x,Vp*y,vv,zcolor=vv,camera=(0,90))


export svdvals_gkl

"""
Compute the largest singular values of a matrix A using the Golub-Kahan-Lanczos
bidiagonalization method `\cite{Golub1965}` implemented in `svdvals_gkl()`.

This implementation uses full one-sided reorthogonalization as described in
`\cite{Simon2000}`.

Inputs
------

- `A`    : The matrix or matrix-like object whose truncated SVD is desired.
           `A` must support `size()`, `A*v` and `A'*v` methods.

- `nvals`: Number of singular values requested.
           Default: 6

- `v0`   : Initial guess vector.
           Default: a randomized unit vector.

Keyword arguments
-----------------

- `maxiter`: Maximum number of iterations.
             Default: the smaller dimension of A.

- `βth`    : The threshold value of β below which an invariant subspace is
             deemed to be found.
             Default: `0.1*√eps(eltype(A))`

- `σth`    : The threshold value below which a Ritz value estimate of the
             singular value is considered to be converged.
             Default: `0.1*√eps(eltype(A))`

Outputs
-------

- `converged_values`   : The requested singular values

- `convergence_history`: A `Dict{Symbol,Any}` containing the following keys:
    - `:isconverged` : Did the calculation converge?
    - `:iters`       : Number of iterations run
    - `:reorth`      : Number of reorthogonalizations
    - `:mvps`        : Number of matrix-vector products computed
    - `:B`           : The `Bidiagonal` matrix computed during the
                       bidiagonalization process
    - `:β`           : Norm of Lanczos iterate
    - `:ω²`          : The Frobenius norm of the difference between A and the
                       low rank approximation to A computable from the
                       currently computed singular values [Simon2000]
    - `:vals`        : Ritz values computed at each iteration
    - `:valerrs`     : Error bounds on Ritz values at each iteration

Side effects
------------

If an invariant subspace is found which smaller than the either dimension of A,
an informational message is printed and only the singular values corresponding
to this subspace are returned.

References
----------

```bibtex
@article{Golub1965,
    author = {Golub, G. and Kahan, W.},
    doi = {10.1137/0702016},
    journal = {Journal of the Society for Industrial and Applied Mathematics
        Series B Numerical Analysis},
    volume = 2,
    number = 2,
    pages = {205--224},
    title = {Calculating the Singular Values and Pseudo-Inverse of a Matrix},
    year = 1965
}

@article{Simon2000,
    author = {Simon, Horst D. and Zha, Hongyuan},
    doi = {10.1137/S1064827597327309},
    journal = {SIAM Journal on Scientific Computing},
    number = 6,
    pages = {2257--2274},
    title = {Low-Rank Matrix Approximation Using the {L}anczos Bidiagonalization
        Process with Applications},
    volume = 21,
    year = 2000
}
```
"""
function svdvals_gkl(A, nvals::Int=6, v0=randn(size(A,2));
    maxiter::Int=minimum(size(A)),
    βth::Real = 0.1*√eps(eltype(A)),
    σth::Real = 0.1*√eps(eltype(A)),
    krestart::Int = 20,
    ktrunc::Int = nvals
    )

    m, n = size(A)
    T = eltype(A)
    Tσ= eltype(svdfact(A[1,1])[:S])
    @assert length(v0) == n
    p = v0

    αs = T[]
    βs = T[]

    α = Inf
    r = Array(T, m)
    u = Array(T, m) #Same as r
    v = Array(T, n) #Same as p

    V = Array(T, min(m, n), 0) #List of converged right vectors

    ω² = ω²₀ = vecnorm(A)^2

    convergence_history = Dict()
    convergence_history[:isconverged] = false
    convergence_history[:ω²] = [ω²]
    convergence_history[:vals] = Vector{Tσ}[]
    convergence_history[:valerrs] = Vector{Tσ}[]
    convergence_history[:mvps] = 0
    convergence_history[:reorth] = 0

    converged_vectors = Int[]
    local β, k, oldσ, S, Δσ, converged_values, converged_values_errors
    fasterror = true
    for k=1:maxiter
        #Reorthogonalize right vectors [Simon2000]
        #Here we use doubly reorthogonalized classical Gram-Schmidt
        if m ≥ n && k>1
            p -= V*(V'p)
            p -= V*(V'p)
            convergence_history[:reorth] += 2size(V,2)
        end

        β = norm(p)

        v[:] = p/β
        #r = A*v
        #k>1 && (r -= β*u)
        A_mul_B!(r, A, v)
        k>1 && (axpy!(-β, u, r))

        #Reorthogonalize left vectors [Simon2000]
        if m < n
            r -= V*(V'r)
            r -= V*(V'r)
            convergence_history[:reorth] += 2size(V,2)
        end

        α = norm(r)
        u[:] = r/α
        #p = A'u - α*v
        Ac_mul_B!(p, A, u)
        axpy!(-α, v, p)

        β = norm(p)
        push!(αs, α)
        push!(βs, β)
        convergence_history[:mvps] += 2

        #Update Simon-Zha approximation error
        ω² -= α^2
        length(βs) > 1 && (ω² -= βs[end-1]^2)
        push!(convergence_history[:ω²], ω²)

        #Compute error bars on singular values
        d = √(α*β)

        #if fasterror
        #    σ = svdvals(Bidiagonal(αs, βs[1:end-1], false))
        #    Δσ = k==1 ? [d] : svdvals_error!(Δσ, oldσ, σ, d)
        #    fasterror = fasterror && all(isfinite(Δσ)) && all(Δσ.>√eps(Tσ))
        #    oldσ = σ
        #end

        #if !fasterror
            S = svdfact(Bidiagonal(αs, βs[1:end-1], false))
            σ = svdvals(S)
            e1= abs(S[:U][end,:])
            e2= abs(S[:Vt][:,end])
            Δσ= Tσ[d*min(e1[i], e2[i]) for i in eachindex(σ)]
        #end

        #Debug svdvals_error!
        #println("Iteration $k: $(√eps())")
        #for i=1:min(k,10)
        #    println(σ[i], '\t', Δσ[i], '\t', Δσ2[i], '\t')
        #end
        #println()
        #end

        push!(convergence_history[:vals], σ)
        push!(convergence_history[:valerrs], Δσ)

        #Check number of converged values
        converged_values = Array(Tσ, 0)
        converged_values_errors = Array(Tσ, 0)
        converged_vectors = Array(Int, 0)
        for i in eachindex(σ)
            if Δσ[i] ≤ σth
                push!(converged_values, σ[i])
                push!(converged_values_errors, Δσ[i])
                push!(converged_vectors, i)
            end
        end

        #true: do complete reorthogonalization
        if true
            V = [V m≥n?v:u]
        end

        #If invariant subspace has been found, stop
        if β ≤ βth
            if k != minimum(size(A))
                #In exact arithmetic, Lanczos is guaranteed to find an
                #invariant subspace corresponding to the entire range of the
                #matrix. For small test matrices it is entirely possible to
                #attain this limit, so that the Krylov subspace is of rank k.
                #For other matrices this may indicate to continue with a
                #different choice of starting vector.
                info("Invariant subspace of dimension $k found")
            end
            convergence_history[:isconverged] = true
            converged_values = σ
            break
        end

        #If at least n converged values have been found, stop
        if length(converged_values) ≥ nvals
            convergence_history[:isconverged] = true
            break
        end
    end

    convergence_history[:iters] = k
    convergence_history[:B] = Bidiagonal(αs, βs[1:end-1], false)
    convergence_history[:β] = βs
    @assert issorted(converged_values, rev=true)
    converged_values, convergence_history
end

"""
Compute approximate errors in the Rayleigh-Ritz approximation to singular
values in the Golub-Kalan-Lanczos bidiagonalization method implemented in
`svdvals_gkl()`.

Warning
-------
The approximations here are only good to ~`√eps()`, since differences smaller
than that cannot be captured reliably in floating point arithmetic.

Inputs
------

- `Δθ`: Old error vector. Will get appended to.
- `τ` : Old Ritz values, sorted either in ascending or descending order.
- `θ` : Current Ritz values, sorted in the same order as `τ`
- `β` : The norm of the current residual vector

Outputs
-------

- `Δθ`: The new error vector, updated in place.

Implementation notes
--------------------

The notation used corresponds to `\cite[2/e, Corollary 7.9.2]{Parlett1980}`.

A given Ritz value ``θ`` has an associated error bound which is the product of
the norm of the residual and last entry of its associated eigenvector
`\cite[Section 13.2]{Parlett1980}`.

`\cite{Hill1992}` provides an approximation to the last component of each
eigenvector given the current Ritz values and the previous Ritz values, using
the Cauchy interlacing theorem.

This function implements the upper bound of Theorem 3 in `\cite{Hill1992}`,
adapted to the computation of error bounds on Ritz values corresponding to
singular values.

The simple modification is to square the Ritz values to get estimates of
(nonzero) eigenvalues of ``A'A`` and apply Theorem 3 of `\cite{Hill1992}` to
the Ritz values of the eigenvalues, thus deriving the final formulae used here:

```math
\Delta\theta_{j}<\beta\times
\begin{cases}
    \sqrt{\frac{\tau_{1}-\theta_{1}}{\theta_{2}-\theta_{1}}
          \frac{\tau_{1}+\theta_{1}}{\theta_{2}+\theta_{1}}}     & i=1\\
    \sqrt{\frac{\theta_{j}-\tau_{j-1}}{\theta_{j}-\theta_{j-1}}
          \frac{\tau_{j}-\theta_{j}}{\theta_{j+1}-\theta_{j}}
          \frac{\theta_{j}+\tau_{j-1}}{\theta_{j}+\theta_{j-1}}
          \frac{\tau_{j}+\theta_{j}}{\theta_{j+1}+\theta_{j}}}   & i=2,\dots,k-1\\
    \sqrt{\frac{\theta_{k}-\tau_{k-1}}{\theta_{k}-\theta_{k-1}}
          \frac{\theta_{k}+\tau_{k-1}}{\theta_{k}+\theta_{k-1}}} & i=k
\end{cases}
```

We also sneak in an `abs()` before taking the outermost square root to avoid
roundoff error when computing very small error bounds. It also has the
advantage of rendering the formulae agnostic to sort order. So long as `τ` and
`θ` are sorted the same way, the same formulae apply since reversing the sort
order merely reverses the indices in the Cauchy interlacing theorem.

References
----------

```bibtex
@book{Parlett1980,
    author = {Parlett, Beresford N},
    address = {Philadelphia, PA},
    doi = {10.1137/1.9781611971163},
    publisher = {SIAM},
    title = {The symmetric eigenvalue problem},
    year = 1980
}

@article{Hill1992,
	Author = {Hill, R O, Jr. and B N Parlett},
	Doi = {10.1137/0613019},
	Journal = {SIAM Journal on Matrix Analysis and Applications},
	Number = 1,
	Pages = {239-247},
	Title = {Refined Interlacing Properties},
	Volume = 13,
	Year = 1992
}
```
"""
function svdvals_error!{T}(
        Δθ::AbstractVector, τ::AbstractVector{T},
        θ::AbstractVector{T}, β::Real=1.0)

    n = length(θ)
    @assert (issorted(τ, rev=true) && issorted(θ, rev=true)) ||
            (issorted(τ) && issorted(θ))
    @assert length(Δθ)==length(τ)==n-1

    #Check Cauchy interlacing property
    #for i in eachindex(τ)
    #    if !(θ[i]+length(θ)^2*eps(T) ≥ τ[i] ≥ θ[i+1]-length(θ)^2*eps(T))
    #        warn("Interlacing violated $i: $(θ[i]) ≥ $(τ[i]) ≥ $(θ[i+1])")
    #        println(min(abs(τ[i] - θ[i])/eps(T), abs(θ[i+1] - τ[i])/eps(T)))
    #    end
    #end

    Δθ[1] = β*√abs((τ[1]-θ[1])*(τ[1]+θ[1])/
               ((θ[2]-θ[1])*(θ[2]+θ[1])))

    for j=2:n-1
        Δθ[j] = β*√abs(((θ[j]-τ[j-1])/(θ[j]-θ[j-1])) *
                    ((θ[j]+τ[j-1])/(θ[j]+θ[j-1])) *
                    ((τ[j]-θ[j])/(θ[j+1]-θ[j])) *
                    ((τ[j]+θ[j])/(θ[j+1]+θ[j])))
    end
    if n>1
        push!(Δθ, β*√abs((θ[n]-τ[n-1])*(θ[n]+τ[n-1])/
                        ((θ[n]-θ[n-1])*(θ[n]+θ[n-1]))))
    end
    Δθ
end


using MovingWeightedLeastSquares
using Base.Random
using Base.Test
using Base.Iterators
using DynamicPolynomials, MultivariatePolynomials

include("sample-generator.jl")
include("cell-linked-list.jl")
include("1d-approximation.jl")
include("single-array-input.jl")
include("1d-derivative.jl")
include("2d-approximation.jl")
include("2d-derivative.jl")


function ReadCSV(name; type=Float32)
    pf = open(name)
    line = readline(pf)
    sep = findfirst(" ", line)[1]
    len = parse(Int, line[1:(sep-1)])
    dim = parse(Int, line[(sep+1):end])
    Points = zeros(type, len, dim)

    for i = 1:len
        line = readline(pf)
        for j = 1:(dim-1)
            sep = findfirst(" ", line)[1]
            val = parse(type, line[1:(sep-1)])
            line = line[(sep+1):end]
            Points[i, j] = val
        end
        val = parse(type, line)
        Points[i, dim] = val
    end

    return  Points

end

function AssignPointsToBlock(Points::Array{Float32,2},
    numOfGrids::Int64, dimensions::Int64)
    grids = [numOfGrids for i = 1:dimensions];
    gridMap1D = grids.^((0:(dimensions-1)));
    gridMap1Doffset = [0; -1*(2:dimensions).^0];
    BlockOfPoint = dropdims(
        mapslices(
        x -> Int32(sum( ((1 .+ div.(min.(x, 0.99999999), 1/numOfGrids)) + gridMap1Doffset ) .*gridMap1D )),
            Points, dims=2),
        dims=2);
    return BlockOfPoint
end

function CountPointsPerBlock(Points::Array{Float32,2},
    numOfGrids::Int64, dimensions::Int64)
    grids = [numOfGrids for i = 1:dimensions];
    gridMap1D = grids.^((0:(dimensions-1)));
    gridMap1Doffset = [0; -1*(2:dimensions).^0];
    PointsPerBlock = zeros(Int32, numOfGrids^dimensions)
    num = size(Points)[1]
    for i = 1:num
        x = Points[i,:]
        x = min.(x, 0.99999999)
        gid = Int64(sum(( (div.(x,1/numOfGrids) .+ 1) + gridMap1Doffset ) .*gridMap1D ))
        PointsPerBlock[gid...] += 1
    end
    IntegralPointsPerBlock = similar(PointsPerBlock)
    IntegralPointsPerBlock[1] = 0#PointsPerBlock[0]
    for i = 2:numOfGrids^dimensions
        IntegralPointsPerBlock[i] = PointsPerBlock[i-1] + IntegralPointsPerBlock[i-1]
    end
    return PointsPerBlock, IntegralPointsPerBlock
end

function ReorderPointsByBlock(Points::Array{Float32,2},
    BlockOfPoint::Array{Int32,1})
    BlockPermutationP = sortperm(BlockOfPoint);
    OrderedPoints = Points[BlockPermutationP,:];
    return OrderedPoints
end


struct PoincareInvariant1st{ET,DT,TT,ΘT}
    equ::ET
    Θ::ΘT
    Δt::TT
    nloop::Int
    ntime::Int
    nsave::Int
    nt::Int
    I::OffsetArray{DT,1,Vector{DT}}
    J::OffsetArray{DT,1,Vector{DT}}
    L::OffsetArray{DT,1,Vector{DT}}
end

function PoincareInvariant1st(f_equ::Function, f_loop::Function, Θ::ΘT, Δt::TT, d::Int, nloop::Int, ntime::Int, nsave::Int=1, DT=Float64) where {TT,ΘT}

    if get_config(:verbosity) > 1
        println()
        println("First Euler-Poincaré Integral Invariant")
        println("=======================================")
        println()
        println(" nloop = ", nloop)
        println(" ntime = ", ntime)
        println(" nsave = ", nsave)
        println(" Δt    = ", Δt)
        println()
    end

    # compute initial conditions
    q₀ = [f_loop(i/nloop) for i in 1:nloop]

    # initialise euation
    equ = f_equ(q₀)

    # create arrays for results
    nt = div(ntime, nsave)

    I = OffsetArray(zeros(DT, nt+1), 0:nt)
    J = OffsetArray(zeros(DT, nt+1), 0:nt)
    L = OffsetArray(zeros(DT, nt+1), 0:nt)

    PoincareInvariant1st{typeof(equ),DT,TT,ΘT}(equ, Θ, Δt, nloop, ntime, nsave, nt, I, J, L)
end

function evaluate_poincare_invariant(pinv::PoincareInvariant1st, sol::Solution)
    p = zeros(size(sol.q[begin],1), size(sol.q,1), size(sol.q,2))
    g = zeros(size(sol.q[begin],1), size(sol.q,1), size(sol.q,2))
    v = zeros(size(sol.q[begin],1), size(sol.q,2))
    γ = zeros(size(sol.q[begin],1), size(sol.q,2))

    compute_one_form(sol.t, sol.q, p, pinv.Θ)

    if isdefined(sol, :λ)
        compute_correction(sol.t, sol.q, sol.λ, g, pinv.equ.g)
    end

    for i in axes(sol.q,1)
        compute_velocity(hcat(sol.q[i,:]...), v)
        pinv.I[i] = compute_loop_integral(hcat(sol.p[i,:]...), v)

        if isdefined(sol, :p)
            pinv.J[i] = compute_loop_integral(hcat(sol.d[i,:]...), v)
        end

        if isdefined(sol, :λ)
            compute_velocity(hcat(sol.λ[i,:]...), γ)
            pinv.L[i] = compute_loop_integral(p[:,i,:] .- pinv.Δt .* g[:,i,:], v .- pinv.Δt .* γ)
        end
    end

    (pinv.I, pinv.J, pinv.L)
end

function write_to_hdf5(pinv::PoincareInvariant1st, sol::Solution, output_file::String)
    # h5open(output_file, isfile(output_file) ? "r+" : "w") do h5
    h5open(output_file, "w") do h5

        write(h5, "t", sol.t)
        write(h5, "I", pinv.I)

        isdefined(sol, :p) ? write(h5, "J", pinv.J) : nothing
        isdefined(sol, :λ) ? write(h5, "L", pinv.L) : nothing

    end
end


"""
FizzBuzz
"""

limit = 100

for counter in 0:1:100;

    if counter % 3 == 0 && counter % 5 == 0
        println("FizzBuzz")
    elseif counter % 3 == 0
        println("Fizz")
    elseif counter % 5 == 0
        println("Buzz")
    else
        println(counter)
    end
end


# Uses new formulation from limbdark paper.

include("s2_stable.jl")

s2diff(b,r) = s2(r,b)-convert(Float64,s2(big(r),big(b)))

# Run some tests:

using PyPlot

nb = 1024; nr = 1024
fig,axes = subplots(3,2)
# Compute r = b+-\epsilon
epsilon = 1e-8
fepsp= zeros(Float64,nb)
fepsbigp= zeros(Float64,nb)
fepsm= zeros(Float64,nb)
fepsbigm= zeros(Float64,nb)
b=linspace(epsilon,2.0,nr)
for ib=1:nb
  fepsbigp[ib] = float(s2(big(b[ib]+epsilon),big(b[ib])))
  fepsp[ib] = float(s2(b[ib]+epsilon,b[ib]))
  fepsbigm[ib] = float(s2(big(b[ib]-epsilon),big(b[ib])))
  fepsm[ib] = float(s2(b[ib]-epsilon,b[ib]))
end
ax = axes[1]
ax[:plot](b,fepsp,label=L"$s_2(b,b+10^{-8})$")
ax[:plot](b,fepsm,label=L"$s_2(b,b-10^{-8})$")
ax[:plot](b,s2.(b,b),label=L"$s_2(b,b)$")
ax[:legend](loc="upper right",fontsize=10)
ax = axes[2]
ax[:plot](b,(fepsp-s2.(b,b))/epsilon,label=L"$s_2(b,b+10^{-8})-s_2(b,b)$")
ax[:plot](b,(fepsbigp-s2.(b,b))/epsilon,linewidth=2,linestyle="--",label=L"$s_2(b,b+10^{-8})-s_2(b,b)$, big")
ax[:plot](b,(fepsm-s2.(b,b))/epsilon,label=L"$s_2(b,b-10^{-8})-s_2(b,b)$")
ax[:plot](b,(fepsbigm-s2.(b,b))/epsilon,linewidth=2,linestyle="--",label=L"$s_2(b,b-10^{-8})-s_2(b,b)$, big")
ax[:legend](loc="right",fontsize=5)
ax = axes[3]
ax[:plot](b,s2diff.(b,b+epsilon),label=L"$s_2(b,b+10^{-8})-$ big")
ax[:plot](b,s2diff.(b,b-epsilon),linewidth=2,linestyle="--",label=L"$s_2(b,b-10^{-8})-$ big")
#ax[:plot](b,fepsp-fepsbigp,label=L"$s_2(b,b+10^{-8})-$ big")
#ax[:plot](b,fepsm-fepsbigm,linewidth=2,linestyle="--",label=L"$s_2(b,b-10^{-8})-$ big")
ax[:legend](loc="lower right",fontsize=6)

# Compute b+r = 1+-\epsilon
epsilon = 1e-8
b=linspace(0.0,2.0,nr)
for ib=1:nr
  fepsbigp[ib] = float(s2(big(1.0-b[ib]+epsilon),big(b[ib])))
  fepsp[ib] = float(s2(1.0-b[ib]+epsilon,b[ib]))
  fepsbigm[ib] = float(s2(big(1.0-b[ib]-epsilon),big(b[ib])))
  fepsm[ib] = float(s2(1.0-b[ib]-epsilon,b[ib]))
end
ax = axes[4]
ax[:plot](b,fepsp,label=L"$s_2(b,1-b+10^{-8})$")
ax[:plot](b,fepsm,label=L"$s_2(b,1-b-10^{-8})$")
ax[:plot](b,s2.(1.0-b,b),label=L"$s_2(b,1-b)$")
ax[:legend](loc="lower right",fontsize=10)
ax = axes[5]
ax[:plot](b,(fepsp-s2.(1.0-b,b))/epsilon,label=L"$s_2(b,1-b+10^{-8})-s_2(b,1-b)$")
ax[:plot](b,(fepsbigp-s2.(1.0-b,b))/epsilon,linewidth=2,linestyle="--",label=L"$s_2(b,1-b+10^{-8})-s_2(b,1-b)$, big")
ax[:plot](b,(fepsm-s2.(1.0-b,b))/epsilon,label=L"$s_2(b,1-b-10^{-8})-s_2(b,1-b)$")
ax[:plot](b,(fepsbigm-s2.(1.0-b,b))/epsilon,linewidth=2,linestyle="--",linestyle="--",label=L"$s_2(b,1-b-10^{-8})-s_2(b,1-b)$, big")
ax[:legend](loc="lower right",fontsize=5)
ax = axes[6]
ax[:plot](b,s2diff.(1.0-b+epsilon,b),label=L"$s_2(b,1-b+10^{-8})-$ big")
ax[:plot](b,s2diff.(1.0-b-epsilon,b),linewidth=2,linestyle="--",label=L"$s_2(b,1-b+10^{-8})-$ big")
ax[:legend](loc="upper right",fontsize=6)
savefig("s2_machine.pdf", bbox_inches="tight")


module aoc_25
    using Test

    export transformer, demo_card_public_key, demo_door_public_key
    export transform_subject, input
    export test_part1
    const transformer = 20201227

    const input = (6270530, 14540258)

#= The card transforms the subject number of 7 according to the card's secret loop size. The result is called the card's public key.
The door transforms the subject number of 7 according to the door's secret loop size. The result is called the door's public key.
The card and door use the wireless RFID signal to transmit the two public keys (your puzzle input) to the other device. Now, the card has the door's public key, and the door has the card's public key. Because you can eavesdrop on the signal, you have both public keys, but neither device's loop size.
The card transforms the subject number of the door's public key according to the card's loop size. The result is the encryption key.
The door transforms the subject number of the card's public key according to the door's loop size. The result is the same encryption key as the card calculated. =#

    const demo_card_public_key = 5764801
    const demo_door_public_key = 17807724

    function transform_subject(subject, max_loop_size, divisor = transformer; value::Int = 1)
        #out = Vector{Int}(undef, max_loop_size )
        for loop_size = 1:max_loop_size
            value *= subject
            value = rem( value, divisor )
            #out[ loop_size ] = value
        end
        value
    end

    export get_loop_numbers
    function get_loop_numbers(public_key::Int; subject = 7)
        loop_size = 0
        value = 1
        n = 1
        while true
            value *= subject
            value = rem( value, transformer )
            if value == public_key
                loop_size = n
                break
            end
            n += 1
        end
        n
    end
    function get_loop_numbers(public_key_1, public_key_2; subject = 7)
        if public_key_1 == public_key_2
            loop_size = get_loop_numbers(public_key_1)
            return loop_size, loop_size
        end
        loop_size_1 = 0
        loop_size_2 = 0
        value = Int(1)
        n = 1
        while loop_size_1 == 0 || loop_size_2 == 0
            value *= subject
            value = rem( value, transformer )
            if value == public_key_1
                loop_size_1 = n
            elseif value == public_key_2
                loop_size_2 = n
            end
            n += 1
        end
        loop_size_1, loop_size_2
    end

    export part1
    function part1()
        loop_sizes = get_loop_numbers.(input)
        encryption_key = transform_subject.(reverse(input), loop_sizes)
        encryption_key[1]
    end

    function test_part1()
        @test 8 == get_loop_numbers( demo_card_public_key )
        @test 11 == get_loop_numbers( demo_door_public_key )
        @test 16311885 == part1()
    end

end


using Statistics

function quantile2(mat::AbstractArray, probs::AbstractVector; byrow::Bool = true)
    (nrow, ncol) = size(mat);
    nprobs = length(probs);

    if byrow
        res = fill(0.0, nrow, nprobs);
        for i = 1:nrow
            res[i, :] = quantile(mat[i, :], probs);
        end
    else
        res = fill(0.0, ncol, nprobs);
        for j = 1:ncol
            res[j, :] = quantile(mat[:, j], probs);
        end
    end
    res;
end

function quantile2_nan(mat::AbstractArray, probs::AbstractVector; byrow::Bool = true)
    (nrow, ncol) = size(mat);
    nprobs = length(probs);

    if byrow
        res = fill(0.0, nrow, nprobs);
        for i = 1:nrow
            res[i, :] = quantile(skipmissing(mat[i, :]), probs);
        end
    else
        res = fill(0.0, ncol, nprobs);
        for j = 1:ncol
            res[j, :] = quantile(skipmissing(mat[:, j]), probs);
        end
    end
    res;
end


@testset "world" begin
    straight_agent = Agent(0.2, 0.0)
    circling_agent = Agent(0.2, 10.0 / 180.0 * pi)
    landmarks =
        [Landmark([2.0, -2.0], 0), Landmark([-1.0, -3.0], 1), Landmark([3.0, 3.0], 2)]
    m = Map()
    camera1 = IdealCamera(landmarks)
    camera2 = IdealCamera(landmarks)
    push!(m, landmarks[1])
    push!(m, landmarks[2])
    push!(m, landmarks[3])

    robot1 = IdealRobot([2.0, 3.0, pi / 6], straight_agent, camera1, 0.05, "blue")
    robot2 = IdealRobot([-2.0, -1.0, pi / 5 * 6], circling_agent, camera2, 0.05, "orange")
    xlim = [-5.5, 10]
    ylim = [-5.5, 10]
    world = World(xlim, ylim)
    push!(world, robot1)
    push!(world, robot2)
    push!(world, m)
    dt = 0.1
    anim = @animate for i = 1:300
        t = dt * i
        annota = "t = $(round(t, sigdigits=3))[s]"
        obsv1 = observations(camera1, robot1.pose_)
        obsv2 = observations(camera2, robot2.pose_)
        p = draw(world, nothing)
        plot(p)
        v, ω = decision(straight_agent, obsv1)
        state_transition(robot1, v, ω, dt)
        v, ω = decision(circling_agent, obsv2)
        state_transition(robot2, v, ω, dt)
    end
    if GUI
        gif(anim, "ch3_robot11.gif", fps = 20)
    end
end

@testset "ch03_robot11" begin
    straight_agent = Agent(0.2, 0.0)
    circling_agent = Agent(0.2, 10.0 / 180.0 * pi)
    landmarks =
        [Landmark([2.0, -2.0], 0), Landmark([-1.0, -3.0], 1), Landmark([3.0, 3.0], 2)]
    m = Map()
    camera1 = IdealCamera(landmarks)
    camera2 = IdealCamera(landmarks)
    push!(m, landmarks)
    robot1 = IdealRobot([2.0, 3.0, pi / 6], straight_agent, camera1, 0.05, "blue")
    robot2 = IdealRobot([-2.0, -1.0, pi / 5 * 6], circling_agent, camera2, 0.05, "orange")
    xlim = [-5.5, 10]
    ylim = [-5.5, 10]
    world = World(xlim, ylim)
    push!(world, robot1)
    push!(world, robot2)
    push!(world, m)
    dt = 0.1
    anim = @animate for i = 1:300
        t = dt * i
        annota = "t = $(round(t, sigdigits=3))[s]"
        obsv1 = observations(camera1, robot1.pose_)
        obsv2 = observations(camera2, robot2.pose_)
        p = draw(world, annota)
        plot(p)
        v, ω = decision(straight_agent, obsv1)
        state_transition(robot1, v, ω, dt)
        v, ω = decision(circling_agent, obsv2)
        state_transition(robot2, v, ω, dt)
    end
    if GUI
        gif(anim, "ch3_robot11.gif", fps = 20)
    end
end


"""
Residual model to pick up the slack between approximation of -∂z(wT) and the truth.
This example uses the ResidualData struct and GP.

    # temperature eq.
    # ∂t(T) + ∂x(uT) + ∂y(vT) + ∂z(wT) = 𝓀∇²T
    # --> horizontal average -->
    #
    #           diffusive  advective
    # ∂t(T) = - ∂z(wT) +  𝓀∂z(∂z(T))
    #
    # GOAL: express wT in terms of only large-scale terms wT = F(T,h,𝓀,μ)
    #  - 𝓀: diffusivity
    #  - μ: coefficient of viscosity

    # T(n+1) = T(n) + Δt * (- ∂z(wT) +  𝓀∂z(∂z(T)))
    # Gp = -∂z(wT)
    # y = T(n+1) - T(n) - Δt * ∂z(∂z(T))

    use GPR to capture the difference between the truth
        Δt(-∂z(wT))
    and the approximation
        T(n+1) - T(n) - Δt * ∂z(∂z(T))

    # target
        approx - truth
"""

# using Statistics, LinearAlgebra, Plots
# include("GP.jl")

struct ResidualData
    truth
    approx
end

# get_truth(𝒟::ProfileData, 𝒢::GP)

δ(ϕ, z) = diff(ϕ) ./ diff(z)

function construct_residual_data(filename, D; N=4)

    # return ProfileData(V, vavg, x, y, x_train, y_train, verification_set, z, zavg, t, Nt, n_train)
end

# filename = "general_strat_16_profiles.jld2"
filename = "dns_profiles.jld2"
D = 16
N = 4

𝒟_wT = construct_profile_data(filename, "wT", D; N=N)
𝒟_T  = construct_profile_data(filename, "T",  D; N=N)
Nt = 𝒟_wT.Nt
t = 𝒟_wT.t
Δt = t[2]-t[1]
κₑ = 𝒟_wT.κₑ

avgd = false
smooth = false

if avgd
    wT = 𝒟_wT.vavg
    T = 𝒟_T.vavg
    z = 𝒟_wT.zavg
    ∂wT∂z = [δ(q, z) for q in wT] # Δt(∂z(wT))
    # Δt * ∂z(∂z(T)) + T(n) - T(n+1)
    Δt∂²T∂z² = [δ( δ(q,z), z[1:end-1] )*Δt for q in T]
    approx = [κₑ*Δt∂²T∂z²[i] .+ T[i][1:end-2] .- T[i+1][1:end-2] for i in 1:(Nt-1)]
else
    wT = 𝒟_wT.v
    if smooth
        wT = smooth_window(wT)
    end

    T  = 𝒟_T.v
    z  = 𝒟_wT.z
    ∂wT∂z = [δ(wT[:,i+1], z)*Δt for i in 1:Nt-1] # Δt(∂z(wT))
    # Δt * ∂z(∂z(T)) + T(n) - T(n+1)
    Δt∂²T∂z² = [δ( δ(T[:,i],z), z[1:end-1] )*Δt for i in 1:Nt]
    approx = [κₑ.*2*Δt∂²T∂z²[i] .+ T[:,i][1:end-2] .- T[:,i+2][1:end-2] for i in 1:(Nt-2)]
end

# Δt(-∂z(wT)) = T(n+1) - T(n) - Δt * ∂z(∂z(T))
# Δt(∂z(wT)) = Δt * ∂z(∂z(T)) + T(n) - T(n+1)

# animation_set = 1:20:(Nt)
animation_set = 1:50
anim = @animate for i in animation_set
    day_string = string(floor(Int, t[i]/86400))
    p1 = scatter(∂wT∂z[i], z, label = "Δt*∂z(wT)", xlims=(-2e-6,2e-6))
    # scatter!(approx[i], z[1:end-2], legend = :topleft, label = "∂z(∂z(T)) + T(n) - T(n+1)", ylabel = "depth", title = "day " * day_string)
    scatter!(approx[i], z[1:end-2], legend = :topright, label = "Δt*κₑ*∂z(∂z(T)) + T(n) - T(n+1)", ylabel = "depth", title = "t = $(𝒟_wT.t[i])")
end
gif(anim, pwd() * "/try2_dns.gif", fps = 2)

function smooth_window(wT)
    # takes in an Nz x Nt array
    Nz, Nt = size(wT)
    smooth_wT = similar(wT)
    for i in 2:(Nz-1)
        for j in 1:Nt
            smooth_wT[i,j] = (wT[i-1,j] + wT[i+1,j])/2
        end
    end
    return smooth_wT
end

smooth_window(wT)


# using Cairo
# using Gadfly
# using DataFrames
# using HypergeometricFunctions
# using SpecialFunctions
# using CategoricalArrays
# using CSV

# function weight(j, s, xi)
#     exp(
#         logabsbinomial(j + s - 2, j - 1)[1] -
#         log(j) +
#         s * log(xi) +
#         (j - 1) * log(1 - xi) +
#         log(_₂F₁(j + s - 1, 1, j + 1, 1 - xi))
#     )
# end

# df_j = DataFrame(j = 1:20)
# df_xi = DataFrame(xi = [0.2, 0.4, 0.6])
# df_s = DataFrame(s = [1, 2, 3, 4])
# df = crossjoin(df_j, df_xi, df_s)
# df[!, :wj] = weight.(df[!, :j], df[!, :s], df[!, :xi])
# df[!, :s] = CategoricalArray(df[!, :s])
# CSV.write("weights.csv", df)

# import Pkg;
# Pkg.activate()
# Pkg.add([
#     "AlgebraOfGraphics",
#     "CategoricalArrays",
#     "CairoMakie",
#     "DataFrames",
#     "HypergeometricFunctions",
#     "LaTeXStrings",
#     "SpecialFunctions"
# ])
# Pkg.activate("")
using AlgebraOfGraphics, CategoricalArrays, CairoMakie, DataFrames
using LaTeXStrings, HypergeometricFunctions, SpecialFunctions
set_aog_theme!()
update_theme!(font = "Arial", fontsize = 28)
theme_minimal
function weight(j, s, xi)
    exp(
        logabsbinomial(j + s - 2, j - 1)[1] -
        log(j) +
        s * log(xi) +
        (j - 1) * log(1 - xi) +
        log(_₂F₁(j + s - 1, 1, j + 1, 1 - xi))
    )
end

df_j = DataFrame(j = 1:20);
df_ϕ = DataFrame(ϕ = [0.1, 0.3, 0.5]);
df_s = DataFrame(s = [1, 2, 3, 4]);
df = crossjoin(df_j, df_ϕ, df_s);
df[!, :w] = weight.(df[!, :j], df[!, :s], df[!, :ϕ]);
df[!, :s] = string.(df[!, :s]);
df[!, :ϕ] = string.(df[!, :ϕ]);
df[!, :ϕ] = LaTeXString.(L"\varphi(x_i \cdot β) = " .* df[!, :ϕ]);
plt = data(df) *
    visual(Lines) *
    mapping(:j, :w, col = :ϕ, color = :s)
axis = (xlabel = L"j", ylabel = L"w_j(x)", color = L"s");
fg = draw(plt; axis, figure = (resolution = (900, 450), px_per_unit = 8))
fg
save("figures/final/lineplot-dgsb-weights.svg", fg, px_per_unit = 3)


instructions = split.(readlines("day8/input.txt"), " ")

visited_lines = Set{Int}()

function execute_line(line_num, instruction, num, acc)
    if instruction == "nop"
        return line_num+1, acc
    elseif instruction == "jmp"
        return line_num+num, acc
    elseif instruction == "acc"
        return line_num+1, acc+num
    end
end

num, acc = execute_line(1, instructions[1][1], parse(Int, instructions[1][2]), 0)
push!(visited_lines, num)
while true
    global acc
    global num, acc = execute_line(num, instructions[num][1], parse(Int, instructions[num][2]), acc)
    if num in visited_lines
        println(acc)
        break
    end
    push!(visited_lines, num)
end


# To be executed from CAGD.jl

todisplay = VERSION <= VersionNumber("1.2") ? true : false

using CAGD
using SparseArrays
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
if todisplay  using ViewerGL; GL = ViewerGL  end

include("./CAGD.jl/paper/examples/cellular_complex.jl");
include("./CAGD.jl/paper/examples/CSG_example.jl");
include("./CAGD.jl/paper/examples/cycles_boundaries.jl");
include("./CAGD.jl/paper/examples/planar_algebra_2D.jl");
include("./CAGD.jl/paper/examples/solid_algebra_3D.jl");
include("./CAGD.jl/paper/examples/space_arrangement_2d.jl");
include("./CAGD.jl/paper/examples/space_arrangement.jl");
include("./CAGD.jl/paper/examples/subdivision_of_2_cells.jl");


# Load Julia packages (libraries) needed.

using StatisticalRethinking
using StatsPlots

ProjDir = @__DIR__

for i in 5:7
  include("$(ProjDir)/m5.$i.jl")
end

# ### snippet 5.29

println()
df = DataFrame!(CSV.File(rel_path("..", "data", "milk.csv"), delim=';'));
df = filter(row -> !(row[:neocortex_perc] == "NA"), df);
df[!, :neocortex_perc] = parse.(Float64, df[:, :neocortex_perc])
df[!, :lmass] = log.(df[:, :mass])
#first(df, 5) |> display

# ### snippet 5.1

scale!(df, [:kcal_per_g, :neocortex_perc, :lmass])
println()

m_5_7 = "
data {
 int < lower = 1 > N; // Sample size
 vector[N] K; // Outcome
 vector[N] NC; // Predictor
 vector[N] M; // Predictor
}

parameters {
 real a; // Intercept
 real bM; // Slope (regression coefficients)
 real bN; // Slope (regression coefficients)
 real < lower = 0 > sigma;    // Error SD
}

model {
  vector[N] mu;               // mu is a vector
  a ~ normal(0, 0.2);           //Priors
  bN ~ normal(0, 0.5);
  bM ~ normal(0, 0.5);
  sigma ~ exponential(1);
  mu = a + bM * M + bN * NC;
  K ~ normal(mu , sigma);     // Likelihood
}
";

# Define the SampleModel and set the output format to :mcmcchains.

m5_7s = SampleModel("m5.7", m_5_7);

# Input data for cmdstan

m5_7_data = Dict("N" => size(df, 1), "M" => df[!, :lmass_s],
    "K" => df[!, :kcal_per_g_s], "NC" => df[!, :neocortex_perc_s]);

# Sample using StanSample

rc = stan_sample(m5_7s, data=m5_7_data);

if success(rc)

  # Describe the draws

  dfa7 = read_samples(m5_7s; output_format=:dataframe)
  println("Normal estimate:")
  p = Particles(dfa7) |> display
  println("Quap estimates:")
  quap(dfa7) |> display

  r1 = plotcoef([m5_5s, m5_6s, m5_7s], [:a, :bN, :bM], "$(ProjDir)/Fig-38a.png",
    "bN & bM Normal estimates")
  r1 |> display

  r2 = plotcoef([m5_5s, m5_6s, m5_7s], [:a, :bN, :bM], "$(ProjDir)/Fig-38b.png",
    "bN & bM Quap estimates", quap)
  r2 |> display
end


"""

This code is based on the original work on the Scilab DoE package by the
following individuals
    Copyright (C) 2012 - 2013 - Michael Baudin
    Copyright (C) 2012 - Maria Christopoulou
    Copyright (C) 2010 - 2011 - INRIA - Michael Baudin
    Copyright (C) 2009 - Yann Collette
    Copyright (C) 2009 - CEA - Jean-Marc Martinez
and
    Copyright (C) 2014, Abraham D. Lee
who converted the package to Python.
"""

"""
$(TYPEDSIGNATURES)

Constructs a Box-Behnken design with size `matrix_size`.

```jldoctest
julia> boxbehnken(4)
27×4 transpose(::Matrix{Float64}) with eltype Float64:
 -1.0  -1.0   0.0   0.0
  1.0  -1.0   0.0   0.0
 -1.0   1.0   0.0   0.0
  1.0   1.0   0.0   0.0
 -1.0   0.0  -1.0   0.0
  1.0   0.0  -1.0   0.0
 -1.0   0.0   1.0   0.0
  1.0   0.0   1.0   0.0
 -1.0   0.0   0.0  -1.0
  1.0   0.0   0.0  -1.0
  ⋮
  0.0  -1.0   0.0   1.0
  0.0   1.0   0.0   1.0
  0.0   0.0  -1.0  -1.0
  0.0   0.0   1.0  -1.0
  0.0   0.0  -1.0   1.0
  0.0   0.0   1.0   1.0
  0.0   0.0   0.0   0.0
  0.0   0.0   0.0   0.0
  0.0   0.0   0.0   0.0

```
"""
function boxbehnken(matrix_size::Int)
    boxbehnken(matrix_size, matrix_size)
end

"""
$(TYPEDSIGNATURES)

Constructs a Box-Behnken design with size `matrix_size` with specified number of center points.

```jldoctest
julia> boxbehnken(4,0)
24×4 transpose(::Matrix{Float64}) with eltype Float64:
 -1.0  -1.0   0.0   0.0
  1.0  -1.0   0.0   0.0
 -1.0   1.0   0.0   0.0
  1.0   1.0   0.0   0.0
 -1.0   0.0  -1.0   0.0
  1.0   0.0  -1.0   0.0
 -1.0   0.0   1.0   0.0
  1.0   0.0   1.0   0.0
 -1.0   0.0   0.0  -1.0
  1.0   0.0   0.0  -1.0
  ⋮
  0.0   1.0   1.0   0.0
  0.0  -1.0   0.0  -1.0
  0.0   1.0   0.0  -1.0
  0.0  -1.0   0.0   1.0
  0.0   1.0   0.0   1.0
  0.0   0.0  -1.0  -1.0
  0.0   0.0   1.0  -1.0
  0.0   0.0  -1.0   1.0
  0.0   0.0   1.0   1.0

```
"""
function boxbehnken(matrix_size::Int, center::Int)
    @assert matrix_size>=3

    A_fact = explicit_fullfactorial(Tuple([-1,1] for i = 1:2))

    rows = floor(Int, (0.5*matrix_size*(matrix_size-1))*size(A_fact)[1])
    A = zeros(rows, matrix_size)

    l = 0
    for i in 1:matrix_size-1
        for j in i+1:matrix_size
            l = l +1
            A[max(0, (l - 1)*size(A_fact)[1])+1:l*size(A_fact)[1], i] = A_fact[:, 1]
            A[max(0, (l - 1)*size(A_fact)[1])+1:l*size(A_fact)[1], j] = A_fact[:, 2]
        end
    end

    if center == matrix_size
        if matrix_size <= 16
            points = [0, 0, 3, 3, 6, 6, 6, 8, 9, 10, 12, 12, 13, 14, 15, 16]
            center = points[matrix_size]
        end
    end

    A = transpose(hcat(transpose(A), transpose(zeros(center, matrix_size))))
end


abstract type DistanceTransform end

struct Chamfer{T} <: DistanceTransform
    dt
end

"""
    Chamfer(img, dt=zeros(Float32, size(img)))

Prepares an array to be `transform`ed using the 3-4 chamfer algorithm
laid out in 'Distance transformations in digital images,
Computer Vision, Graphics, and Image Processing'
[Gunilla Borgefors](https://studentportalen.uu.se/uusp-filearea-tool/download.action?nodeId=214320&toolAttachmentId=64777)

# Arguments
- img: 2D or 3D array to be transformed based on location
    to the nearest background (0) pixel
- tfm: `zeros(Float32, size(img))`
"""
Chamfer(img, dt=zeros(Float32, size(img))) = Chamfer{typeof(dt)}(dt)

"""
    transform(img::AbstractMatrix{T}, tfm::Chamfer)
    transform(img::AbstractArray{T,3}, tfm::Chamfer)

Applies a 3-4 chamfer distance transform to an input image.
Returns an array with spatial information embedded in the
array elements.

# Arguments
- img: 2D or 3D array to be transformed based on location
    to the nearest background (0) pixel

# Citation
'Distance transformations in digital images,
Computer Vision, Graphics, and Image Processing'
[Gunilla Borgefors](DOI: https://doi.org/10.1016/S0734-189X(86)80047-0.)
"""
function transform(img::AbstractMatrix{T}, tfm::Chamfer) where {T}
    dt = tfm.dt
    w, h = size(img)
    # Forward pass
    x = 1
    y = 1
    if img[x, y] == 0
        dt[x, y] = 65535 # some large value
    end
    for x in 1:(w - 1)
        if img[x + 1, y] == 0
            dt[x + 1, y] = 3 + dt[x, y]
        end
    end
    for y in 1:(h - 1)
        x = 1
        if img[x, y + 1] == 0
            dt[x, y + 1] = min(3 + dt[x, y], 4 + dt[x + 1, y])
        end
        for x in 1:(w - 2)
            if img[x + 1, y + 1] == 0
                dt[x + 1, y + 1] = min(
                    4 + dt[x, y], 3 + dt[x + 1, y], 4 + dt[x + 2, y], 3 + dt[x, y + 1]
                )
            end
        end
        x = w

        if img[x, y + 1] == 0
            dt[x, y + 1] = min(4 + dt[x - 1, y], 3 + dt[x, y], 3 + dt[x - 1, y + 1])
        end
    end

    # Backward pass
    for x in (w - 1):-1:1
        y = h
        if img[x, y] == 0
            dt[x, y] = min(dt[x, y], 3 + dt[x + 1, y])
        end
    end
    for y in (h - 1):-1:1
        x = w
        if img[x, y] == 0
            dt[x, y] = min(dt[x, y], 3 + dt[x, y + 1], 4 + dt[x - 1, y + 1])
        end
        for x in 1:(w - 2)
            if img[x + 1, y] == 0
                dt[x + 1, y] = min(
                    dt[x + 1, y],
                    4 + dt[x + 2, y + 1],
                    3 + dt[x + 1, y + 1],
                    4 + dt[x, y + 1],
                    3 + dt[x + 2, y],
                )
            end
        end
        x = 1
        if img[x, y] == 0
            dt[x, y] = min(
                dt[x, y], 4 + dt[x + 1, y + 1], 3 + dt[x, y + 1], 3 + dt[x + 1, y]
            )
        end
    end
    return dt
end

# Helper function for nd arrays
function _transform(img::AbstractMatrix{T}, tfm::Chamfer, dt) where {T}
    w, h = size(img)
    # Forward pass
    x = 1
    y = 1
    if img[x, y] == 0
        dt[x, y] = 65535 # some large value
    end
    for x in 1:(w - 1)
        if img[x + 1, y] == 0
            dt[x + 1, y] = 3 + dt[x, y]
        end
    end
    for y in 1:(h - 1)
        x = 1
        if img[x, y + 1] == 0
            dt[x, y + 1] = min(3 + dt[x, y], 4 + dt[x + 1, y])
        end
        for x in 1:(w - 2)
            if img[x + 1, y + 1] == 0
                dt[x + 1, y + 1] = min(
                    4 + dt[x, y], 3 + dt[x + 1, y], 4 + dt[x + 2, y], 3 + dt[x, y + 1]
                )
            end
        end
        x = w

        if img[x, y + 1] == 0
            dt[x, y + 1] = min(4 + dt[x - 1, y], 3 + dt[x, y], 3 + dt[x - 1, y + 1])
        end
    end

    # Backward pass
    for x in (w - 1):-1:1
        y = h
        if img[x, y] == 0
            dt[x, y] = min(dt[x, y], 3 + dt[x + 1, y])
        end
    end
    for y in (h - 1):-1:1
        x = w
        if img[x, y] == 0
            dt[x, y] = min(dt[x, y], 3 + dt[x, y + 1], 4 + dt[x - 1, y + 1])
        end
        for x in 1:(w - 2)
            if img[x + 1, y] == 0
                dt[x + 1, y] = min(
                    dt[x + 1, y],
                    4 + dt[x + 2, y + 1],
                    3 + dt[x + 1, y + 1],
                    4 + dt[x, y + 1],
                    3 + dt[x + 2, y],
                )
            end
        end
        x = 1
        if img[x, y] == 0
            dt[x, y] = min(
                dt[x, y], 4 + dt[x + 1, y + 1], 3 + dt[x, y + 1], 3 + dt[x + 1, y]
            )
        end
    end
    return dt
end

function transform(img::AbstractArray{T,3}, tfm::Chamfer) where {T}
    dt = tfm.dt
    for z in 1:size(img)[3]
        dt[:, :, z] = _transform(img[:, :, z], tfm, dt[:, :, z])
    end
    return dt
end


"""
        atkinson94(setting, iters, crit)

Runs the Atkinson94 algorithm to find out outliers using LMS method.

# Arguments
- `setting::RegressionSetting`: A regression setting object.
- `iters::Int`: Number of random samples.
- `crit::Float64`: Critical value for residuals

# Description
The algorithm randomly selects initial basic subsets and performs a very robust method, e.g `lms`
to enlarge the basic subset. In each iteration of forward search, the best objective value and parameter
estimates are stored. These values are also used in Atkinson's Stalactite Plot for a visual investigation of
outliers. See `atkinsonstalactiteplot`.

# Output
- `["optimum_index"]`: The iteration number in which the minimum objective is obtained
- `["residuals_matrix"]`: Matrix of residuals obtained in each iteration
- `["outliers"]`: Array of indices of detected outliers
- `["objective"]`: Minimum objective value
- `["coef"]`: Estimated regression coefficients
- `["crit"]`: Critical value given by the user.

# Examples
```julia-repl
julia> reg = createRegressionSetting(@formula(stackloss ~ airflow + watertemp + acidcond), stackloss)
julia> atkinson94(reg)
Dict{Any,Any} with 6 entries:
  "optimum_index"    => 10
  "residuals_matrix" => [0.0286208 0.0620609 … 0.0796249 0.0; 0.0397778 0.120547 … 0.118437 0.0397778; … ; 1.21133 1.80846 … 0.690327 4.14366; 1.61977 0.971592 … 0.616204 3.58098]
  "outliers"         => [1, 3, 4, 21]
  "objective"        => 0.799134
  "coef"             => [-38.3133, 0.745659, 0.432794, 0.0104587]
  "crit"             => 3.0

```
# References
Atkinson, Anthony C. "Fast very robust methods for the detection of multiple outliers."
Journal of the American Statistical Association 89.428 (1994): 1329-1339.
"""
function atkinson94(setting::RegressionSetting; iters=nothing, crit=3.0)
    X, y = @extractRegressionSetting setting
    return atkinson94(X, y, iters=iters, crit=crit)
end

function atkinson94(X::Array{Float64,2}, y::Array{Float64,1}; iters=nothing, crit=3.0)
    n, p = size(X)

    # the median index
    h = Int(floor((n + p + 1) / 2.0))

    if iters === nothing
        iters = minimum([500 * p, 3000])
    end

    bestobjective = Inf
    bestparameters = zeros(Float64, p)
    bestres = zeros(Float64, n, n)
    bestindex = 0
    indices = collect(1:n)

    # store all the sigma values across all iterations
    sigmas = zeros(Float64, iters, n)

    for iter in 1:iters
        m_subset_indices = sample(indices, p, replace=false)

        # stores the n - p sigma values of a forward run
        sigma_tilde = zeros(Float64, n)

        # stores the (n - p) * n residuals during the forward run
        studentized_residuals = zeros(Float64, n, n)
        copy_parameters = false

        @inbounds for m = p:n
            olsreg = ols(X[m_subset_indices,:], y[m_subset_indices])
            betas = coef(olsreg)
            e = (y .- X * betas)
            r = e.^2

            # sigma is the median of the residuals
            sigma_tilde[m] = sqrt(sort(r)[h])
            XXinv = pinv(X[m_subset_indices, :]'X[m_subset_indices, :])

            # scale the residual according to whether it belongs to m_subset_indices or not
            @inbounds for index in 1:n
                if index in m_subset_indices
                    h_i = X[index, :]' * XXinv * X[index, :]
                    studentized_residuals[m, index] = abs(e[index]) / (sigma_tilde[m] * sqrt(abs(1 - h_i)))
                else
                    d_i = X[index, :]' * XXinv * X[index, :]
                    studentized_residuals[m, index] = abs(e[index]) / (sigma_tilde[m] * sqrt(1 + d_i))
                    r[index] = (e[index]^2) / (1 + d_i)
                end
            end

            # increase the size of subset by 1
            ordered_residue_indices = sortperm(r)
            if m != n
                m_subset_indices = ordered_residue_indices[1:m + 1]
            end

            # if sigma was lowest, record the parameters
            if sigma_tilde[m] < bestobjective
                bestobjective = sigma_tilde[m]
                bestparameters = betas
                bestindex = m
                copy_parameters = true
        end
        end
#         sigmas[iter, :] = sigma_tilde

        # copy the entire n^2 residuals for the forward search to generate stalactite plot
        if copy_parameters
            bestres = copy(studentized_residuals)
        end
    end
#     sigma_bar = mean(sigmas, dims=1)
#     for m = p:n
#         bestres[m, :] = bestres[m, :] .* sigma_bar[m]
#     end
    d = Dict()
#     d["sigma_bar"] = sigma_bar
#     d["sigmas"] = sigmas[:, p:end]
    d["coef"] = bestparameters
    d["objective"] = bestobjective
    d["optimum_index"] = bestindex
    d["residuals_matrix"] = bestres[p:end, :]
    d["crit"] = crit
    d["outliers"] = filter(i -> bestres[bestindex, i] > crit, 1:n)
    return d
end


###########################################################################################
# Operators
###########################################################################################
const DerivStencilScratch = Dict{NTuple{2,Int}, Matrix{Float64}}()

function derivStencil(order::Integer, ptsleft::Integer, ptsright::Integer; gridspacing::Real = 1, kw...)

	invM = get!(DerivStencilScratch, (ptsleft, ptsright)) do
		M = zeros(ptsleft + ptsright + 1, ptsleft + ptsright + 1)
		for (m_ind, m) in enumerate(-ptsleft:ptsright), (p_ind, p) in enumerate(0:(ptsleft + ptsright))
			M[p_ind, m_ind] = m^p/factorial(p)
		end
		inv(M)
	end

	M_order = invM[:, order + 1]
	if (ptsleft == ptsright) && isodd(order)
		# Fix loss of numerical precision
		M_order[ptsleft + 1] = 0
	end

	den = gridspacing^order
	M_order ./= den
	return M_order
end

function derivStencil(order::Integer, N::Integer; kwargs...)
	# given N points decide on the number of points to the left and to the right
	if isodd(N)
		return derivStencil(order, div(N, 2), div(N, 2); kwargs...)
	else
		return derivStencil(order, div(N, 2)-1, div(N, 2); kwargs...)
	end
end

function derivStencil!(v, order::Integer, ptsleft::Integer, ptsright::Integer; gridspacing::Real = 1, kw...)

	invM = get!(DerivStencilScratch, (ptsleft, ptsright)) do
		M = zeros(ptsleft + ptsright + 1, ptsleft + ptsright + 1)
		for (m_ind, m) in enumerate(-ptsleft:ptsright), (p_ind, p) in enumerate(0:(ptsleft + ptsright))
			M[p_ind, m_ind] = m^p/factorial(p)
		end
		inv(M)
	end

	v .= @view invM[:, order + 1]
	if (ptsleft == ptsright) && isodd(order)
		# Fix loss of numerical precision
		v[ptsleft + 1] = 0
	end

	den = gridspacing^order
	v ./= den
	return v
end

function derivStencil!(v, order::Integer, N::Integer; kwargs...)
	# given N points decide on the number of points to the left and to the right
	if isodd(N)
		return derivStencil!(v, order, div(N, 2), div(N, 2); kwargs...)
	else
		return derivStencil!(v, order, div(N, 2)-1, div(N, 2); kwargs...)
	end
end

nextodd(n::Integer) = n + 1 + isodd(n)

function ceilsearch(arr, n)
	# assume sorted
	for elem in arr
		if elem>=n
			return elem
		end
	end
	return n
end

function D(N; stencil_gridpts = Dict(6=>3, 42=>5), #= Dictionary of gridpt => order of derivative upto that gridpt =#
	left_edge_npts = 2, left_edge_ghost = false,
	right_edge_npts = 2, right_edge_ghost = false,
	left_Dirichlet = false,
	right_Dirichlet = false,
	kwargs...)

	@assert(N≥2,"Need at least 2 points to compute the derivative")

	S = zeros(N, N)

	D!(S, N; stencil_gridpts = stencil_gridpts,
		left_edge_npts = left_edge_npts,
		right_edge_npts = right_edge_npts,
		left_Dirichlet = left_Dirichlet,
		right_Dirichlet = right_Dirichlet,
		kwargs...)

	return S
end

function D!(S, N; stencil_gridpts = Dict(6=>3, 42=>5), #= Dictionary of gridpt => order of derivative upto that gridpt =#
	left_edge_npts = 2,
	right_edge_npts = 2,
	left_Dirichlet = false,
	right_Dirichlet = false,
	kwargs...)

	@assert(N≥2,"Need at least 2 points to compute the derivative")
	@assert size(S) == (N, N) "Size of matrix must be $((N, N))"

	gridpt_cutoffs_list = sort(collect(keys(stencil_gridpts)))
	maxorder = maximum(values(stencil_gridpts))

	S .= 0

	# Derivatives on the boundary

	v = @view S[1, 1:left_edge_npts]
	derivStencil!(v, 1, 0, left_edge_npts-1; kwargs...)

	v = @view S[end, (N-right_edge_npts + 1):N]
	derivStencil!(v, 1, right_edge_npts-1, 0; kwargs...)

	for gridpt in 2:N-1

		gridpt_cutoff =  ceilsearch(gridpt_cutoffs_list, min(gridpt, N-gridpt + 1))
		if haskey(stencil_gridpts, gridpt_cutoff)
			npts = stencil_gridpts[gridpt_cutoff]
		else
			npts = nextodd(maxorder)
		end

		diagpt = gridpt
		startpt = max(1, diagpt - div(npts, 2)-(diagpt + div(npts, 2) > N ? diagpt + div(npts, 2) - N_cols : 0 ))
		endpt = min(N, startpt + npts -1)
		npts_left = diagpt - startpt
		npts_right = endpt - diagpt

		v = @view S[gridpt, startpt:endpt]
		derivStencil!(v, 1, npts_left, npts_right; kwargs...)
	end

	if left_Dirichlet
		S[1:2, 1] .= 0
	end
	if right_Dirichlet
		S[end-1:end, end] .= 0
	end

	return S
end

function D²(N; stencil_gridpts = Dict(6=>3, 42=>5),
	left_edge_npts = 3, left_edge_ghost = false, right_edge_npts = 3, right_edge_ghost = false, kwargs...)

	@assert(N≥3,"Need at least 3 points to compute the second derivative")

	N_cols = N
	if left_edge_ghost
		N_cols += 1
	end
	if right_edge_ghost
		N_cols += 1
	end

	S = zeros(N, N_cols)

	gridpt_cutoffs_list = sort(collect(keys(stencil_gridpts)))
	maxorder = maximum(values(stencil_gridpts))

	# Derivatives on the boundary

	if left_edge_ghost
		startpt = 2 - div(left_edge_npts, 2)
		endpt = startpt + left_edge_npts-1
		if startpt >=1
			v = @view S[1, startpt:endpt]
			derivStencil!(v, 2, left_edge_npts; kwargs...)
		else
			v = @view S[1, 1:left_edge_npts]
			derivStencil!(v, 2, 1, left_edge_npts-2; kwargs...)
		end
	else
		v = @view S[1, 1:left_edge_npts]
		derivStencil!(v, 2, 0, left_edge_npts-1; kwargs...)
	end

	if right_edge_ghost
		endpt = N +  div(right_edge_npts, 2) + ( left_edge_ghost ? 1 : 0 )
		startpt = endpt - right_edge_npts + 1
		if endpt<=N_cols
			v = @view S[end, startpt:endpt]
			derivStencil!(v, 2, right_edge_npts; kwargs...)
		else
			v = @view S[end, N_cols-right_edge_npts + 1:N_cols]
			derivStencil!(v, 2, right_edge_npts-2, 1; kwargs...)
		end
	else
		v = @view S[end, N_cols-right_edge_npts + 1:N_cols]
		derivStencil!(v, 2, right_edge_npts-1, 0; kwargs...)
	end

	for gridpt in 2:N-1

		gridpt_cutoff =  ceilsearch(gridpt_cutoffs_list, min(gridpt, N-gridpt + 1))
		if haskey(stencil_gridpts, gridpt_cutoff)
			npts = stencil_gridpts[gridpt_cutoff]
		else
			npts = nextodd(maxorder)
		end

		diagpt = gridpt + ( left_edge_ghost ? 1 : 0 )
		startpt = max(1, diagpt - div(npts, 2)-(diagpt + div(npts, 2) > N_cols ? diagpt + div(npts, 2) - N_cols : 0 ))
		endpt = min(N_cols, startpt + npts -1)
		npts_left = diagpt - startpt
		npts_right = endpt - diagpt

		v = @view S[gridpt, startpt:endpt]
		derivStencil!(v, 2, npts_left, npts_right; kwargs...)
	end

	return S
end

function Dⁿ(order, N; stencil_gridpts = Dict(6=>order + 1, 42=>order + 3),
	left_edge_npts = order + 1, left_edge_ghost = false, right_edge_npts = order + 1, right_edge_ghost = false, kwargs...)

	@assert(N≥order + 1,"Need at least $(order + 1) points to compute the derivative")

	N_cols = N
	if left_edge_ghost
		N_cols += 1
	end
	if right_edge_ghost
		N_cols += 1
	end

	S = zeros(N, N_cols)

	gridpt_cutoffs_list = sort(collect(keys(stencil_gridpts)))
	gridpt_cutoff_maxorder = maximum(keys(stencil_gridpts))
	maxorder = stencil_gridpts[gridpt_cutoff_maxorder]

	# Derivatives on the boundary

	if left_edge_ghost
		startpt = 2 - div(left_edge_npts, 2)
		endpt = startpt + left_edge_npts-1
		if startpt >=1
			v = @view S[1, startpt:endpt]
			derivStencil!(v, order, left_edge_npts; kwargs...)
		else
			v = @view S[1, 1:left_edge_npts]
			derivStencil!(v, order, 1, left_edge_npts-2; kwargs...)
		end
	else
		v = @view S[1, 1:left_edge_npts]
		derivStencil!(v, order, 0, left_edge_npts-1; kwargs...)
	end

	if right_edge_ghost
		endpt = N +  div(right_edge_npts, 2) + ( left_edge_ghost ? 1 : 0 )
		startpt = endpt - right_edge_npts + 1
		if endpt<=N_cols
			v = @view S[end, startpt:endpt]
			derivStencil!(v, order, right_edge_npts; kwargs...)
		else
			v = @view S[end, (N_cols-right_edge_npts + 1):N_cols]
			derivStencil!(v, order, right_edge_npts-2, 1; kwargs...)
		end
	else
		v = @view S[end, (N_cols-right_edge_npts + 1):N_cols]
		derivStencil!(v, order, right_edge_npts-1, 0; kwargs...)
	end

	for gridpt in 2:N-1

		gridpt_cutoff =  ceilsearch(gridpt_cutoffs_list, min(gridpt, N-gridpt + 1))
		if haskey(stencil_gridpts, gridpt_cutoff)
			npts = stencil_gridpts[gridpt_cutoff]
		else
			npts = nextodd(maxorder)
		end

		diagpt = gridpt + ( left_edge_ghost ? 1 : 0 )
		startpt = max(1, diagpt - div(npts, 2)-(diagpt + div(npts, 2) > N_cols ? diagpt + div(npts, 2) - N_cols : 0 ))
		endpt = min(N_cols, startpt + npts -1)
		npts_left = diagpt - startpt
		npts_right = endpt - diagpt

		v = @view S[gridpt, startpt:endpt]
		derivStencil!(v, order, npts_left, npts_right; kwargs...)
	end

	return S
end

function dbydr(dr::AbstractArray; stencil_gridpts = Dict(6=>3, 42=>5), kwargs...)
	D(length(dr); stencil_gridpts = stencil_gridpts, kwargs...) ./ dr
end

function dbydr!(S::AbstractMatrix, dr::AbstractArray; stencil_gridpts = Dict(6=>3, 42=>5), kwargs...)
	D!(S, length(dr); stencil_gridpts = stencil_gridpts, kwargs...)
	S ./= dr
	return S
end

function dbydr(N::Integer, dx::Real, stencil_gridpts = Dict(6=>3, 42=>5), kwargs...)
	D(N; stencil_gridpts = stencil_gridpts, gridspacing = dx)
end

function dbydr!(S::AbstractMatrix, N::Integer, dx::Real, stencil_gridpts = Dict(6=>3, 42=>5); kwargs...)
	D!(S, N; stencil_gridpts = stencil_gridpts, gridspacing = dx)
end

function roll(arr, shift)
	# translation with wrap
	# A′ = T(shift) A
	# translation: A′(i) = A(i-shift)
	# translation with wrap: A′(i + 1) = A(mod(i-shift, N)+1), i = 0...N-1

	N = size(arr, ndims(arr))
	shift = mod(shift, N)
	if shift==0
		return arr
	end

	newarr = similar(arr)

	for i in 0:N-1
		newarr[:,:, i + 1] .= @view arr[:,:, mod(i-shift, N)+1]
	end
	return newarr
end


theta = 0.9
mu    = 0.0

data = Matrix{Float64}(1,100)

old_eps = 0.0
for i = 1:length(data)
    r = randn()
    data[i] = old_eps + theta * r
    old_eps = r
end

writecsv("data_ma1.csv", data)


using Test
using LogicCircuits
using ProbabilisticCircuits
using DataFrames: DataFrame
using Suppressor

include("../helper/gpu.jl")

@testset "Likelihood" begin
    # Uses a PC with 4 variables, and tests 3 of the configurations to
    # match with python. Also tests all probabilities sum up to 1.

    prob_circuit = zoo_psdd("little_4var.psdd");
    @test prob_circuit isa ProbCircuit;

    # Step 1. Check Probabilities for 3 samples
    data = DataFrame(BitArray([0 0 0 0; 0 1 1 0; 0 0 1 1]), :auto);
    true_prob = [0.07; 0.03; 0.13999999999999999]

    # Test Sturdel EVI
    samples, _ = sample(prob_circuit, 100000)
    mix, weights, _ = learn_strudel(DataFrame(convert(BitArray, samples), :auto); num_mix = 10,
                                    init_maxiter = 20, em_maxiter = 100, verbose = false)
    mix_calc_prob = exp.(EVI(mix, data, weights))

    @test true_prob ≈ mix_calc_prob atol = 0.1
    mix_calc_prob_all = exp.(EVI(mix, data_all))
    @test 1 ≈ sum(mix_calc_prob_all) atol = 0.1

    cpu_gpu_agree_approx(data_all) do d
        EVI(mix, d, weights)
    end
end


@doc raw"""
    PSFModels.AiryDisk([T=Float64]; fwhm, x=0, y=0, amp=1, maxsize=3, extent=maxsize .* fwhm)
    PSFModels.AiryDisk([T=Float64]; fwhm, pos, amp=1, maxsize=3, extent=maxsize .* fwhm)
    PSFModels.AiryDisk([T=Float64]; fwhm, r, theta, amp=1, maxsize=3, extent=maxsize .* fwhm)

An unnormalized Airy disk. The position can be specified in `(x, y)` coordinates as a `Tuple`, `AbstractVector`, or as separate arguments. By default the model is placed at the origin. The position can also be given as a polar coordinate using `r`/`ρ` and `theta`/`θ`, optionally centered around `origin`.

The `fwhm` can be a scalar (isotropic) or a vector/tuple (diagonal). For efficient calculations, we recommend using [StaticArrays](https://github.com/JuliaArrays/StaticArrays.jl). Here, `maxsize` is a multiple of the fwhm, and can be given as a scalar or as a tuple for each axis. The `extent` defines the bounding box for the model and is used for the default rendering size.

# Functional form

The Airy disk is a distribution over the radius `r` (the square-Euclidean distance)

```
f(x | x̂, FWHM) = [ 2J₁(q) / q ]^2
```
where `J₁` is the first-order Bessel function of the first kind and
```
q ≈ π * r / (0.973 * FWHM)
```
"""
struct AiryDisk{T,FT,VT<:AbstractVector,IT<:Tuple} <: PSFModel{T}
    fwhm::FT
    pos::VT
    amp::T
    indices::IT

    function AiryDisk(::Type{T}, fwhm::FT, pos::VT, amp::T, indices::IT) where {T,VT<:AbstractVector,FT,IT<:Tuple}
        new{T,FT,VT,IT}(fwhm, pos, amp, indices)
    end
end

## constructors
# default type is Float64
AiryDisk(; kwargs...) = AiryDisk(Float64; kwargs...)
function AiryDisk(T; fwhm, amp=one(T), maxsize=3, extent = maxsize .* fwhm, position...)
    # get the position from keyword distpatch
    pos = _position(; position...)
    return AiryDisk(T, fwhm, pos, convert(T, amp), indices_from_extent(pos, extent))
end

Base.size(a::AiryDisk) = map(length, a.indices)
Base.axes(a::AiryDisk) = a.indices

const rz = 3.8317059702075125 / π

function (a::AiryDisk{T})(point::AbstractVector) where T
    radius = a.fwhm * 1.18677
    Δ = euclidean(point, a.pos)
    r = Δ / (radius / rz)
    val = ifelse(iszero(r), a.amp, a.amp * (2 * besselj1(π * r) / (π * r))^2)
    return convert(T, val)
end

function (a::AiryDisk{T,<:Union{AbstractVector,Tuple}})(point::AbstractVector) where T
    weights = SA[(rz / (first(a.fwhm) * 1.18677))^2, (rz / (last(a.fwhm) * 1.18677))^2]
    r = weuclidean(point, a.pos, weights)
    val = ifelse(iszero(r), a.amp, a.amp * (2 * besselj1(π * r) / (π * r))^2)
    return convert(T, val)
end


# Computes graph Stein discrepancy based on the affine objective and a greedy
# 2-spanner of the complete graph on sample points x(i,:) with edges weighted by
# ||x(i,:) - x(i2,:)||
#
# For each j in {1,...,d}, returns a solution g_j to
#   max_{g_j}
#   sum_i q(i) gobjcoeff(i,j) * g_j(x(i,:)) + <gradgobjcoeff(i,:,j), grad g_j(x(i,:)))>
#   subject to, for each spanner edge (i, i2),
#    |g_j(x(i,:))| <= c_{1}, ||grad g_j(x(i,:))||^* <= c_{2},
#    |g_j(x(i,:)) - g_j(x(i2,:))| <= c_{2} ||x(i,:) - x(i2,:)||,
#    ||grad g_j(x(i,:)) - grad g_j(x(i2,:))||^*
#      <= c_3 ||x(i,:) - x(i2,:)||
#    |g_j(x(i,:)) - g_j(x(i2,:)) - <grad g_j(x(i2,:)), x(i,:) - x(i2,:)>|
#      <= (c_3/2) (||x(i,:) - x(i2,:)||)^2
#    |g_j(x(i,:)) - g_j(x(i2,:)) - <grad g_j(x(i,:)), x(i,:) - x(i2,:)>|
#      <= (c_3/2) (||x(i,:) - x(i2,:)||)^2
# where the norm ||.|| is the L1 norm and ||.||^* is its dual.
#
# Args:
#  sample - SteinDiscrete object representing a sample
#  gobjcoefficients - a n x d matrix so jth column has coefficients of
#      the objective for gj
#  gradgobjcoefficients - a n x d x d tensor so [:,:,j] represents
#      the coefficients of grad gj on each row
#  solver - optimization program solver supported by JuMP. Examples:

function affine_graph_discrepancy(sample::SteinDiscrete,
                                  gobjcoefficients::Array{Float64,2},
                                  gradgobjcoefficients::Array{Float64,3},
                                  solver::AbstractMathProgSolver,
                                  operator::AbstractString;
                                  c1::Float64=1.0,
                                  c2::Float64=1.0,
                                  c3::Float64=1.0,
                                  supportlowerbounds=nothing,
                                  supportupperbounds=nothing)

    ## Extract inputs
    points = sample.support
    weights = sample.weights
    n = length(weights)
    d = size(points,2)

    ## reset the support bounds if needed
    if supportlowerbounds == nothing
        supportlowerbounds = repmat([-Inf], d)
    elseif isa(supportlowerbounds, Number)
        supportlowerbounds = rempat([supportlowerbounds], d)
    end
    if supportupperbounds == nothing
        supportupperbounds = repmat([Inf], d)
    elseif isa(supportupperbounds, Number)
        supportupperbounds = rempat([supportupperbounds], d)
    end

    ## Find spanner edge set
    println("[Computing spanner edges]")
    tic(); edges = getspanneredges(points); edgetime = toc()
    println("\tusing $(size(edges,1)) of $(binomial(n,2)) edges");
    if n > 1
        # Compute differences between points connected by an edge
        diffs = points[edges[:,1],:] - points[edges[:,2],:]
        # Compute L1 distances between points connected by an edge
        distances = sum(abs.(diffs),2)
        scaled_squared_distances = (c3/2)*distances.^2
    end

    ## Prepare return values
    objval = Array{Float64}(1, d)
    gopt = Array{Float64}(n, d)
    gradgopt = Array{Float64}(n, d, d)
    solvetime = Array{Float64}(1, d)

    # Distance cutoff for enforcing Lipschitz function constraints
    lipfunccutoff = 2*c1/c2
    # Distance cutoff for enforcing Lipschitz gradient constraints
    lipgradcutoff = 2*c2/c3
    # Distance cutoff for enforcing Taylor compatibility constraints
    taylorcutoff = 4*c2/c3

    ## Solve a different problem for each sample coordinate
    println("[Solving optimization program]")
    for j = 1:d
        tic()
        ## Define optimization problem
        m = Model(solver=solver)
        # Define optimization variables
        @variable(m, -c1 <= g[i=1:n] <= c1)
        @variable(m, -c2 <= gradg[i=1:n,k=1:d] <= c2)
        # Define the optimization objective
        gobj = AffExpr(g[1:n], vec(gobjcoefficients[:,j]), 0.0)
        gradgobj = AffExpr()
        for k = 1:d
            gradgobj += AffExpr(gradg[1:n,k], vec(gradgobjcoefficients[1:n,k,j]), 0.0)
        end
        @objective(m, Max, gobj + gradgobj)
        # Find finite limits of support in dimension j
        limits = [
            supportlowerbounds[j],
            supportupperbounds[j]
        ]
        limits = filter(isfinite, limits)
        # Add boundary constraints if needed
        for i = 1:n, bj = limits
            slackij = points[i,j] - bj
            # Add constraints to ensure gj can vanish on boundary
            # whenever abs(slackij) < lipfunccutoff
            # (otherwise constraints will never be active)
            if abs(slackij) < lipfunccutoff
                @constraint(m, g[i] >= -c2 * abs(slackij))
                @constraint(m, g[i] <= c2 * abs(slackij))
            end
            # Add constraints to ensure \grad gj can vanish in non-j dimension
            # whenever abs(slackij) < lipgradcutoff
            # (otherwise constraints will never be active)
            if abs(slackij) < lipgradcutoff
                constrained_dims = filter(x -> x != j, 1:d)
                for k = constrained_dims
                    @constraint(m, gradg[i,k] >= -c3 * abs(slackij))
                    @constraint(m, gradg[i,k] <= c3 * abs(slackij))
                end
            end
            # Add \grad gj constraints on jth dimension
            # whenever abs(slackij) < taylorcutoff
            # (otherwise constraints will never be active)
            if abs(slackij) < taylorcutoff
                @constraint(m, g[i] - slackij * gradg[i,j] <= (c3/2) * slackij^2)
                @constraint(m, g[i] - slackij * gradg[i,j] >= (-c3/2) * slackij^2)
            end
        end
        # Add pairwise constraints
        if n > 1
            for i = 1:length(distances)
                v1 = edges[i,1]; v2 = edges[i,2]
                # Add Lipschitz function constraints
                # whenever distance < lipgradcutoff
                # (otherwise constraints will never be active)
                if distances[i] < lipfunccutoff
                    @constraint(m, g[v1] - g[v2] <= c2 * distances[i])
                    @constraint(m, g[v1] - g[v2] >= -c2 * distances[i])
                end
                # Add Lipschitz gradient constraints
                # whenever distances[i] < lipgradcutoff
                # (otherwise constraints will never be active)
                if distances[i] < lipgradcutoff
                    for k = 1:d
                        @constraint(m, gradg[v1,k] - gradg[v2,k] <= c3*distances[i])
                        @constraint(m, gradg[v1,k] - gradg[v2,k] >= -c3*distances[i])
                    end
                end
                # Add Taylor compatibility constraints relating g and gradg
                # whenever distance < taylorcutoff
                # (otherwise constraints will never be active)
                if distances[i] < taylorcutoff
                    expr1 = AffExpr(vec(gradg[v1,:]), vec(diffs[i,:]), 0.0)
                    @constraint(m, g[v1] - g[v2] - expr1 <= scaled_squared_distances[i])
                    @constraint(m, g[v1] - g[v2] - expr1 >= -scaled_squared_distances[i])
                    expr2 = AffExpr(vec(gradg[v2,:]), vec(diffs[i,:]), 0.0)
                    @constraint(m, g[v1] - g[v2] - expr2 <= scaled_squared_distances[i])
                    @constraint(m, g[v1] - g[v2] - expr2 >= -scaled_squared_distances[i])
                end
            end
        end
        # Solve the problem
        @time status = JuMP.solve(m)
        solvetime[j] = toc()
        # Package the results
        objval[j] = getobjectivevalue(m)
        gopt[:,j] = getvalue(g)[1:n]
        gradgopt[:,:,j] = getvalue(gradg)[1:n, 1:d]
    end
    # Compute (T_Pg)(x) = {g part} + {grad g part}
    operatorgopt = sum(gopt .* gobjcoefficients,2)
    for j = 1:d
        operatorgopt += sum(gradgopt[:,:,j] .* gradgobjcoefficients[:,:,j],2)
    end

    AffineDiscrepancyResult(
        points,
        weights,
        edges,
        objval,
        gopt,
        gradgopt,
        operatorgopt,
        edgetime,
        solvetime,
        operator
    );
end


basic_params, maneuvering_params = get_KVLCC2_L7_params()

@testset "mmg.jl KVLCC2_L7 turning" begin

    duration = 200  # [s]
    max_δ_rad = 35 * pi / 180.0  # [rad]
    n_const = 17.95  # [rps]

    sampling = duration * 10
    time_list = range(0.00, stop = duration, length = sampling)
    δ_rad_list = max_δ_rad .* ones(Float64, sampling)
    n_p_list = n_const .* ones(Float64, sampling)
    mmg_results = mmg_3dof_simulate(
        basic_params,
        maneuvering_params,
        time_list,
        δ_rad_list,
        n_p_list,
        u0 = 2.29 * 0.512,
        v0 = 0.0,
        r0 = 0.0,
    )
    u, v, r, δ, n_p = mmg_results
    x, y, ψ = calc_position(time_list, u, v, r)
end

@testset "mmg_zigzag_test" begin
    target_δ_rad = 20.0 * π / 180.0
    target_ψ_rad_deviation = 20.0 * π / 180.0
    start_time_second = 0.00
    time_second_interval = 0.01
    end_time_second = 80.00
    time_list = start_time_second:time_second_interval:end_time_second
    n_const = 17.95  # [rps]
    n_p_list = n_const * ones(Float64, length(time_list))
    u_list, v_list, r_list, ψ_list, δ_list = mmg_3dof_zigzag_test(
        basic_params,
        maneuvering_params,
        time_list,
        n_p_list,
        target_δ_rad,
        target_ψ_rad_deviation,
    )
end

@testset "mmg 3dof approx least square method func" begin
    duration = 200  # [s]
    sampling = duration * 10
    time_list = range(0.00, stop = duration, length = sampling)
    max_δ_rad = 35 * pi / 180.0  # [rad]
    n_const = 17.95  # [rps]
    δ_rad_list = max_δ_rad .* ones(Float64, sampling)
    n_p_list = n_const .* ones(Float64, sampling)
    mmg_results = mmg_3dof_simulate(
        basic_params,
        maneuvering_params,
        time_list,
        δ_rad_list,
        n_p_list,
        u0 = 2.29 * 0.512,
        v0 = 0.0,
        r0 = 0.0,
    )
    u, v, r, δ, n_p = mmg_results
    x, y, ψ = calc_position(time_list, u, v, r)
    data = ShipData(time_list, u, v, r, x, y, ψ, δ, n_p)
    estimate_mmg_approx_lsm(
        data,
        basic_params,
        maneuvering_params.k_0,
        maneuvering_params.k_1,
        maneuvering_params.k_2,
    )
end

@testset "mmg 3dof approx time window sampling least square method func" begin
    duration = 200  # [s]
    max_δ_rad = 35 * pi / 180.0  # [rad]
    n_const = 17.95  # [rps]

    sampling = duration * 10
    time_list = range(0.00, stop = duration, length = sampling)
    δ_rad_list = max_δ_rad .* ones(Float64, sampling)
    n_p_list = n_const .* ones(Float64, sampling)
    mmg_results = mmg_3dof_simulate(
        basic_params,
        maneuvering_params,
        time_list,
        δ_rad_list,
        n_p_list,
        u0 = 2.29 * 0.512,
        v0 = 0.0,
        r0 = 0.0,
    )
    u, v, r, δ, n_p = mmg_results
    noize_dist = Normal(0.0, 0.0005)
    u_obs = u + rand(noize_dist, size(u))
    v_obs = v + rand(noize_dist, size(v))
    r_obs = r + rand(noize_dist, size(r))
    x, y, ψ = calc_position(time_list, u_obs, v_obs, r_obs)
    data = ShipData(time_list, u_obs, v_obs, r_obs, x, y, ψ, δ, n_p)
    one_sample_size = 5000
    estimate_mmg_approx_lsm_time_window_sampling(
        data,
        500,
        basic_params,
        maneuvering_params.k_0,
        maneuvering_params.k_1,
        maneuvering_params.k_2,
    )
end

@testset "mmg mcmc sampling func" begin
    duration = 200  # [s]
    max_δ_rad = 35 * pi / 180.0  # [rad]
    n_const = 17.95  # [rps]

    sampling = duration * 10 + 1
    time_list = range(0.00, stop = duration, length = sampling)
    δ_rad_list = max_δ_rad .* ones(Float64, sampling)
    n_p_list = n_const .* ones(Float64, sampling)
    mmg_results = mmg_3dof_simulate(
        basic_params,
        maneuvering_params,
        time_list,
        δ_rad_list,
        n_p_list,
        u0 = 2.29 * 0.512,
        v0 = 0.0,
        r0 = 0.0,
    )
    u, v, r, δ, n_p = mmg_results
    sampling_rate = 10
    time_obs = time_list[1:sampling_rate:end]
    noize_dist = Normal(0.0, 0.01)
    u_obs = u + rand(noize_dist, size(u))
    v_obs = v + rand(noize_dist, size(v))
    r_obs = r + rand(noize_dist, size(r))
    x, y, ψ = calc_position(time_obs, u_obs, v_obs, r_obs)
    data = ShipData(
        time_obs,
        u_obs,
        v_obs,
        r_obs,
        x,
        y,
        ψ,
        δ[1:sampling_rate:end],
        n_p[1:sampling_rate:end],
    )
    n_samples = 10
    n_chains = 1
    model = create_model_for_mcmc_sample_mmg(
        data,
        basic_params,
        maneuvering_params.k_0,
        maneuvering_params.k_1,
        maneuvering_params.k_2,
    )
    chain = nuts_sampling_single_thread(model, n_samples, n_chains)
end


export  MDF,
        lookahead,
        iterative_policy_evaluation,
        policy_evaluation,
        ValueFunctionPolicy,
        greedy,
        PolicyIteration,
        solve,
        backup,
        ValueIteration,
        GaussSeidelValueIteration,
        LinearProgramFormulation,
        tensorform,
        LinearQuadraticProblem

"""
    struct MDP
        γ # discount factor
        𝒮 # state space
        𝒜 # action space
        T # transition function
        R # reward function
        TR # sample transition and reward
    end
"""
struct MDP
    γ # discount factor
    𝒮 # state space
    𝒜 # action space
    T # transition function
    R # reward function
    TR # sample transition and reward
end

"""
    function lookahead(𝒫::MDP, U, s, a)

Calculate state-action value function Q(s,a) from U and valid state transition distribution T:

`Returns`
    Q(s,a) = R(s,a) + γ*sum(T(s,a,s′)*U(s′) for s′ in 𝒮)

"""
function lookahead(𝒫::MDP, U, s, a)
    𝒮, T, R, γ = 𝒫.𝒮, 𝒫.T, 𝒫.R, 𝒫.γ
    return R(s,a) + γ*sum(T(s,a,s′)*U(s′) for s′ in 𝒮)
end

"""
    function lookahead(𝒫::MDP, U::Vector, s, a)

Calculate state-action value function Q(s,a) from U and valid state transition distribution T:

    return R(s,a) + γ*sum(T(s,a,s′)*U[i] for (i,s′) in enumerate(𝒮))

Complexity: 𝒪(|𝒮|)
"""
function lookahead(𝒫::MDP, U::Vector, s, a)
    𝒮, T, R, γ = 𝒫.𝒮, 𝒫.T, 𝒫.R, 𝒫.γ
    return R(s,a) + γ*sum(T(s,a,s′)*U[i] for (i,s′) in enumerate(𝒮))
end

"""
    function iterative_policy_evaluation(𝒫::MDP, π, k_max)

 Approximate policy evaluation for policy `π`, using k_max iterations of 1-step lookahead, starting from U(s) = 0.

Complexity: 𝒪(|𝒮|² × k_max)
"""
function iterative_policy_evaluation(𝒫::MDP, π, k_max)
    𝒮, T, R, γ = 𝒫.𝒮, 𝒫.T, 𝒫.R, 𝒫.γ
    U = [0.0 for s in 𝒮]
    for k in 1:k_max
        U = [lookahead(𝒫, U, s, π(s)) for s in 𝒮]
    end
    return U
end

"""
    function policy_evaluation(𝒫::MDP, π)

Exact policy evaluation using matrix-based Moore-Penrose pseudoinverse equation.
Requires O(|𝒮|³) time.

`Returns:` [ U(s) ∀ s ∈ {𝒮} ]
"""
function policy_evaluation(𝒫::MDP, π)
    𝒮, R, T, γ = 𝒫.𝒮, 𝒫.R, 𝒫.T, 𝒫.γ

    # R(s, π(s)) ∀ s ∈ {𝒮}
    R′ = [R(s, π(s)) for s in 𝒮] # nx1 = |𝒮|x1
    T′ = [T(s, π(s), s′) for s in 𝒮, s′ in 𝒮] # nxm = |𝒮|x|𝒮|
    return (I - γ*T′)\R′
end

"""
    struct ValueFunctionPolicy
        𝒫 # problem
        U # utility function
    end
"""
struct ValueFunctionPolicy
    𝒫 # problem
    U # utility function
end

"""
    function greedy(𝒫::MDP, U, s)

Find the greedy action and its expected utility starting from state `s` and using estimated optimal (greedy) state value function `U`.  Calculates estimate as maximum over action space of 1-step lookahead:

    u, a = _findmax(a->lookahead(𝒫, U, s, a), 𝒫.𝒜)

Complexity: 𝒪(|𝒜|×|𝒮|)
"""
function greedy(𝒫::MDP, U, s)
    u, a = _findmax(a->lookahead(𝒫, U, s, a), 𝒫.𝒜)
    return (a=a, u=u)
end

"""
    function (π::ValueFunctionPolicy)(s)

Greedy policy function mapping states to greedy action.
"""
function (π::ValueFunctionPolicy)(s)
    return greedy(π.𝒫, π.U, s).a
end

# ---------------- Policy Iteration ----------------

"""
    struct PolicyIteration
        π # initial policy
        k_max # maximum number of iterations
    end
"""
struct PolicyIteration
    π # initial policy
    k_max # maximum number of iterations
end

"""
    function solve(M::PolicyIteration, 𝒫::MDP)

Iteratively improves an initial policy `π` to obtain an optimal policy for an MDP `𝒫` with discrete state and action spaces.  Returns improved policy.
"""
function solve(M::PolicyIteration, 𝒫::MDP)
    π, 𝒮 = M.π, 𝒫.𝒮
    converged = false

    for k = 1:M.k_max
        # exact solution for U_π(s)
        U = policy_evaluation(𝒫, π)
        # Use new utility function for finding greedy action.
        π′ = ValueFunctionPolicy(𝒫, U)
        # same action given for every state, => convergence (current π is best)
        if all(π(s) == π′(s) for s in 𝒮)
            break
        end
        # Not converged: update policy with new greedy policy.
        π = π′
    end
    return π
end

# -------------------------- Value Iteration --------------------------

"""
    function backup(𝒫::MDP, U, s)

Calculate maxₐ Q(s,a).
Complexity: O(|𝒮|×|𝒜|).

    return maximum(lookahead(𝒫, U, s, a) for a in 𝒫.𝒜)
"""
function backup(𝒫::MDP, U, s)
    return maximum(lookahead(𝒫, U, s, a) for a in 𝒫.𝒜)
end

"""
    struct ValueIteration
        k_max # maximum number of iterations
    end
"""
struct ValueIteration
    k_max # maximum number of iterations
end

"""
    function solve(M::ValueIteration, 𝒫::MDP)

Iteratively improves a value function `U` to obtain an optimal policy for an MDP `𝒫` with discrete state and action spaces.  Terminates after `k_max` iterations.

Order O(k_max × |𝒮| × |𝒜| × |𝒮|)
# Returns

Optimal policy `π`
"""
function solve(M::ValueIteration, 𝒫::MDP)
    U = [0.0 for s in 𝒫.𝒮]
    for k = 1:M.k_max
        U = [backup(𝒫, U, s) for s in 𝒫.𝒮]
    end
    return ValueFunctionPolicy(𝒫, U)
end

#
# @doc raw"""
#     function solve(M::ValueIteration, 𝒫::MDP, δ::Float64)
#
# Iteratively improves a value function `U` to obtain an optimal policy for an MDP `𝒫` with discrete state and action spaces.  Terminates after Bellman residual falls below threshold δ:
#
# $\|\|U\_{k+1} - U\_k\|\|\_{\infty} < \delta$
#
# # Returns
#
# Optimal policy `π`
# """
# function solve(M::ValueIteration, 𝒫::MDP, δ::Float64)
#     U = [0.0 for s in 𝒫.𝒮]
#     for k = 1:M.k_max
#         U = [backup(𝒫, U, s) for s in 𝒫.𝒮]
#     end
#     return ValueFunctionPolicy(𝒫, U)
# end

# ----------------- Asyncronous Value Iteration --------------------
#
# Similar to Value Iteration, but only subset of states updated every iteration.
#
# Common asyncronous value iteration method: Gauss-Seidel.
#   Sweeps through an ordering of states and applies Bellman update in place.
"""
    struct GaussSeidelValueIteration
        k_max # maximum number of iterations
    end
"""
struct GaussSeidelValueIteration
    k_max # maximum number of iterations
end

"""
    function solve(M::GaussSeidelValueIteration, 𝒫::MDP)

Same as Value Iteration but update U one state at a time.  Convergence speed depends on ordering of states in `𝒮`.
"""
function solve(M::GaussSeidelValueIteration, 𝒫::MDP)
    𝒮, 𝒜, T, R, γ = 𝒫.𝒮, 𝒫.𝒜, 𝒫.T, 𝒫.R, 𝒫.γ
    U = [0.0 for s in 𝒮]
    for k = 1:M.k_max
        for (i, s) in enumerate(𝒮)
            u = backup(𝒫, U, s)
            U[i] = u
        end
    end
    return ValueFunctionPolicy(𝒫, U)
end

# ------------- Linear Programming --------------

"""
    struct LinearProgramFormulation end
"""
struct LinearProgramFormulation end

"""
    function tensorform(𝒫::MDP)

Convert an MDP into its tensor form, where the states and actions consist of integer indicies, the reward function is a matrix, and the transition function is a three-dimensional tensor.
"""
function tensorform(𝒫::MDP)
    𝒮, 𝒜, R, T = 𝒫.𝒮, 𝒫.𝒜, 𝒫.R, 𝒫.T
    𝒮′ = eachindex(𝒮)
    𝒜′ = eachindex(𝒜)
    R′ = [R(s,a) for s in 𝒮, a in 𝒜]
    T′ = [T(s,a,s′) for s in 𝒮, a in 𝒜, s′ in 𝒮]
    return 𝒮′, 𝒜′, R′, T′
end

"""
    solve(𝒫::MDP) = solve(LinearProgramFormulation(), 𝒫)

Default MDP solver is the LinearProgramFormulation solver.
"""
solve(𝒫::MDP) = solve(LinearProgramFormulation(), 𝒫)

"""
    function solve(M::LinearProgramFormulation, 𝒫::MDP)
"""
function solve(M::LinearProgramFormulation, 𝒫::MDP)
    𝒮, 𝒜, R, T = tensorform(𝒫)
    model = Model(GLPK.Optimizer)
    @variable(model, U[𝒮])
    @objective(model, Min, sum(U))
    @constraint(model, [s=𝒮,a=𝒜], U[s] ≥ R[s,a] + 𝒫.γ*T[s,a,:]⋅U)
    optimize!(model)
    return ValueFunctionPolicy(𝒫, value.(U))
end

# ---------- Linear Quadratic Problem -----------

"""
    struct LinearQuadraticProblem
        Ts # transition matrix with respect to state
        Ta # transition matrix with respect to action
        Rs # reward matrix with respect to state (negative semidefinite)
        Ra # reward matrix with respect to action (negative definite)
        h_max # horizon
    end
"""
struct LinearQuadraticProblem
    Ts # transition matrix with respect to state
    Ta # transition matrix with respect to action
    Rs # reward matrix with respect to state (negative semidefinite)
    Ra # reward matrix with respect to action (negative definite)
    h_max # horizon
end

"""
    function solve(𝒫::LinearQuadraticProblem)
"""
function solve(𝒫::LinearQuadraticProblem)
    Ts, Ta, Rs, Ra, h_max = 𝒫.Ts, 𝒫.Ta, 𝒫.Rs, 𝒫.Ra, 𝒫.h_max
    V = zeros(size(Rs))
    πs = Any[s -> zeros(size(Ta, 2))]
    for h in 2:h_max
        V = Ts'*(V - V*Ta*((Ta'*V*Ta + Ra) \ Ta'*V))*Ts + Rs
        L = -(Ta'*V*Ta + Ra) \ Ta' * V * Ts
        push!(πs, s -> L*s)
    end
    return πs
end


################################################################################
#                           TEST: Compose Rotations
################################################################################

for i = 1:samples
    # DCMs
    # ==========================================================================

    # Sample 4 DCMs.
    for sym in [:D1, :D2, :D3, :D4]
        @eval $sym = angle_to_dcm(EulerAngles(-pi + 2*pi*rand(),
                                              -pi + 2*pi*rand(),
                                              -pi + 2*pi*rand(),
                                              rand($rot_seq_array)))
    end

    # Test the function `compose_rotation`.
    @test          D1 === compose_rotation(D1)
    @test       D2*D1  ≈  compose_rotation(D1,D2)       atol=1e-14
    @test    D3*D2*D1  ≈  compose_rotation(D1,D2,D3)    atol=1e-14
    @test D4*D3*D2*D1  ≈  compose_rotation(D1,D2,D3,D4) atol=1e-14

    # Euler angle and axis
    # ==========================================================================

    # Sample 4 Euler angle and axis.
    for sym in [:ea1, :ea2, :ea3, :ea4]
        α   = -π + 2π*rand()
        v   = randn(3)
        v   = v/norm(v)

        @eval $sym = EulerAngleAxis($α,$v)
    end

    ear1 = ea2*ea1
    eac1 = compose_rotation(ea1,ea2)

    ear2 = ea3*ea2*ea1
    eac2 = compose_rotation(ea1,ea2,ea3)

    ear3 = ea4*ea3*ea2*ea1
    eac3 = compose_rotation(ea1,ea2,ea3,ea4)

    # Test the function `compose_rotation`.
    @test ea1 === compose_rotation(ea1)
    @test ear1 ≈ eac1 atol = 1e-14
    @test ear2 ≈ eac2 atol = 1e-14
    @test ear3 ≈ eac3 atol = 1e-14

    # Euler angles
    # ==========================================================================

    # Sample 4 Euler angles.
    for sym in [:Θ1, :Θ2, :Θ3, :Θ4]
        θx      = -π + 2π*rand()
        θy      = -π + 2π*rand()
        θz      = -π + 2π*rand()
        rot_seq = Meta.quot(rand([:XYZ,:ZYX]))

        @eval $sym = EulerAngles($θx, $θy, $θz, $rot_seq)
    end

    Θr1 = Θ2*Θ1
    Θc1 = compose_rotation(Θ1,Θ2)

    Θr2 = Θ3*Θ2*Θ1
    Θc2 = compose_rotation(Θ1,Θ2,Θ3)

    Θr3 = Θ4*Θ3*Θ2*Θ1
    Θc3 = compose_rotation(Θ1,Θ2,Θ3,Θ4)

    # Test the function `compose_rotation`.
    @test Θ1 === compose_rotation(Θ1)
    @test Θr1 ≈ Θc1 atol = 1e-14
    @test Θr2 ≈ Θc2 atol = 1e-14
    @test Θr3 ≈ Θc3 atol = 1e-14

    # Quaternions
    # ==========================================================================

    # Sample 4 quaternions.
    for sym in [:q1, :q2, :q3, :q4]
        @eval $sym = angle_to_quat(EulerAngles(-pi + 2*pi*rand(),
                                               -pi + 2*pi*rand(),
                                               -pi + 2*pi*rand(),
                                               rand($rot_seq_array)))
    end

    # Test the function `compose_rotation`.
    @test           q1     === compose_rotation(q1)
    @test       (q1*q2)[:]  ≈  compose_rotation(q1,q2)[:]       atol=1e-14
    @test    (q1*q2*q3)[:]  ≈  compose_rotation(q1,q2,q3)[:]    atol=1e-14
    @test (q1*q2*q3*q4)[:]  ≈  compose_rotation(q1,q2,q3,q4)[:] atol=1e-14
end