using DifferentialEquations
using Test
using MRIgeneralizedBloch
using MRIgeneralizedBloch: apply_hamiltonian_freeprecession!, apply_hamiltonian_gbloch!

##
max_error = 5e-2

## set parameters
ω0 = 200.0
m0s = 0.15
R1f = 0.3
R2f = 1 / 65e-3
R1s = 2
T2s = 10e-6
Rx = 30.0
TE = 3.5e-3
mfun(p, t; idxs=nothing) = typeof(idxs) <: Number ? 0.0 : zeros(30)
alg = Tsit5()

## random initialization
ms = m0s * rand()
mf = (1 - m0s) * rand()
ϑ = rand() * π / 2
φ = rand() * 2π
m0 = [mf * sin(ϑ) * cos(φ), mf * sin(ϑ) * sin(φ), mf * cos(ϑ), ms, 1]

## baseline ODE solution
FP_sol = solve(ODEProblem(apply_hamiltonian_freeprecession!, m0, (0, TE), (ω0, m0s, R1f, R2f, Rx, R1s)), alg)

## Analytical gradients (using ApproxFun)
grad_list = [grad_m0s(), grad_R1f(), grad_R2f(), grad_Rx(), grad_R1s(), grad_T2s(), grad_ω0(), grad_B1()]

FP_sol_grad = solve(ODEProblem(apply_hamiltonian_freeprecession!, [m0; zeros(5 * (length(grad_list)),1)], (0, TE), (ω0, m0s, R1f, R2f, Rx, R1s, grad_list)), alg)

## FD derivative wrt. m0s
dm0s = 1e-9
FP_sol_dm0s = solve(ODEProblem(apply_hamiltonian_freeprecession!, m0, (0, TE), (ω0, (m0s + dm0s), R1f, R2f, Rx, R1s)), alg)

t = 0 : 1e-5 : TE
dxf = similar(t)
dyf = similar(t)
dzf = similar(t)
dzs = similar(t)
dxf_fd = similar(t)
dyf_fd = similar(t)
dzf_fd = similar(t)
dzs_fd = similar(t)

for i = 1 : length(t)
        dxf_fd[i] = (FP_sol_dm0s(t[i])[1] - FP_sol(t[i])[1]) /dm0s
        dyf_fd[i] = (FP_sol_dm0s(t[i])[2] - FP_sol(t[i])[2]) /dm0s
        dzf_fd[i] = (FP_sol_dm0s(t[i])[3] - FP_sol(t[i])[3]) /dm0s
        dzs_fd[i] = (FP_sol_dm0s(t[i])[4] - FP_sol(t[i])[4]) /dm0s

        dxf[i] = FP_sol_grad(t[i])[6]
        dyf[i] = FP_sol_grad(t[i])[7]
        dzf[i] = FP_sol_grad(t[i])[8]
        dzs[i] = FP_sol_grad(t[i])[9]
end

@test dxf ≈ dxf_fd atol = 1e-3
@test dyf ≈ dyf_fd atol = 1e-3
@test dzf ≈ dzf_fd rtol = max_error
@test dzs ≈ dzs_fd rtol = max_error

## test derivative wrt. R1f
dR1f = 1e-9
FP_sol_dR1f = solve(ODEProblem(apply_hamiltonian_freeprecession!, m0, (0, TE), (ω0, m0s, (R1f+ dR1f), R2f, Rx, R1s)), alg)

for i = 1: length(t)
        dxf_fd[i] = (FP_sol_dR1f(t[i])[1] - FP_sol(t[i])[1]) /dR1f
        dyf_fd[i] = (FP_sol_dR1f(t[i])[2] - FP_sol(t[i])[2]) /dR1f
        dzf_fd[i] = (FP_sol_dR1f(t[i])[3] - FP_sol(t[i])[3]) /dR1f
        dzs_fd[i] = (FP_sol_dR1f(t[i])[4] - FP_sol(t[i])[4]) /dR1f

        dxf[i] = FP_sol_grad(t[i])[11]
        dyf[i] = FP_sol_grad(t[i])[12]
        dzf[i] = FP_sol_grad(t[i])[13]
        dzs[i] = FP_sol_grad(t[i])[14]
end

@test dxf ≈ dxf_fd atol = 1e-3
@test dyf ≈ dyf_fd atol = 1e-3
@test dzf ≈ dzf_fd rtol = max_error
@test dzs ≈ dzs_fd rtol = max_error

## test derivative wrt. R2f
dR2f = 1e-6
FP_sol_dR2f = solve(ODEProblem(apply_hamiltonian_freeprecession!, m0, (0, TE), (ω0, m0s, R1f, (R2f+dR2f), Rx, R1s)), alg)

for i = 1: length(t)
        dxf_fd[i] = (FP_sol_dR2f(t[i])[1] - FP_sol(t[i])[1]) /dR2f
        dyf_fd[i] = (FP_sol_dR2f(t[i])[2] - FP_sol(t[i])[2]) /dR2f
        dzf_fd[i] = (FP_sol_dR2f(t[i])[3] - FP_sol(t[i])[3]) /dR2f
        dzs_fd[i] = (FP_sol_dR2f(t[i])[4] - FP_sol(t[i])[4]) /dR2f

        dxf[i] = FP_sol_grad(t[i])[16]
        dyf[i] = FP_sol_grad(t[i])[17]
        dzf[i] = FP_sol_grad(t[i])[18]
        dzs[i] = FP_sol_grad(t[i])[19]
end

@test dxf ≈ dxf_fd rtol = max_error
@test dyf ≈ dyf_fd rtol = max_error
@test dzf ≈ dzf_fd atol = 1e-4
@test dzs ≈ dzs_fd atol = 1e-4

## test derivative wrt. Rx
dRx = 1e-9
FP_sol_dRx = solve(ODEProblem(apply_hamiltonian_freeprecession!, m0, (0, TE), (ω0, m0s, R1f, R2f, (Rx+dRx), R1s)), alg)

for i = 1: length(t)
        dxf_fd[i] = (FP_sol_dRx(t[i])[1] - FP_sol(t[i])[1]) /dRx
        dyf_fd[i] = (FP_sol_dRx(t[i])[2] - FP_sol(t[i])[2]) /dRx
        dzf_fd[i] = (FP_sol_dRx(t[i])[3] - FP_sol(t[i])[3]) /dRx
        dzs_fd[i] = (FP_sol_dRx(t[i])[4] - FP_sol(t[i])[4]) /dRx

        dxf[i] = FP_sol_grad(t[i])[21]
        dyf[i] = FP_sol_grad(t[i])[22]
        dzf[i] = FP_sol_grad(t[i])[23]
        dzs[i] = FP_sol_grad(t[i])[24]
end

@test dxf ≈ dxf_fd atol = 1e-4
@test dyf ≈ dyf_fd atol = 1e-4
@test dzf ≈ dzf_fd rtol = max_error
@test dzs ≈ dzs_fd rtol = max_error

## test derivative wrt. R1s
dR1s = 1e-9
FP_sol_dR1s = solve(ODEProblem(apply_hamiltonian_freeprecession!, m0, (0, TE), (ω0, m0s, R1f, R2f, Rx, (R1s+ dR1s))), alg)

for i = 1: length(t)
        dxf_fd[i] = (FP_sol_dR1s(t[i])[1] - FP_sol(t[i])[1]) /dR1s
        dyf_fd[i] = (FP_sol_dR1s(t[i])[2] - FP_sol(t[i])[2]) /dR1s
        dzf_fd[i] = (FP_sol_dR1s(t[i])[3] - FP_sol(t[i])[3]) /dR1s
        dzs_fd[i] = (FP_sol_dR1s(t[i])[4] - FP_sol(t[i])[4]) /dR1s

        dxf[i] = FP_sol_grad(t[i])[26]
        dyf[i] = FP_sol_grad(t[i])[27]
        dzf[i] = FP_sol_grad(t[i])[28]
        dzs[i] = FP_sol_grad(t[i])[29]
end

@test dxf ≈ dxf_fd atol = 1e-3
@test dyf ≈ dyf_fd atol = 1e-3
@test dzf ≈ dzf_fd rtol = max_error
@test dzs ≈ dzs_fd rtol = max_error

## test derivative wrt. T2s
dT2s = 1e-14
FP_sol_dT2s = solve(ODEProblem(apply_hamiltonian_freeprecession!, m0, (0, TE), (ω0, m0s, R1f, R2f, Rx, R1s)), alg)

for i = 1: length(t)
        dxf_fd[i] = (FP_sol_dT2s(t[i])[1] - FP_sol(t[i])[1]) /dT2s
        dyf_fd[i] = (FP_sol_dT2s(t[i])[2] - FP_sol(t[i])[2]) /dT2s
        dzf_fd[i] = (FP_sol_dT2s(t[i])[3] - FP_sol(t[i])[3]) /dT2s
        dzs_fd[i] = (FP_sol_dT2s(t[i])[4] - FP_sol(t[i])[4]) /dT2s

        dxf[i] = FP_sol_grad(t[i])[31]
        dyf[i] = FP_sol_grad(t[i])[32]
        dzf[i] = FP_sol_grad(t[i])[33]
        dzs[i] = FP_sol_grad(t[i])[34]
end

@test dxf ≈ dxf_fd rtol = max_error
@test dyf ≈ dyf_fd rtol = max_error
@test dzf ≈ dzf_fd rtol = max_error
@test dzs ≈ dzs_fd rtol = max_error

## test derivative wrt. ω0
dω0 = 1
FP_sol_dω0 = solve(ODEProblem(apply_hamiltonian_freeprecession!, m0, (0, TE), ((ω0 + dω0), m0s, R1f, R2f, Rx, R1s)), alg)

for i = 1: length(t)
        dxf_fd[i] = (FP_sol_dω0(t[i])[1] - FP_sol(t[i])[1]) /dω0
        dyf_fd[i] = (FP_sol_dω0(t[i])[2] - FP_sol(t[i])[2]) /dω0
        dzf_fd[i] = (FP_sol_dω0(t[i])[3] - FP_sol(t[i])[3]) /dω0
        dzs_fd[i] = (FP_sol_dω0(t[i])[4] - FP_sol(t[i])[4]) /dω0

        dxf[i] = FP_sol_grad(t[i])[36]
        dyf[i] = FP_sol_grad(t[i])[37]
        dzf[i] = FP_sol_grad(t[i])[38]
        dzs[i] = FP_sol_grad(t[i])[39]
end

@test dxf ≈ dxf_fd rtol = max_error
@test dyf ≈ dyf_fd rtol = max_error
@test dzf ≈ dzf_fd atol = 1e-4
@test dzf ≈ dzf_fd atol = 1e-4

## test derivative wrt. B1
dB1 = 1e-9
FP_sol_dB1 = solve(ODEProblem(apply_hamiltonian_freeprecession!, m0, (0, TE), (ω0, m0s, R1f, R2f, Rx, R1s)), alg)

for i = 1: length(t)
        dxf_fd[i] = (FP_sol_dB1(t[i])[1] - FP_sol(t[i])[1]) /dB1
        dyf_fd[i] = (FP_sol_dB1(t[i])[2] - FP_sol(t[i])[2]) /dB1
        dzf_fd[i] = (FP_sol_dB1(t[i])[3] - FP_sol(t[i])[3]) /dB1
        dzs_fd[i] = (FP_sol_dB1(t[i])[4] - FP_sol(t[i])[4]) /dB1

        dxf[i] = FP_sol_grad(t[i])[41]
        dyf[i] = FP_sol_grad(t[i])[42]
        dzf[i] = FP_sol_grad(t[i])[43]
        dzs[i] = FP_sol_grad(t[i])[44]
end

@test dxf ≈ dxf_fd rtol = max_error
@test dyf ≈ dyf_fd rtol = max_error
@test dzf ≈ dzf_fd rtol = max_error
@test dzs ≈ dzs_fd rtol = max_error

#############################################################
## R1a
#############################################################
R1a = 1
grad_list = [grad_R1a()]

## baseline ODE solution
FP_sol = solve(ODEProblem(apply_hamiltonian_freeprecession!, m0, (0, TE), (ω0, m0s, R1a, R2f, Rx, R1a)), alg)

FP_sol_grad = solve(ODEProblem(apply_hamiltonian_freeprecession!, [m0; zeros(5 * (length(grad_list)),1)], (0, TE), (ω0, m0s, R1f, R2f, Rx, R1s, grad_list)), alg)

## test derivative wrt. R1a
dR1a = 1e-9
FP_sol_dR1f = solve(ODEProblem(apply_hamiltonian_freeprecession!, m0, (0, TE), (ω0, m0s, (R1a + dR1a), R2f, Rx, (R1a + dR1a))), alg)

for i = 1: length(t)
        dxf_fd[i] = (FP_sol_dR1f(t[i])[1] - FP_sol(t[i])[1]) /dR1a
        dyf_fd[i] = (FP_sol_dR1f(t[i])[2] - FP_sol(t[i])[2]) /dR1a
        dzf_fd[i] = (FP_sol_dR1f(t[i])[3] - FP_sol(t[i])[3]) /dR1a
        dzs_fd[i] = (FP_sol_dR1f(t[i])[4] - FP_sol(t[i])[4]) /dR1a

        dxf[i] = FP_sol_grad(t[i])[6]
        dyf[i] = FP_sol_grad(t[i])[7]
        dzf[i] = FP_sol_grad(t[i])[8]
        dzs[i] = FP_sol_grad(t[i])[9]
end

@test dxf ≈ dxf_fd atol = 1e-3
@test dyf ≈ dyf_fd atol = 1e-3
@test dzf ≈ dzf_fd rtol = max_error
@test dzs ≈ dzs_fd rtol = max_error

## plot the functions for debug purposes
# using Plots
# plotlyjs()
# plot(t, dxf, ticks=:native, label="xf", legend=:bottomleft)
# plot!(t, dxf_fd, label="xf_fd")
# plot!(t, dyf, ticks=:native, label="yf")
# plot!(t, dyf_fd, label="yf_fd")
# plot!(t, dzf, ticks=:native, label="zf")
# plot!(t, dzf_fd, label="zf_fd")
# plot!(t, dzs, label="zs")
# plot!(t, dzs_fd, label="zs_fd")


struct ConTab{T1<:AbstractMatrix, T2<:AbstractVector}
    X::T1
    mx::T2
    my::T2
