"""
    Egg(Bot)
		q = (x, z, t)
			x - lateral position
			z - vertical position
			t - body orientation
"""
struct Egg{T}
    # dimensions
    nq::Int # generalized coordinates
    nu::Int # controls
    nw::Int # parameters

    # parameters
    mass_body::T
    inertia_body::T
    radius_a::T
    radius_b::T

    # environment
    friction_body_world::T
	gravity::T
end

# Lagrangian
function lagrangian(model::Egg, q, q̇)
    L = 0.0

    # body kinetic energy
    v_body = q̇[1:2]
    ω_body = q̇[3]
    L += 0.5 * model.mass_body * dot(v_body, v_body)
    L += 0.5 * model.inertia_body * ω_body^2.0

    # body potential energy
    z_body = q[2]
    L -= model.mass_body * model.gravity * z_body

    return L
end

# signed distance
function signed_distance(model::Egg, q)
    [
     q[2] - max(model.radius_a, model.radius_b)  # body-world distance
    ]
end

function signed_distance(model::Egg, q, pt1, pt2)
    # error between egg and point on nearest surface
    e1 = pt1 - pt2
    [
     sqrt(dot(e1, e1)) # body-world distance
    ]
end

# input matrix
function input_jacobian(model, q)
	[
     0.0;
     0.0;
     1.0
    ]
end

# dimensions
nq = 3
nu = 1
nw = 0

# parameters
mass_body = 1.0
inertia_body = 0.1
radius_a = 0.1
radius_b = 0.05

# environment
friction_body_world = 0.5
gravity = 9.81

# model
model = Egg(nq, nu, nw,
            mass_body, inertia_body,
            radius_a, radius_b,
            friction_body_world,
            gravity)

# code-gen

nq = model.nq
nu = model.nu
nw = model.nw

# variables
@variables q[1:nq] q̇[1:nq]

# Lagrangian
L = lagrangian(model, q, q̇)

ddL = Symbolics.hessian(L, [q; q̇])#, simplify=true)
dLq = Symbolics.gradient(L, q)#, simplify=true)
ddLq̇q = ddL[nq .+ (1:nq), 1:nq]

M = ddL[nq .+ (1:nq), nq .+ (1:nq)]
C = ddLq̇q * q̇ - dLq

mass_matrix_sym = eval(Symbolics.build_function(M, q)[1])
dynamics_bias_sym = eval(Symbolics.build_function(C, q, q̇)[1])

ϕ = signed_distance(model, q)
dϕ = Symbolics.jacobian(ϕ, q)
signed_distance_jacobian = eval(Symbolics.build_function(dϕ, q)[1])

function contact_jacobian(model, q)
    [
     1.0 0.0 0.0; # note that body torque due to friction is added in residual
     0.0 1.0 0.0
    ]
end

function lagrangian_derivatives(model, q, v)
	D1L = -1.0 * dynamics_bias_sym(q, v)
    D2L = mass_matrix_sym(q) * v
	return D1L, D2L
end

function dynamics(model, h, q0, q1, u1, w1, λ1, q2)
	# evalutate at midpoint
	qm1 = 0.5 * (q0 + q1)
    vm1 = (q1 - q0) / h[1]
    qm2 = 0.5 * (q1 + q2)
    vm2 = (q2 - q1) / h[1]

	D1L1, D2L1 = lagrangian_derivatives(model, qm1, vm1)
	D1L2, D2L2 = lagrangian_derivatives(model, qm2, vm2)

	return (0.5 * h[1] * D1L1 + D2L1 + 0.5 * h[1] * D1L2 - D2L2
		+ input_jacobian(model, qm2) * u1[1]
        + λ1)
end

@variables h[1:1] q0[1:nq] q1[1:nq] u1[1:nu] w1[1:nw] λ1[1:3] q2[1:nq]
dyn = dynamics(model, h, q0, q1, u1, w1, λ1, q2)

function residual(model, z, θ, η, μ)
    # dimensions
    nq = model.nq
    nu = model.nu
    nw = model.nw

    # unpack data
    q0 = θ[1:nq]
    q1 = θ[nq .+ (1:nq)]
    u1 = θ[2nq .+ (1:nu)]
    w1 = θ[2nq + nu .+ (1:nw)]
    friction_body_world = θ[2nq + nu + nw .+ (1:1)]
    h = θ[2nq + nu + nw + 1 .+ (1:1)]

    # unpack contact points and directions
    pt_contact_body = η[1:2]
    pt_contact_surface = η[3:4]
    dir_contact = η[5:6]

    # unpack variables
    q2 = z[1:nq]
    γ1 = z[nq .+ (1:1)]
    sγ1 = z[nq + 1 .+ (1:1)]
    ψ1 = z[nq + 1 + 1 .+ (1:1)]
    b1 = z[nq + 1 + 1 + 1 .+ (1:1)]
    sψ1 = z[nq + 1 + 1 + 1 + 1 .+ (1:1)]
    sb1 = z[nq + 1 + 1 + 1 + 1 + 1 .+ (1:1)]

    rotation_body = rotation(dir_contact)
    r_contact = [pt_contact_body - q2[1:2]; 0.0]

    # tangential velocity
    v1 = (q2 - q1) / h[1]
    vT_body = (rotation_body * (v1[1:2] + cross([0.0; 0.0; v1[3]], r_contact)[1:2]))[1]

    # contact forces
    J = contact_jacobian(model, q2)
    λ1 = transpose(J) * transpose(rotation_body) * [b1[1]; γ1[1]]

    λ1[3] += cross(r_contact, [transpose(rotation_body) * [b1[1]; 0.0]; 0.0])[3] # friction on body creates a moment

	[
     dynamics(model, h, q0, q1, u1, w1, λ1, q2);
	 signed_distance(model, q2, pt_contact_body, pt_contact_surface) - sγ1;
     vT_body - sb1[1];
     friction_body_world[1] * γ1[1] - ψ1[1];
     γ1 .* sγ1 .- μ[1];
	 cone_product([ψ1[1]; b1[1]], [sψ1[1]; sb1[1]]) - [μ[1]; 0.0];
    ]
end

nz = nq + 1 + 1 + 1 + 1 + 1 + 1
ny = 1 + 1 + 1 + 1 + 1 + 1
nθ = 2nq + nu + nw + 1 + 1
nη = 2 + 2 + 2

@variables z[1:nz] θ[1:nθ] η[1:nη] μ[1:1]
r = residual(model, z, θ, η, μ)
rz = Symbolics.jacobian(r, z)
rη = Symbolics.jacobian(r, η)
rθ = Symbolics.jacobian(r, θ)
rz_sp = Symbolics.sparsejacobian(r, z)
rη_sp = Symbolics.sparsejacobian(r, η)
rθ_sp = Symbolics.sparsejacobian(r, θ)

r_Egg! = eval(Symbolics.build_function(r, z, θ, η, μ)[2])
rz_Egg! = eval(Symbolics.build_function(rz, z, θ, η)[2])
rη_Egg! = eval(Symbolics.build_function(rη, z, θ, η)[2])
rθ_Egg! = eval(Symbolics.build_function(rθ, z, θ, η)[2])

r_ = eval(Symbolics.build_function(r, z, θ, η, μ)[1])
rz_ = eval(Symbolics.build_function(rz, z, θ, η)[1])
rη_ = eval(Symbolics.build_function(rη, z, θ, η)[1])
rθ_ = eval(Symbolics.build_function(rθ, z, θ, η)[1])
rz_sp! = eval(Symbolics.build_function(rz_sp.nzval, z, θ, η)[2])
rη_sp! = eval(Symbolics.build_function(rη_sp.nzval, z, θ, η)[2])
rθ_sp! = eval(Symbolics.build_function(rθ_sp.nzval, z, θ, η)[2])

r0 = zeros(length(r))
rz0 = zeros(size(rz))
rθ0 = zeros(size(rθ))
rη0 = zeros(size(rη))

rz0_sp = zeros(size(rz_sp.nzval))
rθ0_sp = zeros(size(rθ_sp.nzval))
rη0_sp = zeros(size(rη_sp.nzval))

z0 = rand(nz)
θ0 = rand(nθ)
η0 = rand(nη)

μ0 = [1.0]

# using BenchmarkTools
# @benchmark r!($r0, $z0, $θ0, $η0, $μ0)
# @benchmark rz!($rz0, $z0, $θ0, $η0)
# @benchmark rη!($rη0, $z0, $θ0, $η0)
# @benchmark rθ!($rθ0, $z0, $θ0, $η0)

# @benchmark $r0 .= r_($z0, $θ0, $η0, $μ0)
# @benchmark $rz0 .= rz_($z0, $θ0, $η0)
# @benchmark $rη0 .= rη_($z0, $θ0, $η0)
# @benchmark $rθ0 .= rθ_($z0, $θ0, $η0)

# @benchmark rz_sp!($rz0_sp, $z0, $θ0, $η0)
# @benchmark rη_sp!($rη0_sp, $z0, $θ0, $η0)
# @benchmark rθ_sp!($rθ0_sp, $z0, $θ0, $η0)

##
q0 = [0.0, 1.0, 0.0 * π]
z0[1:nq] = copy(q0)
z0[nq .+ (1:1)] .= 1.0
z0[nq + 1 .+ (1:1)] .= 1.0
z0[nq + 1 + 1 .+ (1:1)] .= 1.0
z0[nq + 1 + 1 + 1 .+ (1:1)] .= 0.1
z0[nq + 1 + 1 + 1 + 1 .+ (1:1)] .= 1.0
z0[nq + 1 + 1 + 1 + 1 + 1 .+ (1:1)] .= 0.1

k_body = kinematics_body(model, q0)
k_foot = kinematics_foot(model, q0)

p_body = contact_point(surf, k_body, r=2.0 * model.body_radius)
p_foot = contact_point(surf, k_foot, r=2.0 * model.body_radius)

η0 = [p_body[1]; p_foot[1]; p_body[3]; p_foot[3]]
rz0 = rz_(z0, θ0, η0)

r1 = rη_(z0, θ0, η0)[:, 1:2] * p_body[6] * kinematics_body_jacobian(model, q0)
r2 = rη_(z0, θ0, η0)[:, 2 .+ (1:2)] * p_foot[6] * kinematics_foot_jacobian(model, q0)
r3 = rη_(z0, θ0, η0)[:, 4 .+ (1:2)] * p_body[7] * p_body[6] * kinematics_body_jacobian(model, q0)

rη_update = r1 + r2 + r3

rz0[:, 1:nq] += rη_update

function res!(r, z, θ, μ)
    q2 = z[1:nq]
    k_body = kinematics_body(model, q2)
    k_foot = kinematics_foot(model, q2)

    p_body = contact_point(surf, k_body, r=2.0 * model.body_radius)
    p_foot = contact_point(surf, k_foot, r=2.0 * model.foot_radius)

    η = [p_body[1]; p_foot[1]; p_body[3]; p_foot[3]]

    r_Egg!(r, z, θ, η, μ)
end

function res_z!(rz, z, θ)
    q2 = z[1:nq]
    k_body = kinematics_body(model, q2)
    k_foot = kinematics_foot(model, q2)

    p_body = contact_point(surf, k_body, r=2.0 * model.body_radius)
    p_foot = contact_point(surf, k_foot, r=2.0 * model.foot_radius)

    # @show p_body[1]
    # @show p_foot[1]

    η = [p_body[1]; p_foot[1]; p_body[3]; p_foot[3]]

    rz_Egg!(rz, z, θ, η)
    rη = rη_(z, θ, η)

    r1 = rη[:, 1:2] * p_body[6] * kinematics_body_jacobian(model, q2)
    r2 = rη[:, 2 .+ (1:2)] * p_foot[6] * kinematics_foot_jacobian(model, q2)
    r3 = rη[:, 4 .+ (1:2)] * p_body[7] * p_body[6] * kinematics_body_jacobian(model, q2)
    r4 = rη[:, 6 .+ (1:2)] * p_foot[7] * p_foot[6] * kinematics_foot_jacobian(model, q2)

    rη_update = r1 + r2 + r3 + r4

    rz[:, 1:nq] += rη_update
end

# res(z0, θ0, μ0)
# res_z(z0, θ0, μ0)

# res!(r0, z0, θ0, μ0)
# res_z!(rz0, z0, θ0)

# r0


"""
    Newton{N,S,T,X}

A [Newton](https://en.wikipedia.org/wiki/Newton_polynomial)
interpolating polynomial uses a basis `1`, `(x-x_0)`,
`(x-x_0)(x-x_1)`, ..., `(x-x0)(x-x1)⋅⋅⋅(x-x_{n-1})` and coefficients
(in forward form) `f[x_0]`, `f[x_0,x_1]`, ...,`f[x_0,...,x_n]`. The
Newton class stores the nodes (after sorting) and the Newton tableau
used to generate the coefficients on fitting.

The easiest way to construct an instance is with `fit`, as in:

```jldoctest Newton
julia> using Polynomials, SpecialPolynomials

julia> xs = [1,2,3,4]; f(x)= x^3 - 2x + 1;

julia> p = fit(Newton, xs, f)
Newton(5.0⋅p_1(x) + 6.0⋅p_2(x) + 1.0⋅p_3(x))

julia> p.(xs) == f.(xs)  # p interpolates
true

julia> convert(Polynomial, p)
Polynomials.Polynomial(1.0 - 2.0*x + 1.0*x^3)
```

"""
struct Newton{N, S<:Number, T <: Number, X, V} <: AbstractInterpolatingPolynomial{T, X}
    xs::Vector{S}
    tableau::Matrix{T}
    coeffs::V
    function Newton(xs::Vector{S}, nt::AbstractMatrix{T}, var::Symbol=:x) where {S,T}
        N = length(xs)
        xs = sort(unique(xs))
        N != length(xs) && throw(ArgumentError("xs must be unique"))
        v = view(nt, 1, :)
        new{N,S,T, Symbol(var), typeof(v)}(xs, nt, v)
    end
end

export Newton

basis_symbol(::Type{<:Newton}) = "p"

## Boilerplate code reproduced here, as there are two type parameters
Base.convert(::Type{P}, p::P) where {P <: Newton} =  p
Base.convert(P::Type{Newton{N,S,T}}, p::Newton) where {N,S,T} = Newton{N,S,T,Polynomial.indeterminate(P,p)}(p.xs, p.tableau)
Base.promote_rule(::Type{Newton{N, S, T}}, ::Type{U}) where {N, S, T, U <: Number} = Newton{N, S, promote_type(T,U)}

Polynomials.domain(::Type{<:Newton}) = Polynomials.Interval(-Inf, Inf)
Polynomials.variable(p::Newton{N, S, T}, var::Polynomials.SymbolLike=:x) where {N,S, T} = fit(Newton, p.xs, x -> x, var=var)
Polynomials.zero(p::Newton{N, S, T}) where {N,S,T} = 0*p
function Polynomials.one(p::Newton)
    xs = sort(p.xs)
    nt = newton_tableau(one, xs)
    Newton(xs, nt, Polynomials.indeterminate(p))
end

"""
    newton_tableau(f, xs::Vector)
    newton_tableau(xs, ys)

Construct the
[matrix](https://en.wikipedia.org/wiki/Divided_differences#Matrix_form)
form of Newton's divided differences. This assumes the xs are unique.

"""
newton_tableau(f, xs::AbstractVector{S}) where {S} = newton_tableau(xs, f.(xs)/one(S))
function newton_tableau(xs::AbstractVector{S}, ys::Vector{T}) where {S,T}
    xs = sort(unique(xs))
    n = length(xs)
    M = diagm(0=>ys/one(S))
    for i in 2:n
        for j in 1:n-i+1
            k = i + j -1
            M[j, k] = (M[j+1,k] - M[j,k-1]) / (xs[j-1+1+(i-1)] - xs[j-1+1])
        end
    end
    M