end

function contab(x, y; normalize = true, correction = 0.5, verbose = true)
    t = [count((i .=== x) .& (j .=== y)) for i in sort(unique(skipmissing(x))), j in sort(unique(skipmissing(y)))]
    t = Matrix{Float64}(t)
    t[t .== 0] .= correction / sum(t)
    if normalize
        t = t ./ sum(t)
        if any(t .< eps())
            @warn "Some normalized frequency are smaller than `eps()`. Use larger `correction` value."
        end
        if any(t .< 0.01) && verbose
            @warn "Probabilities of some elements in the contigency tabel, $(findall(x -> x < 0.01, t)) are less than 0.01. The estimation of polychoric correlation may be unstable. For `snpolyc`, use a smaller number of `step` to solve it."
        end
    end
    return ConTab(t, sum(t; dims = 2)[:], sum(t; dims = 1)[:])
end

function contab(t; normalize = true, correction = 0.5, verbose = true)
    t = Matrix(t)
    t[t .== 0] .= correction / sum(t)
    if normalize
        t = t ./ sum(t)
        if any(t .< eps())
            @warn "Some normalized frequency are smaller than `eps()`. Use more large `correction` value."
        end
        if any(t .< 0.01) && verbose
            @warn "Some elements in the contigency tabel, $(findall(x -> x < 0.01, t)). The estimation of polychoric correlation may be unstable. For `snpolyc`, use a larger number of `step` to solve it."
        end
    end
    return contab(t, sum(t; dims = 2)[:], sum(t; dims = 1)[:])
end

# Calculate all mass, not density.

"""
    ppdf(f, max, min, step)

Compute the volume of the rectangular area ∈ [min, max] by using trapezoidal rule.
"""
function ppdf(f, max, min, step)
    x = range(min[1], stop = max[1], step = step)
    y = range(min[2], stop = max[2], step = step)
    return trapz((x, y), [f([i, j]) for i in x, j in y])
    # return hcubature(f, max, min)[1]
end

"""
Compute, more efficiently than using the trapezoidal rule, the volume of the bivariate normal distributions.
"""
function ppdf(f, m)
    g(h, k, θ) = exp(-0.5 * (h^2 + k^2 - 2*h*k*θ)/(1-θ^2)) / sqrt(1-θ^2 )
    Φ(x) = cdf(Normal(0, 1), x)
    return quadgk(t->g(m[1], m[2], t), 0, f.Σ[1,2])[1] / 2π + Φ(m[1])*Φ(m[2])
    # v = range(0, f.Σ[1,2], length = 20)
    # V = @trapz v x g(m[1], m[2], x)
    # return V / 2π + Φ(m[1])*Φ(m[2])
end

# Numerical integration with sophisticated 1dim version
function H(a, b, c, d, mvd::MvNormal)
    A1 = a == -Inf || c == -Inf ?  0.0 : ppdf(mvd, [a, c])
    A2 = a == -Inf ? 0.0 : ppdf(mvd, [a, d])
    A3 = c == -Inf ? 0.0 : ppdf(mvd, [b, c])
    A4 = d == Inf && b == Inf ? 1.0 : ppdf(mvd, [b, d])
    return A1 - A2 - A3 + A4
end

# Loss function
function loss(t::ConTab, ξx, ξy, d::MultivariateDistribution)
    I, J = size(t.X)
    h = zeros(Float64, I, J)
    @fastmath for i in 1:I, j in 1:J
        if i == I && j == J && sum(h) < 1.0
            h[i, j] = 1.0 - sum(h)
        elseif i == I
            h[i, j] = t.my[j] - sum(h[1:end-1, j])[1]
        elseif j == J
            h[i, j] = t.mx[i] - sum(h[i, 1:end-1])[1]
        else
            h[i, j] = H(ξx[i], ξx[i+1], ξy[j], ξy[j+1], d)
        end
    end
    if any(h .< 0.0)
        h = abs.(h)
    end
    h = h ./ sum(h)
    return - sum(t.X .* log.(h))
end

function polyc(t::ConTab)
    # Fix marginal freq
    cumx = [0.0; cumsum(t.mx)[1:end-1]; 1.0]
    cumy = [0.0; cumsum(t.my)[1:end-1]; 1.0]
    ξx = quantile.(Normal(0, 1), cumx)
    ξy = quantile.(Normal(0, 1), cumy)
    opt = optimize(ρ -> loss(t, ξx, ξy, MvNormal([0, 0], [1.0 ρ; ρ 1.0])), -1.0, 1.0, Brent())
    return (ρ = opt.minimizer, τ₁ = ξx, τ₂ = ξy)
end

function polyc(x, y; verbose = true)
    tab = contab(x, y; verbose = verbose)
    # Fix marginal freq
    polyc(tab)
end

"""
    polycor(X::Union{AbstractDataFrame, AbstractMatrix}; lower_tri = false)
Estimate polychoric correlation matrix from X.

# Arguments
- `X` Matrix contains categorical vectors.
- `lower_tri` logical.

# Examples
```julia

using Distributions
# Support function
function discretize(x, th)
    sum(th .< x)
end
# Gen sim data.
raw = rand(MvNormal([0, 0], [1 .5; .5 1]), 10000)'
τ1 = sort!(rand(Uniform(-2, 2), 3))
τ2 = sort!(rand(Uniform(-2, 2), 4))
obs = [discretize.(raw[:, 1], Ref(τ1)) discretize.(raw[:, 2], Ref(τ2)) ]

polycor(obs)
```
"""
function polycor(X)
    J = size(X, 2)
    r = Matrix{AbstractFloat}(undef, J, J)
    r[diagind(r)] .= 1.0
    for i in 1:J
        x = @view X[:, i]
        for j in i:J
            y = @view X[:, j]
            r[j, i] = polyc(x, y; verbose = false).ρ
            r[i, j] = r[j, i]
        end
    end
    return r
end

function replace_diagonal!(M)
    offdiagind(A) = (ι for ι in CartesianIndices(A) if ι[1] ≠ ι[2])
    L = maximum(M[collect(offdiagind(M))])
    M[diagind(M)] .= L
end


using Singular

function add(l1 :: Vector{spoly{A}}, l2 :: Vector{spoly{A}}) where A
  length(l1) == length(l2) || error("vector dimensions not equal")
  ln = length(l1)
  R = parent(l1[1])
  carry = zero(R)
  l3 = [zero(R) for _ in 1:ln]
  for ix in 1:ln
    l3[ix] = carry+l1[ix]+l2[ix]
    #note that xy ∨ xz ∨ yz == xy + xz + yz
    carry = carry*l1[ix]+carry*l2[ix]+l1[ix]*l2[ix]
  end
  return l3
end

function sub(l1 :: Vector{spoly{A}}, l2 :: Vector{spoly{A}}) where A
  length(l1) == length(l2) || error("vector dimensions not equal")
  ln = length(l1)
  R = parent(l1[1])
  carry = one(R)
  l3 = [zero(R) for _ in 1:ln]
  for ix in 1:ln
    l3[ix] = carry+l1[ix]+l2[ix]+one(R)
    #note that xy ∨ xz ∨ yz == xy + xz + yz
    carry = carry*l1[ix]+carry*(l2[ix]+one(R))+l1[ix]*(l2[ix]+one(R))
  end
  return l3
end

function shr(u :: Vector{spoly{A}}; step :: Int = 1) where A
  length(u) > 0 || error("vector dimension equal to zero")
  R = parent(u[1])
  v = [zero(R) for _ in 1:length(u)]
  for ix in 1:length(u)-step
    v[ix] = u[ix+step]
  end
  return v
end

function and(l1 :: Vector{spoly{A}}, l2 :: Vector{spoly{A}}) where A
  length(l1) == length(l2) || error("vector dimensions not equal")
  ln = length(l1)
  R = parent(l1[1])
  l3 = [zero(R) for _ in 1:ln]
  for ix in 1:ln
    l3[ix] = l1[ix]*l2[ix]
  end
  return l3
end

#note: returns a vector of length `length(u)+step`
function shl_ext(u :: Vector{spoly{A}}; step :: Int = 1) where A
  length(u) > 0 || error("vector dimension equal to zero")
  R = parent(u[1])
  v = [zero(R) for _ in 1:(length(u)+step)]
  for ix in step+1:length(u)+step
    v[ix] = u[ix-step]
  end
  return v
end

#note: returns a vector of length `length(l1)+1`
function add_ext(l1 :: Vector{spoly{A}}, l2 :: Vector{spoly{A}}) where A
  length(l1) == length(l2) || error("vector dimensions not equal")
  ln = length(l1)
  R = parent(l1[1])
  carry = zero(R)
  l3 = [zero(R) for _ in 1:ln+1]
  for ix in 1:ln
    l3[ix] = carry+l1[ix]+l2[ix]
    #note that xy ∨ xz ∨ yz == xy + xz + yz
    carry = carry*l1[ix]+carry*l2[ix]+l1[ix]*l2[ix]
  end
  l3[ln+1] = carry
  return l3
end

function mapreduce(l::Vector{spoly{A}}, I::sideal{spoly{A}}) where A
  return map(x->reduce(x,I),l)
end

vars = ["b$ix" for ix in 0:15]
R, b = PolynomialRing(Fp(2), vars)
I = Ideal(R, [b[ix]^2-b[ix] for ix in 1:16])
I = std(I) #compute Gröbner basis of I
(b0,b1,b2,b3,b4,b5,b6,b7,b8,b9,b10,b11,b12,b13,b14,b15) = b
reduce(b0^2+b1^3, I)

#0f = 00001111
#01 = 00000001
_5555 = [one(R),zero(R),one(R),zero(R),one(R),zero(R),one(R),zero(R),
         one(R),zero(R),one(R),zero(R),one(R),zero(R),one(R),zero(R)]
_3333 = [one(R),one(R),zero(R),zero(R),one(R),one(R),zero(R),zero(R),
         one(R),one(R),zero(R),zero(R),one(R),one(R),zero(R),zero(R)]
_0f0f = [one(R),one(R),one(R),one(R),zero(R),zero(R),zero(R),zero(R),
         one(R),one(R),one(R),one(R),zero(R),zero(R),zero(R),zero(R)]
g = [b0,b1,b2,b3,b4,b5,b6,b7,b8,b9,b10,b11,b12,b13,b14,b15]

di1 = deepcopy(g);
dx1 = shr(di1);
dx2 = and(dx1,_5555);

#di1 - dx2

di2 = sub(di1,dx2);
di3 = mapreduce(di2,I);
ax1 = deepcopy(di3);
di4 = shr(di3; step=2);
ax2 = and(ax1, _3333);
di5 = and(di4, _3333);
ax3 = mapreduce(ax2, I);
di6 = mapreduce(di5, I);
di7 = add(di6, ax3);
di8 = mapreduce(di7, I);
ax4 = deepcopy(di8);
ax5 = shr(ax4; step=4);
ax6 = add(ax5, di8);
ax7 = mapreduce(ax6, I);
ax8 = and(ax7, _0f0f);
ax9 = mapreduce(ax8, I);
#imul   ax,ax,0x0101
#0x0101 = 0b0000000100000001
ax10 = [ax9[1],ax9[2],ax9[3],ax9[4],ax9[5],ax9[6],ax9[7],ax9[8],
         ax9[9],ax9[10],ax9[11],ax9[12],ax9[13],ax9[14],ax9[15],ax9[16],
         zero(R),zero(R),zero(R),zero(R),zero(R),zero(R),zero(R),zero(R)];
ax11 = shl_ext(ax9; step=8);
ax12 = add_ext(ax10,ax11);
ax13 = mapreduce(ax12, I)
#shr ax, 0x8
ax14 = [ax13[9],ax13[10],ax13[11],ax13[12],
        ax13[13],ax13[14],ax13[15],ax13[16],
        zero(R),zero(R),zero(R),zero(R),
        zero(R),zero(R),zero(R),zero(R)]

function equiv(u::BitVector)
  n = sum(a)
  v = [0 for _ in 1:32]
  v[1] = n % 2
  ix = 1
  while 2^ix <= n
    v[ix+1] = binomial(n,2^ix) % 2
    ix += 1
  end
  return v
end


# continuous 1-d; different scalings should lead to different sizes, different boundary condtions shouldn't
@testset "Construction Types" begin
    xSizes = (33, 67, 128)
    boundaries = (DEFAULT_BOUNDARY, padded, NaivePer)
    sVals = [1, 2, 3.5, 8, 16]
    βs = [1, 1.5,4.0]
    waveTypes =(ContinuousWavelets.morl, dog0, paul4, cDb2, ContinuousWavelets.morse)
    @testset "xSz=$xSize, b=$boundary, s=$s, β=$β, wfc=$(wave)" for xSize in xSizes, boundary in boundaries, s in sVals, β in βs, wave in waveTypes
        wfc = ContinuousWavelets.wavelet(wave, s=s, boundary=boundary,β=β)
        xr = rand(Float64, xSize);
        x32 = rand(Float32, xSize);
        xc = rand(ComplexF64, xSize);
        # the sizes are of course broken at this size, so no warnings needed
        yc = 3; yr = 3; y32 = 3
        with_logger(ConsoleLogger(stderr, Logging.Error)) do
            yr = ContinuousWavelets.cwt(xr, wfc);
            y32 = ContinuousWavelets.cwt(x32, wfc);
            yc = ContinuousWavelets.cwt(xc, wfc);
        end
        if isAnalytic(wfc.waveType)
            @test eltype(yr) == ComplexF64
            @test eltype(y32) == ComplexF32
            @test eltype(yc) == ComplexF64
        else
            @test eltype(yr) == Float64
            @test eltype(y32) == Float32
            @test eltype(yc) == ComplexF64
        end
        nOctaves, totalWavelets, sRanges, sWidths = ContinuousWavelets.getNWavelets(xSize, wfc);
        @test totalWavelets == size(yc, 2)
    end
end
# xSize = xSizes[3]; boundary = boundaries[3]; s = sVals[end]; β = βs[3]; wave = waveTypes[4]
# xSize = 33; boundary = SymBoundary(); s = 1; β = 1; wave = dog2


"""
$(TYPEDEF)

The supertype for Lattices in `D`-dimensional space. Type `T` is used to represent coordinates. The boundary conditions can be either periodic (`PB=true`) or free (`PB=false`).
"""
abstract type AbstractLattice{D, T, PB} <: AbstractNodeCollection{D, T} end