end

"""
    update_newton_tableau(x, fx, xs, nt)

Update the newton tableau, `nt`, generated from the `xs` for some
function `f`. The point `(x,fx)` is a new `x` value (from the `xs`)
and `fx` the function value.

"""
function update_newton_tableau(x::T, fx::S, xs, nt::Matrix{R}) where {T, S, R}
    n = size(nt, 1)
    M = zeros(R, n+1, n+1)
    for i in 1:n
        for j in i:n
            M[i,j] = nt[i,j]
        end
    end
    M[end,end] = fx
    xxs = vcat(xs, x)
    for i  in n:-1:1
        M[i, n+1] = (M[i+1, n+1] - M[i,n]) / (xxs[end] - xxs[end-(n-i+1)])
    end

    xxs, M
end

## We use the forward difference formula
## the coefficients are the top row of the upper
## triangular tableau
Polynomials.coeffs(p::Newton) = p.tableau[1,:]

## Evaluation
## p(x) = f[x0] + f[x0,x1](x-x0) + ... + f[x0,x1,...,x_n](x-x0)(x-x1)...(x-x_{n-1})
## use nested evaluation (http://pages.cs.wisc.edu/~amos/412/lecture-notes/lecture08.pdf)
(p::Newton)(x)  = evalpoly(x,p)
function Polynomials.evalpoly(x, p::Newton{N}) where {N}
    xs, cs = p.xs, coeffs(p)
    tot = cs[end]
    for j in N-1:-1:1
        tot = cs[j] + (x - xs[j])*tot
    end
    tot
end

## compute (p,dp)  at x:
## http://pages.cs.wisc.edu/~amos/412/lecture-notes/lecture08.pdf
function _dual(p::Newton{N, S,T}, x) where {N,S,T}
    xs, cs = p.xs, coeffs(p)
    tot = cs[end]
    totp = zero(S)
    for j in N-1:-1:1
        delta = (x - xs[j])
        tot, totp = cs[j] + delta * tot, tot + delta * totp
    end
    tot, totp
end

function Base.:+(p1::Newton{N,S,T,X}, p2::Newton{M,S1,T1,Y}) where {N,S,T,X,M,S1,T1,Y}
    X == Y || throw(ArgumentError("p1 and p2 must share the same variable"))
    if N == M && p1.xs == p2.xs
        tableau = p1.tableau + p2.tableau
        return Newton(p1.xs, tableau, X)
    else
        p,q = N >= M ? (p1, p2) : (p2, p1)
        qq = fit(Newton, p.xs, q) # stretch than add
        ## https://en.wikipedia.org/wiki/Divided_differences
        ## T_{f+g}x = T_fx + T_gx
        p + qq
    end

end

function Base.:*(p1::Newton{N,S,T,X}, p2::Newton{M,S1,T1,Y}) where {N,S,T,X,M,S1,T1,Y}
    ## Basic idea
    ## find larger of xs, expand to xxs
    ## refit p1 and p2 on this
    ## use  product of tableaus
    Polynomials.assert_same_variable(p1, p2)
    p,q = N >= M ? (p1, p2) : (p2, p1)
    xs = p.xs
    new_xs = _new_nodes(xs, q.xs)

    xxs = vcat(xs, new_xs)
    pp = fit(Newton, xxs, p)
    qq = fit(Newton, xxs, q)

    ## https://en.wikipedia.org/wiki/Divided_differences
    ## T_{f*g}(x) = T_f(x) * T_g(x)

    Newton(pp.xs, pp.tableau * qq.tableau, X)

end

## Scalar
function Base.:+(p::Newton{N,T,S,X}, c::Number) where {N,T,S,X}
    tableau = copy(p.tableau)
    for i in 1:N
        tableau[i,i] += c
    end
    Newton(p.xs, tableau, X)
end

Base.:*(c::Number, p::Newton) = p*c
function Base.:*(p::P, c::Number) where {P <: Newton}
    Newton(p.xs, c*p.tableau, Polynomials.indeterminate(p))
end

Base.:-(p::P) where {P<:Newton} = p*(-1)

function  Base.convert(Q::Type{<:Polynomial},  p::Newton{N, S,T})  where  {N,S, T}
    p(variable(Polynomial{S}))
end

# return derivative
# to evaluate derivative at a value
# use  _dual
# These are not efficient due to the need to  compute a tableau
function Polynomials.derivative(p::Newton{N,S,T}, order::Int=1) where {N,S,T}
    order < 0 && throw(ArgumentError("order must be positive"))
    order == 0 && return p
    N==1 &&  return fit(Newton, p.xs, zero)

    x = variable(Polynomial{S})
    q,qq = _dual(p, x)
    dp = fit(Newton, p.xs[1:end-1], qq)
    if order > 1
        dp = derivative(dp, order-1)
    end
    dp

end

# integrate
function Polynomials.integrate(p::Newton{N,S,T}, C::Number=0) where {N,S,T}
    q = convert(Polynomial, p)
    Q = integrate(q, C)
    if length(p.xs) >= 2
        new_x = p.xs[end] + (p.xs[end] - p.xs[end-1])
        xxs = vcat(p.xs, new_x)
    else
        xxs = vcat(p.xs, p.xs[1]+1)
    end
    fit(Newton, xxs, Q, var=Polynomial.indeterminate(p))

end

function Polynomials.fit(P::Type{<:Newton},
                         xs::Vector{S},
                         ys::AbstractVector{T};
                         var = :x,) where {S, T}
    ind = sortperm(xs)
    xs = xs[ind]
    ys = ys[ind]
    nt = newton_tableau(xs, ys)

    Newton(xs, nt, var)
end

function Polynomials.fit(P::Type{<:Newton}, xs::Vector{S}, f; var=:x) where {S}
    xs = sort(xs)
    nt = newton_tableau(f, xs)
    Newton(xs, nt, var)
end


"""
    trunk_build(trunk,M)

Generate a dictionary of trunk functions evaluated at ``t = 0``

"""
function trunk_build(trunk,M)
    return Dict(i => function (x) return trunk(vcat(0,x))[i] end for i in 1:M)
end