"""
$(TYPEDEF)
$(TYPEDFIELDS)

This type describes a regular arrangment of nodes in `D`-dimensional space with boundary condition `PB` (`PB=true` for periodic boundary conditions, `PB=false` for free periodic boundary conditions). A general lattice consists of identicall cells (combinations of nodes) arranged as a Bravais lattice.
"""
struct RegularLattice{D, T, PB, CT, L<:Union{Symbol,Nothing}} <: AbstractLattice{D, T, PB}
    """
    The number of basis cells along each of the basis directions.
    """
    lattice_dims::NTuple{D, Int}
    """
    Coordinates of the primitive vectors of the underlying Bravais lattice. `primitive_vecs[:, i]` gives the ``i``-th primitive vector.
    """
    primitive_vecs::SMatrix{D,D, T}
    """
    Repeated basis cell of the lattice.
    """
    basis_cell::CT
    """
    Label of the Lattice.
    """
    label::L
    """
    Total number of cells.
    """
    num_of_cells::Int
    """
    Total number of nodes.
    """
    num_of_nodes::Int
    """
    Cartesian index of the central cell (in the case of even length among of dimensions it is an approximation).
    """
    central_cell::CartesianIndex{D}
    """
    $(SIGNATURES)

    Constructs a Regular Lattice enforcing the condition that `typeof(basis_cell)<:AbstractCell{D,T}`.

    #Arguments
    - `lattice_dims::NTuple{D,Int}`: specifies the number of basis cells along each of the basis directions.
    - `primitive_vecs::SMatrix{D,D,T}`: stores as columns the primitive vectors of the underlying Bravais lattice.
    - `basis_cell::AbstractCell{D,T}`: specifies the repeated basis cell of the lattice; if the basis cell does not have any special structure, an instance of `TrivialCell` should be used.
    - `label::Union{Symbol,Nothing}`: the label of the lattice.
    - `PB::Bool`: specifies if periodic or free boundary conditions are applied; default value is `true`.
    """
	function RegularLattice(lattice_dims::NTuple{D, Int}, primitive_vecs::SMatrix{D,D,T}, basis_cell::AbstractCell{D,T}, label::Union{Symbol,Nothing}, PB::Bool) where {D,T}
        @assert D>=1
        num_of_cells = prod(lattice_dims)
        central_cell = CartesianIndex(div.(lattice_dims,2).+1)
        new{D, T, PB, typeof(basis_cell), typeof(label)}(lattice_dims, primitive_vecs, basis_cell, label, num_of_cells, length(basis_cell)*num_of_cells, central_cell)
    end
end

### Outer constructors.

function RegularLattice(lattice_dims::NTuple{D, Int}, primitive_vecs::SMatrix{D,D,T1}, basis_cell::AbstractCell{D,T2}, label, PB) where {D,T1,T2}
    T = promote_type(T1,T2)
    return RegularLattice(lattice_dims, SMatrix{D,D,T}(primitive_vecs), switch_coord_type(basis_cell, T), label, PB)
end

"""
`RegularLattice(lattice_dims::NTuple{D, Int}, a::T=1; label=:cubic, periodic=true)`

Constructs hypercubic lattice with lattice parameter `a` and trivial unit cell.
"""
function RegularLattice(lattice_dims::NTuple{D, Int}, a::T=1; periodic=true, label=:cubic) where {D, T<:Number}
    primitive_vecs = one(SMatrix{D,D,Int})*a
    basis_cell = TrivialCell{D,T}()
    return RegularLattice(lattice_dims, primitive_vecs, basis_cell, label, periodic)
end
function RegularLattice(lattice_dims::NTuple{D,Int}, primitive_vecs::SMatrix{D,D,T}; label=:simple, periodic = true) where {D,T}
    basis_cell = TrivialCell{D,T}()
    return RegularLattice(lattice_dims, primitive_vecs, basis_cell, label, periodic)
end
"""
`RegularLattice(lattice_dims::NTuple{D, Int}, primitive_vecs::SMatrix{D,D}, basis_cell::AbstractCell{D}; label=nothing, periodic=true)`

Convenient constructor which allows to specify the label and boundary condition as keyword arguments.
"""
RegularLattice(lattice_dims::NTuple{D, Int}, primitive_vecs::SMatrix{D,D}, basis_cell::AbstractCell{D}; label=nothing, periodic=true) where {D} =
    RegularLattice(lattice_dims, primitive_vecs, basis_cell, label, periodic)

"""
$(TYPEDSIGNATURES)

Return the number of lattice separate groups.
"""
@inline num_of_groups(lattice::RegularLattice) = num_of_groups(lattice.basis_cell)

Base.@propagate_inbounds group_size(lattice::RegularLattice, ig) = lattice.num_of_cells*group_size(lattice.basis_cell, ig)

Base.length(lattice::RegularLattice) = lattice.num_of_nodes

##
## Indexing interface.

"""
`getindex(lattice::RegularLattice{D}, Index...)`

Returns the coordinate of the node with index `Index`. The configurations of indices for different types of unit cell are the following:
- `TrivialCell`: `I::CartesianIndex{D}`.
- `HomogeneousCell`: `I::CartesianIndex{D}, ic::Int`.
- `InhomogeneousCell`: `I::CartesianIndex{D}, ic::Int, ig::Int`.
When iterating over lattice, the index at the left is iterated faster. For example, in the case of `InhomogeneousCell` I is iterated first, then `ic`, then `ig`.

In the case of periodic lattice, if `CartesianIndex` is outside of the lattice dims, it is simply translated back inside.
"""

Base.@propagate_inbounds function Base.getindex(lattice::RegularLattice{D,T}, I::CartesianIndex{D}, Ic...) where {D,T}
    @boundscheck check_lattice_index(lattice, I)
	@boundscheck check_cell_index(lattice.basis_cell, Ic...)
    @inbounds lattice.primitive_vecs*SVector{D}(I.I) + lattice.basis_cell[Ic...]
end
Base.@propagate_inbounds function Base.getindex(lattice::RegularLattice{D,T,false,<:TrivialCell}, I::CartesianIndex{D}) where {D,T}
    @boundscheck check_lattice_index(lattice, I)
    @inbounds lattice.primitive_vecs*SVector{D}(I.I)
end
Base.@propagate_inbounds function Base.getindex(lattice::RegularLattice{D,T,true}, I::CartesianIndex{D}, Ic...) where {D,T}
    @boundscheck check_cell_index(lattice.basis_cell, Ic...)
    @inbounds lattice.primitive_vecs*SVector{D}(mod1.(I.I, lattice.lattice_dims)) + lattice.basis_cell[Ic...]
end
Base.@propagate_inbounds function Base.getindex(lattice::RegularLattice{D,T,true,<:TrivialCell}, I::CartesianIndex{D}) where {D,T}
    @inbounds lattice.primitive_vecs*SVector{D}(mod1.(I.I, lattice.lattice_dims))
end

@inline check_lattice_index(lattice::RegularLattice{D}, I::CartesianIndex{D}) where {D} = _check_cartesian_index(true, I.I, lattice.lattice_dims) || throw(BoundsError(lattice, I))
@inline _check_cartesian_index(b, i, stop) = _check_cartesian_index(b & (1 <= i[1] <= stop[1]), Base.tail(i), Base.tail(stop))
@inline _check_cartesian_index(b, i::Tuple, ::Tuple{}) = false
@inline _check_cartesian_index(b, i::Tuple{}, stop::Tuple) = false
@inline _check_cartesian_index(b, i::Tuple{}, stop::Tuple{}) = b

## Relative coordinate

Base.@propagate_inbounds function relative_coordinate(lattice::RegularLattice{D,T,false}, I1::TI, I2::TI) where {D,T, TI<:Union{CartesianIndex{D}, Tuple{CartesianIndex{D},Vararg{Int}}}}
    return lattice[I1] - lattice[I2]
end
"""
$(TYPEDSIGNATURES)

For periodic lattice, the "shortest" relative coordinate is calculated instead.

For, that the following heuristic is used. Cartesian indices of the two nodes are shifted by the same amount, so that the cartesian index of the second node corresponds to the central cell of the lattice. Then, the cartesian index of the first node is translated back inside the lattice. The relative coordinate is computed using the resulting indices.

This heuristic guarantees that
`relative_coordinate(lattice, I1, I2) == -relative_coordinate(lattice, I2, I1)`.
"""
Base.@propagate_inbounds function relative_coordinate(lattice::RegularLattice{D,T,true,<:TrivialCell}, I1::TI, I2::TI) where {D,T, TI<:CartesianIndex{D}}
    return lattice[I1 + lattice.central_cell - I2] - lattice[lattice.central_cell]
end
Base.@propagate_inbounds function relative_coordinate(lattice::RegularLattice{D,T,true}, I1::TI, I2::TI) where {D,T, TI<:Tuple{CartesianIndex{D},Vararg{Int}}}
    return lattice[first(I1) + lattice.central_cell - first(I2), Base.tail(I1)...] - lattice[lattice.central_cell, Base.tail(I2)...]
end

#Base.@propagate_inbounds function translate_indices(lattice::RegularLattice{D,T,PB,<:TrivialCell}, Is) where {D,T,PB}
#    @boundscheck map(x -> check_lattice_index(lattice, x), Is)
#    return Is
#end
#Base.@propagate_inbounds function translate_indices(lattice::RegularLattice{D,T,PB,<:HomogeneousCell}, Is) where {D,T,PB}
#    @boundscheck map(x -> check_lattice_index(lattice, x), Is)
#    return Base.product(Is, Base.OneTo(length(lattice.unit_cell)))
#end
#Base.@propagate_inbounds function translate_indices(lattice::RegularLattice{D,T,PB,<:InhomogeneousCell}, Is, ig::Int) where {D,T,PB}
#    @boundscheck map(x -> check_lattice_index(lattice, x), Is)
#    return Base.product(Is, Base.OneTo(group_size(lattice.unit_cell, ig)), ig)
#end

#@inline to_default(lattice::RegularLattice{D,T,PB,<:TrivialCell}, x::Int, y::Int, z::Int) = CartesianIndex


# Model
include_model("double_integrator")

# Horizon
T = 11

tf = 1.0
h0 = tf / (T-1) # timestep

# Initial and final states
x1 = [1.0; 0.0]
xT = [0.0; 0.0]

xl, xu = state_bounds(model, T, x1 = x1, xT = xT)

# Objective
obj = quadratic_tracking_objective(
        [t < T ? Diagonal(ones(model.n)) : Diagonal(ones(model.n)) for t = 1:T],
        [Diagonal(1.0 * ones(model.m)) for t = 1:T-1],
        [xT for t = 1:T],
		[zeros(model.m) for t = 1:T])

# Problem
prob = trajectory_optimization_problem(model,
			   obj,
			   T,
               xl = xl,
               xu = xu)

# Initialization
x0 = linear_interpolation(x1, xT, T) # linear interpolation for states
u0 = [ones(model.m) for t = 1:T-1]

# Pack trajectories into vector
z0 = pack(x0, u0, prob)

# Solve
# include_snopt() # set path
# @time z, info = solve(prob, copy(z0), nlp = :SNOPT7)
@time z, info = solve(prob, copy(z0), nlp = :ipopt)

# Visualize
using Plots
x, u = unpack(z, prob)
plot(hcat(x...)', width = 2.0)
plot(hcat(u...)', width = 2.0, linetype = :steppost)


immutable PolyharmonicSpline{K,S} <: AbstractInterpolation
end

PolyharmonicSpline() = PolyharmonicSpline{2,Val{0}}
PolyharmonicSpline(K) = PolyharmonicSpline{K,Val{0}}
PolyharmonicSpline(K,S) = PolyharmonicSpline{K,Val{S}}

typealias ThinPlateSpline{S} PolyharmonicSpline{2,S}
ThinPlateSpline() = ThinPlateSpline{Val{0}}
ThinPlateSpline(S) = ThinPlateSpline{Val{S}}

function default_settings{D,T,K,S}(::Type{PolyharmonicSpline{K,S}},axes::NTuple{D,T})
    N = length(axes[1])
    return Irregular{D,N}, None
end

type PolyharmonicSplineInterpolation{T,D,N}
  coeff::Vector{T}
  values::Vector{T}
  centers::Array{T,2}
end

@inline function polyharmonicK(r,K)
    if iseven(K)
        p = r < 1.0 ? (r^(K-1))*log(r^r) : (r^K)*log(r)
    else
        p = r^K
    end
    return p
end

function preprocess{K,S}(phs::Type{PolyharmonicSpline{K,S}}, axes, f)
    values = convert(Array, f)
    centers = hcat(axes...)'
    n, m = size(centers)
    m != length(values) && throw(DimensionMismatch())

    M = zeros(m,m)
    N = zeros(m,n+1)

    @inbounds for i=1:m
        N[i,1] = 1
        N[i,2:end] = centers[:,i]
        for j=i:m
            M[i,j] = polyharmonicK(norm(centers[:,i] .- centers[:,j]),K)
            M[j,i] = M[i,j]
        end
    end
    s = S.parameters[1]
    M = M .+ s.*eye(m)
    L = [[M N]; [N' zeros(n+1,n+1)]]

    w = pinv(L)*[values; zeros(n+1)]

    w, centers, values = promote(w, centers, values)
    return PolyharmonicSplineInterpolation{eltype(w),n,length(w)}(w,values,centers)
end

function apply_bc{K,S}(::Type{PolyharmonicSpline{K,S}},BC::Type{None}, I::Symbol, x::Symbol)
    apply_none()
end

@inline function get_indices{K,S,D,N}(::Type{PolyharmonicSpline{K,S}}, G::Type{Irregular{D,N}}, I, x)
    return ()
end

@inline function get_value{K,S,T,D,N}(::Type{PolyharmonicSpline{K,S}}, I::PolyharmonicSplineInterpolation{T,D,N}, x, inds)
    v = zero(Float64)
    l = N - (D+1)
    centers = I.centers
    coeff = I.coeff
    @inbounds for i=1:l
        v = v + coeff[i]*polyharmonicK(norm(x .- centers[:,i]),K)
    end
    v = v + coeff[l+1]
    @inbounds for i=2:D+1
        v = v + coeff[l+i]*x[i-1]
    end
    return v
end


using SpecialFunctions: loggamma, digamma
import Distributions: entropy
using StaticArrays

export Dirichlet, DirichletClassifier, Categorical

################################################################################
##### Dirichlet
################################################################################
"""
    Dirichlet{N,T}

Statically sized Dirichlet distributions. In contrast to say `Distributions.Dirichlet` the
the number of categories is encoded as a type parameter.
"""
struct Dirichlet{N,T}
    α::NTuple{N,T}
    Dirichlet(α::NTuple{N,T}) where{N,T} = new{N,T}(α)
end

Dirichlet(args...) = Dirichlet(promote(args...))
Dirichlet(t::Tuple) = Dirichlet(t...)
Dirichlet(d::Dirichlet) = d
function Base.show(io::IO, d::Dirichlet)
    print(io, "Dirichlet", d.α)
end

function Statistics.var(d::Dirichlet)
    # var is the diagonal of the cov matrix
    # this is the convention from Distributions.jl
    α = SVector(d.α)
    S = sum(α)
    m = α./S
    α .* (S .- α) ./ (S .^2 .* (S .+ 1))
end

to_NTuple(::Val{n}, x::NTuple{n}) where {n} = x
to_NTuple(::Val{n}, x::StaticVector{n}) where {n} = Tuple(x)
to_NTuple(N::Val, itr) = to_NTuple_dynamic(N, itr)

to_NTuple_dynamic(::Val{0}, itr) = ()
to_NTuple_dynamic(::Val{1}, itr) = (first(itr),)
function to_NTuple_dynamic(::Val{2}, itr)
    x1,x2 = itr
end
function to_NTuple_dynamic(::Val{3}, itr)
    x1,x2,x3 = itr
end
function to_NTuple_dynamic(::Val{4}, itr)
    x1,x2,x3,x4 = itr
end
to_NTuple_dynamic(::Val{n}, itr) where {n} = NTuple{n}(itr)::NTuple{n}

function Distributions.logpdf(d::Dirichlet{N}, x) where {N}
    x = to_NTuple(Val(N), x)
    α = d.α
    A = sum(loggamma.(α)) - loggamma(sum(α))
    sum(log.(x) .* (α .- 1)) - A
end

"""

    DirichletClassifier(Val(k))

Layer that takes a tensor of floats of size (k*n, b) and turns it into
a tensor of Dirichlet distributions of size (n,b).
"""
struct DirichletClassifier{k}
    K::Val{k}
end

val(::Val{k}) where {k} = k
NClasses(::Type{DirichletClassifier{k}}) where {k} = Val(k)
NClasses(o) = NClasses(typeof(o))
function NClasses(T::Type)
    error("""
          NClasses(::Type{$T}) not implemented.
          """)
end
nclasses(o) = val(NClasses(o))

function (o::DirichletClassifier)(arr::AbstractMatrix)
    @argcheck size(arr, 1) == nclasses(o)
    K = NClasses(o)
    rows = ntuple(K) do i
        view(arr, i, :)
    end
    map(rows...) do args...
        Dirichlet(softplus.(args)...)
    end
end

"""

    kl_uniform(d::Dirichlet)

Given `P = Dirichlet(α)` compute the Kullback Leibler divergence with respect to the uniform Dirichlet distribution:

    `D_KL(P || Dirichlet(1,1,1...))`
"""
function kl_uniform(d::Dirichlet)
    # For any distribution P = p(x)dx
    # And a uniform distribution U = u*dx where u = 1/vol(Ω)
    # We have
    # D_KL(P||U) = ∫ log(u/p(x)) p(x) dx
    #            = ∫ log(u) p(x) dx - ∫log(p(x)) p(x) dx
    #            = entropy(U)       - entropy(P)
    #
    # This also makes sense intuitively:
    # D_KL(P||U) measures the extra price we pay for using the uniform coding scheme
    # on P distributed data.
    # In the uniform scheme, we pay always a flat rate of entropy(U) for any sample,
    # while with the optimal code we only pay entropy(P) on average per sample
    entropy_uniform = -loggamma(nclasses(d))
    entropy_uniform - entropy(d)
end

NClasses(::Type{<:Dirichlet{k}}) where {k} = Val(k)
function Distributions.entropy(d::Dirichlet)
    k = nclasses(d)
    α = d.α
    α₀ = sum(α)
    ψ = digamma
    logBα = log_multivariate_beta(α)
    logBα + (α₀ - k)* ψ(α₀) - sum((α .- 1) .* ψ.(α))
end

function log_multivariate_beta(α)
    sum(loggamma, α) - loggamma(sum(α))
end

function Statistics.mean(d::Dirichlet)
    a = inv(sum(d.α))
    return a * d.α
end

predict(d::Dirichlet) = argmax(d.α)
evidence(d::Dirichlet) = sum(d.α) - nclasses(d) # is this correct?

################################################################################
##### Categorical
################################################################################
struct Categorical{N,T}
    p::NTuple{N,T} # Σpᵢ = 1 must hold
end

function posterior_predictive(o::Dirichlet)
    Categorical(o.α ./ sum(o.α))
end
Base.@propagate_inbounds Distributions.logpdf(o::Categorical, i) = log(o.p[i])
Base.@propagate_inbounds Distributions.pdf(o::Categorical, i)    = o.p[i]

# conversion
Distributions.Dirichlet(d::Dirichlet) = Distributions.Dirichlet(collect(d.α))
Distributions.Categorical(d::Categorical) = Distributions.Categorical(d.p)

for (D_EDL, D_Dist) in [
                        (:Dirichlet  , :(Distributions.Dirichlet  )),
                        (:Categorical, :(Distributions.Categorical)),
                        (:Normal     , :(Distributions.Normal     )),
                        ]

    @eval Base.convert(::Type{Distributions.Distribution}, d::$D_EDL) = Base.convert($D_Dist, d)
    @eval Base.convert(::Type{$D_Dist}, d::$D_EDL) = $D_Dist(d)
end

sampletype(::Type{Dirichlet{N,T}}) where {N,T} = T
sampletype(::Type{Categorical{N,T}}) where {N,T} = T


#test time evolution (DifferentialEquations 4.0)
using ProjectedGPE

function timeevolution2()
sinfo = Params()

#== start template ==#
#fundamental constants/units
  ħ = 1.0545718e-34
  kB = 1.38064852e-23
  amu = 1.660339040e-27
  Bohr = 5.29e-11
#Rb87 mass and scattering length
  m = 86.909180527*amu
  as = 100*Bohr
#trap frequencies
  ωx = 2π
  ωy = 4ωx
  ωz = 0.
#choice of time, length, energy units
  t0 = 1.0/ωy
  x0 = sqrt(ħ*t0/m)
  E0 = ħ/t0
#interactions
  #g  = (4*pi*ħ^2*as/m)*x0^3/E0 #dimensionless 3D
  g = 0.1 #test 2D
#damping parameters (dimensionless)
  γ = 0.05
  ℳ  = 0.0
#chemical potential (dimensionless)
  μ  = 12.0
#time evolution parameters
  ti = 0.0
  tf = 3.0/γ  #evolve for 2 damping times
  Nt = 50
  t = collect(linspace(ti,tf,Nt))
  dt = 0.01π/μ #integrate step size [ - should have dt ≪ 2π/μ]
#== end template ==#

@pack sinfo = ωx,ωy,ωz,γ,ℳ,g,t0,x0,E0,μ,ti,tf,Nt,t,dt

#Initialize CField (dimensionless)
  basis = "Hermite"
  ecut = 30*ħ*ωy/E0
  Ω = [ωx; ωy]*t0
  cinfo = makecinfo(ecut,Ω,basis)
  tinfo = maketinfo(cinfo,4)
  @unpack en,P,M = cinfo;
  #x,wx,Tx,y,wy,Ty = makealltrans(M,4,Ω)
  #W = wx.*wy'
#test transform
  c0  = randn(M...)+im*randn(M...); c0=P.*c0
  ψ0  = c2x(c0,tinfo) #initial condition
  ψ = similar(ψ0)  #a field to write to in place

#equation of motion
  function nlin!(dc,c,tinfo)
      c2x!(ψ,c,tinfo)
      x2c!(dc,abs2.(ψ).*ψ,tinfo)
  end

  #in place
  function Lgp!(dc,c,p,t)
      nlin!(dc,c,tinfo)
      dc .= P.*(-im*(1-im*γ)*((en - μ).*c .+ g*dc))
  end

tspan = (t[1],t[end])
prob = ODEProblem(Lgp!,c0,tspan)
alg =    DP5() #abstol = 1e-3, reltol = 1e-3

println("Started evolution ...")
#@time sol = solve(prob,alg,dt=dt,saveat=t,alg_hints=[:stiff],save_everystep=false,dense=false);
@time sol = solve(prob,alg,dt=dt,saveat=t);
println("... Finished.")
return sinfo,cinfo,sol
end

sinfo,cinfo,sol = timeevolution2()

#plot solution for 2D
## Transform to cartesian grid
@unpack ħ,m,ωx,ωy,ωz,γ,ℳ,g,x0,t0,E0,μ,ti,tf,Nt,t,dt = sinfo
#@unpack M,Ω,ecut,P,en = cinfo

#examine solution
using PyPlot

#lets be explicit about units:
Rx = sqrt(2μ*E0/m/ωx^2)
Ry = sqrt(2μ*E0/m/ωy^2)
yMax=1.5Ry
xMax=1.5Rx

Nx = 400
Ny = Nx
x = collect(linspace(-xMax,xMax,Nx))
y = collect(linspace(-yMax,yMax,Ny))
tinfop = maketinfoplot(cinfo,x/x0,y/x0)

#Plot
#using PyPlot
f=figure(figsize=(12,3))
#@manipulate for i=1:length(t) withfig(f,clear=true) do
    ψ = c2x(sol[end],tinfop);
    pcolormesh(x/x0,y/x0,g*abs2.(ψ'))
    xlabel("x/x0")
    ylabel("y/x0")
    colorbar()
    #end
#end


module DomainIntegrals

using QuadGK, LinearAlgebra
using GaussQuadrature, FastGaussQuadrature, HCubature

using StaticArrays
using IntervalSets, DomainSets
using CompositeTypes.Display

import Base: convert, checkbounds, length,
    *, ∘

import DomainSets:
    domaintype,
    codomaintype,
    prectype,
    numtype,
    forward_map,
    components

using DomainSets: ×,
    EmptySpace, FullSpace,
    euclideandimension

include("extra/common.jl")

include("gauss.jl")
include("functional/measure.jl")
include("weights.jl")
include("singularity.jl")
include("strategy.jl")
include("integrand.jl")
include("processing/properties.jl")
include("processing/measure.jl")
include("processing/domain.jl")
include("integral.jl")
include("complexplane.jl")

include("functional/forms.jl")

end # module


## test_fem.jl : test GRF defined on FEM

@testset "finite element mesh      " begin

    ## Load node/element table ##
    p, t = star()
    m = Matern(0.2,2.0)

    ## Cholesky ##
    grf = GaussianRandomField(CovarianceFunction(2,m),Cholesky(),p,t)
    @test isa(grf,GaussianRandomField)
    @test isa(grf.cov,CovarianceFunction)
    @test isa(grf.cov.cov,Matern)
    @test ndims(grf.cov) == 2
    @test isa(grf,GaussianRandomField{Cholesky})
    @test length(grf.pts) == 2
    @test size(grf.pts[1],1) == 2
    @test size(grf.pts[1],2) == size(p,1)
    @test size(grf.pts[2],1) == 3
    @test size(grf.pts[2],2) == size(t,1)
    #plot_trisurf(grf); close()
    #tricontourf(grf); close()
    #plot(grf); close()

    ## Spectral ##
    grf = GaussianRandomField(CovarianceFunction(2,m),Spectral(),p,t)
    @test isa(grf,GaussianRandomField{Spectral})

    ## CirculantEmbedding ##
    @test_throws ArgumentError GaussianRandomField(CovarianceFunction(2,m),CirculantEmbedding(),p,t)

    ## KarhunenLoeve ##
    grf = GaussianRandomField(CovarianceFunction(2,m),KarhunenLoeve(100),p,t,mode="nodes")
    @test isa(grf,GaussianRandomField{KarhunenLoeve{100}})
    #plot_eigenvalues(grf); close()
    #plot_eigenfunction(grf,1); close()

    ## Anisotropic random field ##
    a = AnisotropicExponential([1000 0; 0 1000])
    cov = CovarianceFunction(2,a)
    grf = GaussianRandomField(cov,KarhunenLoeve(500),p,t,quad=GaussLegendre())
    @test isa(grf.cov.cov,AnisotropicExponential)

    ## Center mode ##
    grf = GaussianRandomField(CovarianceFunction(2,m),KarhunenLoeve(100),p,t,mode="center")
    @test isa(grf,GaussianRandomField{KarhunenLoeve{100}})
    @test_throws ArgumentError GaussianRandomField(CovarianceFunction(2,m),KarhunenLoeve(100),p,t,mode="cente")
    @test_throws ArgumentError GaussianRandomField(1,CovarianceFunction(2,m),KarhunenLoeve(100),p,t,mode="cente")
    @test_throws ArgumentError GaussianRandomField(zeros(size(t,1)),CovarianceFunction(2,m),KarhunenLoeve(100),p,t,mode="cente")

    ## GaussianRandomField with mean value ##
    grf = GaussianRandomField(1,CovarianceFunction(2,m),Cholesky(),p,t,mode="center")
    @test isa(grf,GaussianRandomField{Cholesky})
    @test length(grf.mean) == size(t,1)
    @test all(grf.mean.==1)

    grf = GaussianRandomField(ones(size(p,1)),CovarianceFunction(2,m),Spectral(),p,t)
    @test isa(grf,GaussianRandomField{Spectral})
    @test length(grf.mean) == size(p,1)
    @test all(grf.mean.==1)

    ## Non-bounding box method ##
    p, t = Lshape()
    grf = grf = GaussianRandomField(CovarianceFunction(2,Matern(0.2,1.0)),Spectral(),p,t,n=10)
    @test size(grf.data.eigenfunc,2) == 10

end


# Finite-core model
function bound_leg_velocity(a, b, Γ, ε)
    na, nb, σ = norm(a), norm(b), dot(a, b)
    term_1    = (na^2 - σ) / √(na^2 + ε^2) + (nb^2 - σ) / √(nb^2 + ε^2)
    term_2    = a × b / (na^2 * nb^2 - σ^2 + ε^2 * (na^2 + nb^2 - 2 * na * nb))

    Γ/4π * term_1 * term_2
end

trailing_leg_velocity(r, Γ, u, ε) = Γ/4π * normalize(r) × u / (norm(r) - dot(r, u) + ε^2 / (norm(r) + dot(r, u)))
trailing_legs_velocities(a, b, Γ, u, ε) = trailing_leg_velocity(a, Γ, u, ε) - trailing_leg_velocity(b, Γ, u, ε)
total_horseshoe_velocity(a, b, Γ, u, ε) = bound_leg_velocity(a, b, Γ, ε) + trailing_legs_velocities(a, b, Γ, u, ε)

horseshoe_velocity(r, line :: Line, Γ, direction, ε) = total_horseshoe_velocity(r1(r, line), r2(r, line), Γ, direction, ε)


module How
    export howmuch

    function howmuch(m::Int, n::Int)::Array{Tuple{String,String,String}}
        min_value, max_value = minmax(m,n)
        # Only values "37+63idx" are possible
        idx_to_value(idx) = 37 + 63idx
        value_to_idx(v) = (v - 37) / 63
        min_idx = ceil(Int,value_to_idx(min_value))
        max_idx = floor(Int,value_to_idx(max_value))
        result = []
        for idx = min_idx:max_idx
            money = idx_to_value(idx)
            boat_cost = money ÷ 7
            car_cost = money ÷ 9
            push!(result, ("M: $money","B: $boat_cost","C: $car_cost") )
        end
        result
    end

end


module ViscFoil

include("VorticityStream.jl")
# function integral_momentum_equation(dx, x, p, s)
#     θ, dθ = x                       # State vector
#     ρ_e, u_e, a_e, τ_w, δ_st = p    # Parameters
#     H = δ_st/θ                      # Shape parameter
#     c_f = τ_w/(1/2 * ρ_e * u_e^2)   # Skin-friction coefficient
#     M_e, H = u_e/a_e                # Mach number
#     dx[1] = dθ                      # LHS: dθ/ds
#     dx[2] = c_f / 2  - (H + 2 - M_e^2) * θ/u_e * ??? # RHS: c_f/2  - (H + 2 - Mₑ²) * θ/U_e * dU_e/ds
# end

function edge_velocity(σs, panels, wake_panels, u)
    all_panels          = [panels; wake_panels]
    foil_vortices       = vortex_matrix(panels, panels)
    all_sources         = source_matrix(panels, all_panels)
    wake_wake_vortices  = vortex_matrix(wake_panels, wake_panels)
    wake_foil_vortices  = vortex_matrix(wake_panels, panels)
    wake_boundary       = boundary_vector(colpoints.(wake_panels), u)

    P       = -foil_vortices^(-1) * all_sources
    P_s     =
    # D1  = visc_visc_vortices # PLACEHOLDER
    D1      = wake_foil_vortices * P + wake_wake_vortices
    U_inv   = [ (zeros ∘ length)(panels) ;
                    wake_boundary  ]
    D       = [ P
                D1 ]

    U_e     = U_inv + D * σs
end

end


const branches = 4
const nmax = 500

const rooted = zeros(BigInt, nmax + 1)
const unrooted = zeros(BigInt, nmax + 1)
rooted[1] = rooted[2] = unrooted[1] = unrooted[2] = 1
const c = zeros(BigInt, branches)

function tree(br, n, l, sum, cnt)
    for b in br+1:branches
        sum += n
        if (sum > nmax) || (l * 2 >= sum && b >= branches)
            return
        elseif b == br + 1
            c[br + 1] = rooted[n + 1] * cnt
        else
            c[br + 1] *= rooted[n + 1] + b - br - 1
            c[br + 1] = div(c[br + 1], b - br)
        end
        if l*2 < sum
            unrooted[sum + 1] += c[br + 1]
        end
        if b < branches
            rooted[sum + 1] += c[br + 1]
        end
        for m in n-1:-1:1
            tree(b, m, l, sum, c[br + 1])
        end
    end
end

bicenter(n) = if iseven(n) unrooted[n + 1] += div(rooted[div(n, 2) + 1] * (rooted[div(n, 2) + 1] + 1), 2) end

function paraffins()
    for n in 1:nmax
        tree(0, n, n, 1, one(BigInt))
        bicenter(n)
        println("$n: $(unrooted[n + 1])")
    end
end

paraffins()


include("../benchmark.jl")

using JuMP #, Ipopt

m = Model(solver=OnePhaseSolver(init!mu_scale=0.01))
@variable(m, x >= 0)
@variable(m, y >= 0)
@objective(m, Min, x-y)
@NLconstraint(m, x^2 + y^2 <= 1)

solve(m)

getobjectivevalue(m)


@testset "Cartesian" begin
  x = [π, exp(1), √2]
  for ct ∈ (ToCartesian(x->x, 3), FromCartesian(x->x, 3))
    @test J(ct, x) == 1
    m = zeros(3,3) + I
    gᵢⱼ(ct, x) ≈ m
    gⁱʲ(ct, x) ≈ m
    dl(ct, x, Covariant) ≈ ones(3)
    dl(ct, x, Contravariant) ≈ ones(3)

    v = CovariantVector(x->[x[1], x[2], x[3]], ct, true)
    @test div(v)(x) ≈ 3
    @test curl(v)(x) ≈ zeros(3)
    v = CovariantVector(x->[x[3], x[1], x[2]], ct, true)
    @test div(v)(x) ≈ 0
    @test curl(v)(x) ≈ ones(3)
    @test div(curl(v))(x) ≈ 0
    @test div(curl(v))(x) ≈ 0
    v = CovariantVector(x->[-x[2], -x[3], -x[1]], ct, true)
    @test div(v)(x) ≈ 0
    @test curl(v)(x) ≈ ones(3)
    @test div(curl(v))(x) ≈ 0
    @test div(curl(v))(x) ≈ 0
  end
end


import Base: muladd, fma

"""
    two_sum(a, b)
Computes `hi = fl(a+b)` and `lo = err(a+b)`.
"""
@inline function two_sum(a::T, b::T) where {T<:AbstractFloat}
    hi = a + b
    a1 = hi - b
    b1 = hi - a1
    lo = (a - a1) + (b - b1)
    return hi, lo
end

"""
    two_hilo_sum(a, b)
*unchecked* requirement `|a| ≥ |b|`
Computes `s = fl(a+b)` and `e = err(a+b)`.
"""
@inline function two_hilo_sum(a::T, b::T) where {T<:AbstractFloat}
    s = a + b
    e = b - (s - a)
    return s, e
end

@inline function two_prod(a::T, b::T) where {T<:AbstractFloat}
    s = a * b
    t = fma(a, b, -s)
    return s, t
end

function fma(xhi::T, xlo::T, yhi::T, ylo::T, zhi::T, zlo::T) where {T<:AbstractFloat}
   chi, c1 = two_prod(xhi, yhi)
   t0 = xlo * ylo
   t1 = fma(xhi, ylo, t0)
   c2 = fma(xlo, yhi, t1)
   c3 = c1 + c2
   dhi, dlo = two_hilo_sum(chi, c3)

   shi, slo = two_sum(zhi, dhi)
   thi, tlo = two_sum(zlo, dlo)
   c = slo + thi
   vhi, vlo = two_hilo_sum(shi, c)
   w = tlo + vlo
   hi, lo = two_hilo_sum(vhi, w)
   return DoubleFloat{T}(hi, lo)
end

@inline function fma(x::DoubleFloat{T}, y::DoubleFloat{T}, z::DoubleFloat{T}) where {T<:AbstractFloat}
   return fma(x.hi, x.lo, y.hi, y.lo, z.hi, z.lo)
end

@inline function fma(x::DoubleFloat{T}, y::DoubleFloat{T}, z::Tuple{T,T}) where {T<:AbstractFloat}
   return fma(x.hi, x.lo, y.hi, y.lo, z[1], z[2])
end
@inline function fma(x::DoubleFloat{T}, y::Tuple{T,T}, z::DoubleFloat{T}) where {T<:AbstractFloat}
   return fma(x.hi, x.lo, y[1], y[2], z.hi, z.lo)
end
@inline function fma(x::Tuple{T,T}, y::DoubleFloat{T}, z::DoubleFloat{T}) where {T<:AbstractFloat}
   return fma(x[1], x[2], y.hi, y.lo, z.hi, z.lo)
end
@inline function fma(x::DoubleFloat{T}, y::Tuple{T,T}, z::Tuple{T,T}) where {T<:AbstractFloat}
   return fma(x.hi, x.lo, y[1], y[2], z[1], z[2])
end
@inline function fma(x::Tuple{T,T}, y::Tuple{T,T}, z::DoubleFloat{T}) where {T<:AbstractFloat}
   return fma(x[1], x[2], y[1], y[2], z.hi, z.lo)
end
@inline function fma(x::Tuple{T,T}, y::DoubleFloat{T}, z::Tuple{T,T}) where {T<:AbstractFloat}
   return fma(x[1], x[2], y.hi, y.lo, z[1], z[2])
end
@inline function fma(x::Tuple{T,T}, y::Tuple{T,T}, z::Tuple{T,T}) where {T<:AbstractFloat}
   return fma(x[1], x[2], y[1], y[2], z[1], z[2])
end

function muladd(xhi::T, xlo::T, yhi::T, ylo::T, zhi::T, zlo::T) where {T<:AbstractFloat}
   chi, c1 = two_prod(xhi, yhi)
   t0 = xlo * ylo
   c2 = fma(xlo, yhi, t0)
   c3 = c1 + c2
   dhi, dlo = two_hilo_sum(chi, c3)

   shi, slo = two_sum(zhi, dhi)
   thi, tlo = two_sum(zlo, dlo)
   c = slo + thi
   vhi, vlo = two_hilo_sum(shi, c)
   w = tlo + vlo
   hi, lo = two_hilo_sum(vhi, w)
   return DoubleFloat{T}(hi, lo)
end

@inline function muladd(x::DoubleFloat{T}, y::DoubleFloat{T}, z::DoubleFloat{T}) where {T<:AbstractFloat}
   return muladd(x.hi, x.lo, y.hi, y.lo, z.hi, z.lo)
end

@inline function muladd(x::DoubleFloat{T}, y::DoubleFloat{T}, z::Tuple{T,T}) where {T<:AbstractFloat}
   return muladd(x.hi, x.lo, y.hi, y.lo, z[1], z[2])
end
@inline function muladd(x::DoubleFloat{T}, y::Tuple{T,T}, z::DoubleFloat{T}) where {T<:AbstractFloat}
   return muladd(x.hi, x.lo, y[1], y[2], z.hi, z.lo)
end
@inline function muladd(x::Tuple{T,T}, y::DoubleFloat{T}, z::DoubleFloat{T}) where {T<:AbstractFloat}
   return muladd(x[1], x[2], y.hi, y.lo, z.hi, z.lo)
end
@inline function muladd(x::DoubleFloat{T}, y::Tuple{T,T}, z::Tuple{T,T}) where {T<:AbstractFloat}
   return muladd(x.hi, x.lo, y[1], y[2], z[1], z[2])
end
@inline function muladd(x::Tuple{T,T}, y::Tuple{T,T}, z::DoubleFloat{T}) where {T<:AbstractFloat}
   return muladd(x[1], x[2], y[1], y[2], z.hi, z.lo)
end
@inline function muladd(x::Tuple{T,T}, y::DoubleFloat{T}, z::Tuple{T,T}) where {T<:AbstractFloat}
   return muladd(x[1], x[2], y.hi, y.lo, z[1], z[2])
end
@inline function muladd(x::Tuple{T,T}, y::Tuple{T,T}, z::Tuple{T,T}) where {T<:AbstractFloat}
   return muladd(x[1], x[2], y[1], y[2], z[1], z[2])
end


# v0.6.0

rad = π / 4
deg = 45.0

@show rad deg
@show sin(rad) sin(deg2rad(deg))
@show cos(rad) cos(deg2rad(deg))
@show tan(rad) tan(deg2rad(deg))
@show asin(sin(rad)) asin(sin(rad)) |> rad2deg
@show acos(cos(rad)) acos(cos(rad)) |> rad2deg
@show atan(tan(rad)) atan(tan(rad)) |> rad2deg


"""
```
carter_kohn_smoother(y, T, R, C, Q, Z, D, E, s_0, P_0;
    Nt0 = 0, draw_states = true)

carter_kohn_smoother(regime_indices, y, Ts, Rs, Cs, Qs,
    Zs, Ds, Es, s_0, P_0; Nt0 = 0, draw_states = true)
```
This program is a simulation smoother based on Carter and Kohn's \"On Gibbs
Sampling for State Space Models\" (Biometrika, 1994). It recursively samples
from the conditional distribution of time-t states given the full set of
observables and states from time t+1 to time T. Unlike the Durbin-Koopman
simulation smoother, this one does rely on inverting potentially singular
matrices using the Moore-Penrose pseudoinverse.

The state space is augmented with shocks (see `?augment_states_with_shocks`),
and the augmented state space is Kalman filtered and smoothed. Finally, the
smoothed states and shocks are indexed out of the augmented state vectors.

The original state space (before augmenting with shocks) is given by:

```
s_{t+1} = C + T*s_t + R*ϵ_t    (transition equation)
y_t     = D + Z*s_t + u_t      (measurement equation)

ϵ_t ∼ N(0, Q)
u_t ∼ N(0, E)
Cov(ϵ_t, u_t) = 0
```

### Inputs

- `y`: `Ny` x `Nt` matrix containing data `y_1, ... , y_T`
- `s_0`: `Ns` x 1 initial state vector
- `P_0`: `Ns` x `Ns` initial state covariance matrix

**Method 1 only:** state-space system matrices `T`, `R`, `C`, `Q`, `Z`,
`D`, `E`. See `?kalman_filter`
pp
**Method 2 only:** `regime_indices` and system matrices for each regime `Ts`,
`Rs`, `Cs`, `Qs`, `Zs`, `Ds`, `Es`. See `?kalman_filter`

where:

- `Nt`: number of periods for which we have data
- `Ns`: number of states
- `Ne`: number of shocks
- `Ny`: number of observables

### Keyword Arguments

- `Nt0`: if greater than 0, the returned smoothed states and shocks matrices
  will be shorter by that number of columns (taken from the beginning)
- `draw_states`: indicates whether to draw smoothed states from the distribution
  `N(s_{t|T}, P_{t|T})` or to use the mean `s_{t|T}` (reducing to the Hamilton
  smoother)

### Outputs

- `s_smth`: `Ns` x `Nt` matrix of smoothed states `s_{t|T}`
- `ϵ_smth`: `Ne` x `Nt` matrix of smoothed shocks `ϵ_{t|T}`
"""
function carter_kohn_smoother(y::AbstractMatrix,
    T::Matrix{S}, R::Matrix{S}, C::Vector{S},
    Q::Matrix{S}, Z::Matrix{S}, D::Vector{S}, E::Matrix{S},
    s_0::Vector{S}, P_0::Matrix{S};
    Nt0::Int = 0, draw_states::Bool = true) where {S<:AbstractFloat}

    Nt = size(y, 2)
    carter_kohn_smoother(UnitRange{Int}[1:Nt], y, Matrix{S}[T], Matrix{S}[R], Vector{S}[C],
        Matrix{S}[Q], Matrix{S}[Z], Vector{S}[D], Matrix{S}[E], s_0, P_0;
        Nt0 = Nt0, draw_states = draw_states)
end

function carter_kohn_smoother(regime_indices::Vector{UnitRange{Int}}, y::AbstractMatrix,
    Ts::Vector{Matrix{S}}, Rs::Vector{Matrix{S}}, Cs::Vector{Vector{S}}, Qs::Vector{Matrix{S}},
    Zs::Vector{Matrix{S}}, Ds::Vector{Vector{S}}, Es::Vector{Matrix{S}},
    s_0::Vector{S}, P_0::Matrix{S};
    Nt0::Int = 0, draw_states::Bool = true) where {S<:AbstractFloat}

    # Dimensions
    Nt = size(y,     2) # number of periods of data
    Ns = size(Ts[1], 1) # number of states
    Ne = size(Rs[1], 2) # number of shocks

    # Augment state space with shocks
    Ts, Rs, Cs, Zs, s_0, P_0 =
        augment_states_with_shocks(regime_indices, Ts, Rs, Cs, Qs, Zs, s_0, P_0)

    # Kalman filter stacked states and shocks stil_t
    _, stil_pred, Ptil_pred, stil_filt, Ptil_filt, _, _, _, _ =
        kalman_filter(regime_indices, y, Ts, Rs, Cs, Qs, Zs, Ds, Es, s_0, P_0,
                      outputs = [:pred, :filt])

    # Smooth the stacked states and shocks recursively, starting at t = T-1 and
    # going backwards
    stil_smth = copy(stil_filt)

    for i = length(regime_indices):-1:1
        # Get state-space system matrices for this regime
        T = Ts[i]

        for t in reverse(regime_indices[i])
            if t == Nt
                μ = stil_filt[:, end]
                Σ = Ptil_filt[:, :, end]
            else
                J = Ptil_filt[:, :, t] * T' * pinv(Ptil_pred[:, :, t+1])
                μ = stil_filt[:, t] + J*(stil_smth[:, t+1] - stil_pred[:, t+1]) # stil_{t|T}
                Σ = Ptil_filt[:, :, t] - J*T*Ptil_filt[:, :, t]                 # Ptil_{t|T}
            end

            # Draw stil_t ∼ N(stil_{t|T}, Ptil_{t|T})
            stil_smth[:, t] = if draw_states
                U, eig, _ = svd(Σ)
                μ + U * diagm(0 => (sqrt.(eig))) * randn(Ns+Ne)
            else
                μ
            end
        end
    end

    # Index out states and shocks
    s_smth = stil_smth[1:Ns,     :] # s_{t|T}
    ϵ_smth = stil_smth[Ns+1:end, :] # ϵ_{t|T}

    # Trim the presample if needed
    if Nt0 > 0
        insample = Nt0+1:Nt
        s_smth = s_smth[:, insample]
        ϵ_smth = ϵ_smth[:, insample]
    end

    return s_smth, ϵ_smth
end


using DifferentialEquations, Plots

function fhn!(du, u, p, t)
	du[1] = (u[1]-u[1]*u[1])*(u[1]-p[2]) - u[2]
	du[2] = p[3]*(p[1]*u[1] - u[2])
	return nothing
end

function lor!(du, u, p, t)
	du[1] = p[1]*(u[2]-u[1])
	du[2] = u[1]*(p[2]-u[3]) - u[2]
	du[3] = u[1]*u[2] - p[3]*u[3]
	du .= du .* 0.005
	return nothing
end

function f!(du, u, p, t)
	fhn!(view(du,1:2), view(u,1:2), view(p,1:3), t)
	lor!(view(du,3:5), view(u,3:5), view(p,4:6), t)
	#du[1] = du[1] + 0.05*(u[5]-27.0)
	return nothing
end

u0 = randn(Float64,2)
u1 = randn(Float64,3)

p0 = [0.37; 0.1; 0.01]
p1 = [10.0;28.0; 8/3]

tspan = (0.0, 5000.0)

prob0 = ODEProblem(fhn!, u0, tspan, p0)
prob1 = ODEProblem(lor!, u1, tspan, p1)
prob  = ODEProblem(f!, [u0;u1], tspan, [p0;p1])

sol0  = solve(prob0, Tsit5())
sol1  = solve(prob1, Tsit5())
sol   = solve(prob,  Tsit5())

plot(sol0)
plot!(sol1)
plot!(sol, linestyle=:dash, linewidth=2)

function condition(u,t,integrator)
	u[3]
end
function affect!(integrator)
  integrator.u[1] = integrator.u[1] + rand(Float64)
end
cb = ContinuousCallback(condition,affect!)
prob = remake(prob, tspan=(0.0,10000.0))
sol   = solve(prob, Tsit5(); callback=cb)
plot(sol, vars=(0,1))


const C = CartesianIndex{2}

function day25()
    part = [0, 0]

    lines = filter(s -> !isempty(strip(s)), strip.(readlines("AoCdata/AoC_2021_day25.txt")))
    nrows, ncols = length(lines), length(first(lines))
    mat = fill('.', nrows, ncols)
    for i in 1:nrows
        mat[i, :] .= collect(lines[i])
    end
    tmat, moved, steps = deepcopy(mat), 1, 0
    while moved > 0
        moved = 0
        for c in CartesianIndices(mat)
            if mat[c] == '>'
                dest = C(c[1], mod1(c[2] + 1, ncols))
                if mat[dest] == '.'
                    tmat[c] = '.'
                    tmat[dest] = '>'
                    moved += 1
                end
            end
        end
        mat = deepcopy(tmat)
        for c in CartesianIndices(mat)
            if mat[c] == 'v'
                dest = C(mod1(c[1] + 1, nrows), c[2])
                if mat[dest] == '.'
                    tmat[c] = '.'
                    tmat[dest] = 'v'
                    moved += 1
                end
            end
        end
        mat = deepcopy(tmat)
        steps += 1
    end

    part[1] = steps

    return part
end

part = day25()
println("Part 1: ", part[1])


"""
    za = agg(z,nas,nat)

Aggregate field `z` to an array `za` of size `(nas,nas,nat)`
"""
function agg(zi,nas,nat)
        nss=size(zi);
        if(length(nss)>=3)
           ns=nss[1]; nt=nss[3];
        else
           ns=nss[1]; nt=1;
        end
        sdim=div(ns,nas);
        tdim=div(nt,nat);
        ir=1:sdim; jr=ir; kr=1:tdim;
        rs=0:sdim:(ns-sdim);
        rt=0:tdim:(nt-tdim);
	z=zeros(Float64,nas,nas,nat);
        n=zeros(Int64,nas,nas,nat);
        for i=1:sdim
           for j=1:sdim
              for k=1:tdim
		 n=n+.~isnan.(zi[i.+rs,j.+rs,k.+rt]);
		 z=z+zi[i.+rs,j.+rs,k.+rt].*.~isnan.(zi[i.+rs,j.+rs,k.+rt]) ;
              end
           end
        end
        z=z./n;
end


"""
    pred_vgp(x_pred, gp_params)
Obtain predictions from variational Gaussian process at locations x_pred.
Returns tuple with predicted means and predicted standard deviations
"""
function pred_vgp(x_pred, gp_params::SVGP_params)
    opt_xi    = gp_params.inducing_locs
    opt_m     = gp_params.inducing_mean
    opt_cov   = Hermitian(gp_params.inducing_L * transpose(gp_params.inducing_L))

    ni, nd    = size(opt_xi)

    rho        = exp.(gp_params.log_rho)
    kappa      = exp(gp_params.log_kappa[1])
    sigsq      = exp(gp_params.log_sigma[1])^2

    nugg_scale = 1.e-6 * kappa
    if nugg_scale < 1.e-8
        nugg_scale = 1.e-8
    end

    cov_pred   = covfn(x_pred, x_pred, gp_params) + I * nugg_scale
    cov_mat_i  = covfn(gp_params.inducing_locs, gp_params.inducing_locs, gp_params) + I * nugg_scale
    cross_mat  = covfn(gp_params.inducing_locs, x_pred, gp_params)
    cross_inv  = cov_mat_i \ cross_mat

    p_mean = cross_inv' * (opt_m .- gp_params.const_mean[1]) .+ gp_params.const_mean[1]
    p_cov  = diag( cov_pred + cross_inv' * ( opt_cov - cov_mat_i) * cross_inv);

    return p_mean, sqrt.(p_cov)
end

"""
    pred_vgp(x_pred, gp_obj)
Obtain predictions from variational Gaussian process at locations x_pred.
Returns tuple with predicted means and predicted standard deviations.
Included for backward compatibility, but only part of gp_obj used is the params
"""

function pred_vgp(x_pred, gp_obj::SVGP_obj)
    gp_params = gp_obj.params
    opt_xi    = gp_params.inducing_locs
    opt_m     = gp_params.inducing_mean
    opt_cov   = Hermitian(gp_params.inducing_L * transpose(gp_params.inducing_L))

    ni, nd    = size(opt_xi)

    rho        = exp.(gp_params.log_rho)
    kappa      = exp(gp_params.log_kappa[1])
    sigsq      = exp(gp_params.log_sigma[1])^2

    nugg_scale = 1.e-6 * kappa
    if nugg_scale < 1.e-8
        nugg_scale = 1.e-8
    end

    cov_pred   = covfn(x_pred, x_pred, gp_params) + I * nugg_scale
    cov_mat_i  = covfn(gp_params.inducing_locs, gp_params.inducing_locs, gp_params) + I * nugg_scale
    cross_mat  = covfn(gp_params.inducing_locs, x_pred, gp_params)
    cross_inv  = cov_mat_i \ cross_mat

    p_mean = cross_inv' * (opt_m .- gp_params.const_mean[1]) .+ gp_params.const_mean[1]
    p_cov  = diag( cov_pred + cross_inv' * ( opt_cov - cov_mat_i) * cross_inv);

    return p_mean, sqrt.(p_cov)
end


using LightGraphs
using MatrixDepot
function symmetrize(A)
    println("Symmetrizing ")
    tic()
    if !issymmetric(A)
        println(STDERR, "the matrix is not symmetric using A+A'")
        A = A + A'
    end
    if isa(A, Base.LinAlg.Symmetric)
        A.data.nzval = abs(A.data.nzval)
    else
        A.nzval = abs(A.nzval)
    end
    #= spA = abs(sparse(A)) =#
    toc()
    return A
end

function loadmat(matname)
    println("Reading MTX for $matname")
    tic()
    try
        info = matrixdepot(matname, :get)
    end
    A = matrixdepot(matname, :read)
    A = symmetrize(A)
    g = Graph(A)
    toc()
    println(STDERR, "size(A) = $(size(A))")
    return g
end

function vec2tree(v::Vector{Int})
    nv = length(v)
    I = Vector{Int}()
    J = Vector{Int}()
    for i in 1:nv
        if v[i] > 0
            push!(I, i)
            push!(J, v[i])
        end
    end
    #= @assert length(I) == length(J) =#
    ncolumns = max(maximum(J), maximum(I))
    nentries = length(J)
    n = ncolumns
    return sparse(I,J, fill(1, nentries), n,n)
end

names = ["Newman/football" , "Newman/cond-mat-2003", "SNAP/amazon0302", "SNAP/roadNet-CA"]
for matname in names
    g = loadmat(matname)
    seed = 1
    println("bfs_tree original")
    # woah bfs_tree allocates a lot of memory
    @time tdg = LightGraphs.bfs_tree(g, seed)
    println(tdg)
    # preallocating the output tree reduces the number of allocations.
    # much faster
    println("bfs_tree!(vector)")
    visitor = LightGraphs.TreeBFSVisitorVector(zeros(Int, nv(g)))
    @time tvec = LightGraphs.bfs_tree!(visitor, g, seed)
    println("converting to Sparse")
    @time m = vec2tree(visitor.tree)
    println("converting to DiGraph")
    @time h = DiGraph(m)
end


using GZip

"Where to download mnist from"
mnisturl = "http://yann.lecun.com/exdb/mnist"

"Where to download mnist to"
mnistdir = "../data/"

"""

This utility loads the [MNIST](http://yann.lecun.com/exdb/mnist)
hand-written digits dataset.  There are 60000 training and 10000 test
examples. Each input x consists of 784 pixels representing a 28x28
grayscale image.  The pixel values are converted to Float32 and
normalized to [0,1].  Each output y is a UInt8 indicating the correct
class.  10 is used to represent the digit 0.

```
# Usage:
include(Pkg.dir("Knet/data/mnist.jl"))
xtrn, ytrn, xtst, ytst = mnist()
# xtrn: 28×28×1×60000 Array{Float32,4}
# ytrn: 60000-element Array{UInt8,1}
# xtst: 28×28×1×10000 Array{Float32,4}
# ytst: 10000-element Array{UInt8,1}
```

"""
function mnist()
    global _mnist_xtrn,_mnist_ytrn,_mnist_xtst,_mnist_ytst
    if !@isdefined _mnist_xtrn
        @info "Loading MNIST..."
        _mnist_xtrn = _mnist_xdata("train-images-idx3-ubyte.gz")
        _mnist_xtst = _mnist_xdata("t10k-images-idx3-ubyte.gz")
        _mnist_ytrn = _mnist_ydata("train-labels-idx1-ubyte.gz")
        _mnist_ytst = _mnist_ydata("t10k-labels-idx1-ubyte.gz")
    end
    return _mnist_xtrn,_mnist_ytrn,_mnist_xtst,_mnist_ytst
end

"Utility to view a MNIST image, requires the Images package"
mnistview(x,i)=colorview(Gray,permutedims(x[:,:,1,i],(2,1)))

function _mnist_xdata(file)
    a = _mnist_gzload(file)[17:end]
    reshape(a ./ 255f0, (28,28,1,div(length(a),784)))
end

function _mnist_ydata(file)
    a = _mnist_gzload(file)[9:end]
    a[a.==0] .= 10
    # full(sparse(a,1:length(a),1f0,10,length(a)))
    return a
end

function _mnist_gzload(file)
    if !isdir(mnistdir)
        mkpath(mnistdir)
    end
    path = joinpath(mnistdir,file)
    if !isfile(path)
        url = "$mnisturl/$file"
        download(url, path)
    end
    f = gzopen(path)
    a = read(f)
    close(f)
    return(a)
end

nothing


# test g-distribution
@testset "g-dist" begin
    a = 3.5
    samples2 = [KissMCMC.sample_g(a) for i=1:50000];

    @test all(1/a .<= samples2 .<= a)
    @test KissMCMC.cdf_g_inv(1, a) ≈ a
    @test KissMCMC.cdf_g_inv(0, a) ≈ 1/a
    z = 1/a:0.01:a
    meang = sum(z.*KissMCMC.g_pdf.(z, a)) *step(z)
    @test isapprox(mean(samples2), meang, atol=1e-2)
    stdg = sqrt(sum((meang.-z).^2 .* KissMCMC.g_pdf.(z, a)) *step(z))
    @test isapprox(std(samples2), stdg, atol=1e-2)
end

# test MCMC
@testset "emcee" begin
    for tc in testcases
        pdf = make_lpdf(tc.dist)
        @inferred pdf(tc.theta0)
        theta0s = make_theta0s(tc.theta0, tc.ball_radius,
                               pdf, tc.nwalkers,
                               hasblob=tc.hasblob)
        samples = emcee(pdf, theta0s;
                        niter=tc.niter,
                        hasblob=tc.hasblob,
                        use_progress_meter=false,
                        tc.otherkws...);
        @test length.(samples[1:3]) == (tc.nwalkers, tc.nwalkers, tc.nwalkers)
        if tc.hasblob
            @test length(samples[4]) == tc.nwalkers
        else
            @test samples[4]==nothing
        end
        @test length(samples[1][1]) == tc.niter÷tc.nwalkers÷2 # ÷2 because of burnin
        thetas, accept_ratio, logdensities, blobs = squash_walkers(samples...;
                                                                   verbose=false,
                                                                   tc.squash_walkers_kws...)

        !tc.hasblob && @test blobs==nothing
        @test length(thetas) == tc.niter÷2
        @test length(logdensities) == tc.niter÷2
        @test accept_ratio>0.1
        test_mean_std(thetas, tc, tc.tole)
        test_blobs(tc.blob_truths, blobs)
        #test_logdensities(tc.truths, logdensities, pdf)
    end
end


if visualtests
  @testset "Plot recipes" begin
    Random.seed!(2018)
    data₁ = rand(LogNormal(2, 2), 5000)
    data₂ = rand(Pareto(.9, 10), 5000)

    Random.seed!(2018)
    xs = rand(LogNormal(0, 1), 5000)

    @testset "Return levels" begin
      @test_reference "data/ReturnLevels.png" returnplot(xs,label="log-normal")
    end

    @testset "Mean excess" begin
      plt1 = excessplot(data₁, label="Lognormal")
      plt2 = excessplot(data₂, label="Pareto")
      plt = plot(plt1, plt2, size=(600,800), layout=(2,1))
      @test_reference "data/MeanExcess.png" plt
    end

    @testset "Pareto quantile" begin
      plt1 = paretoplot(data₁, label="Lognormal")
      plt2 = paretoplot(data₂, label="Pareto")
      plt = plot(plt1, plt2, size=(600,800), layout=(2,1))
      @test_reference "data/ParetoQuantile.png" plt
    end
  end
end


struct BudgetedElementaryPathInstance{T} <: CombinatorialInstance
  graph::AbstractGraph{T}
  rewards::Dict{Edge{T}, Float64}
  weights::Dict{Edge{T}, Int}
  src::T
  dst::T

  budget::Int # weights * solution >= budget

  function BudgetedElementaryPathInstance(graph::AbstractGraph{T}, rewards::Dict{Edge{T}, Float64},
                                          weights::Dict{Edge{T}, Int}, src::T, dst::T;
                                          budget::Union{Nothing, Int}=nothing,
                                          max_weight::Union{Nothing, Int}=nothing) where T
    # Error checking.
    if src == dst
      error("Source node is the same as destination node.")
    end

    if ! isnothing(budget) && budget < 0
      error("Budget is present and negative.")
    end

    if any(collect(values(weights)) .< 0)
      error("At least a weight is negative.")
    end

    if length(rewards) != length(weights)
      error("Not the same number of values and weights; these two vectors must have the same size.")
    end

    # Complete the optional parameters.
    if isnothing(max_weight) # Not useful to store, it is just used to compute the max budget if it is not given.
      max_weight = maximum(collect(values(weights)))
    end

    if isnothing(budget)
      d = length(values) # Dimension of the problem.
      budget = d * max_weight
    end

    # Return a new instance.
    new{T}(graph, rewards, weights, src, dst, budget)
  end
end

graph(i::BudgetedElementaryPathInstance{T}) where T = i.graph
rewards(i::BudgetedElementaryPathInstance{T}) where T = i.rewards
weights(i::BudgetedElementaryPathInstance{T}) where T = i.weights
src(i::BudgetedElementaryPathInstance{T}) where T = i.src
dst(i::BudgetedElementaryPathInstance{T}) where T = i.dst
budget(i::BudgetedElementaryPathInstance{T}) where T = i.budget

# dimension(i::BudgetedElementaryPathInstance{T}) where T = ne(graph(i))
reward(i::BudgetedElementaryPathInstance{T}, u::T, v::T) where T = rewards(i)[Edge(u, v)]
weight(i::BudgetedElementaryPathInstance{T}, u::T, v::T) where T = weights(i)[Edge(u, v)]

struct BudgetedElementaryPathSolution{T} <: CombinatorialSolution
  instance::BudgetedElementaryPathInstance{T}
  path::Vector{Edge{T}}
  states::Dict{Tuple{T, Int}, Float64}
  solutions::Dict{Tuple{T, Int}, Vector{Edge{T}}}
end

function paths_all_budgets(s::BudgetedElementaryPathSolution{T}, max_budget::Int) where T
  if max_budget > budget(s.instance)
    @warn "The asked maximum budget $max_budget is higher than the instance budget $(budget(s.instance)). Therefore, some values have not been computed and are unavailable."
  end

  mb = min(max_budget, budget(s.instance))
  return Dict{Int, Vector{Edge{T}}}(
    budget => s.solutions[s.instance.dst, budget] for budget in 0:mb)
end

function paths_all_budgets_as_tuples(s::BudgetedElementaryPathSolution{T}, max_budget::Int) where T
  if max_budget > budget(s.instance)
    @warn "The asked maximum budget $max_budget is higher than the instance budget $(budget(s.instance)). Therefore, some values have not been computed and are unavailable."
  end

  mb = min(max_budget, budget(s.instance))
  return Dict{Int, Vector{Tuple{T, T}}}(
    budget => [(src(e), dst(e)) for e in s.solutions[s.instance.dst, budget]]
    for budget in 0:mb)
end


"""
$(SIGNATURES)

Calculate the Direction Of Arrival (DOA) with respect to the manifold
`manifold`. The calculated DOA is returned as `Spherical` object, where the
azimuth is counted counter-clockwise from the x axis and the elevation is
counted from the xy-plane towards the zenith.
If the initial azimuth `init_az` and elevation `init_el` are provided, this
will use the hill climbing algorithm.
The hill climbing algorithm should be a lot faster than plain maximum search,
but may only find a local maximum, which might be desired.

# Examples
Using the signal subspace:
```julia-repl
julia> est_doa(manifold, a -> norm(signal_subspace' * a) / norm(a))
```

Using the noise subspace (MUSIC):
```julia-repl
julia> est_doa(manifold, a -> norm(a) / norm(noise_subspace' * a))
```

# Arguments:
- `manifold::AbstractManifold`: Manifold of the antenna array
- `reduction_function`: Function that reduces the manifold vector to a scalar
"""
function est_doa(
    manifold::AbstractManifold,
    reduction_function;
    init_az = nothing,
    init_el = nothing,
    kwargs...
)
    p = Pattern(manifold, reduction_function; kwargs...)
    if isnothing(init_az) || isnothing(init_el)
        _, idx = findmax(p.values)
    else
        init_az_idx = floor(Int, mod2pi(init_az) / Float64(p.azs.step) + 1)
        init_el_idx = floor(Int, (π / 2 - init_el) / Float64(p.els.step) + 1)
        _, idx = hill_climbing(p.values, init_el_idx, init_az_idx)
    end
    az = p.azs[idx.I[2]]
    el = p.els[idx.I[1]]
    return Spherical(1, az, π / 2 - el)
end

"""
$(SIGNATURES)

This is a convenience function of `est_doa`.
See `est_doa` for more details.
"""
function est_doa_by_noise_subspace(
    manifold::AbstractManifold{N},
    noise_subspace;
    init_az = nothing,
    init_el = nothing,
    kwargs...
) where N
    size(noise_subspace, 1) == N || ArgumentError("Number of antenna channels must match")
    est_doa(
        manifold,
        a -> norm(a) / norm(noise_subspace' * a);
        init_az = init_az,
        init_el = init_el,
        kwargs...
    )
end

"""
$(SIGNATURES)

This is a convenience function of `est_doa`.
See `est_doa` for more details.
"""
function est_doa_by_signal_subspace(
    manifold::AbstractManifold{N},
    signal_subspace;
    init_az = nothing,
    init_el = nothing,
    kwargs...
) where N
    size(signal_subspace, 1) == N || ArgumentError("Number of antenna channels must match")
    est_doa(
        manifold,
        a -> norm(signal_subspace' * a) / norm(a);
        init_az = init_az,
        init_el = init_el,
        kwargs...
    )
end

"""
$(SIGNATURES)

This is a convenience function of `est_doa`.
See `est_doa` for more details.
"""
est_doa_by_music(m, ns; init_az = nothing, init_el = nothing, k...) =
    est_doa_by_noise_subspace(m, ns, init_az = init_az, init_el = init_el, k...)

"""
$(SIGNATURES)

Multi-dimensionial hill climbing.

# Arguments:
- `values::AbstractArray`: Some multi-dimensional array
- `inits`: Starting indices
"""
function hill_climbing(values, inits...)
    ndims(values) == length(inits) || ArgumentError("Number of inits does not fit dimensions of values")
    indices = CartesianIndices(values)
    lower_index_bounds, upper_index_bounds = first(indices), last(indices)
    step_size = oneunit(lower_index_bounds)
    current_index = CartesianIndex(inits...)
    current_value = values[current_index]
    while true
        next_index = current_index
        for J in max(lower_index_bounds, current_index - step_size):min(upper_index_bounds, current_index + step_size)
            if values[J] > current_value
                next_index = J
                current_value = values[J]
            end
        end
        if current_index == next_index
            return current_value, current_index
        else
            current_index = next_index
        end
    end
end


using DelayEmbeddings, SparseArrays
include("GroupSlices.jl")

export
    TransferOperator, # the probabilities estimator
    InvariantMeasure, invariantmeasure,
    transfermatrix

"""
    TransferOperator(ϵ::RectangularBinning) <: BinningProbabilitiesEstimator

A probability estimator based on binning data into rectangular boxes dictated by
the binning scheme `ϵ`, then approxmating the transfer (Perron-Frobenius) operator
over the bins, then taking the invariant measure associated with that transfer operator
as the bin probabilities. Assumes that the input data are sequential (time-ordered).

This implementation follows the grid estimator approach in Diego et al. (2019)[^Diego2019].

## Description

The transfer operator ``P^{N}``is computed as an `N`-by-`N` matrix of transition
probabilities between the states defined by the partition elements, where `N` is the
number of boxes in the partition that is visited by the orbit/points.

If  ``\\{x_t^{(D)} \\}_{n=1}^L`` are the ``L`` different ``D``-dimensional points over
which the transfer operator is approximated, ``\\{ C_{k=1}^N \\}`` are the ``N`` different
partition elements (as dictated by `ϵ`) that gets visited by the points, and
 ``\\phi(x_t) = x_{t+1}``, then

```math
P_{ij} = \\dfrac
{\\#\\{ x_n | \\phi(x_n) \\in C_j \\cap x_n \\in C_i \\}}
{\\#\\{ x_m | x_m \\in C_i \\}},
```

where ``\\#`` denotes the cardinal. The element ``P_{ij}`` thus indicates how many points
that are initially in box ``C_i`` end up in box ``C_j`` when the points in ``C_i`` are
projected one step forward in time. Thus, the row ``P_{ik}^N`` where
``k \\in \\{1, 2, \\ldots, N \\}`` gives the probability
of jumping from the state defined by box ``C_i`` to any of the other ``N`` states. It
follows that ``\\sum_{k=1}^{N} P_{ik} = 1`` for all ``i``. Thus, ``P^N`` is a row/right
stochastic matrix.

### Invariant measure estimation from transfer operator

The left invariant distribution ``\\mathbf{\\rho}^N`` is a row vector, where
``\\mathbf{\\rho}^N P^{N} = \\mathbf{\\rho}^N``. Hence, ``\\mathbf{\\rho}^N`` is a row
eigenvector of the transfer matrix ``P^{N}`` associated with eigenvalue 1. The distribution
``\\mathbf{\\rho}^N`` approximates the invariant density of the system subject to the
partition `ϵ`, and can be taken as a probability distribution over the partition elements.

In practice, the invariant measure ``\\mathbf{\\rho}^N`` is computed using
[`invariantmeasure`](@ref), which also approximates the transfer matrix. The invariant distribution
is initialized as a length-`N` random distribution which is then applied to ``P^{N}``.
The resulting length-`N` distribution is then applied to ``P^{N}`` again. This process
repeats until the difference between the distributions over consecutive iterations is
below some threshold.

## Probability and entropy estimation

- `probabilities(x::AbstractDataset, est::TransferOperator{RectangularBinning})` estimates
    probabilities for the bins defined by the provided binning (`est.ϵ`)
- `genentropy(x::AbstractDataset, est::TransferOperator{RectangularBinning})` does the same,
    but computes generalized entropy using the probabilities.

See also: [`RectangularBinning`](@ref), [`invariantmeasure`](@ref).

[^Diego2019]: Diego, D., Haaga, K. A., & Hannisdal, B. (2019). Transfer entropy computation using the Perron-Frobenius operator. Physical Review E, 99(4), 042212.
"""
struct TransferOperator{R} <: BinningProbabilitiesEstimator
    ϵ::R
end

# If x is not sorted, we need to look at all pairwise comparisons
function inds_in_terms_of_unique(x)
    U = unique(x)
    N = length(x)
    Nu = length(U)
    inds = zeros(Int, N)

    for j = 1:N
        xⱼ = view(x, j)
        for i = 1:Nu
            # using views doesn't allocate
            @inbounds if xⱼ == view(U, i)
                inds[j] = i
            end
        end
    end

    return inds
end

# Taking advantage of the fact that x is sorted reduces runtime by 1.5 orders of magnitude
# for datasets of >100 000+ points
function inds_in_terms_of_unique_sorted(x) # assumes sorted
    @assert issorted(x)
    U = unique(x)
    N, Nu = length(x), length(U)
    prev = view(x, 1)
    inds = zeros(Int, N)
    uidx = 1
    @inbounds for j = 1:N
        xⱼ = view(x, j)
        # if the current value has changed, then we know that the corresponding index
        # for the unique point must be incremented by 1
        if xⱼ != prev
            prev = xⱼ
            uidx += 1
        end
        inds[j] = uidx
    end

    return inds
end

function inds_in_terms_of_unique(x, sorted::Bool)
    if sorted
        return inds_in_terms_of_unique_sorted(x)
    else
        return inds_in_terms_of_unique(x)
    end
end

inds_in_terms_of_unique(x::AbstractDataset) = inds_in_terms_of_unique(x.data)

"""
    TransferOperatorApproximationRectangular(to, ϵ::RectangularBinning, mini, edgelengths,
        bins, sort_idxs)

The `N`-by-`N` matrix `to` is an approximation to the transfer operator, subject to the
partition `ϵ`, computed over some set of sequentially ordered points.

For convenience, `mini` and `edgelengths` provide the minima and box edge lengths along
each coordinate axis, as determined by applying `ϵ` to the points. The coordinates of
the (leftmost, if axis is ordered low-high) box corners are given in `bins`.

Only bins actually visited by the points are considered, and `bins` give the coordinates
of these bins. The element `bins[i]` correspond to the `i`-th state of the system, which
corresponds to the `i`-th column/row of the transfer operator `to`.

`sort_idxs` contains the indices that would sort the input points. `visitors` is a
vector of vectors, where `visitors[i]` contains the indices of the (sorted)
points that visits `bins[i]`.
See also: [`RectangularBinning`](@ref).
"""
struct TransferOperatorApproximationRectangular{T<:Real}
    transfermatrix::AbstractArray{T, 2}
    ϵ::RectangularBinning
    mini
    edgelengths
    bins
    sort_idxs::Vector{Int}
    visitors::Vector{Vector{Int}}
end
# TODO: The above is type unstable!

"""
    transferoperator(pts::AbstractDataset{D, T}, ϵ::RectangularBinning) → TransferOperatorApproximationRectangular

Estimate the transfer operator given a set of sequentially ordered points subject to a
rectangular partition given by `ϵ`.

## Example

```julia
using DynamicalSystems, Plots, Entropy
D = 4
ds = Systems.lorenz96(D; F = 32.0)
N, dt = 20000, 0.1
orbit = trajectory(ds, N*dt; dt = dt, Ttr = 10.0)

# Estimate transfer operator over some coarse graining of the orbit.
transferoperator(orbit, RectangularBinning(10))
```

See also: [`RectangularBinning`](@ref).
"""
function transferoperator(pts::AbstractDataset{D, T}, ϵ::RectangularBinning;
        boundary_condition = :circular) where {D, T<:Real}

    L = length(pts)
    mini, edgelengths = Entropies.minima_edgelengths(pts, ϵ)

    # The L points visits a total of L bins, which are the following bins:
    visited_bins = Entropies.encode_as_bin(pts, mini, edgelengths)
    sort_idxs = sortperm(visited_bins)

    # TODO: fix re-indexing after sorting. Sorting is much faster, so we want to do so.
    #sort!(visited_bins)

    # There are N=length(unique(visited_bins)) unique bins.
    # Which of the unqiue bins does each of the L points visit?
    visits_whichbin = inds_in_terms_of_unique(visited_bins, false) # set to true when sorting is fixed

    # `visitors` lists the indices of the points visiting each of the N unique bins.
    slices = GroupSlices.groupslices(visited_bins)
    visitors = GroupSlices.groupinds(slices)

    # first_visited_by == [x[1] for x in visitors]
    first_visited_by = GroupSlices.firstinds(slices)
    L = length(first_visited_by)

    I = Int32[]
    J = Int32[]
    P = Float64[]

    # Preallocate target index for the case where there is only
    # one point of the orbit visiting a bin.
    target_bin_j::Int = 0
    n_visitsᵢ::Int = 0

    if boundary_condition == :circular
        #warn("Using circular boundary condition")
        append!(visits_whichbin, [1])
    elseif boundary_condition == :random
        #warn("Using random circular boundary condition")
        append!(visits_whichbin, [rand(1:length(visits_whichbin))])
    else
        error("Boundary condition $(boundary_condition) not implemented")
    end

    # Loop over the visited bins bᵢ
    for i in 1:L
        # How many times is this bin visited?
        n_visitsᵢ = length(visitors[i])

        # If both conditions below are true, then there is just one
        # point visiting the i-th bin. If there is only one visiting point and
        # it happens to be the last, we skip it, because we don't know its
        # image.
        if n_visitsᵢ == 1 && !(i == visits_whichbin[end])
            # To which bin does the single point visiting bᵢ jump if we
            # shift it one time step ahead along its orbit?
            target_bin_j = visits_whichbin[visitors[i][1] + 1][1]

            # We now know that exactly one point (the i-th) does the
            # transition from i to the target j.
            push!(I, i)
            push!(J, target_bin_j)
            push!(P, 1.0)
        end

        # If more than one point of the orbit visits the i-th bin, we
        # identify the visiting points and track which bins bⱼ they end up
        # in after the forward linear map of the points.
        if n_visitsᵢ > 1
            timeindices_visiting_pts = visitors[i]

            # TODO: Introduce circular boundary condition. Simply excluding
            # might lead to a cascade of loosing points.

            # If bᵢ is the bin visited by the last point in the orbit, then
            # the last entry of `visiting_pts` will be the time index of the
            # last point of the orbit. In the next time step, that point will
            # have nowhere to go along its orbit (precisely because it is the
            # last data point). Thus, we exclude it.
            if i == visits_whichbin[end]
                #warn("Removing last point")
                n_visitsᵢ = length(timeindices_visiting_pts) - 1
                timeindices_visiting_pts = timeindices_visiting_pts[1:(end - 1)]
            end

            # To which boxes do each of the visitors to bᵢ jump in the next
            # time step?
            target_bins = visits_whichbin[timeindices_visiting_pts .+ 1]
            unique_target_bins = unique(target_bins)

            # Count how many points jump from the i-th bin to each of
            # the unique target bins, and use that to calculate the transition
            # probability from bᵢ to bⱼ.
            #for j in 1:length(unique_target_bins)
            for (j, bᵤ) in enumerate(unique(target_bins))
                n_transitions_i_to_j = sum(target_bins .== bᵤ)

                push!(I, i)
                push!(J, bᵤ)
                push!(P, n_transitions_i_to_j / n_visitsᵢ)
            end
        end
    end

    # Transfer operator is just the normalized transition probabilities between the boxes.
    TO = sparse(I, J, P)

    # Compute the coordinates of the visited bins. bins[i] corresponds to the i-th
    # row/column of the transfer operator
    unique!(visited_bins)
    bins = [β .* edgelengths .+ mini for β in visited_bins]

    TransferOperatorApproximationRectangular(TO, ϵ, mini, edgelengths, bins,
        sort_idxs, visitors)
end

"""
    InvariantMeasure(to, ρ)

Minimal return struct for [`invariantmeasure`](@ref) that contains the estimated invariant
measure `ρ`, as well as the transfer operator `to` from which it is computed (including
bin information).

See also: [`invariantmeasure`](@ref).
"""
struct InvariantMeasure{T}
    to::T
    ρ::Probabilities
end

function invariantmeasure(iv::InvariantMeasure)
    return iv.ρ, iv.to.bins
end

import LinearAlgebra: norm
"""
    invariantmeasure(x::AbstractDataset, ϵ::RectangularBinning) → iv::InvariantMeasure

Estimate an invariant measure over the points in `x` based on binning the data into
rectangular boxes dictated by the binning scheme `ϵ`, then approximate the transfer
(Perron-Frobenius) operator over the bins. From the approximation to the transfer operator,
compute an invariant distribution over the bins. Assumes that the input data are sequential.

Details on the estimation procedure is found the [`TransferOperator`](@ref) docstring.

## Example

```julia
using DynamicalSystems, Plots, Entropies
D = 4
ds = Systems.lorenz96(D; F = 32.0)
N, dt = 20000, 0.1
orbit = trajectory(ds, N*dt; dt = dt, Ttr = 10.0)

# Estimate the invariant measure over some coarse graining of the orbit.
iv = invariantmeasure(orbit, RectangularBinning(15))

# Get the probabilities and bins
invariantmeasure(iv)
```

## Probabilities and bin information

    invariantmeasure(iv::InvariantMeasure) → (ρ::Probabilities, bins::Vector{<:SVector})

From a pre-computed invariant measure, return the probabilities and associated bins.
The element `ρ[i]` is the probability of visitation to the box `bins[i]`. Analogous to
[`binhist`](@ref).

!!! hint "Transfer operator approach vs. naive histogram approach"

    Why bother with the transfer operator instead of using regular histograms to obtain
    probabilities?

    In fact, the naive histogram approach and the
    transfer operator approach are equivalent in the limit of long enough time series
    (as ``n \\to \\intfy``), which is guaranteed by the ergodic theorem. There is a crucial
    difference, however:

    The naive histogram approach only gives the long-term probabilities that
    orbits visit a certain region of the state space. The transfer operator encodes that
    information too, but comes with the added benefit of knowing the *transition
    probabilities* between states (see [`transfermatrix`](@ref)).

See also: [`InvariantMeasure`](@ref).
"""
function invariantmeasure(to::TransferOperatorApproximationRectangular;
        N::Int = 200, tolerance::Float64 = 1e-8, delta::Float64 = 1e-8)

    TO = to.transfermatrix
    #=
    # Start with a random distribution `Ρ` (big rho). Normalise it so that it
    # sums to 1 and forms a true probability distribution over the partition elements.
    =#
    Ρ = rand(Float64, 1, size(to.transfermatrix, 1))
    Ρ = Ρ ./ sum(Ρ, dims = 2)

    #=
    # Start estimating the invariant distribution. We could either do this by
    # finding the left-eigenvector of M, or by repeated application of M on Ρ
    # until the distribution converges. Here, we use the latter approach,
    # meaning that we iterate until Ρ doesn't change substantially between
    # iterations.
    =#
    distribution = Ρ * to.transfermatrix

    distance = norm(distribution - Ρ) / norm(Ρ)

    check = floor(Int, 1 / delta)
    check_pts = floor.(Int, transpose(collect(1:N)) ./ check) .* transpose(collect(1:N))
    check_pts = check_pts[check_pts .> 0]
    num_checkpts = size(check_pts, 1)
    check_pts_counter = 1

    counter = 1
    while counter <= N && distance >= tolerance
        counter += 1
        Ρ = distribution

        # Apply the Markov matrix to the current state of the distribution
        distribution = Ρ * to.transfermatrix

        if (check_pts_counter <= num_checkpts &&
           counter == check_pts[check_pts_counter])

            check_pts_counter += 1
            colsum_distribution = sum(distribution, dims = 2)[1]
            if abs(colsum_distribution - 1) > delta
                distribution = distribution ./ colsum_distribution
            end
        end

        distance = norm(distribution - Ρ) / norm(Ρ)
    end
    distribution = dropdims(distribution, dims = 1)

    # Do the last normalisation and check
    colsum_distribution = sum(distribution)

    if abs(colsum_distribution - 1) > delta
        distribution = distribution ./ colsum_distribution
    end

    # Find partition elements with strictly positive measure.
    δ = tolerance/size(to.transfermatrix, 1)
    inds_nonzero = findall(distribution .> δ)

    # Extract the elements of the invariant measure corresponding to these indices
    return InvariantMeasure(to, Probabilities(distribution))
end

function invariantmeasure(x::AbstractDataset, ϵ::RectangularBinning)
    to = transferoperator(x, ϵ)
    invariantmeasure(to)
end

function probabilities(x::AbstractDataset, est::TransferOperator)
    to = transferoperator(x, est.ϵ)
    iv = invariantmeasure(to)
    return iv.ρ
end

"""
    transfermatrix(iv::InvariantMeasure) → (M::AbstractArray{<:Real, 2}, bins::Vector{<:SVector})

Return the transfer matrix/operator and corresponding bins. Here, `bins[i]` corresponds
to the i-th row/column of the transfer matrix. Thus, the entry `M[i, j]` is the
probability of jumping from the state defined by `bins[i]` to the state defined by
`bins[j]`.

See also: [`TransferOperator`](@ref).
"""
function transfermatrix(iv::InvariantMeasure)
    return iv.to.transfermatrix, iv.to.bins
end