"""
    trunk_ortho_build(utilde,L1,L2,nodes,weights)

Generate a dictionary of orthonormal custom basis functions which use an orthogonal polynomial expansion for evaluation away from the specified quadrature nodes.

"""
function trunk_ortho_build(utilde,L1,L2,nodes,weights)
    trunk_ortho = [];
    W = diagm(0 => weights);
    N = size(utilde,2);
    for i in 1:N
        coeffs = zeros(size(utilde,1));
        coeffs[1] = legendre_norm.(nodes,L1,L2,0)'*W*utilde[:,i];
        for j in 1:(size(utilde,1)-1)
            coeffs[j+1] = legendre_norm.(nodes,L1,L2,j)'*W*utilde[:,i]
        end
        push!(trunk_ortho, (x) -> coeffs'*legendre_norm_collect.(x,L1,L2,(size(utilde,1)-1)));
    end
    return trunk_ortho
end

"""
    build_basis(trunk,L1,L2,M,nodes,weights)

Generate a dictionary of orthonormal custom basis functions based on a specified quadrature rule. An orthogonal polynomial expansion is utilized for the evaluation at points away from the specified quadrature nodes.

"""
function build_basis(trunk,L1,L2,M,nodes,weights)
    trunk_func = trunk_build(trunk,M);

    W = diagm(0 => weights);

    A = zeros(size(nodes,1),M);
    for i in 1:size(nodes,1)
        for k in 1:M
            A[i,k] = trunk_func[k](nodes[i]);
        end
    end

    B = W^(1/2)*A;
    F = svd(B);
    utilde = W^(-1/2)*F.U;

    ortho_trunk_func = trunk_ortho_build(utilde,L1,L2,nodes,weights);

    # Check if the functions are orthonormal to a specified precision at the specified nodes
    orthonormal_check(basis_eval(ortho_trunk_func,nodes),weights;tol=1e-12)

    return ortho_trunk_func, F.S
end

"""
    basis_eval(basis,nodes)

Evaluate the basis functions for a specified spatial grid or at a specified point. Evaluation for a specified grid returns an ``N x M`` matrix with ``N`` respresenting the number of nodes and ``M`` representing the number of custom basis functions. Evalaution for a specified point returns a ``1 x M`` vector.

"""
function basis_eval(basis,nodes::Array)
    apply(f,x) = f(x);
    return permutedims(apply.(basis,permutedims(nodes)));
end
function basis_eval(basis,node)
    apply(f,x) = f(x);
    return permutedims(apply.(basis,node));
end

"""
    expansion_coefficients(basis,fnc,nodes,weights)

Compute the expansion coefficients...

"""
function expansion_coefficients(basis::Array{Any,1},fnc::Function,nodes,weights)
    return basis_eval(basis,nodes)'*diagm(0 => weights)*fnc.(nodes)
end
function expansion_coefficients(basis::Array{Any,1},fnc,nodes,weights)
    return basis_eval(basis,nodes)'*diagm(0 => weights)*fnc
end
function expansion_coefficients(basis::Array{Float64,2},fnc::Function,nodes,weights)
    return basis'*diagm(0 => weights)*fnc.(nodes)
end
function expansion_coefficients(basis::Array{Float64,2},fnc,nodes,weights)
    return basis'*diagm(0 => weights)*fnc
end

"""
    expansion_approximation(basis,coefficients,nodes)

Compute the expansion approximation...

"""
function expansion_approximation(basis::Array{Any,1},coefficients,nodes)
    return basis_eval(basis,nodes)*coefficients
end
function expansion_approximation(basis::Array{Float64,2},coefficients,nodes)
    return basis*coefficients
end

"""
    basis_derivative(basis,nodes)

Compute the derivatives of each basis function using auto differentiation.

"""
function basis_derivative(basis)
  Dbasis = []
  for i in 1:length(basis);
    push!(Dbasis,(x)->ForwardDiff.derivative.(basis[i],x))
  end
  return Dbasis
end

"""
    save_basis(basis,M,pde_function)

"""
function save_basis(basis,M,pde_function)
    @save @sprintf("basis_%i_M_%s.bson",M,pde_function) basis
end

"""
    load_basis(basis,M,pde_function)

"""
function load_basis(M,pde_function)
    @load @sprintf("basis_%i_M_%s.bson",M,pde_function) basis
    return basis
end


@testset "ACWD" begin
  x₁ = randn(8)
  y₁ = acwt(x₁, wavelet(WT.db4))
  @test iacwt(y₁) ≈ x₁

  x₂ = randn(8,8)
  y₂ = acwt(x₂, wavelet(WT.db4))
  @test iacwt(y₂) ≈ x₂
end

@testset "ACWPT" begin
  x₁ = randn(8)
  y₁ = acwt(x₁, wavelet(WT.db4))
  y₂ = acwpt(x₁, wavelet(WT.db4))

  @test y₂[:,8] ≈ y₁[:,1]
  @test iacwpt(y₂, maketree(x₁, :full)) ≈ x₁
end


using GeometricStructures
using Base.Test

import GeometricStructures.intersection
import GeometricStructures.convert
import GeometricStructures.convexhull

include("points.jl")
include("linesegments.jl")
include("convexhull.jl")


export DiaSymSemiseparable

struct DiaSymSemiseparable <: SymSemiseparableMatrix
    n::Int64
    p::Int64
    U::AbstractArray
    V::AbstractArray
    d::AbstractArray
end

# Constuctors
function DiaSymSemiseparable(U::AbstractArray, V::AbstractArray, d::AbstractArray)
    if size(U,1) == size(V,1) && size(U,2) == size(V,2) && length(d) == size(U,1)
        return DiaSymSemiseparable(size(U,1),size(U,2),U,V,d);
    else
        error("Dimension mismatch between the generators U, V and d")
    end
end
# Mappings
mul!(y::AbstractArray, L::DiaSymSemiseparable, 		            b::AbstractArray) = dss_mul_mat!(y, L.U, L.V, L.d, b);
mul!(y::AbstractArray, L::AdjointOperator{DiaSymSemiseparable}, b::AbstractArray) = dss_mul_mat!(y, L.A.U, L.A.V, L.A.d, b);
function inv!(y, K::DiaSymSemiseparable, b::AbstractArray)
	L = DiaSymSemiseparableChol(K)
	y[:,:] = L'\(L\b)
end
function inv!(y, K::AdjointOperator{DiaSymSemiseparable}, b::AbstractArray)
	L = DiaSymSemiseparableChol(K.A)
	y[:,:] = L'\(L\b)
end

#################################################
#### Cholesky factoriaztion of:              ####
####    Higher-order quasiseparable matrices ####
#################################################

#### Matrix-matrix product ####
function dss_mul_mat!(Y::Array, U::Array, V::Array, d::Array, X::Array)
    n, m = size(U);
    mx = size(X,2);
    Vbar = zeros(m,mx);
    Ubar = U'*X;
    for i = 1:n
        tmpV = V[i,:];
        tmpU = U[i,:];
        tmpX = X[i:i,:];
        Ubar -= tmpU .* tmpX;
        Vbar += tmpV .* tmpX;
        Y[i,:] = tmpU'*Vbar + tmpV'*Ubar + d[i]*tmpX;
    end
end
#### Going from the Cholesky factorization back to matrix Σ ####
function dss_create_vd(U::Array, W::Array, dbar::Array)
    n,m = size(U)
    d = zeros(n)
    V = zeros(n,m)
    P = zeros(m,m)
    for i = 1:n
        tmpU = U[i,:]
        tmpW = W[i,:]
        d[i] = dbar[i]^2 - tmpU'*tmpW
        V[i,:] = P*tmpU
        P += tmpW*tmpW';
    end
    return V, d
end


### A Pluto.jl notebook ###
# v0.14.8

using Markdown
using InteractiveUtils

# ╔═╡ 837951f4-68ec-11eb-048d-6b27b9055428
using Pkg, DrWatson

# ╔═╡ 91489170-68ec-11eb-3611-150d5febf223
begin
	using Distributions
	using StatsPlots
	using StatsBase
	using LaTeXStrings
	using CSV
	using DataFrames
	using LinearAlgebra
	using Random
	using ParetoSmoothedImportanceSampling
	using StanQuap
	using StatisticalRethinking
	using StatisticalRethinkingPlots
end

# ╔═╡ 7046ea12-68ec-11eb-18ae-a724ca79a21f
md" ## Clip-07-35s.jl"

# ╔═╡ ace5ed4c-68ec-11eb-05ed-1f1ff2d72e1b
begin
	df = CSV.read(sr_datadir("WaffleDivorce.csv"), DataFrame);
	scale!(df, [:Marriage, :MedianAgeMarriage, :Divorce])
end;

# ╔═╡ 4c47f312-68ed-11eb-37de-ff782118d258
stan5_1_t = "
data {
 int < lower = 1 > N; // Sample size
 vector[N] D; // Outcome
 vector[N] A; // Predictor
}

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

transformed parameters {
	vector[N] mu;
	mu = a + + bA * A;
}

model {
  a ~ normal( 0 , 0.2 );
  bA ~ normal( 0 , 0.5 );
  sigma ~ exponential( 1 );
  D ~ student_t( 2, mu , sigma );
}
generated quantities{
	vector[N] log_lik;
	for (i in 1:N)
		log_lik[i] = student_t_lpdf(D[i] | 2, mu[i], sigma);
}
";

# ╔═╡ 4c482dfa-68ed-11eb-2db9-a54ab816beb7
begin
	data = (N=size(df, 1), D=df.Divorce_s, A=df.MedianAgeMarriage_s,
		M=df.Marriage_s)
	m5_1s_t = SampleModel("m5.1s_t", stan5_1_t)
	rc5_1s_t = stan_sample(m5_1s_t; data)

	if success(rc5_1s_t)
		post5_1s_t_df = read_samples(m5_1s_t, :dataframe)
		PRECIS(post5_1s_t_df[:, [:a, :bA, :sigma]])
	end
end

# ╔═╡ 806d4352-69b7-11eb-00a5-791917fa9d45
if success(rc5_1s_t)
	nt5_1s_t = read_samples(m5_1s_t, :namedtuple)
	log_lik_1_t = nt5_1s_t.log_lik'
	waic_m5_1s_t = waic(log_lik_1_t)
end

# ╔═╡ 6210f1ba-68ef-11eb-1d55-653d1c8a147c
begin
    b5_1s_t = post5_1s_t_df[:, [:a, :bA, :sigma]]
    mu5_1s_t = b5_1s_t.a .+ b5_1s_t.bA * df.MedianAgeMarriage_s'
	lp5_1s_t = logpdf.(TDist(2), mu5_1s_t)
	waic(lp5_1s_t)
end

# ╔═╡ 3a9ff05e-6a16-11eb-2ba0-530fb4bc18f4
size(mu5_1s_t)

# ╔═╡ 8f107504-68ed-11eb-002b-372884f598b0
stan5_2_t = "
data {
  int N;
  vector[N] D;
  vector[N] M;
}
parameters {
  real a;
  real bM;
  real<lower=0> sigma;
}
transformed parameters {
	vector[N] mu;
	mu = a + bM * M;
}
model {
  a ~ normal( 0 , 0.2 );
  bM ~ normal( 0 , 0.5 );
  sigma ~ exponential( 1 );
  D ~ student_t( 2, mu , sigma );
}
generated quantities{
	vector[N] log_lik;
	for (i in 1:N)
		log_lik[i] = student_t_lpdf(D[i] | 2, mu[i], sigma);
}
";

# ╔═╡ cbb9ef56-68ed-11eb-387b-252ff6dc5b21
begin
	m5_2s_t = SampleModel("m5.2_t", stan5_2_t);
	rc5_2s_t = stan_sample(m5_2s_t; data)
	if success(rc5_2s_t)
		post5_2s_t_df = read_samples(m5_2s_t, :dataframe)
		PRECIS(post5_2s_t_df[:, [:a, :bM, :sigma]])
	end
end

# ╔═╡ 5d216004-69b7-11eb-34d4-8be76d4ae47c
if success(rc5_2s_t)
	nt5_2s_t = read_samples(m5_2s_t, :namedtuple)
	log_lik_2_t = nt5_2s_t.log_lik'
	waic_m5_2s_t = waic(log_lik_2_t)
end

# ╔═╡ 87cd0b48-68f2-11eb-3704-af5e98a9c463
begin
    b5_2s_t = post5_2s_t_df[:, [:a, :bM, :sigma]]
    mu5_2s_t = b5_2s_t.a .+ b5_2s_t.bM * df.Marriage_s'
	lp5_2s_t = logpdf.(TDist.(2), mu5_2s_t)
	waic(lp5_2s_t)
end

# ╔═╡ 71a48f7a-68ee-11eb-1b0e-331393e95551
stan5_3_t = "
data {
  int N;
  vector[N] D;
  vector[N] M;
  vector[N] A;
}
parameters {
  real a;
  real bA;
  real bM;
  real<lower=0> sigma;
}
transformed parameters {
	vector[N] mu;
	mu = a + + bA * A + bM * M;
}
model {
  a ~ normal( 0 , 0.2 );
  bA ~ normal( 0 , 0.5 );
  bM ~ normal( 0 , 0.5 );
  sigma ~ exponential( 1 );
  D ~ student_t( 2, mu , sigma );
}
generated quantities{
	vector[N] log_lik;
	for (i in 1:N)
		log_lik[i] = student_t_lpdf(D[i] | 2, mu[i], sigma);
}
";

# ╔═╡ b5b487ec-68ee-11eb-3ee7-0d98367028ff
begin
	m5_3s_t = SampleModel("m5.3_t", stan5_3_t);
	rc5_3s_t = stan_sample(m5_3s_t; data);

	if success(rc5_3s_t)
		post5_3s_t_df = read_samples(m5_3s_t, :dataframe)
		PRECIS(post5_3s_t_df[:, [:a, :bA, :bM, :sigma]])
	end
end

# ╔═╡ 4e7a00e8-69b4-11eb-15ef-df1a53282396
if success(rc5_3s_t)
	nt5_3s_t = read_samples(m5_3s_t, :namedtuple)
	log_lik_3_t = nt5_3s_t.log_lik'
	waic_m5_3s_t = waic(log_lik_3_t)
end

# ╔═╡ 51d570db-4ff0-49b4-b4c1-8e709a6b1564
nt5_3s_t

# ╔═╡ c2ba84b0-68f2-11eb-3c5a-17e03541dcb1
begin
    b5_3s_t = post5_3s_t_df[:, [:a, :bA, :bM, :sigma]]
    mu5_3s_t = b5_3s_t.a .+ b5_3s_t.bM * df.Marriage_s' +
		b5_3s_t.bA * df.MedianAgeMarriage_s'
	lp5_3s_t = logpdf.(TDist.(2), mu5_3s_t)
	waic(lp5_3s_t)
end

# ╔═╡ 7b051a2e-68f3-11eb-3ee8-e77c639fc390
[waic_m5_1s_t.WAIC, waic_m5_2s_t.WAIC, waic_m5_3s_t.WAIC]

# ╔═╡ ed4469dc-68f3-11eb-3b0f-f54305283306
begin
	loo5_1s_t, loos5_1s_t, pk5_1s_t = psisloo(log_lik_1_t)
	loo5_2s_t, loos5_2s_t, pk5_2s_t = psisloo(log_lik_2_t)
	loo5_3s_t, loos5_3s_t, pk5_3s_t = psisloo(log_lik_3_t)
	[-2loo5_1s_t, -2loo5_2s_t, -2loo5_3s_t]
end

# ╔═╡ 2c4b5ffc-68f4-11eb-21d6-7d2db834d3f8
begin
	pk_plot(pk5_1s_t)
	annotate!([(13 + 1, pk5_1s_t[13] + 0.02, Plots.text(df[13, :Loc],
		6, :red, :right))])
end

# ╔═╡ 3766d148-695b-11eb-19dc-0b25e6704b3d
begin
	waic_5_1s_pw_t = waic(lp5_1s_t; pointwise=true)
	scatter(pk5_1s_t, waic_5_1s_pw_t.penalty,
		xlab="PSIS Pareto k", ylab="WAIC penalty", leg=false)
	vline!([0.5])
	for state in [13, 20, 34, 44, 50]
		annotate!([([pk5_1s_t[state] + 0.02], [waic_5_1s_pw_t.penalty[state]],
			Plots.text(df[state, :Loc], 6, :red, :right))])
	end
	plot!()
end

# ╔═╡ aedb425c-68f4-11eb-319d-a958c2aacde1
begin
	pk_plot(pk5_3s_t)
	annotate!([(13 + 1, pk5_3s_t[13] + 0.02, Plots.text(df[13, :Loc],
		6, :red, :right))])
end

# ╔═╡ 75ce4e64-6958-11eb-333e-dbee637ecdf2
begin
	waic_5_3s_pw_t = waic(lp5_3s_t; pointwise=true)
	scatter(pk5_3s_t, waic_5_3s_pw_t.penalty,
		xlab="PSIS Pareto k", ylab="WAIC penalty", leg=false)
	vline!([0.5])
	for state in [13, 20, 34, 44, 50]
		annotate!([([pk5_3s_t[state] + 0.02], [waic_5_3s_pw_t.penalty[state]],
			Plots.text(df[state, :Loc], 6, :red, :right))])
	end
	plot!()
end

# ╔═╡ 8ec864f0-697d-11eb-029b-93e11b2e247b
waic_5_3s_pw_t

# ╔═╡ eb6b3ee8-697d-11eb-3380-a99e3c132720
pk5_3s_t

# ╔═╡ 5eae5f80-68f4-11eb-2d8e-5db775fbc1c9
pk_plot(pk5_2s_t)

# ╔═╡ 56735994-695b-11eb-28dd-21ff7dfa325d
begin
	waic_5_2s_pw_t = waic(lp5_2s_t; pointwise=true)
	scatter(pk5_2s_t, waic_5_2s_pw_t.penalty,
		xlab="PSIS Pareto k", ylab="WAIC penalty", leg=false)
	vline!([0.5])
	for state in [13, 20, 34, 44, 50]
		annotate!([([pk5_2s_t[state] + 0.01], [waic_5_2s_pw_t.penalty[state] + 0.01],
			Plots.text(df[state, :Loc], 6, :red, :right))])
	end
	plot!()
end

# ╔═╡ c3169ed8-694e-11eb-0616-97d25d8b56e1
md" ## End of clip-07-35s.jl"

# ╔═╡ Cell order:
# ╟─7046ea12-68ec-11eb-18ae-a724ca79a21f
# ╠═837951f4-68ec-11eb-048d-6b27b9055428
# ╠═91489170-68ec-11eb-3611-150d5febf223
# ╠═ace5ed4c-68ec-11eb-05ed-1f1ff2d72e1b
# ╠═4c47f312-68ed-11eb-37de-ff782118d258
# ╠═4c482dfa-68ed-11eb-2db9-a54ab816beb7
# ╠═806d4352-69b7-11eb-00a5-791917fa9d45
# ╠═6210f1ba-68ef-11eb-1d55-653d1c8a147c
# ╠═3a9ff05e-6a16-11eb-2ba0-530fb4bc18f4
# ╠═8f107504-68ed-11eb-002b-372884f598b0
# ╠═cbb9ef56-68ed-11eb-387b-252ff6dc5b21
# ╠═5d216004-69b7-11eb-34d4-8be76d4ae47c
# ╠═87cd0b48-68f2-11eb-3704-af5e98a9c463
# ╠═71a48f7a-68ee-11eb-1b0e-331393e95551
# ╠═b5b487ec-68ee-11eb-3ee7-0d98367028ff
# ╠═4e7a00e8-69b4-11eb-15ef-df1a53282396
# ╠═51d570db-4ff0-49b4-b4c1-8e709a6b1564
# ╠═c2ba84b0-68f2-11eb-3c5a-17e03541dcb1
# ╠═7b051a2e-68f3-11eb-3ee8-e77c639fc390
# ╠═ed4469dc-68f3-11eb-3b0f-f54305283306
# ╠═2c4b5ffc-68f4-11eb-21d6-7d2db834d3f8
# ╠═3766d148-695b-11eb-19dc-0b25e6704b3d
# ╠═aedb425c-68f4-11eb-319d-a958c2aacde1
# ╠═75ce4e64-6958-11eb-333e-dbee637ecdf2
# ╠═8ec864f0-697d-11eb-029b-93e11b2e247b
# ╠═eb6b3ee8-697d-11eb-3380-a99e3c132720
# ╠═5eae5f80-68f4-11eb-2d8e-5db775fbc1c9
# ╠═56735994-695b-11eb-28dd-21ff7dfa325d
# ╟─c3169ed8-694e-11eb-0616-97d25d8b56e1


using Plexifications
using Base.Test: @test, @test_throws

@test_throws ErrorException begin
    inv(BiNilplex(0, 0.0))
end

@test begin
    !iszerodivisor(BiNilplex(1,2,3,4))
end

@test begin
    iszerodivisor(BiNilplex(0,0,3,4))
end

@test begin
    iszerodivisor(BiNilplex(0,3,0,4))
end

@test begin
    a = Plexifications.randomBigFloat()
    isreal(BiNilplex(a))
end

@test begin
    a = 1
    !isreal(BiNilplex(a, a))
end

@test begin
    z = random(BiNilplex{Int})
    z == +(z)
end

@test begin
    a = Plexifications.randomBigFloat()
    real(BiNilplex(a)) == a
end

@test begin
    io = IOBuffer()
    show(io, BiNilplex(1, 2, 3, 4))
    l = String(take!(io))
    r = "[1: 1, A: 2, B: 3, AB: 4]"
    l == r
end

@test begin
    x = random(BiNilplex{BigFloat})
    y = random(BiNilplex{BigFloat})
    x != y
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(BiNilplex{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(BiNilplex{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(Nlplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = x * y
    r = y * x
    l == r
end

@test begin
    x = random(Nlplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = x * y
    r = y * x
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(Nlplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    z = random(BiNilplex{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    z = random(BiNilplex{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    barA(barA(x)) == x
end

@test begin
    x = random(BiNilplex{BigInt})
    barB(barB(x)) == x
end

@test begin
    x = random(BiNilplex{BigInt})
    barB(barA(x)) == barA(barB(x))
end

@test begin
    x = random(BiNilplex{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = barA(x * y)
    r = barA(x) * barA(y)
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = barA(x + y)
    r = barA(x) + barA(y)
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = barA(x - y)
    r = barA(x) - barA(y)
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = barB(x * y)
    r = barB(x) * barB(y)
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = barB(x + y)
    r = barB(x) + barB(y)
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = barB(x - y)
    r = barB(x) - barB(y)
    l == r
end

@test begin
    x = random(BiNilplex{Rational{BigInt}})
    y = random(BiNilplex{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(BiNilplex{Rational{BigInt}})
    y = random(BiNilplex{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(BiNilplex{Rational{BigInt}})
    y = random(BiNilplex{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    z = random(BiNilplex{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(BiNilplex{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    z = random(BiNilplex{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    z = random(BiNilplex{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    z = random(BiNilplex{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    n = norm(x)
    n > big(0) || n == big(0)
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = norm(x * y)
    r = norm(x) * norm(y)
    l == r
end

@test_throws ErrorException begin
    inv(BiNilplex(0, 0.0))
end

@test_throws ErrorException begin
    random(BiNilplex{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(BiNilplex{Int})
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(BiNilplex{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(BiNilplex{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(BiNilplex{BigInt})
    z == z + zero(BiNilplex{BigInt})
end

@test begin
    z = random(BiNilplex{BigInt})
    z == z * one(BiNilplex{BigInt})
end

@test begin
    z = random(BiNilplex{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(BiNilplex{BigInt})
    z == z * one(z)
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = commutator(x, y)
    r = zero(BiNilplex{BigInt})
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = alternatorL(x, y)
    r = zero(BiNilplex{BigInt})
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = alternatorR(x, y)
    r = zero(BiNilplex{BigInt})
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    l = flexator(x, y)
    r = zero(BiNilplex{BigInt})
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    z = random(BiNilplex{BigInt})
    l = associator(x, y, z)
    r = zero(BiNilplex{BigInt})
    l == r
end

@test begin
    x = random(BiNilplex{BigInt})
    y = random(BiNilplex{BigInt})
    z = random(BiNilplex{BigInt})
    l = jacobiator(x, y, z)
    r = 3 * x * y * z
    l == r
end

@test begin
    l = BiNilplex(Nlplex(1,2))
    r = BiNilplex(1,2,0)
    l == r
end

@test begin
    l = BiNilplex(1, 2, 3.0)
    r = BiNilplex(1.0, 2.0, 3.0)
    l == r
end

@test begin
    l = BiNilplex(1, 2, 3, 4.0)
    r = BiNilplex(1.0, 2.0, 3.0, 4.0)
    l == r
end


mutable struct Sigmoid <: DiffableFunction
    grad_field :: GradField
end

function _sigmoid(x::T) where {T <: Real}
    one(T) / (one(T) + exp(-x))
end

@inline forward(::Sigmoid, x) = _sigmoid(x)

function backward(f::Sigmoid, gy::Variable{T}) where {T <: Real}
    y = f.grad_field.outputs[1]
    return (y * (1 - y)) * gy
end

function backward(f::Sigmoid, gy::Variable{T}) where {T <: AbstractArray}
    y = f.grad_field.outputs[1]
    @. return (y * (1 - y)) * gy
end

sigmoid(x::Variable) = Sigmoid(GradField())(x)

Base.broadcasted(::typeof(sigmoid), x::Variable)  = Broadcasting(_sigmoid, Sigmoid(GradField()))(x)


#typealias MultiindexSet{T} Array{T}
#typ MultiindexSet{T} <: AbstractArray{T, 2}

function multiindex{T<:Integer}(m::Integer,p::T)
    if m==0
        I = zeros(T,1,0)
        return I
    end
    I = zeros(T,0,m)
    I1 = multiindex(m-1, p)
    ord = reshape(sum(I1,2), size(I1,1))
    for q=0:p
        ind=map(o -> (o<=q), ord)
        n=sum(ind)
        Iq = [I1[ind,:] (q-sum(I1[ind,:],2))]
        I = [I; Iq]
    end
    return I
end

function multiindex_order{T<:Integer}(I::AbstractArray{T})
  return vec(sum(I,2))
end

export
    multiindex,
    multiindex_order


# static calculations used by the full model (static and dynamic)
function L_tilde(g, z_hat, Ω, E, S, parameters)
    @unpack N, θ, κ, χ, ζ = parameters
    return Ω * ((N-1) * z_hat^(-θ)*κ + ζ*(S + E / χ)) # (33) (nests the stationary case with E = δ)
end

function L_tilde_x(z_hat, Ω, parameters)
    @unpack N, θ, κ = parameters
    return Ω * (N - 1) * z_hat^(-θ) * κ # (34)
end

L_tilde_E(Ω, E, parameters) = Ω * parameters.ζ * E / parameters.χ # (35) (nests the stationary case with E = δ)
L_tilde_a(Ω, S, parameters) = Ω * parameters.ζ * S # (36)

function S(g, parameters)
    @unpack θ, μ, υ = parameters
    return θ * (g - μ - θ * υ^2 /2) # (32)
end

w(z_bar, parameters) = (parameters.σ - 1)/(parameters.σ) * z_bar # (C.13)

function x(w, parameters)
    @unpack Theta, ζ, η = parameters
    return ζ * (1 - η + η * Theta / w) # (C.14)
end

function λ_ii(z_hat, parameters)
    @unpack N, σ, θ, d = parameters
    return  1/(1 + (N-1)*z_hat^(σ-1-θ)*d^(1-σ)) # (C.79)
end

c(L_tilde, Ω, z_bar) = (1 - L_tilde)*z_bar # (C.15)
function entry_residual(v_1, Ξ₁, parameters)
    @assert parameters.η ≈ 0.
    Ξ₁*v_1 - parameters.ζ*(1- parameters.χ)/parameters.χ # (56)
end
π_min(L_tilde, z_bar, parameters) = (1 - L_tilde) / ((parameters.σ-1)*z_bar^(parameters.σ-1)) # (38)

function π_rat(z_hat, parameters)
    @unpack θ, σ, d, κ, N = parameters
    return θ/(1 + θ - σ) + (N - 1) * d^(1-σ) * (σ - 1)*z_hat^(σ - 1 - θ)/(1 + θ - σ) # (C.17)
end

function z_bar(z_hat, Ω, parameters)
    @unpack θ, σ, N, d = parameters
    return (Ω*θ/(1 + θ - σ) * (1 + (N - 1)*d^(1 - σ)*z_hat^(σ - 1 - θ)))^(1/(σ-1)) # (C.11)
end

# subset of the above used in the dynamic model
function static_equilibrium(Ξ₁, v_1, g, z_hat, E, Ω, z, parameters)
    @unpack θ, σ, N, d, χ, ζ, κ = parameters
    S_t = S(g, parameters)
    L_tilde_t = L_tilde(g, z_hat, Ω, E, S_t, parameters)
    z_bar_t = z_bar(z_hat, Ω, parameters)
    w_t = w(z_bar_t, parameters)
    π_min_t = (1 - L_tilde_t) / ((σ-1)*z_bar_t^(σ-1))  # (C.20)

    i_vectorized = z .>= log(z_hat) # Vectorized indicator function
    π_t = π_min_t * (1.0.+(N-1)*d^(1-σ)*i_vectorized) - (N-1)*κ*exp.(-(σ-1).*z).*i_vectorized  # (39)
    entry_residual_t = entry_residual(v_1, Ξ₁, parameters)

    return (S = S_t, L_tilde = L_tilde_t, z_bar = z_bar_t, π_min = π_min_t, π = π_t, entry_residual = entry_residual_t,
            w = w_t)
end

consumption_equivalent(U, U_old, parameters) = isapprox(parameters.γ, 1.0) ?  exp(parameters.ρ*(U - U_old)) - 1 : (U/U_old)^(1.0/(1.0 - parameters.γ)) -1# H.23 main paper
ACR(new_λ, old_λ, params) = (-1/params.θ)*log(new_λ/old_λ)


# Packages required
using Hwloc
using .Threads
using Parameters
using LinearAlgebra
using Calculus
using DifferentialEquations
using ForwardDiff
using Distributions
using StatsBase
using Random
using SymPy
using PyPlot
using DataFrames
using CSV


include("death_rate_zero.jl")
include("birth_rate_zero.jl")
include("equal_rates.jl")
include("unequal_rates.jl")

"""
    gradient_hessian(η, i, j, t)

Compute the gradient and Hessian of the log-transition probability of a simple
birth and death process evaluated at the point ``η = (λ, μ)^{\\prime}``.
Variable `i` is the initial population size, `j` is the final population size,
and `t` is the elapsed time.
"""
function gradient_hessian(
  η::Vector{F},
  i::I,
  j::I,
  t::F
)::Tuple{Vector{F}, Symmetric{F, Matrix{F}}} where {
  F <: AbstractFloat,
  I <: Integer
}
  ϵ = floatmin(F)

  if t < ϵ
    fill(zero(F), 2), Symmetric(fill(zero(F), (2, 2)))
  elseif (η[1] < ϵ) && (η[2] < ϵ)
    if i == j
      fill(- i * t, 2), Symmetric([zero(F) (i * t)^2; (i * t)^2 zero(F)])
    else
      fill(F(NaN), 2), Symmetric(fill(F(NaN), (2, 2)))
    end
  elseif η[1] < ϵ
    gradient_hessian_λ_zero(η[2], i, j, t)
  elseif η[2] < ϵ
    gradient_hessian_μ_zero(η[1], i, j, t)
  elseif η[2] ≈ η[1]
    # TODO: Improve numerical accuracy of formulas
    if (F === BigFloat) || (η[1] * t > 0.001)
      gradient_hessian_equal_rates(η[1], F(i), F(j), t)
    else
      setprecision(BigFloat, 256) do
        l = BigFloat(η[1])
        a = BigFloat(i)
        b = BigFloat(j)
        s = BigFloat(t)

        g, H = gradient_hessian_equal_rates(l, a, b, s)

        F[g[1], g[2]], Symmetric(F[H[1, 1] H[2, 1]; H[1, 2] H[2, 2]])
      end
    end
  else
    if (F === BigFloat) || (log(η[1] / η[2]) / ((η[1] - η[2]) * t) < 1_000)
      gradient_hessian_unequal_rates(η, F(i), F(j), t)
    else
      setprecision(BigFloat, 256) do
        l = BigFloat(η[1])
        m = BigFloat(η[2])
        a = BigFloat(i)
        b = BigFloat(j)
        s = BigFloat(t)

        g, H = gradient_hessian_unequal_rates([l, m], a, b, s)

        F[g[1], g[2]], Symmetric(F[H[1, 1] H[2, 1]; H[1, 2] H[2, 2]])
      end
    end
  end
end

function gradient_hessian(
  η::Vector{F},
  x::ObservationDiscreteTimeEven
)::Tuple{Vector{F}, Symmetric{F, Matrix{F}}} where {
  F <: AbstractFloat
}
  ∇ = fill(zero(Float64), 2)
  H = fill(zero(Float64), (2, 2))

  for n = 1:size(x.state, 2), s = 2:size(x.state, 1)
    a, b = gradient_hessian(η, x.state[s - 1, n], x.state[s, n], x.u)
    ∇ .+= a
    H .+= b
  end

  ∇, Symmetric(H)
end


using FileIO
using LinearAlgebra
using Parameters
include("network_sims.jl")
include("random_matrices.jl")

output_dir = "results2-analysis"

N = 3000
g = 1.5
etas = [-0.3, 0.0, 0.3]
s = 0.0
J_seed = 5000
h0_seed = 8000

dt = 0.01
t_bound = 300.0
wait = 200.0
savedunits = 10
timesteps = round(Int64, (t_bound - wait) / dt)
ts = dt .* collect(1:timesteps)

activity = Array{Float64,3}(undef, timesteps, savedunits, length(etas))
eigval_parts = Array{Float64,3}(undef, N, 2, length(etas))

for (i, eta) in enumerate(etas)
    J = generate_J(N, eta, seed=J_seed)
    h0 = generate_nv(N, Float64, seed=h0_seed)
    A = sim_activity(J, g, h0, s=s, dt=dt, t_bound=t_bound)
    activity[:,:,i] = view(A, 1:savedunits,
            round(Int, wait / dt) + 1:round(Int, t_bound / dt))'

    eigval_parts[:,1,i] = real(eigvals(J))
    eigval_parts[:,2,i] = imag(eigvals(J))
end

fn = "example_activity_and_eigvals_results2.jld2"
out = Dict{String,Any}()
@pack! out = ts, activity, eigval_parts
save(joinpath(output_dir, fn), out)


#=
Annealed importance sampling (AIS) to estimate the partition function (and hence
the log-likelihood).
A nice explanation of AIS in general can be found in Goodfellow et al Deep Learning book.
Salakhutdinov et al (10.1145/1390156.1390266, http://www.cs.utoronto.ca/~rsalakhu/papers/bm.pdf)
discusses AIS for RBMs specifically.

AIS tends to understimate the log of the partition function (in probability).
In contrast, Reverse AIS estimator (RAISE) can be used to obtain a stochastic upper bound.
See http://proceedings.mlr.press/v38/burda15.html.
Combining the two we can "sandwiches" the true value to have an idea if the Monte Carlo
chains have converged.

Addendum: I think Burda's paper has a typo. The correct expression for the weights In
reverse AIS (which I use here) can be found in Upadhya et al 2015, Equation 10
(https://link.springer.com/chapter/10.1007/978-3-319-26535-3_62).

Bonus: A discussion of estimating partition function in RBMs, comparing several algorithms:

https://www.sciencedirect.com/science/article/pii/S0004370219301948

For a variant or RAISE: https://arxiv.org/abs/1511.02543
=#

"""
    ais(rbm0, rbm1, v0, βs)

Provided `v0` is an equilibrated sample from `rbm0`, returns `F` such that `mean(exp.(F))` is
an unbiased estimator of `Z1/Z0`, the ratio of partition functions of `rbm1` and `rbm0`.

!!! tip Use [`logmeanexp`](@ref)
    `logmeanexp(F)`, using the function `logmeanexp`[@ref] provided in this package,
    tends to give a better approximation of `log(Z1) - log(Z0)` than `mean(F)`.
"""
function ais(rbm0::RBM, rbm1::RBM, v::AbstractArray, βs::AbstractVector)
    @assert issorted(βs) && 0 == first(βs) ≤ last(βs) == 1
    F = free_energy(rbm0, v)
    for β in βs
        if iszero(β) || isone(β)
            continue
        else
            rbm = anneal(rbm0, rbm1; β)
            F -= free_energy(rbm, v)
            v = sample_v_from_v(rbm, v)
            F += free_energy(rbm, v)
        end
    end
    F -= free_energy(rbm1, v)
    return F
end

function ais(rbm0::RBM, rbm1::RBM, v0::AbstractArray; nbetas::Int=2)
    βs = range(0, 1, nbetas)
    return ais(rbm0, rbm1, v0, βs)
end

"""
    aise(rbm, [βs]; [nbetas], init=rbm.visible, nsamples=1)

AIS estimator of the log-partition function of `rbm`. It is recommended to fit `init` to
the single-site statistics of `rbm` (or the data).

!!! tip Use large `nbetas`
    For more accurate estimates, use larger `nbetas`. It is usually better to have
    large `nbetas` and small `nsamples`, rather than large `nsamples` and small `nbetas`.
"""
function aise(rbm::RBM, βs::AbstractVector{<:Real}; init::AbstractLayer=rbm.visible, nsamples::Int=1)
    rbm0 = anneal_zero(init, rbm)
    v0 = sample_from_inputs(init, Falses(size(init)..., nsamples))
    F = ais(rbm0, rbm, v0, βs)
    return F .+ log_partition_zero_weight(rbm0)
end

"""
    raise(rbm::RBM, βs; v, init)

Reverse AIS estimator of the log-partition function of `rbm`.
While `aise` tends to understimate the log of the partition function, `raise` tends to
overstimate it. `v` must be an equilibrated sample from `rbm`.

!!! tip Use [`logmeanexp`](@ref)
    If `F = raise(...)`, then `-logmeanexp(-F)`, using the function `logmeanexp`[@ref]
    provided in this package, tends to give a better approximation of `log(Z)` than `mean(F)`.

!!! tip Sandwiching the log-partition function
    If `Rf = aise(...)`, `Rr = raise(...)` are the AIS and reverse AIS estimators, we have the
    stochastic bounds `logmeanexp(Rf) ≤ log(Z) ≤ -logmeanexp(-Rr)`.
"""
function raise(rbm::RBM, βs::AbstractVector; v::AbstractArray, init::AbstractLayer=rbm.visible)
    rbm0 = anneal_zero(init, rbm)
    F = ais(rbm, rbm0, v, βs)
    return log_partition_zero_weight(rbm0) .- F
end

aise(rbm::RBM; nbetas::Int=10000, kw...) = aise(rbm, range(0, 1, nbetas); kw...)
raise(rbm::RBM; nbetas::Int=10000, kw...) = raise(rbm, range(0, 1, nbetas); kw...)

"""
    anneal(rbm0, rbm1; β)

Returns an RBM that interpolates between `rbm0` and `rbm1`.
Denoting by `E0(v, h)` and `E1(v, h)` the energies assigned by `rbm0` and `rbm1`,
respectively, the returned RBM assigns energies given by:

    E(v,h) = (1 - β) * E0(v) + β * E1(v, h)
"""
function anneal(rbm0::RBM, rbm1::RBM; β::Real)
    vis = anneal(rbm0.visible, rbm1.visible; β)
    hid = anneal(rbm0.hidden, rbm1.hidden; β)
    w = (1 - β) * rbm0.w + β * rbm1.w
    return RBM(vis, hid, w)
end

anneal(init::Binary, final::Binary; β::Real) = Binary(
    oftype(final.θ, (1 - β) * init.θ + β * final.θ)
)
anneal(init::Spin, final::Spin; β::Real) = Spin(
    oftype(final.θ, (1 - β) * init.θ + β * final.θ)
)
anneal(init::Potts, final::Potts; β::Real) = Potts(
    oftype(final.θ, (1 - β) * init.θ + β * final.θ)
)
anneal(init::Gaussian, final::Gaussian; β::Real) = Gaussian(
    oftype(final.θ, (1 - β) * init.θ + β * final.θ),
    oftype(final.γ, (1 - β) * init.γ + β * final.γ)
)
anneal(init::ReLU, final::ReLU; β::Real) = ReLU(
    oftype(final.θ, (1 - β) * init.θ + β * final.θ),
    oftype(final.γ, (1 - β) * init.γ + β * final.γ)
)
anneal(init::dReLU, final::dReLU; β::Real) = dReLU(
    oftype(final.θp, (1 - β) * init.θp + β * final.θp),
    oftype(final.θn, (1 - β) * init.θn + β * final.θn),
    oftype(final.γp, (1 - β) * init.γp + β * final.γp),
    oftype(final.γn, (1 - β) * init.γn + β * final.γn)
)
anneal(init::pReLU, final::pReLU; β::Real) = pReLU(
    oftype(final.θ, (1 - β) * init.θ + β * final.θ),
    oftype(final.γ, (1 - β) * init.γ + β * final.γ),
    oftype(final.Δ, (1 - β) * init.Δ + β * final.Δ),
    oftype(final.η, (1 - β) * init.η + β * final.η)
)
anneal(init::xReLU, final::xReLU; β::Real) = xReLU(
    oftype(final.θ, (1 - β) * init.θ + β * final.θ),
    oftype(final.γ, (1 - β) * init.γ + β * final.γ),
    oftype(final.Δ, (1 - β) * init.Δ + β * final.Δ),
    oftype(final.ξ, (1 - β) * init.ξ + β * final.ξ)
)

anneal_zero(init::AbstractLayer, rbm1::RBM) = RBM(init, anneal_zero(rbm1.hidden), Zeros(rbm1.w))

anneal_zero(l::Binary) = Binary(zero(l.θ))
anneal_zero(l::Spin) = Spin(zero(l.θ))
anneal_zero(l::Potts) = Potts(zero(l.θ))
anneal_zero(l::Gaussian) = Gaussian(zero(l.θ), l.γ)
anneal_zero(l::ReLU) = ReLU(zero(l.θ), l.γ)
anneal_zero(l::dReLU) = dReLU(zero(l.θp), zero(l.θn), l.γp, l.γn)
anneal_zero(l::pReLU) = pReLU(zero(l.θ), l.γ, zero(l.Δ), l.η)
anneal_zero(l::xReLU) = xReLU(zero(l.θ), l.γ, zero(l.Δ), l.ξ)

"""
    log_partition_zero_weight(rbm)

Log-partition function of a zero-weight version of `rbm`.
"""
log_partition_zero_weight(rbm::RBM) = -free_energy(rbm.visible) - free_energy(rbm.hidden)

"""
    logmeanexp(A; dims=:)

Computes `log.(mean(exp.(A); dims))`, in a numerically stable way.
"""
function logmeanexp(A::AbstractArray; dims=:)
    R = logsumexp(A; dims)
    N = length(A) ÷ length(R)
    return R .- log(N)
end

"""
    logvarexp(A; dims=:)

Computes `log.(var(exp.(A); dims))`, in a numerically stable way.
"""
function logvarexp(
    A::AbstractArray; dims=:, corrected::Bool=true, logmean=logmeanexp(A; dims)
)
    R = logsumexp(2logsubexp.(A, logmean); dims)
    N = length(A) ÷ length(R)
	if corrected
		return R .- log(N - 1)
    else
        return R .- log(N)
    end
end

"""
    logstdexp(A; dims=:)

Computes `log.(std(exp.(A); dims))`, in a numerically stable way.
"""
function logstdexp(
    A::AbstractArray; dims=:, corrected::Bool=true, logmean=logmeanexp(A; dims)
)
    return logvarexp(A; dims, corrected, logmean) / 2
end


struct HailstoneSeq{T<:Integer}
	start::T
end

Base.eltype(::HailstoneSeq{T}) where T = T

Base.start(hs::HailstoneSeq) = (-1, hs.start)
Base.done(::HailstoneSeq, state) = state == (1, 4)
function Base.next(::HailstoneSeq, state)
	_, s2 = state
	s1 = s2
	if iseven(s2)
		s2 = s2 ÷ 2
	else
		s2 = 3s2 + 1
	end
	return s1, (s1, s2)
end

function Base.length(hs::HailstoneSeq)
	r = 0
	for _ in hs
		r += 1
	end
	return r
end

function Base.show(io::IO, hs::HailstoneSeq)
	f5 = collect(Iterators.take(hs, 5))
	print(io, "HailstoneSeq(", join(f5, ", "), "...)")
end

hs = HailstoneSeq(27)
println("Collection of the Hailstone sequence from 27: $hs")
cl = collect(hs)
println("First 5 elements: ", join(cl[1:5], ", "))
println("Last 5 elements: ", join(cl[end-4:end], ", "))

Base.isless(h::HailstoneSeq, s::HailstoneSeq) = length(h) < length(s)
println("The number with the longest sequence under 100,000 is: ", maximum(HailstoneSeq.(1:100_000)))


cd(@__DIR__); include("setups/simpletree.jl")
gr(dpi = 200)

##
ind_eig = [5, 6, 11, 18, 21, 25, 38]
plot(layout = Plots.grid(1, 7), size = (1400, 500))
for i = 1:length(ind_eig)
    l = ind_eig[i]
    w = 𝚽[:, l]
    scatter!(X[:, 1], X[:, 2], marker_z = w, ms = 5, c = :viridis, subplot = i, mswidth = 0)
    plot!(frame = :none, cbar = false, grid = false, legend = false, subplot = i)
    plot!(title = latexstring("\\mathbf{\\phi}_{", l - 1, "}"), titlefont = 18, subplot = i)
end
plot!(subplot = 10, grid = false, frame = :none)
plt = current()
savefig(plt, "../figs/simpletree_f_rngwf_top21_ref_eigenvectors.png")


module Utils

export batched_outer, batched_kronecker, add_dim, gram_matrix

using Flux

"""
# Parameters

- `A`: array of shape (`m`, `num_batch`)
- `B`: array of shape (`n`, `num_batch`)

# Return

- `C`: array of shape (`m`, `n`, `num_batch`), where `C[:, :, i] = A[:, i] * B[:, i]'`

"""
function batched_outer(A, B)
    @assert ndims(A) == 2
    @assert ndims(B) == 2
    size(A, 2) == size(B, 2) || throw(DimensionMismatch("batch size mismatch"))
    num_batch = size(A, 2)
    A′ = reshape(A, :, 1, num_batch)
    B′ = reshape(B, 1, :, num_batch)
    Flux.batched_mul(A′, B′)
end

"""
mat1: (m, n, batch)
mat2: (j, k, batch)
"""
function batched_kronecker(mat1, mat2)
    @assert ndims(mat1) == 3
    @assert ndims(mat2) == 3
    @assert size(mat1, 3) == size(mat2, 3)
    num_batch = size(mat1, 3)
    mat1_flatten = reshape(mat1, :, 1, num_batch)
    mat2_flatten = reshape(mat2, 1, :, num_batch)
    curried_reshape(new_shape...) = (m)-> reshape(m, new_shape...)
    curried_permdims(perms...) = (m) -> permutedims(m, perms)
    Flux.batched_mul(
        mat1_flatten,
        mat2_flatten,
        ) |>
    curried_reshape(size(mat1)[1:end-1]..., size(mat2)[1:end-1]..., num_batch) |>
    curried_permdims(3, 1, 4, 2, 5) |>
    curried_reshape(size(mat1, 1)*size(mat2, 1), size(mat1, 2)*size(mat2, 2), num_batch)
end

"""
https://stackoverflow.com/a/59849291/3908797
"""
add_dim(x::AbstractArray{T, N}) where {T, N} = reshape(x, Val(N+1))

"""
- `A`: (d, m, batch)
- `B`: (d, n, batch)

return C: (m, n, batch)
"""
gram_matrix(A, B) = Flux.batched_mul(permutedims(A, (2, 1, 3)), B)

"""
`labels`: (m, batch)
"""
function onehot(labels)
    num_batch = size(labels, 2)
    encoded = []
    for i in 1:num_batch
        uniq_label = unique(labels[:, i])
        encoder(l) = Flux.onehot(l, uniq_label)
        push!(
            encoded,
            vcat(encoder.(labels[:, i])...)
        )
    end
    hcat(encoded...)
end

end


using RunTests
using Base.Test

@testmodule ExtraStatisticsTest begin
  using MLLibrary

  function test_meanColumn()
    A = [1 2 3 4;5 6 7 8]
    @test MLLibrary.Statistics.meanColumn(A)==[3.0 4.0 5.0 6.0]
    # @test 1==0
  end

  function test_stdColumn()
    A = [1 2 3 4;5 6 7 8]
    @test MLLibrary.Statistics.stdColumn(A)==[2.8284271247461903 2.8284271247461903 2.8284271247461903 2.8284271247461903]
    # @test 1==0
  end

  function test_sumColumn()
    A = [1 2 3 4;5 6 7 8]
    @test MLLibrary.Statistics.sumColumn(A)==[6 8 10 12]
    # @test 1==0
  end

  function test_getMinIndices()
    A=[1;2;-1;5;3;2;9;-3; 10]
    @test MLLibrary.Statistics.getMinIndices(A)==[8]
    @test MLLibrary.Statistics.getMinIndices(A,2)==[8, 3]
    @test MLLibrary.Statistics.getMinIndices(A,3)==[8, 3, 1]
    @test MLLibrary.Statistics.getMinIndices(A,4)==[8, 3, 1, 2]
  end

  function test_getMajority()
    A=[1;2;-1;5;3;2;9;-3;10]
    @test MLLibrary.Statistics.getMajority(A)==2
    A=[1;2;1;5;3;2;9;-3;10]
    maj=MLLibrary.Statistics.getMajority(A)
    @test maj==1 || maj==2
    A=[1]
    @test MLLibrary.Statistics.getMajority(A)==1

  end
end


# # "Point exponential decay" time-stepper convergence test

using PyPlot
using Oceananigans

# Define a few utilities for running tests and unpacking and plotting results
include("ConvergenceTests/ConvergenceTests.jl")

using .ConvergenceTests
using .ConvergenceTests.PointExponentialDecay: run_test

defaultcolors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
removespine(side) = gca().spines[side].set_visible(false)
removespines(sides...) = [removespine(side) for side in sides]

""" Run advection-diffusion test for all Nx in resolutions. """
function run_convergence_test(stop_time, proposal_Δt, timestepper)

    # Adjust time-step
    stop_iterations = [round(Int, stop_time / dt) for dt in proposal_Δt]
                 Δt = [stop_time / stop_iter for stop_iter in stop_iterations]

    # Run the tests
    results = [run_test(timestepper=timestepper, Δt=dt, stop_iteration=stop_iter)
               for (dt, stop_iter) in zip(Δt, stop_iterations)]

    return results, Δt
end

unpack_errors(results) = map(r -> r.L₁, results)

stop_time = 3
Δt = 10 .^ range(-3, 0, length=30)  # Equally spaced in log space.
ab2_results, Δt = run_convergence_test(stop_time, Δt, :QuasiAdamsBashforth2)
rk3_results, Δt = run_convergence_test(stop_time, Δt, :RungeKutta3)

ab2_L₁ = unpack_errors(ab2_results)
rk3_L₁ = unpack_errors(rk3_results)

fig, axs = subplots()

loglog(Δt, ab2_L₁, "o--", alpha=0.8, linewidth=1, label="Measured error, Quasi-second-order Adams-Bashforth")
loglog(Δt, rk3_L₁, "^-.", alpha=0.8, linewidth=1, label="Measured error, Third-order Runge-Kutta")

# Guide line to confirm second-order scaling
loglog(Δt, ab2_L₁[1] .* Δt / Δt[1], "k-", alpha=0.4, linewidth=2, label=L"\sim \Delta t")
loglog(Δt, rk3_L₁[1] .* (Δt / Δt[1]).^3, "r-", alpha=0.4, linewidth=2, label=L"\sim \Delta t^3")

xlabel(L"\Delta t")
ylabel("\$L_1\$-norm of \$ | c_\\mathrm{simulated}(t) - \\mathrm{e}^{-t} | \$")
title("Oceananigans time-stepper convergence for \$ c(t) = \\mathrm{e}^{-t} \$")
removespines("top", "right")
legend()

filepath = joinpath(@__DIR__, "figs", "point_exponential_decay_time_stepper_convergence.png")
savefig(filepath, dpi=480)

test_rate_of_convergence(ab2_L₁, Δt, data_points=9, expected=1.0, atol=0.01)
test_rate_of_convergence(rk3_L₁, Δt, data_points=9, expected=3.0, atol=0.01)


abstract type Quadrature end
abstract type Euler <: Quadrature end
struct euler <: Euler end
abstract type RK4 <: Quadrature end
struct rk4 <: RK4 end

function integrate_dynamics(model, x0, us, ts, quadrature = euler())
    xs = Array{typeof(x0[1]),1}(undef, length(x0) * length(ts))
    integrate_dynamics!(xs, model, x0, us, ts, quadrature)
    return xs
end

"""
    integrate_dyanmics!(xs, model, x0, us, ts, quadrature)

Integrates the trajectory provided control inputs. Modifies `xs` in place.

# Inputs

* `xs::Vector{Float64}` - preallocated memory consisting of concatenated state vectors to be populated with simulation states for each element of `ts`
* `model<:AbstractModel` - model to be simulated
* `x0::Vector{Float64}` - initial state
* `us::Vector{Float64}` - concatenated control inputs prescribed for the simulation for each element of `ts`
* `ts::Vector{Float64}` - vector of times at which dynamics are evaluated
* `quadrature<:Quadrature` - what kind of quadrature is to be used for integration; `euler()` and `rk4()` currently supported

"""
function integrate_dynamics!(xs, model, x0, us, ts, quadrature::Euler)
    nx = model.nx # number of states
    @assert nx == length(x0) "length of x0 and model.nx inconsistent"
    nu = model.nu # number of controls
    @assert nu == Int(length(us) / length(ts)) "length of us and length of ts and model.nu inconsistent"
    @assert length(xs) == length(x0) * length(ts) "length of states x not consistent with timesteps and x0"

    # initial knot is the initial guess
    xs[1:nx] .= x0

    # loop through time
    for i=1:length(ts)-1
        Δt = ts[i+1] - ts[i]
        x = xs[(i-1) * nx + 1 : i * nx]
        u = us[(i-1) * nu + 1 : i * nu]
        xdot = dynamics(model, x, u)
        xs[i * nx + 1 : (i+1) * nx] .= x + xdot * Δt
    end

    return nothing
end

function integrate_dyanmics!(xs, model, x0, us, ts, quadrature::RK4)
    nx = model.nx # number of states
    @assert nx == length(x0) "length of x0 and model.nx inconsistent"
    nu = model.nu # number of controls
    @assert nu == Int(length(us) / length(ts)) "length of us and length of ts and model.nu inconsistent"
    @assert length(xs) == length(x0) * length(ts) "length of states x not consistent with timesteps and x0"

    # initial knot is the initial guess
    xs[1:nx] .= x0

    # loop through time
    for i=1:length(ts)-1
        Δt = ts[i+1] - ts[i]
        x = xs[(i-1) * nx + 1 : i * nx]
        u = us[(i-1) * nu + 1 : i * nu]
        u_next = us[i * nu + 1 : (i+1) * nu]
        F1 = dynamics(model, x, u)
        F2 = dynamics(model, x + .5 * Δt .* F1, u)
        F3 = dynamics(model, x + .5 * Δt .* F2, u_next)
        F4 = dynamics(model, x +  Δt .* F3, u_next)
        xs[i * nx + 1 : (i+1) * nx] .= x + Δt/6 .* (F1 + 2 .* F2 + 2 .* F3 + F4)
    end

    return nothing
end


# Step 1
struct Fib
    n::Int
end

# Step 2
function Base.iterate(fib::Fib, state=(i=1, vals=(1, 1)))
    state.i > fib.n && return nothing

    if state.i <= 2
        return 1, (i=state.i + 1, vals=(1, 1))
    end

    aₙ₋₂, aₙ₋₁ = state.vals
    aₙ = aₙ₋₂ + aₙ₋₁

    aₙ, (i=state.i + 1, vals=(aₙ₋₁, aₙ))
end

# Step 3
Base.length(fib::Fib) = fib.n

# Step 4
Base.eltype(::Type{Fib}) = Int


### A Pluto.jl notebook ###
# v0.11.4

using Markdown
using InteractiveUtils

# ╔═╡ 96694b30-f034-11ea-1e07-f10a958db49f
md""" So far we have discussed limits in Calculus of various types, continuity, mean value theorem, etc..

In this unit we will discuss derivatives and their applications.

Given two points $(x_1,y_1), (x_2,y_2)$ we can calculate the slope $m$ between them:

$$m = \frac{y_2-y_1}{x_2-x_1}$$
"""

# ╔═╡ 6271a790-f035-11ea-116b-11e6dcef81b1
md"""We can calculate slope of linear functions, ever since we took our first Algebra course.

But what about "non-linear" fuctions?

Let's consider our friendly quadradic fuction $f(x) = x^2$.

Here's a graph:"""

# ╔═╡ 3283a680-f050-11ea-13da-6f5150506375
md"What about slope between two points on the graph?"

# ╔═╡ db616880-f052-11ea-25e3-1bc2914021fb
md"""Hmm. That doesn't look like that'll convey useful information.

What about if we look at two closer points?"""

# ╔═╡ e76835ee-f039-11ea-002c-c5eea484adf4
using Gadfly

# ╔═╡ 621a3160-f038-11ea-007d-c7b61f010b89
begin
	p = plot((z)->z^2, -1.0,1.0)
	#push!(p, layer(x=[slider], y =[slider^2], size=[.05],color=[colorant"gold"]))
	push!(p, Coord.cartesian(fixed=true))
end

# ╔═╡ 8d3b5ff0-f050-11ea-1d4f-173996efe68b
begin
	p1 = plot((z)->z^2, -1.0,1.0)
	push!(p1,layer((z)->(.416*z+.4583),-.5, .92))
	push!(p1, Coord.cartesian(fixed=true))
	push!(p1,Guide.title("Slope between (-0.5,.25) and (.92,.85)"))
end

# ╔═╡ 333c9510-f036-11ea-3e76-d1005102754d
begin
	p2 = plot((z)->z^2, -1.5,1.5)
	push!(p2,layer((z)->(1.6911*z-.7058),.5, 1.5))
	push!(p2, Coord.cartesian(fixed=true))
	push!(p2,Guide.title("Slope between (.75,.5625) and (.92,.85)"))
end

# ╔═╡ Cell order:
# ╟─e76835ee-f039-11ea-002c-c5eea484adf4
# ╟─96694b30-f034-11ea-1e07-f10a958db49f
# ╟─6271a790-f035-11ea-116b-11e6dcef81b1
# ╟─621a3160-f038-11ea-007d-c7b61f010b89
# ╟─3283a680-f050-11ea-13da-6f5150506375
# ╟─8d3b5ff0-f050-11ea-1d4f-173996efe68b
# ╟─db616880-f052-11ea-25e3-1bc2914021fb
# ╟─333c9510-f036-11ea-3e76-d1005102754d


### A Pluto.jl notebook ###
# v0.17.3

using Markdown
using InteractiveUtils

# ╔═╡ 0c9c514e-598c-418f-90e5-4a74ffb5bdc0
begin
    import Pkg
    # activate a temporary environment
    Pkg.activate(mktempdir())
    Pkg.add([
        Pkg.PackageSpec(name="Plots", version="1"),
        Pkg.PackageSpec(name="PlutoUI", version="0.7"),
        Pkg.PackageSpec(name="Distributions"),
        Pkg.PackageSpec(name="DataFrames", version="1"),
        Pkg.PackageSpec(name="StanSample", version="5.2"),
        Pkg.PackageSpec(name="StatisticalRethinking", version="4.4"),
    ])
    using Plots, PlutoUI
	using Distributions, DataFrames
	using StanSample, StatisticalRethinking
end

# ╔═╡ f2224d36-7972-4e21-badc-3a401bf6f0bb
pwd()

# ╔═╡ 7d6d3864-3bff-414d-b077-ab59673f6d25
@__DIR__

# ╔═╡ eb07b03c-4617-4e07-b68a-476b903460e0
Pkg.pkg"status"

# ╔═╡ ed6fe9e3-df1d-4c60-b3f0-9cab2c81eb78
stan0_0 = "
data {
	int N;
	vector[N] D;
}
parameters {
	real a;
	real sigma;
}
model {
	a ~ normal(0, 2);
	sigma ~ exponential(1);
	D ~ normal(a, sigma);
}";

# ╔═╡ a022fd36-61aa-4e85-92e0-0ab55339468b
data = (N=1000, D=rand(Normal(0,1), 1000));

# ╔═╡ 3104104b-b90b-4279-a79e-a94502869172
begin
	m0_0s = SampleModel("m0_0s", stan0_0)
	rc0_0s = stan_sample(m0_0s; data)
end;

# ╔═╡ 3dab9dbd-2955-4a11-816f-d986ce0529bd
if success(rc0_0s)
	df0_0s = read_samples(m0_0s, :dataframe)
	PRECIS(df0_0s)
end

# ╔═╡ Cell order:
# ╠═f2224d36-7972-4e21-badc-3a401bf6f0bb
# ╠═7d6d3864-3bff-414d-b077-ab59673f6d25
# ╠═0c9c514e-598c-418f-90e5-4a74ffb5bdc0
# ╠═eb07b03c-4617-4e07-b68a-476b903460e0
# ╠═ed6fe9e3-df1d-4c60-b3f0-9cab2c81eb78
# ╠═a022fd36-61aa-4e85-92e0-0ab55339468b
# ╠═3104104b-b90b-4279-a79e-a94502869172
# ╠═3dab9dbd-2955-4a11-816f-d986ce0529bd


# test index notation using @tensor macro
#-----------------------------------------
withblas = TensorContractionsXD.use_blas() ? "with" : "without"
withcache = TensorContractionsXD.use_cache() ? "with" : "without"
@testset "Index Notation $withblas BLAS and $withcache cache" begin
    t0 = time()
    A = randn(Float64, (3,5,4,6))
    p = (4,1,3,2)
    C1 = permutedims(A, p)
    @tensor C2[4,1,3,2] := A[1,2,3,4]
    @test C1 ≈ C2
    @test_throws TensorContractionsXD.IndexError begin
        @tensor C[1,2,3,4] := A[1,2,3]
    end
    @test_throws TensorContractionsXD.IndexError begin
        @tensor C[1,2,3,4] := A[1,2,2,4]
    end
    println("tensorcopy: $(time()-t0) seconds")
    t0 = time()

    B=randn(Float64, (5,6,3,4))
    p=[3,1,4,2]
    @tensor C1[3,1,4,2] := A[3,1,4,2] + B[1,2,3,4]
    C2=A+permutedims(B, p)
    @test C1 ≈ C2
    @test_throws DimensionMismatch begin
        @tensor C[1,2,3,4] := A[1,2,3,4] + B[1,2,3,4]
    end
    println("tensoradd: $(time()-t0) seconds")
    t0 = time()

    A=randn(Float64, (50,100,100))
    @tensor C1[a] := A[a, b', b']
    C2=zeros(50)
    for i=1:50
        for j=1:100
            C2[i]+=A[i, j, j]
        end
    end
    @test C1 ≈ C2
    A=randn(Float64, (3,20,5,3,20,4,5))
    @tensor C1[e, a, d] := A[a, b, c, d, b, e, c]
    C2=zeros(4,3,3)
    for i1=1:4, i2=1:3, i3=1:3
        for j1=1:20, j2=1:5
            C2[i1, i2, i3]+=A[i2, j1, j2, i3, j1, i1, j2]
        end
    end
    @test C1 ≈ C2
    println("tensortrace: $(time()-t0) seconds")
    t0 = time()

    A=randn(Float64, (3,20,5,3,4))
    B=randn(Float64, (5,6,20,3))
    @tensor C1[a, g, e, d, f] := A[a, b, c, d, e] * B[c, f, b, g]
    C2=zeros(3,3,4,3,6)
    for a=1:3, b=1:20, c=1:5, d=1:3, e=1:4, f=1:6, g=1:3
        C2[a, g, e, d, f] += A[a, b, c, d, e] * B[c, f, b, g]
    end
    @test C1 ≈ C2
    @test_throws TensorContractionsXD.IndexError begin
        @tensor A[a, b, c, d] * B[c, f, b, g]
    end
    println("tensorcontract 1: $(time()-t0) seconds")
    t0 = time()

    A=randn(Float64, (5,5,5,5))
    B=rand(ComplexF64, (5,5,5,5))
    @tensor C1[1,2,5,6,3,4,7,8] := A[1,2,3,4] * B[5,6,7,8]
    C2=reshape(kron(reshape(B, (25,25)), reshape(A, (25,25))), (5,5,5,5,5,5,5,5))
    @test C1 ≈ C2
    @test_throws TensorContractionsXD.IndexError begin
        @tensor C[a, b, c, d, e, f, g, i] := A[a, b, c, d] * B[e, f, g, h]
    end
    println("tensorcontract 2: $(time()-t0) seconds")
    t0 = time()

    Da=10
    Db=15
    Dc=4
    Dd=8
    De=6
    Df=7
    Dg=3
    Dh=2
    A=rand(ComplexF64, (Da, Dc, Df, Da, De, Db, Db, Dg))
    B=rand(ComplexF64, (Dc, Dh, Dg, De, Dd))
    C=rand(ComplexF64, (Dd, Dh, Df))
    @tensor D1[d, f, h] := A[a, c, f, a, e, b, b, g] * B[c, h, g, e, d]+0.5 * C[d, h, f]
    D2=zeros(ComplexF64, (Dd, Df, Dh))
    for d=1:Dd, f=1:Df, h=1:Dh
        D2[d, f, h] += 0.5 * C[d, h, f]
        for a=1:Da, b=1:Db, c=1:Dc, e=1:De, g=1:Dg
            D2[d, f, h] += A[a, c, f, a, e, b, b, g] * B[c, h, g, e, d]
        end
    end
    @test D1 ≈ D2
    @test norm(vec(D1)) ≈ sqrt(abs((@tensor scalar(D1[d, f, h] * conj(D1[d, f, h])))))
    println("tensorcontract 3: $(time()-t0) seconds")
    t0 = time()

    Abig=randn(Float64, (30,30,30,30))
    A=view(Abig,1 .+ 3 .* (0:9),2 .+ 2 .* (0:6),5 .+ 4 .* (0:6),4 .+ 3 .* (0:8))
    p=[3,1,4,2]
    Cbig=zeros(ComplexF64, (50,50,50,50))
    C=view(Cbig,13 .+ (0:6),11 .+ 4 .* (0:9),15 .+ 4 .* (0:8),4 .+ 3 .* (0:6))
    Acopy = copy(A)
    Ccopy = copy(C)
    @tensor C[3,1,4,2] = A[1,2,3,4]
    @tensor Ccopy[3,1,4,2] = Acopy[1,2,3,4]
    @test C ≈ Ccopy
    @test_throws TensorContractionsXD.IndexError begin
        @tensor C[3,1,4,2] = A[1,2,3]
    end
    @test_throws DimensionMismatch begin
        @tensor C[3,1,4,2] = A[3,1,4,2]
    end
    @test_throws TensorContractionsXD.IndexError begin
        @tensor C[1,1,2,3] = A[1,2,3,4]
    end
    println("views: $(time()-t0) seconds")
    t0 = time()

    Cbig=zeros(ComplexF64, (50,50,50,50))
    C=view(Cbig, 13 .+ (0:6), 11 .+ 4 .* (0:9), 15 .+ 4 .* (0:8), 4 .+ 3 .* (0:6))
    Acopy=permutedims(copy(A), p)
    Ccopy=copy(C)
    α=randn(Float64)
    β=randn(Float64)
    @tensor C[3,1,4,2] = β * C[3,1,4,2] + α * A[1,2,3,4]
    Ccopy=β * Ccopy+α * Acopy
    @test C ≈ Ccopy
    @test_throws TensorContractionsXD.IndexError begin
        @tensor C[3,1,4,2] = 0.5 * C[3,1,4,2] + 1.2 * A[1,2,3]
    end
    @test_throws DimensionMismatch  begin
        @tensor C[3,1,4,2] = 0.5 * C[3,1,4,2] + 1.2 * A[3,1,2,4]
    end
    @test_throws TensorContractionsXD.IndexError  begin
        @tensor C[1,1,2,3] = 0.5 * C[1,1,2,3] + 1.2 * A[1,2,3,4]
    end
    println("more views: $(time()-t0) seconds")
    t0 = time()

    Abig=rand(Float64, (30,30,30,30))
    A=view(Abig,1 .+ 3 .* (0:8),2 .+ 2 .* (0:14),5 .+ 4 .* (0:6),7 .+ 2 .* (0:8))
    Bbig=rand(ComplexF64, (50,50))
    B=view(Bbig,13 .+ (0:14),3 .+ 5 .* (0:6))
    Acopy=copy(A)
    Bcopy=copy(B)
    α=randn(Float64)
    @tensor B[b, c] += α * A[a, b, c, a]
    for i=1 .+ (0:8)
        Bcopy += α * view(A, i, :, :, i)
    end
    @test B ≈ Bcopy
    @test_throws TensorContractionsXD.IndexError begin
        @tensor B[b, c] += α * A[a, b, c]
    end
    @test_throws DimensionMismatch begin
        @tensor B[c, b] += α * A[a, b, c, a]
    end
    @test_throws TensorContractionsXD.IndexError begin
        @tensor B[c, b] += α * A[a, b, a, a]
    end
    @test_throws DimensionMismatch begin
        @tensor B[c, b] += α * A[a, b, a, c]
    end
    println("even more views: $(time()-t0) seconds")
    t0 = time()

    Abig=rand(Float64, (30,30,30,30))
    A=view(Abig, 1 .+ 3 .* (0:8), 2 .+ 2 .* (0:14), 5 .+ 4 .* (0:6), 7 .+ 2 .* (0:8))
    Bbig=rand(ComplexF64, (50,50,50))
    B=view(Bbig, 3 .+ 5 .* (0:6), 7 .+ 2 .* (0:7), 13 .+ (0:14))
    Cbig=rand(ComplexF32, (40,40,40))
    C=view(Cbig,3 .+ 2 .* (0:8),13 .+ (0:8),7 .+ 3 .* (0:7))
    Acopy=copy(A)
    Bcopy=copy(B)
    Ccopy=copy(C)
    α=randn(Float64)
    for d=1 .+ (0:8), a=1 .+ (0:8), e=1 .+ (0:7)
        for b=1 .+ (0:14), c=1 .+ (0:6)
            Ccopy[d, a, e] -=α * A[a, b, c, d] * conj(B[c, e, b])
        end
    end
    @tensor C[d, a, e] -= α * A[a, b, c, d] * conj(B[c, e, b])
    @test C ≈ Ccopy
    println("and some more views: $(time()-t0) seconds")
    t0 = time()

    Cbig=rand(ComplexF32, (40,40,40))
    C=view(Cbig,3 .+ 2 .* (0:8),13 .+ (0:8),7 .+ 3 .* (0:7))
    Ccopy=copy(C)
    for d=1 .+ (0:8), a=1 .+ (0:8), e=1 .+ (0:7)
        for b=1 .+ (0:14), c=1 .+ (0:6)
            Ccopy[d, a, e] += α * A[a, b, c, d] * conj(B[c, e, b])
        end
    end
    @tensor C[d, a, e] += α * A[a, b, c, d] * conj(B[c, e, b])
    @test C ≈ Ccopy
    @test_throws TensorContractionsXD.IndexError begin
        @tensor C[d, a, e] += α * A[a, b, c, a] * B[c, e, b]
    end
    @test_throws TensorContractionsXD.IndexError begin
        @tensor C[d, a, e] += α * A[a, b, c, d] * B[c, b]
    end
    @test_throws TensorContractionsXD.IndexError begin
        @tensor C[d, e] += α * A[a, b, c, d] * B[c, e, b]
    end
    @test_throws DimensionMismatch begin
        @tensor C[d, e, a] += α * A[a, b, c, d] * B[c, e, b]
    end
    println("Float32 views: $(time()-t0) seconds")
    t0 = time()

    # Simple function example
    @tensor function f(A, b)
        w[x] := (1//2)*A[x,y]*b[y]
        return w
    end
    for T in (Float32, Float64, ComplexF32, ComplexF64, BigFloat)
        A = rand(T, 10, 10)
        b = rand(T, 10)
        @test f(A,b) ≈ (1//2)*A*b
    end

    # Example from README.md
    α=randn()
    A=randn(5,5,5,5,5,5)
    B=randn(5,5,5)
    C=randn(5,5,5)
    D=zeros(5,5,5)
    @tensor begin
        D[a, b, c] = A[a, e, f, c, f, g] * B[g, b, e] + α * C[c, a, b]
        E[a, b, c] := A[a, e, f, c, f, g] * B[g, b, e] + α * C[c, a, b]
    end
    @test D == E
    println("readme example: $(time()-t0) seconds")
    t0 = time()

    # Some tensor network examples
    @testset for T in (Float32, Float64, ComplexF32, ComplexF64, BigFloat)
        D1, D2, D3 = 30, 40, 20
        d1, d2 = 2, 3
        A1 = rand(T, D1, d1, D2) .- 1//2
        A2 = rand(T, D2, d2, D3) .- 1//2
        rhoL = rand(T, D1, D1) .- 1//2
        rhoR = rand(T, D3, D3) .- 1//2
        H = rand(T, d1, d2, d1, d2) .- 1//2
        A12 = reshape(reshape(A1, D1 * d1, D2) * reshape(A2, D2, d2 * D3), (D1, d1, d2, D3))
        rA12 = reshape(reshape(rhoL * reshape(A12, (D1, d1*d2*D3)), (D1*d1*d2, D3)) * rhoR, (D1, d1, d2, D3))
        HrA12 = permutedims(reshape(reshape(H, (d1 * d2, d1*d2)) * reshape(permutedims(rA12, (2,3,1,4)), (d1 * d2, D1 * D3)), (d1, d2, D1, D3)), (3,1,2,4))
        E = dot(A12, HrA12)
        @tensor HrA12′[a, s1, s2, c] := rhoL[a, a'] * A1[a', t1, b] * A2[b, t2, c'] * rhoR[c', c] * H[s1, s2, t1, t2]
        @tensor HrA12′′[:] := rhoL[-1, 1] * H[-2, -3, 4, 5] * A2[2, 5, 3] * rhoR[3, -4] * A1[1, 4, 2] # should be contracted in exactly same order
        @tensor HrA12′′′[a, s1, s2, c] := rhoL[a, a'] * H[s1, s2, t1, t2] * A2[b, t2, c'] * rhoR[c', c] * A1[a', t1, b] order=(a',b,c',t1,t2)# should be contracted in exactly same order
        @tensoropt HrA12′′′′[:] := rhoL[-1, 1] * H[-2, -3, 4, 5] * A2[2, 5, 3] * rhoR[3, -4] * A1[1, 4, 2]

        @test HrA12′ == HrA12′′ == HrA12′′′ # should be exactly equal
        @test HrA12 ≈ HrA12′
        @test HrA12 ≈ HrA12′′′′
        @test HrA12′′ == ncon([rhoL, H, A2, rhoR, A1],
                                [[-1,1],[-2,-3,4,5],[2,5,3],[3,-4],[1,4,2]])
        @test HrA12′′ == @ncon([rhoL, H, A2, rhoR, A1],
                                [[-1,1],[-2,-3,4,5],[2,5,3],[3,-4],[1,4,2]];
                                order = [1,2,3,4,5], output=[-1,-2,-3,-4])
        @test E ≈ @tensor scalar(rhoL[a', a] * A1[a, s, b] * A2[b, s', c] * rhoR[c, c'] * H[t, t', s, s'] * conj(A1[a', t, b']) * conj(A2[b', t', c']))
    end
    println("tensor network examples: $(time()-t0) seconds")
    t0 = time()

    @testset "tensoropt" begin
        A=randn(5,5,5,5)
        B=randn(5,5,5)
        C=randn(5,5,5)
        @tensoropt (a=>χ,b=>χ^2,c=>2*χ,d=>χ,e=>5,f=>2*χ) D1[a,b,c,d] := A[a,e,c,f]*B[g,d,e]*C[g,f,b]
        @tensoropt (a=χ,b=χ^2,c=2*χ,d=χ,e=5,f=2*χ) D2[a,b,c,d] := A[a,e,c,f]*B[g,d,e]*C[g,f,b]
        @tensoropt ((a,d)=>χ,b=>χ^2,(c,f)=>2*χ,e=>5) D3[a,b,c,d] := A[a,e,c,f]*B[g,d,e]*C[g,f,b]
        @tensoropt ((a,d)=χ,b=χ^2,(c,f)=2*χ,e=5) D4[a,b,c,d] := A[a,e,c,f]*B[g,d,e]*C[g,f,b]
        @test D1 == D2 == D3 == D4
        _optdata = optex -> TensorContractionsXD.optdata(optex, :(D1[a,b,c,d] := A[a,e,c,f]*B[g,d,e]*C[g,f,b]))
        optex1 = :((a=>χ,b=>χ^2,c=>2*χ,d=>χ,e=>5,f=>2*χ))
        optex2 = :((a=χ,b=χ^2,c=2*χ,d=χ,e=5,f=2*χ))
        optex3 = :(((a,d)=>χ,b=>χ^2,(c,f)=>2*χ,e=>5))
        optex4 = :(((a,d)=χ,b=χ^2,(c,f)=2*χ,e=5))
        optex5 = :(((a,)=>χ,b=>χ^2,(c,)=>2χ,d=>χ,e=>5,f=>χ*2,()=>12345))
        @test _optdata(optex1) == _optdata(optex2) == _optdata(optex3) == _optdata(optex4) == _optdata(optex5)
        optex6 = :(((a,b,c)=χ,))
        optex7 = :((a,b,c))
        @test _optdata(optex6) == _optdata(optex7)
        optex8 = :(((a,b,c)=1,(d,e,f,g)=χ))
        optex9 = :(!(a,b,c))
        @test _optdata(optex8) == _optdata(optex9)
        optex10 = :((a=>χ,b=>χ^2,c=2*χ,d=>χ,e=>5,f=2*χ))
        optex11 = :((a=χ,b=χ^2,c=2*χ,d,e=5,f))
        @test_throws ErrorException _optdata(optex10)
        @test_throws ErrorException _optdata(optex11)
    end

    @testset "Issue 83" begin
        op1 = randn(2, 2)
        op2 = randn(2, 2)
        op3 = randn(2, 2)

        f83(op,op3) = @ncon((op, op3), ([-1 3], [3 -3]))

        b = f83(op1,op3)
        bcopy = deepcopy(b)
        c = f83(op2,op3)
        @test b == bcopy
        @test b != c
    end

    # diagonal matrices
    @testset for T in (Float32, Float64, ComplexF32, ComplexF64)
        A = randn(T, 10, 10, 10, 10)
        @tensor C[3,1,4,2] := A[1,2,3,4]
        U,S,V = svd(reshape(C, (100, 100)))
        U2 = reshape(view(U, 1:2:100, :), (5,10,100))
        S2 = Diagonal(S)
        V2 = reshape(view(V, 1:2:100, :), (5, 10, 100))
        @tensor A2[1,2,3,4] := U2[3,1,a]*S2[a,a']*conj(V2[4,2,a'])
        @test A2 ≈ view(A, :, :, 1:2:10, 1:2:10)

        B = randn(T, 10)
        @tensor C1[3,1,4,2] := A[a,2,3,4]*Diagonal(B)[a,1]
        @tensor C2[3,1,4,2] := A[1,a,3,4]*conj(Diagonal(B)[a,2])
        @tensor C3[3,1,4,2] := conj(A[1,2,a,4])*Diagonal(B)[a,3]
        @tensor C4[3,1,4,2] := conj(A[1,2,3,a])*conj(Diagonal(B)[a,4])
        @tensor C1′[3,1,4,2] := Diagonal(B)[a,1]*A[a,2,3,4]
        @tensor C2′[3,1,4,2] := conj(Diagonal(B)[a,2])*A[1,a,3,4]
        @tensor C3′[3,1,4,2] := Diagonal(B)[a,3]*conj(A[1,2,a,4])
        @tensor C4′[3,1,4,2] := conj(Diagonal(B)[a,4])*conj(A[1,2,3,a])
        @test C1 == C1′
        @test C2 == C2′
        @test C3 == C3′
        @test C4 == C4′
        for i = 1:10
            @test C1[:,i,:,:] ≈ permutedims(A[i,:,:,:],(2,3,1))*B[i]
            @test C2[:,:,:,i] ≈ permutedims(A[:,i,:,:],(2,1,3))*conj(B[i])
            @test C3[i,:,:,:] ≈ conj(permutedims(A[:,:,i,:],(1,3,2)))*B[i]
            @test C4[:,:,i,:] ≈ conj(permutedims(A[:,:,:,i],(3,1,2)))*conj(B[i])
        end

        @tensor D1[1,2] := A[1,2,a,b]*Diagonal(B)[a,b]
        @tensor D2[1,2] := A[1,b,2,a]*conj(Diagonal(B)[a,b])
        @tensor D3[1,2] := conj(A[a,2,1,b])*Diagonal(B)[a,b]
        @tensor D4[1,2] := conj(A[a,1,b,2])*conj(Diagonal(B)[a,b])
        @tensor D1′[1,2] := Diagonal(B)[a,b]*A[1,2,a,b]
        @tensor D2′[1,2] := conj(Diagonal(B)[a,b])*A[1,b,2,a]
        @tensor D3′[1,2] := Diagonal(B)[a,b]*conj(A[a,2,1,b])
        @tensor D4′[1,2] := conj(Diagonal(B)[a,b])*conj(A[a,1,b,2])
        @test D1 == D1′
        @test D2 == D2′
        @test D3 == D3′
        @test D4 == D4′

        E1 = zero(D1)
        E2 = zero(D2)
        E3 = zero(D3)
        E4 = zero(D4)
        for i = 1:10
            E1[:,:] += A[:,:,i,i]*B[i]
            E2[:,:] += A[:,i,:,i]*conj(B[i])
            E3[:,:] += A[i,:,:,i]'*B[i]
            E4[:,:] += conj(A[i,:,i,:])*conj(B[i])
        end
        @test D1 ≈ E1
        @test D2 ≈ E2
        @test D3 ≈ E3
        @test D4 ≈ E4

        F = randn(T,(10,10))
        @tensor G[a,c,b,d] := F[a,b]*Diagonal(B)[c,d]
        @test reshape(G,(100,100)) ≈ kron(Diagonal(B), F)
    end
    println("diagonal examples: $(time()-t0) seconds")
    t0 = time()

end


using FEMfunctions

import Gmsh: gmsh
using LinearAlgebra

using AbstractPlotting
using SparseArrays
using SpecialMatrices
using Polynomials
using Printf
using WriteVTK

include("makeMesh.jl")
include("twodAdvectionDiffusion.jl")

κ = 1.0
ω = -1.0    # advection velocity parameter

### define the N parameters that describe the inner boundary
N = 160      # number of BSplines used to represent the inner boundary

### describe the inner boundary using BSplines
r = ones(N,1)

x,eConn,eConn2, innerNodes,innerX, outerNodes,outerX = makeMesh(r)

sort!(innerNodes);
sort!(outerNodes);

#   Let's look at the mesh...
nNodes = size(x,2)
nElements = size(eConn,2)

xT = zeros(Float64,nNodes,2)   # xT = transpose(x[1:2,:])
eC = zeros(Int64,nElements,6)  # eC = transpose(eConn), converted to Int64
for i=1:nNodes
  xT[i,1] = x[1,i]
  xT[i,2] = x[2,i]
end
for i=1:nElements#-nElementsSensor # the sensor elements are orientated correctly (why?)
  eC[i,1] = convert(Int64,eConn[1,i])
  eC[i,2] = convert(Int64,eConn[3,i])
  eC[i,3] = convert(Int64,eConn[2,i])
  eC[i,4] = convert(Int64,eConn[6,i])
  eC[i,5] = convert(Int64,eConn[5,i])
  eC[i,6] = convert(Int64,eConn[4,i])
end

velocity = Array{Float64,2}(undef,nNodes,2)
for i=1:nNodes
  velocity[i,1] = -ω*xT[i,2]
  velocity[i,2] = ω*xT[i,1]
end

temperature,exactTemperature = twodAdvectionDiffusion(xT,eC,innerNodes,outerNodes,velocity,κ,ω)

M = twodMassMatrix(xT,eC)
error = temperature-exactTemperature
#  The inner boundary isn't exactly a circle of radius 1, so we have a weaker error
@test sqrt( dot(error,(M*error)) ) < 2e-1

#  To manually check the function saveFEMasVTK, enter the following line
#  and use ParaView:
#  saveFEMasVTK("advDiff",xT,eC,["temperature"],temperature,["velocity"],velocity)


using Images, TestImages, StaticArrays, ImageTracking, ImageView, LinearAlgebra, CoordinateTransformations, Gtk.ShortNames
#=Image Credit:  C. Liu. Beyond Pixels: Exploring New Representations and
#Applications for Motion Analysis. Doctoral Thesis. Massachusetts Institute of
#Technology. May 2009. =#
img1 = load("demo/table1.jpg")
img2 = load("demo/table2.jpg")

corners = imcorner(img1, method=shi_tomasi)
I = findall(!iszero, corners)
r, c = (getindex.(I, 1), getindex.(I, 2))
points = map((ri, ci) -> SVector{2}(Float64(ri), Float64(ci)), r, c)

algorithm = LucasKanade(20, window_size = 11,
                            pyramid_levels = 4,
                            eigenvalue_threshold = 0.000001)
flow, indicator = optical_flow(Gray{Float32}.(img1), Gray{Float32}.(img2),points, algorithm)

# Keep the subset of points that were successfully tracked and determine
# correspondences.
valid_points = points[indicator]
valid_flow = flow[indicator]
valid_correspondence = map((x,Δx)-> x+Δx, valid_points, valid_flow)

# Convert (row,columns) to (x,y) convention and round to nearest integer.
pts0 = map(x-> round.(Int,(last(x),first(x))), points)
pts1 = map(x-> round.(Int,(last(x),first(x))), valid_points)
pts2 = map(x-> round.(Int,(last(x),first(x))), valid_correspondence)
lines = map((p1, p2) -> (p1,p2), pts1, pts2)

# Visualise the optical flow. Red lines demarcate optical flow on the keypoints
# that were successfully tracked.
guidict = imshow(img1)
idx2 = annotate!(guidict, AnnotationLines(lines, linewidth=2.0, color=RGB(1,0,0), coord_order="xyxy"))


"""
    Sobel-sample x+y in [0,10]x[0,10],
    then minimize it on Section([NaN,10.0]),
    and verify that the minimum is on x,y=(0,10)
    rather than in (0,0)
"""

using Surrogates
using Test

lb = [ 0.0, 0.0, 0.0]
ub = [10.0,10.0,10.0]
x = Surrogates.sample(10,lb,ub,LatinHypercubeSample())
f = x -> x[1]+x[2]+x[3]
y = f.(x)
f([0,0,0]) == 0

f_hat = Kriging(x,y,lb,ub)

f_hat([0,0,0])

isapprox(f([0,0,0]), f_hat([0,0,0]))

""" The global minimum is at (0,0) """

(xy_min, f_hat_min) = surrogate_optimize(
    f,
    DYCORS(), lb, ub,
    f_hat,
    SobolSample())

isapprox(xy_min[1], 0.0, atol=1e-3)

""" The minimum on the (0,10) section is around (0,10) """

section_sampler_z_is_10 = SectionSample(
    [NaN64, NaN64, 10.0],
    Surrogates.UniformSample())

@test [3] == Surrogates.fixed_dimensions(section_sampler_z_is_10)
@test [1,2] == Surrogates.free_dimensions(section_sampler_z_is_10)

Surrogates.sample(5, lb, ub, section_sampler_z_is_10)

(xy_min, f_hat_min) = surrogate_optimize(
    f,
    EI(), lb, ub,
    f_hat,
    section_sampler_z_is_10, maxiters=1000)

isapprox(xy_min[1],  0.0, atol=0.1)
isapprox(xy_min[2],  0.0, atol=0.1)
isapprox(xy_min[3], 10.0, atol=0.1)


module TestDescentUpdates
using Test
using LinearAlgebra
using OptimizationAlgorithms
using OptimizationAlgorithms: Gradient, ScaledGradient, StoppingCriterion, fixedpoint!
using OptimizationAlgorithms: PolyakMomentum, Nesterov, ADAM, AMSGRAD
using ForwardDiff: gradient

ε = 1e-4
@testset "PolyakMomentum" begin
    n = 2
    A = randn(n, n)
    A = A'A + I
    μ = randn(n)
    f = x -> (x-μ)'A*(x-μ)
    x = randn(n)
    G = Gradient(f, x)
    α = .1
    β = .5
    G = PolyakMomentum(G, copy(x), α, β)
    fixedpoint!(G, x, StoppingCriterion(x, dx = 1e-2ε))
    @test norm(gradient(f, x)) < ε
end

@testset "Nesterov" begin
    # TODO: better test for nesterov
    n = 2
    A = randn(n, n)
    A = A'A + I
    μ = randn(n)
    f = x -> (x-μ)'A*(x-μ)
    x = randn(n)
    α = 5e-2
    β = .5
    G = Gradient(f, x)
    G = Nesterov(G, x, α, β)
    fixedpoint!(G, x, StoppingCriterion(x, dx = 1e-2ε))
    @test norm(gradient(f, x)) < ε
end

@testset "ADAM" begin
    n = 2
    A = randn(n, n)
    A = A'A + I
    μ = randn(n)
    f = x -> (x-μ)'A*(x-μ)
    x = randn(n)
    G = Gradient(f, x)
    α = .1
    β₁ = 0.5
    β₂ = 0.999 # this should be large
    G = ADAM(G, x, α, β₁, β₂)
    fixedpoint!(G, x, StoppingCriterion(x, dx = 1e-2ε, maxiter = 1024))
    @test norm(gradient(f, x)) < ε
end

@testset "AMSGRAD" begin
    n = 2
    A = randn(n, n)
    A = A'A + I
    μ = randn(n)
    f = x -> (x-μ)'A*(x-μ)
    x = randn(n)
    G = Gradient(f, x)
    α = .1
    β₁ = 0.5
    β₂ = 0.999 # this should be large
    G = AMSGRAD(G, x, α, β₁, β₂)
    fixedpoint!(G, x, StoppingCriterion(x, dx = 1e-2ε, maxiter = 1024))
    @test norm(gradient(f, x)) < ε
end

end # TestDescentUpdates


using Parameters
using PyPlot

include("cr_model.jl")

function retrieve_RC(vectordata, RC)
    newarray = zeros(size(vectordata)[1],size(vectordata)[2])

    for i in 1:size(vectordata)[1]
        for j in 1:size(vectordata)[2]
            newarray[i,j] = vectordata[i,j][RC]
        end
    end
    return newarray
end

## CR exp resource, type I consumer
function expI_resiso(p)
    @unpack r, a = p
    return r / a
end

function expI_coniso(p)
    @unpack m, a, e = p
    return m / ( a * e )
end

let
    par = CRPar(a=1.6, m=0.7)
    resconrange = 0.0:0.01:3.0
    vector_data = [cr_expI_model([R, C], par, NaN) for C in resconrange, R in resconrange]
    U = retrieve_RC(vector_data, 1)
    V = retrieve_RC(vector_data, 2)
    speed = sqrt.(U.^2 .+ V.^2)
    lw = 5 .* speed ./ maximum(speed) # Line Widths

    expI_plot = figure()
    hlines(expI_resiso(par), 0.0, 3.0, colors="blue")
    vlines(expI_coniso(par), 0.0, 3.0, colors="orange")
    streamplot(collect(resconrange), collect(resconrange), U, V, density = 0.6, color = "k", linewidth = lw)
    xlabel("Resource")
    ylabel("Consumer")
    xlim(0,3)
    ylim(0,3)
    return expI_plot
end

## CR log resource, type I consumer
function logI_resiso(R, p)
    @unpack r, k, a = p
    return r * (- R + k) / (a * k)
end

function logI_coniso(p)
    @unpack m, a, e = p
    return m / (a * e)
end

let
    par = CRPar()
    Rrange = 0.0:0.1:3.0
    resconrange = 0.0:0.005:3.0
    vector_data = [cr_logI_model([R, C], par, NaN) for C in resconrange, R in resconrange]
    U = retrieve_RC(vector_data, 1)
    V = retrieve_RC(vector_data, 2)
    speed = sqrt.(U.^2 .+ V.^2)
    lw = 5 .* speed ./ maximum(speed)

    logI_plot = figure()
    plot(Rrange, [logI_resiso(R, par) for R in Rrange], color="blue")
    vlines(logI_coniso(par), 0.0, 3.0, colors="orange")
    streamplot(collect(resconrange), collect(resconrange), U, V, density = 0.6, color = "k", linewidth = lw)
    xlabel("Resource")
    ylabel("Consumer")
    xlim(0,3)
    ylim(0,3)
    return logI_plot
end

## CR log resource, type II consumer

function logII_resiso(R, p)
    @unpack r, k, a, h = p
    return r * (-R^2 * a * h + R * a * h * k - R + k) / (a * k)
end

function logII_coniso(p)
    @unpack m, a, e, h = p
    return m / (a * (e - h * m))
end

let
    par = CRPar(r = 1.5, e = 0.45)
    Rrange = 0.0:0.1:3.0
    resconrange = 0.0:0.005:3.0
    vector_data = [cr_logII_model([R, C], par, NaN) for C in resconrange, R in resconrange]
    U = retrieve_RC(vector_data, 1)
    V = retrieve_RC(vector_data, 2)
    speed = sqrt.(U.^2 .+ V.^2)
    lw = 5 .* speed ./ maximum(speed)

    logII_plot = figure()
    plot(collect(Rrange), [logII_resiso(R, par) for R in Rrange], color="blue")
    vlines(logII_coniso(par), 0.0, 3.0, colors="orange")
    streamplot(collect(resconrange), collect(resconrange), U, V, density = 0.6, color = "k", linewidth = lw)

    xlabel("Resource")
    ylabel("Consumer")
    xlim(0,3)
    ylim(0,3)
    return logII_plot
end

## CR exp resource, type II consumer

let
    par = CRPar(r=0.5)
    Rrange = 0.0:0.1:100.0
    resconrange = 0.0:0.1:100.0
    vector_data = [cr_expII_model([R, C], par, NaN) for C in resconrange, R in resconrange]
    U = retrieve_RC(vector_data, 1)
    V = retrieve_RC(vector_data, 2)
    speed = sqrt.(U.^2 .+ V.^2)
    lw = 5 .* speed ./ maximum(speed)

    expII_plot = figure()
    plot(collect(Rrange), [expII_resiso(R, par) for R in Rrange], color="blue")
    vlines(expII_coniso(par), 0.0, 100.0, colors="orange")
    streamplot(collect(resconrange), collect(resconrange), U, V, density = 0.6, color = "k", linewidth = lw)

    xlabel("Resource")
    ylabel("Consumer")
    xlim(0,30)
    ylim(0,30)
    return expII_plot
end


"""
    modularity(g, c)

Return a value representing Newman's modularity `Q` for the undirected graph
`g` given the partitioning vector `c`.
"""
function modularity end
@traitfn function modularity(g::::(!IsDirected), c::Vector)
    m = 2*ne(g)
    m == 0 && return 0.
    nc = maximum(c)
    a = zeros(Int,nc)
    Q = 0
    for u in vertices(g)
        for v in neighbors(g,u)
            if u <= v
                c1 = c[u]
                c2 = c[v]
                if c1 == c2
                    Q += 2
                end
                a[c1] += 1
                a[c2] += 1
            end
        end
    end
    Q = Q*m
    @inbounds for i=1:nc
        Q -= a[i]*a[i]
    end
    return Q/m/m
end


@testset "CauchyProblem" begin

    @testset "constant solution" begin
    prob = CauchyProblem(identity, b=1, d_dϕb=0)

    θ = solve(prob)

    ϕ = range(0, 20, length=100)

    @test all(@inferred θ.(ϕ) .≈ θ.i)
    @test all(@inferred θ.(ϕ) .≈ θ.b)
    @test all(@inferred d_dϕ.(θ, ϕ) .≈ 0)
    @test θ.iterations == 0
    @test isnan(@inferred θ(-1))
    end

    @testset "exact" begin
    # Reference: Philip (1960) Table 1, No. 13
    # https://doi.org/10.1071/PH600001
    D = θ -> 0.5*(1 - log(θ))

    prob1 = DirichletProblem(D, i=0, b=1)
    θ1 = solve(prob1)

    prob2 = CauchyProblem(D, b=θ1.b, d_dϕb=θ1.d_dϕb)
    θ2 = solve(prob2)

    @test θ2.i == θ1.i
    @test θ2.b == θ1.b
    @test θ2.d_dϕb == θ1.d_dϕb
    @test θ2.ϕi == θ1.ϕi
    @test θ2.ϕb == θ1.ϕb

    @test θ2.iterations == 0

    @test isnan(@inferred θ2(-1))

    end
end


#=

  Killer Sudoku in Julia ConstraintSolver.jl

  http://en.wikipedia.org/wiki/Killer_Sudoku
  """
  Killer sudoku (also killer su doku, sumdoku, sum doku, addoku, or
  samunamupure) is a puzzle that combines elements of sudoku and kakuro.
  Despite the name, the simpler killer sudokus can be easier to solve
  than regular sudokus, depending on the solver's skill at mental arithmetic;
  the hardest ones, however, can take hours to crack.

  ...

  The objective is to fill the grid with numbers from 1 to 9 in a way that
  the following conditions are met:

    * Each row, column, and nonet contains each number exactly once.
    * The sum of all numbers in a cage must match the small number printed
      in its corner.
    * No number appears more than once in a cage. (This is the standard rule
      for killer sudokus, and implies that no cage can include more
      than 9 cells.)

  In 'Killer X', an additional rule is that each of the long diagonals
  contains each number once.
  """

  Here we solve the problem from the Wikipedia page, also shown here
  http://en.wikipedia.org/wiki/File:Killersudoku_color.svg

  The output is:
    2 1 5 6 4 7 3 9 8
    3 6 8 9 5 2 1 7 4
    7 9 4 3 8 1 6 5 2
    5 8 6 2 7 4 9 3 1
    1 4 2 5 9 3 8 6 7
    9 7 3 8 1 6 4 2 5
    8 2 1 7 3 9 5 4 6
    6 5 9 4 2 8 7 1 3
    4 3 7 1 6 5 2 8 9

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

=#

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

#
# Basic Sudoku model
#
function sudoku(model,x,n)
    # Row and column constraints
    for rc = 1:n
        @constraint(model, x[rc,:] in CS.AllDifferent())
        @constraint(model, x[:,rc] in CS.AllDifferent())
    end

    # Cell constraints
    ns = round(Int, sqrt(n))
    n2 = ns-1
    for i in 1:ns:n, j in 1:ns:n
        @constraint(model, vec([x[i+k,j+l] for k in 0:n2, l in 0:n2]) in CS.AllDifferent())
    end

end

function killer_sudoku(problem,print_solutions=true,all_solutions=true)

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

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

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

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

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            "activity.decay" => 0.999, # default 0.999
                                                            "activity.max_probes" => 100, # default, 10
                                                            "activity.max_confidence_deviation" => 20, # default 20

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

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))
    n = 9
    # define the variables
    @variable(model, 1 <= x[1:n,1:n] <= n, Int)

    sudoku(model, x, n) # plain Sudoku model

    # The Killer Sudoku hints
    for hint in problem
        s,lst = hint
        xs = [x[t[1],t[2]] for t in lst]
        @constraint(model,s == sum(xs))
        @constraint(model, xs in CS.AllDifferent())
    end

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                display(x_val)

            end
        end
    else
        println("status:$status")
    end

    return status
end

#
# hints
#
killer_sudoku_problems = Dict(
   :1 => [
            #  [Sum, [list of indices in X]]
            [ 3, [[1,1], [1,2]]],
            [15, [[1,3], [1,4], [1,5]]],
            [22, [[1,6], [2,5], [2,6], [3,5]]],
            [ 4, [[1,7], [2,7]]],
            [16, [[1,8], [2,8]]],
            [15, [[1,9], [2,9], [3,9], [4,9]]],
            [25, [[2,1], [2,2], [3,1], [3,2]]],
            [17, [[2,3], [2,4]]],
            [ 9, [[3,3], [3,4], [4,4]]],
            [ 8, [[3,6], [4,6],[5,6]]],
            [20, [[3,7], [3,8],[4,7]]],
            [ 6, [[4,1], [5,1]]],
            [14, [[4,2], [4,3]]],
            [17, [[4,5], [5,5],[6,5]]],
            [17, [[4,8], [5,7],[5,8]]],
            [13, [[5,2], [5,3],[6,2]]],
            [20, [[5,4], [6,4],[7,4]]],
            [12, [[5,9], [6,9]]],
            [27, [[6,1], [7,1],[8,1],[9,1]]],
            [ 6, [[6,3], [7,2],[7,3]]],
            [20, [[6,6], [7,6], [7,7]]],
            [ 6, [[6,7], [6,8]]],
            [10, [[7,5], [8,4],[8,5],[9,4]]],
            [14, [[7,8], [7,9],[8,8],[8,9]]],
            [ 8, [[8,2], [9,2]]],
            [16, [[8,3], [9,3]]],
            [15, [[8,6], [8,7]]],
            [13, [[9,5], [9,6],[9,7]]],
            [17, [[9,8], [9,9]]]
            ]

)

@time killer_sudoku(killer_sudoku_problems[:1])