# helper functions

"""
    r_lattice(prob)

Return a 2d array of size (3, N_r) where the columns are the lattice coordinates of a uniform r grid in one unit cell with size_r(prob)[i] points in direction i.
"""
function r_lattice(prob::AbstractBSEProblem)
    N_rs = size_r(prob)
    rl_1d_1 = range(0.0, stop = 1.0, length = N_rs[1] + 1)[1:N_rs[1]]
    rl_1d_2 = range(0.0, stop = 1.0, length = N_rs[2] + 1)[1:N_rs[2]]
    rl_1d_3 = range(0.0, stop = 1.0, length = N_rs[3] + 1)[1:N_rs[3]]
    rl_1 = kron(ones(N_rs[3]), ones(N_rs[2]), rl_1d_1)
    rl_2 = kron(ones(N_rs[3]), rl_1d_2, ones(N_rs[1]))
    rl_3 = kron(rl_1d_3, ones(N_rs[2]), ones(N_rs[1]))
    rl = transpose(hcat(rl_1, rl_2, rl_3))
    return rl
end

"""
    r_cartesian(prob)

Return a 2d array of size (3, N_r) where the columns are the cartesian coordinates of the r grid generated by r_lattice.
"""
function r_cartesian(prob::AbstractBSEProblem)
    a_mat = lattice_matrix(prob)
    rl = r_lattice(prob)
    rc = similar(rl)
    for ir in 1:size(rl, 2)
        rc[:, ir] = a_mat * rl[:, ir]
    end
    return rc
end

"""
    k_lattice(prob[, shift])

Return a 2d array of size `(3, N_k)` where the columns are the lattice coordinates of a uniform ``k`` grid in one reciprocal unit cell with `size_k(prob)[i]` points in direction `i`. The grid can optionally be shifted by a vector `shift`.
"""
function k_lattice(prob::AbstractBSEProblem, shift = [0.0, 0.0, 0.0])
    N_ks = size_k(prob)
    kl_1d_1 = mod.(range(0.0, stop = 1.0, length = N_ks[1] + 1)[1:N_ks[1]] .+ shift[1], 1.0)
    kl_1d_2 = mod.(range(0.0, stop = 1.0, length = N_ks[2] + 1)[1:N_ks[2]] .+ shift[2], 1.0)
    kl_1d_3 = mod.(range(0.0, stop = 1.0, length = N_ks[3] + 1)[1:N_ks[3]] .+ shift[3], 1.0)
    kl_1 = kron(ones(N_ks[3]), ones(N_ks[2]), kl_1d_1)
    kl_2 = kron(ones(N_ks[3]), kl_1d_2, ones(N_ks[1]))
    kl_3 = kron(kl_1d_3, ones(N_ks[2]), ones(N_ks[1]))
    kl = transpose(hcat(kl_1, kl_2, kl_3))
    return kl
end

"""
    k_cartesian(prob[, shift])

Return a 2d array of size `(3, N_k)` where the columns are the
lattice coordinates of a uniform ``k`` grid in the first Brillouin zone
with `size_k(prob)[i]` points in direction `i`. The grid can
optionally be shifted by a vector `shift`.
The cartesian representation is chosen to lie in the first Brillouin
zone, i.e. to have minimal distance from the origin.
"""
function k_cartesian(prob::AbstractBSEProblem, shift = [0.0, 0.0, 0.0])
    a_mat = lattice_matrix(prob)
    b_mat = 2 * π * inv(a_mat)
    kl = k_lattice(prob, b_mat \ shift)
    kc = similar(kl)
    for ik in 1:size(kl, 2)
        kc[:, ik] = k_lattice_to_reduced_bz(kl[:, ik], b_mat)
    end
    return kc
end

"""
    k_lattice_to_reduced_bz(k_l, b_mat)

Helper function to compute cartesian coordinate of reciprocal point given in lattice coordinates.
"""
function k_lattice_to_reduced_bz(k_l, b_mat)
    shifted_norms = zeros(2, 2, 2)
    for i1 in 1:2, i2 in 1:2, i3 in 1:2
        shifted_norms[i1, i2, i3] = norm(b_mat * (k_l - [i1 - 1, i2 - 1, i3 - 1]))
    end
    m, i = findmin(shifted_norms)
    return b_mat * (k_l - [i[1] - 1, i[2]  - 1, i[3] - 1])
end

"""
    size_q(prob)

Number of differences of k-points in the different directions.
"""
function size_q(prob::AbstractBSEProblem)
    size_q(size_k(prob))
end
function size_q(N_ks::Tuple{Int,Int,Int})
    return (n -> n == 1 ? 1 : 2 * n).(N_ks) # double size unless size is 1
end

"""
    q_lattice(prob)

Return a 2d array of size `(3, N_q)` where the columns are the lattice coordinates of the uniform ``q`` grid generated by a uniform ``k`` grid in one reciprocal unit cell with `size_k(prob)[i]` points in direction `i`.
"""

function q_lattice(prob::AbstractBSEProblem)
    return q_lattice(size_k(prob))
end
function q_lattice(N_ks::Tuple{Int,Int,Int})
    N_qs = size_q(N_ks)
    ql_1d_1 = vcat(range(0, stop = 1, length = N_ks[1] + 1)[1:(end - 1)], range(-1, stop = 0, length = N_ks[1] + 1)[(end - (N_qs[1] - N_ks[1])):(end - 1)])
    ql_1d_2 = vcat(range(0, stop = 1, length = N_ks[2] + 1)[1:(end - 1)], range(-1, stop = 0, length = N_ks[2] + 1)[(end - (N_qs[2] - N_ks[2])):(end - 1)])
    ql_1d_3 = vcat(range(0, stop = 1, length = N_ks[3] + 1)[1:(end - 1)], range(-1, stop = 0, length = N_ks[3] + 1)[(end - (N_qs[3] - N_ks[3])):(end - 1)])
    ql_1 = kron(ones(N_qs[3]), ones(N_qs[2]), ql_1d_1)
    ql_2 = kron(ones(N_qs[3]), ql_1d_2, ones(N_qs[1]))
    ql_3 = kron(ql_1d_3, ones(N_qs[2]), ones(N_qs[1]))
    ql = transpose(hcat(ql_1, ql_2, ql_3))
    return ql
end

"""
    Ω0_volume(prob)

Return the volume of the unit cell.
"""
function Ω0_volume(prob::AbstractBSEProblem)
    return abs(det(lattice_matrix(prob)))
end

"""
    BZ_volume(prob)

Return the volume of the first Brillouin zone. The convention is to use
the normalization b_mat^* * a_mat = 2π with the lattice vectors a_mat
and the reciprocal vectors b_mat.
"""
function BZ_volume(prob::AbstractBSEProblem)
    return (2 * π)^3 / Ω0_volume(prob)
end

"""
    compute_v_hat(prob[, cutoff])

Compute the Fourier coefficients of the Coulomb potential. The frequencies are stored in the order compatible with the 3d FFT.
"""
function compute_v_hat(prob::AbstractBSEProblem, cutoff = Inf)
    N_rs = size_r(prob)
    b_mat = 2 * π * inv(lattice_matrix(prob))
    v_hat = 4 * pi * vec(mapslices(x -> norm(x) < 1e-10 || norm(b_mat * x) > cutoff ? 0.0 : 1 / norm(b_mat * x)^2, fftfreq(N_rs...); dims = 1))
    return v_hat
end

"""
    fftfreq(n1, n2, n3)

Generate an 3d array of Fourier frequencies in the order that they appear in the 3d FFT.
"""
function fftfreq(n1, n2, n3)
    f1 = vcat(collect(0:(div(n1, 2) - (1 - n1 % 2))), collect(-div(n1, 2):-1))
    f2 = vcat(collect(0:(div(n2, 2) - (1 - n2 % 2))), collect(-div(n2, 2):-1))
    f3 = vcat(collect(0:(div(n3, 2) - (1 - n3 % 2))), collect(-div(n3, 2):-1))

    F1 = kron(ones(n3), ones(n2), f1)
    F2 = kron(ones(n3), f2, ones(n1))
    F3 = kron(f3, ones(n2), ones(n1))

    return vcat(F1', F2', F3')
end

"""
    G_vector_to_index(G, N_rs)

Return the index `i` such that `fftfreq(N_rs...)[:, i] == G`.
"""
function G_vector_to_index(G, N_rs) # TODO: make fast for all dimensions
    if length(G) == 3
        ind = mod1(G[1] + 1, N_rs[1]) + N_rs[1] * (mod1(G[2] + 1, N_rs[2]) - 1) + N_rs[1] * N_rs[2] * (mod1(G[3] + 1, N_rs[3]) - 1)
    else
        ind = mod1(G[1] + 1, N_rs[1])
        for id in 2:length(G)
            ind += prod(N_rs[1:(id - 1)]) * (mod1(G[id] + 1, N_rs[id]) - 1)
        end
    end
    return ind
end


# SplineSpace represents a Spline, right now piecewise constant HeavisideSpace is only implemented case
struct SplineSpace{order,T,R} <: Space{PiecewiseSegment{T},R}
    domain::PiecewiseSegment{T}
end

SplineSpace{m,T}(d::PiecewiseSegment{T}) where {m,T} = SplineSpace{m,T,real(eltype(T))}(d)
SplineSpace{m,T}(d::AbstractVector) where {m,T} = SplineSpace{m}(PiecewiseSegment(sort(d)))

SplineSpace{m}(d::PiecewiseSegment{T}) where {m,T} = SplineSpace{m,T,real(eltype(T))}(d)
SplineSpace{m}(d::AbstractVector) where {m} = SplineSpace{m}(PiecewiseSegment(sort(d)))

const HeavisideSpace{T,R} = SplineSpace{0,T,R}
dimension(h::SplineSpace{λ}) where {λ} = length(h.domain.points)+λ-1

convert(::Type{HeavisideSpace},d::PiecewiseSegment) = HeavisideSpace{eltype(d)}(d)

convert(::Type{HeavisideSpace},d::AbstractVector) =
    HeavisideSpace(PiecewiseSegment(sort(d)))

spacescompatible(a::SplineSpace{λ},b::SplineSpace{λ}) where {λ} = domainscompatible(a,b)

function evaluate(c::AbstractVector{T}, s::HeavisideSpace{<:Real}, x::Real) where T
    p = domain(s).points
    for k=1:length(c)
        if p[k] ≤ x ≤ p[k+1]
            return c[k]
        end
    end
    return zero(T)
end

function evaluate(c::AbstractVector{T}, s::SplineSpace{1,<:Real}, x::Real) where T
    p = domain(f).points
    c = f.coefficients
    for k=1:length(p)-1
        if p[k] ≤ x ≤ p[k+1]
            return (x-p[k])*c[k+1]/(p[k+1]-p[k]) + (p[k+1]-x)*c[k]/(p[k+1]-p[k])
        end
    end
    return zero(T)
end

function points(sp::HeavisideSpace,n)
    x=sp.domain.points
    (x[1:end-1] + diff(x)/2)[1:n]
end

points(sp::SplineSpace{1},n) = sp.domain.points[1:n]

for λ = [0,1]
    @eval begin
        function transform(S::SplineSpace{$λ},vals::AbstractVector,plan...)
            @assert length(vals) ≤ dimension(S)
            vals
        end
        itransform(S::SplineSpace{$λ},cfs::AbstractVector,plan...) = pad(cfs,dimension(S))
    end
end

bandwidths(D::ConcreteDerivative{H}) where {H<:HeavisideSpace} = (0,1)
rangespace(D::ConcreteDerivative{H}) where {H<:HeavisideSpace} = DiracSpace(domain(D).points[2:end-1])

function getindex(D::ConcreteDerivative{H,<:Any,T},k::Integer,j::Integer) where {H<:HeavisideSpace,T}
    if k==j
        -one(T)
    elseif j==k+1
        one(T)
    else
        zero(T)
    end
end

Base.sum(f::Fun{HS}) where {HS<:HeavisideSpace} = dotu(f.coefficients,diff(space(f).domain.points))
function Base.sum(f::Fun{SplineSpace{1,T,R}}) where {T,R}
    vals=pad(f.coefficients,dimension(space(f)))
    dfs=diff(space(f).domain.points)
    ret=vals[1]*dfs[1]/2
    for k=2:length(vals)-1
        ret+=vals[k]*(dfs[k]+dfs[k+1])/2
    end
    ret+=vals[end]*dfs[end]/2
    ret
end

#diffentiate HeavisideSpace
function differentiate(f::Fun{<:HeavisideSpace})
    dp=domain(f).points
    cfs=f.coefficients
    diff=0.0
    for n=1:length(cfs)
        diff=diff+cfs[n]*(DiracDelta(dp[n])-DiracDelta(dp[n+1]))
    end
    return diff
end

#Derivative Operator for HeavisideSpace
function Derivative(H::HeavisideSpace, k::Int)
    @assert k == 1
    ConcreteDerivative(H)
end

differentiate(f::Fun{SplineSpace{1,T,R}}) where {T,R} =
    Fun(HeavisideSpace(space(f).domain),
        diff(pad(f.coefficients,dimension(space(f))))./diff(space(f).domain.points))

integrate(f::Fun{HeavisideSpace{T,R}}) where {T,R} =
    Fun(SplineSpace{1,T,R}(space(f).domain),
        [0;cumsum(f.coefficients).*diff(space(f).domain.points)])


## DotProduct is not following the PreMetric rules since d(x, x) != 0 and d(x, y) >= 0 for all x, y
struct DotProduct <: Distances.UnionPreMetric end

@inline function Distances._evaluate(::DotProduct, a::AbstractVector, b::AbstractVector)
    @boundscheck if length(a) != length(b)
        throw(
            DimensionMismatch(
                "first array has length $(length(a)) which does not match the length of the second, $(length(b)).",
            ),
        )
    end
    return dot(a, b)
end

Distances.result_type(::DotProduct, Ta::Type, Tb::Type) = promote_type(Ta, Tb)

@inline Distances.eval_op(::DotProduct, a::Real, b::Real) = a * b
@inline function (dist::DotProduct)(a::AbstractArray, b::AbstractArray)
    return Distances._evaluate(dist, a, b)
end
@inline (dist::DotProduct)(a::Number, b::Number) = a * b


# ---
# title: 1423. Maximum Points You Can Obtain from Cards
# id: problem1423
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array, Dynamic Programming, Sliding Window
# link: <https://leetcode.com/problems/maximum-points-you-can-obtain-from-cards/description/>
# hidden: true
# ---
#
# There are several cards  **arranged in a row** , and each card has an
# associated number of points The points are given in the integer array
# `cardPoints`.
#
# In one step, you can take one card from the beginning or from the end of the
# row. You have to take exactly `k` cards.
#
# Your score is the sum of the points of the cards you have taken.
#
# Given the integer array `cardPoints` and the integer `k`, return the _maximum
# score_ you can obtain.
#
#
#
# **Example 1:**
#
#
#
#     Input: cardPoints = [1,2,3,4,5,6,1], k = 3
#     Output: 12
#     Explanation: After the first step, your score will always be 1. However, choosing the rightmost card first will maximize your total score. The optimal strategy is to take the three cards on the right, giving a final score of 1 + 6 + 5 = 12.
#
#
# **Example 2:**
#
#
#
#     Input: cardPoints = [2,2,2], k = 2
#     Output: 4
#     Explanation: Regardless of which two cards you take, your score will always be 4.
#
#
# **Example 3:**
#
#
#
#     Input: cardPoints = [9,7,7,9,7,7,9], k = 7
#     Output: 55
#     Explanation: You have to take all the cards. Your score is the sum of points of all cards.
#
#
# **Example 4:**
#
#
#
#     Input: cardPoints = [1,1000,1], k = 1
#     Output: 1
#     Explanation: You cannot take the card in the middle. Your best score is 1.
#
#
# **Example 5:**
#
#
#
#     Input: cardPoints = [1,79,80,1,1,1,200,1], k = 3
#     Output: 202
#
#
#
#
# **Constraints:**
#
#   * `1 <= cardPoints.length <= 10^5`
#   * `1 <= cardPoints[i] <= 10^4`
#   * `1 <= k <= cardPoints.length`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using LinearAlgebra
using Distributions: InverseWishart
using Random
using HighFrequencyCovariance

twister = MersenneTwister(10)
IW_dist(n) = InverseWishart(n, Matrix(Float64.(I(n))))

# slightly non psd non diagonal version.
A = Hermitian([1.02 0.5 0.9; 0.5 0.98 0.9; 0.9 0.9 1.01])
# Diagonal weighting.
W_root = Diagonal(rand(twister, size(A)[1]))
invW = inv(W_root^2)
U_ed = project_to_U(A, invW)
all( abs.(diag(U_ed) .- 1.0) .< 100*eps())
isa(U_ed, Hermitian)
S_ed = project_to_S(A, W_root)
all(eigen(S_ed).values .> -100*eps())

# hermitian weighting.
W_root = Hermitian(rand(twister, IW_dist(3)  ))
invW = inv(W_root^2)
U_ed = project_to_U(A, invW)
all( abs.(diag(U_ed) .- 1.0) .< 100*eps())
isa(U_ed, Hermitian)
S_ed = project_to_S(A, W_root)
all(eigen(S_ed).values .> -100*eps())

# Really non psd non diagonal version.
A = Hermitian([1.02 0.5 0.9; 0.5 0.98 -0.9; 0.9 -0.9 1.01])
# Diagonal weighting.
W_root = Diagonal(rand(twister, size(A)[1]))
invW = inv(W_root^2)
U_ed = project_to_U(A, invW)
all( abs.(diag(U_ed) .- 1.0) .< 100*eps())
isa(U_ed, Hermitian)
S_ed = project_to_S(A, W_root)
all(eigen(S_ed).values .> -100*eps())

# hermitian weighting.
W_root = Hermitian(rand(twister, IW_dist(3)  ))
invW = inv(W_root^2)
U_ed = project_to_U(A, invW)
all( abs.(diag(U_ed) .- 1.0) .< 100*eps())
isa(U_ed, Hermitian)
S_ed = project_to_S(A, W_root)
all(eigen(S_ed).values .> -100000*eps())

# Diagonal A
A = Diagonal([1.02 0 0; 0 0.98 0; 0 0 1.05])
# Diagonal weighting.
W_root = Diagonal(rand(twister, size(A)[1]))
invW = inv(W_root^2)
U_ed = project_to_U(A, invW)
all( abs.(diag(U_ed) .- 1.0) .< 100*eps())
isa(U_ed, Diagonal)
S_ed = project_to_S(A, W_root)
all(eigen(S_ed).values .> -100*eps())

# hermitian weighting.
W_root = Hermitian(rand(twister, IW_dist(3)  ))
invW = inv(W_root^2)
U_ed = project_to_U(A, invW)
all( abs.(diag(U_ed) .- 1.0) .< 100*eps())
isa(U_ed, Hermitian)
S_ed = project_to_S(A, W_root)
all(eigen(S_ed).values .> -100000*eps())

# Copypasting this here to avoid exporting it
function sqrt_psd(A::Hermitian)
    eigenvalues, eigenvectors = eigen(A)
    return Hermitian(construct_matrix_from_eigen(sqrt.(eigenvalues), eigenvectors))
end
sqrt_psd(A::Diagonal) = sqrt(A)

# A identify matrix case.
N = 4
A = Hermitian(I(N))
W_root = Hermitian(rand(twister, IW_dist(N)  ))
Dykstra = Hermitian(zeros(N,N))
W_inv = inv(W_root^2)
W_inv_sqrt = sqrt_psd(W_inv)
Y = A

Y, Dykstra = iterate_higham(Y, Dykstra, W_root, W_inv, W_inv_sqrt)
all(abs.(Y .- I(N)) .< 100*eps())
Y, Dykstra = iterate_higham(Y, Dykstra, W_root, W_inv, W_inv_sqrt)
all(abs.(Y .- I(N)) .< 100*eps())

# If A were input as diagonal
A = Diagonal(I(N))
Dykstra = Diagonal(zeros(N,N))
Y = A
Y, Dykstra = iterate_higham(Y, Dykstra, W_root, W_inv, W_inv_sqrt)
all(abs.(Y .- I(N)) .< 100*eps())
Y, Dykstra = iterate_higham(Y, Dykstra, W_root, W_inv, W_inv_sqrt)
all(abs.(Y .- I(N)) .< 100*eps())

# If A and E_root were diagonal
A = Diagonal(Float64.(I(N)))
W_root = Diagonal(rand(twister, N))
Dykstra = Diagonal(zeros(N,N))
W_inv = inv(W_root^2)
W_inv_sqrt = sqrt_psd(W_inv)
Y = A
Y, Dykstra = iterate_higham(Y, Dykstra, W_root, W_inv, W_inv_sqrt)
all(abs.(Y .- I(N)) .< 100*eps())
Y, Dykstra = iterate_higham(Y, Dykstra, W_root, W_inv, W_inv_sqrt)
all(abs.(Y .- I(N)) .< 100*eps())

# A valid correlation matrix case
N = 3
A, _ = cov2cor(Hermitian(rand(twister, IW_dist(N)  )))
W_root = Hermitian(rand(twister, IW_dist(N)  ))
Dykstra = Hermitian(zeros(N,N))
W_inv = inv(W_root^2)
W_inv_sqrt = sqrt_psd(W_inv)
Y = A

Y, Dykstra = iterate_higham(Y, Dykstra, W_root, W_inv, W_inv_sqrt)
all(abs.(Y .- A) .< 100*eps())
Y, Dykstra = iterate_higham(Y, Dykstra, W_root, W_inv, W_inv_sqrt)
all(abs.(Y .- A) .< 100*eps())

# A Slightly non psd case
N = 3
A = Hermitian([1.02 0.5 0.9; 0.5 0.98 0.9; 0.9 0.9 1.01])
W = Hermitian(abs.(cov2cor(Hermitian(rand(twister, IW_dist(N)  )))[1]))
updated_matrix, counter, convergence = nearest_correlation_matrix(A, W; doDykstra = true, stop_at_first_correlation_matrix = true, max_iterates = 1000)
updated_matrix2, counter2, convergence2 = nearest_correlation_matrix(A, W; doDykstra = true, stop_at_first_correlation_matrix = false, max_iterates = 10000)
updated_matrix3, counter3, convergence3 = nearest_correlation_matrix(A, W; doDykstra = false, stop_at_first_correlation_matrix = true, max_iterates = 10000)
all(abs.(updated_matrix .- updated_matrix2) .< 100*eps()) # These should be the same as once it is psd the matrix stops updating.
any(abs.(updated_matrix .- updated_matrix3) .> 0.0001) # These should be different as without Dykstra one is biased
# Testing convergences.
minimum(eigen(updated_matrix).values) > -10*eps()
minimum(eigen(updated_matrix2).values) > -10*eps()
minimum(eigen(updated_matrix3).values) > -10*eps()

# With diagonal weighting.
N = 3
A = Hermitian([1.02 0.5 0.9; 0.5 0.98 0.9; 0.9 0.9 1.01])
W = Diagonal(rand(twister,N))
updated_matrix, counter, convergence = nearest_correlation_matrix(A, W; doDykstra = true, stop_at_first_correlation_matrix = true, max_iterates = 1000)
updated_matrix2, counter2, convergence2 = nearest_correlation_matrix(A, W; doDykstra = true, stop_at_first_correlation_matrix = false, max_iterates = 10000)
updated_matrix3, counter3, convergence3 = nearest_correlation_matrix(A, W; doDykstra = false, stop_at_first_correlation_matrix = true, max_iterates = 10000)
all(abs.(updated_matrix .- updated_matrix2) .< 100*eps()) # These should be the same as once it is psd the matrix stops updating.
any(abs.(updated_matrix .- updated_matrix3) .> 0.0001) # These should be different as without Dykstra one is biased
# Testing convergences.
minimum(eigen(updated_matrix).values) > -10*eps()
minimum(eigen(updated_matrix2).values) > -10*eps()
minimum(eigen(updated_matrix3).values) > -10*eps()

# A really non psd case
A = Hermitian([1.02 0.5 0.9; 0.5 0.98 -0.9; 0.9 -0.9 1.01])
W = Hermitian(abs.(cov2cor(Hermitian(rand(twister, IW_dist(N)  )))[1]))
updated_matrix, counter, convergence = nearest_correlation_matrix(A, W; doDykstra = true, stop_at_first_correlation_matrix = true, max_iterates = 100000)
updated_matrix2, counter2, convergence2 = nearest_correlation_matrix(A, W; doDykstra = true, stop_at_first_correlation_matrix = false, max_iterates = 10000)
updated_matrix3, counter3, convergence3 = nearest_correlation_matrix(A, W; doDykstra = false, stop_at_first_correlation_matrix = true, max_iterates = 10000)
all(abs.(updated_matrix .- updated_matrix2) .< 1000*eps()) # These should be the same as once it is psd the matrix stops updating.
any(abs.(updated_matrix .- updated_matrix3) .> 0.0001) # These should be different as without Dykstra one is biased
# Testing convergences.
minimum(eigen(updated_matrix).values) > -10*eps()
minimum(eigen(updated_matrix2).values) > -10*eps()


"""
    Line(p1, p2)

Construct a `Line` geometric object (representing a segment) with given beginning and end
points, which are generally 2-vectors or 3-vectors.
"""
struct Line{N, T <: Real}
    p1::SVector{N, T}
    p2::SVector{N, T}
end

function Line(p1::StaticVector{N, T}, p2::StaticVector{N, T}) where {N, T <: Real}
    Line{N, T}(p1, p2)
end

function Line{N}(p1::AbstractVector{T}, p2::AbstractVector{T}) where {N, T <: Real}
    Line{N, T}(p1, p2)
end

# Note: This doesn't compare the ordering of points
function Base.:(==)(l1::Line{N}, l2::Line{N}) where N
    return l1.p1 == l2.p1 && l1.p2 == l2.p2
end

function Base.isequal(l1::Line{N}, l2::Line{N}) where N
    isequal(l1.p1, l2.p1) && isequal(l1.p2, l2.p2)
end

function Base.hash(l::Line, h::UInt)
    hash(l.p1, hash(l.p2, hash(UInt === UInt64 ? 0x627c5acc5b1e3d3d : 0x94c690be, h)))
end

# Use getindex to get points along the line
length(l::Line) = norm(l.p1 - l.p2)
function getindex(l::Line, t)
    t = t/length(l)
    (1-t)*l.p1 + t*l.p2
end
if VERSION < v"0.7"
    endof(l::Line) = length(l) # e.g. l[end - 1.0] is the point 1 metre from end of line
else
    lastindex(l::Line) = length(l) # e.g. l[end - 1.0] is the point 1 metre from end of line
end

eltype(::Line{N, T}) where {N, T} = SVector{N,T}
eltype(::Type{Line{N, T}}) where {N, T} = SVector{N,T}

getindex(l::Line, v::AbstractVector) = map(x -> l[x], v)

# Transformations (Affine only, for the moment)
(trans::Translation{V})(l::Line) where {V} = Line(trans(l.p1), trans(l.p2))
(trans::LinearMap{M})(l::Line) where {M} = Line(trans(l.p1), trans(l.p2))
(trans::AffineMap{M,V})(l::Line) where {M,V} = Line(trans(l.p1), trans(l.p2))

# Geometric operations
convert2d(l::Line{2}) = l
convert2d(l::Line{3}) = Line(convert2d(l.p1), convert2d(l.p2))
convert3d(l::Line{2}) = Line(convert3d(l.p1), convert3d(l.p2))
convert3d(l::Line{2}, z::Real) = Line(convert3d(l.p1, z), convert3d(l.p2, z))
convert3d(l::Line{3}) = l
convert3d(l::Line{3}, z::Real) = Line(convert3d(l.p1, z), convert3d(l.p2, z))

function intersects(l1::Line, l2::Line)
    return distance(l1, l2) == 0
end

function show(io::IO, l::Line{N}) where N
    print(io, "Line{$N}(")
    print(io, l.p1)
    print(io, ", ")
    print(io, l.p2)
    print(io, ")")
end

lines(l::Line) = [l]


function _scale_value(k::T1, from::NTuple{2,T2}, to::NTuple{2,T3}) where {T1 <: Number, T2 <: Number, T3 <: Number}
    m, M = from
    l, U = to
    inter = (k-m)/(M-m)
    return inter*(U-l)+l
end

function finish_layout!(L; link=false)
    x = getfield.(values(L), :x)
    y = getfield.(values(L), :y)
    range_x = (minimum(x), maximum(x))
    range_y = (minimum(y), maximum(y))
    range_total = (minimum([range_x..., range_y...]), maximum([range_x..., range_y...]))
    if link
        range_x = range_total
        range_y = range_total
    end
    distribute_layout!(L, range_x, range_y)
end

function distribute_layout!(L, X::NTuple{2,T}, Y::NTuple{2,T}) where {T <: Number}
    x = getfield.(values(L), :x)
    y = getfield.(values(L), :y)
    range_x = (minimum(x), maximum(x))
    range_y = (minimum(y), maximum(y))
    for k in keys(L)
        L[k].x = EcologicalNetworksPlots._scale_value(L[k].x, range_x, X)
        L[k].y = EcologicalNetworksPlots._scale_value(L[k].y, range_y, Y)
    end
end

function _spread_levels!(L; ratio::Float64=1.0)
    for s in keys(L)
        L[s].y *= ratio
    end
end


module m_stokes_driven_tht6_veclap
using Test
using LinearAlgebra
using StaticArrays
using MeshCore: nrelations, nentities
using MeshSteward: T6block, T6toT3
using MeshSteward: Mesh, attach!, baseincrel, boundary
using MeshSteward: vselect, geometry, summary
using MeshSteward: vtkwrite
using Elfel.RefShapes: manifdim, manifdimv
using Elfel.FElements: FEH1_T6, FEH1_T3, refshape, Jacobian
using Elfel.FESpaces: FESpace, ndofs, setebc!, nunknowns, doftype
using Elfel.FESpaces: numberfreedofs!, numberdatadofs!, numberdofs!
using Elfel.FESpaces: scattersysvec!, makeattribute, gathersysvec!, edofcompnt
using Elfel.FESpaces: highestfreedofnum, highestdatadofnum
using Elfel.FEIterators: FEIterator, ndofsperel, elnodes, eldofs
using Elfel.FEIterators: jacjac
using Elfel.QPIterators: QPIterator, bfun, bfungrad, weight
using Elfel.Assemblers: SysmatAssemblerSparse, start!, finish!, assemble!
using Elfel.Assemblers: SysvecAssembler
using Elfel.LocalAssemblers: LocalMatrixAssembler, LocalVectorAssembler, init!
# using UnicodePlots

mu = 0.25 # dynamic viscosity
A = 1.0 # length of the side of the square
N = 100;# number of subdivisions along the sides of the square domain

function genmesh()
    # Hood-Taylor pair of meshes is needed
    # This mesh will be for the velocities
    vmesh = Mesh()
    attach!(vmesh, T6block(A, A, N, N), "velocity")
    # This mesh will be used for the pressures. Notice that it needs to be
        # "compatible" with the velocity mesh in the sense that they need to share
        # the nodes at the corners of the triangles.
        pmesh = Mesh()
    attach!(pmesh, T6toT3(baseincrel(vmesh, "velocity")), "pressure")
    return vmesh, pmesh
end

function assembleK(uxfesp, uyfesp, pfesp, tndof, mu)
    function integrateK!(ass, elits, qpits, mu)
        uxnedof, uynedof, pnedof = ndofsperel.(elits)
        kuxux = LocalMatrixAssembler(uxnedof, uxnedof, 0.0)
        kuyuy = LocalMatrixAssembler(uynedof, uynedof, 0.0)
        kuxp = LocalMatrixAssembler(uxnedof, pnedof, 0.0)
        kuyp = LocalMatrixAssembler(uynedof, pnedof, 0.0)
        for el in zip(elits...)
            uxel, uyel, pel = el
            init!(kuxux, eldofs(uxel), eldofs(uxel))
            init!(kuyuy, eldofs(uyel), eldofs(uyel))
            init!(kuxp, eldofs(uxel), eldofs(pel))
            init!(kuyp, eldofs(uyel), eldofs(pel))
            for qp in zip(qpits...)
                uxqp, uyqp, pqp = qp
                Jac, J = jacjac(pel, pqp)
                JxW = J * weight(pqp)
                gradNp = bfungrad(pqp, Jac)
                gradNux = bfungrad(uxqp, Jac)
                gradNuy = bfungrad(uyqp, Jac)
                Np = bfun(pqp)
                for j in 1:uxnedof, i in 1:uxnedof
                    kuxux[i, j] += (mu * JxW) * (2 * gradNux[i][1] * gradNux[j][1] + gradNux[i][2] * gradNux[j][2])
                end
                for j in 1:uynedof, i in 1:uynedof
                    kuyuy[i, j] += (mu * JxW) * (gradNuy[i][1] * gradNuy[j][1] + 2 * gradNuy[i][2] * gradNuy[j][2])
                end
                for j in 1:pnedof, i in 1:uxnedof
                    kuxp[i, j] += (-JxW) * (gradNux[i][1] * Np[j])
                end
                for j in 1:pnedof, i in 1:uynedof
                    kuyp[i, j] += (-JxW) * (gradNuy[i][2] * Np[j])
                end
            end
            assemble!(ass, kuxux)
            assemble!(ass, kuyuy)
            assemble!(ass, kuxp)
            assemble!(ass, transpose(kuxp))
            assemble!(ass, kuyp)
            assemble!(ass, transpose(kuyp))
        end
        return ass
    end

    elits = (FEIterator(uxfesp), FEIterator(uyfesp), FEIterator(pfesp))
    qargs = (kind = :default, npts = 3,)
    qpits = (QPIterator(uxfesp, qargs), QPIterator(uyfesp, qargs), QPIterator(pfesp, qargs))
    ass = SysmatAssemblerSparse(0.0)
    start!(ass, tndof, tndof)
    integrateK!(ass, elits, qpits, mu)
    return finish!(ass)
end

function solve!(U, K, F, nu)
    KT = K * U
    U[1:nu] = K[1:nu, 1:nu] \ (F[1:nu] - KT[1:nu])
end

function test()
    vmesh, pmesh = genmesh()
    # Velocity spaces
    uxfesp = FESpace(Float64, vmesh, FEH1_T6(), 1)
    uyfesp = FESpace(Float64, vmesh, FEH1_T6(), 1)
    locs = geometry(vmesh)
    inflate = A / N / 100
    # Part of the boundary that is immovable
    boxes = [[0.0 A 0.0 0.0], [0.0 0.0 0.0 A], [A A 0.0 A]]
    for box in boxes
        vl = vselect(locs; box = box, inflate = inflate)
        for i in vl
            setebc!(uxfesp, 0, i, 1, 0.0)
            setebc!(uyfesp, 0, i, 1, 0.0)
        end
    end
    # The lid: driven in the X direction
    uxbar = 1.0
    box = [0.0 A A A]
    vl = vselect(locs; box = box, inflate = inflate)
    for i in vl
        setebc!(uxfesp, 0, i, 1, uxbar)
        setebc!(uyfesp, 0, i, 1, 0.0)
    end
    # Pressure space
    pfesp = FESpace(Float64, pmesh, FEH1_T3(), 1)
    setebc!(pfesp, 0, 1, 1, 0.0)
    # Number the degrees of freedom
    numberdofs!([uxfesp, uyfesp, pfesp])
    tndof = ndofs(uxfesp) + ndofs(uyfesp) + ndofs(pfesp)
    tnunk = nunknowns(uxfesp) + nunknowns(uyfesp) + nunknowns(pfesp)
    @test (tndof, tnunk) == (91003, 89402)
    # Assemble the coefficient matrix
    K = assembleK(uxfesp, uyfesp, pfesp, tndof, mu)
    # p = spy(K, canvas = DotCanvas)
    # display(p)
    # Solve the system
    U = fill(0.0, tndof)
    gathersysvec!(U, [uxfesp, uyfesp, pfesp])
    F = fill(0.0, tndof)
    solve!(U, K, F, tnunk)
    scattersysvec!([uxfesp, uyfesp, pfesp], U)
    # Postprocessing
    makeattribute(pfesp, "p", 1)
    makeattribute(uxfesp, "ux", 1)
    makeattribute(uyfesp, "uy", 1)
    vtkwrite("m_stokes_driven_tht6_veclap-p", baseincrel(pmesh), [(name = "p",), ])
    vtkwrite("m_stokes_driven_tht6_veclap-v", baseincrel(vmesh), [(name = "ux",), (name = "uy",)])
    try rm("m_stokes_driven_tht6_veclap-p.vtu"); catch end
    try rm("m_stokes_driven_tht6_veclap-v.vtu"); catch end
    true
end

end
using .m_stokes_driven_tht6_veclap
m_stokes_driven_tht6_veclap.test()

module m_stokes_driven_t3b

using LinearAlgebra
using StaticArrays
using MeshCore: nrelations, nentities, ir_identity
using MeshSteward: T3block
using MeshSteward: Mesh, attach!, baseincrel, boundary
using MeshSteward: vselect, geometry, summary
using MeshSteward: vtkwrite
using Elfel.RefShapes: manifdim, manifdimv
using Elfel.FElements: FEH1_T3_BUBBLE, FEH1_T3, refshape, Jacobian
using Elfel.FESpaces: FESpace, ndofs, setebc!, nunknowns, doftype
using Elfel.FESpaces: numberfreedofs!, numberdatadofs!, numberdofs!
using Elfel.FESpaces: scattersysvec!, makeattribute, gathersysvec!, edofcompnt
using Elfel.FESpaces: highestfreedofnum, highestdatadofnum
using Elfel.FEIterators: FEIterator, ndofsperel, elnodes, eldofs
using Elfel.FEIterators: jacjac
using Elfel.QPIterators: QPIterator, bfun, bfungrad, weight
using Elfel.Assemblers: SysmatAssemblerSparse, start!, finish!, assemble!
using Elfel.Assemblers: SysvecAssembler
using Elfel.LocalAssemblers: LocalMatrixAssembler, LocalVectorAssembler, init!
# using UnicodePlots

mu = 0.25 # dynamic viscosity
A = 1.0 # length of the side of the square
N = 100;# number of subdivisions along the sides of the square domain

function genmesh()
    # This mesh will be both for the velocities and for the pressure
    mesh = Mesh()
    attach!(mesh, T3block(A, A, N, N), "velocity+pressure")
    ir = baseincrel(mesh)
    eidir = ir_identity(ir)
    attach!(mesh, eidir)
    return mesh
end

function assembleK(uxfesp, uyfesp, pfesp, tndof, mu)
    function integrateK!(ass, elits, qpits, mu)
        uxnedof, uynedof, pnedof = ndofsperel.(elits)
        kuxux = LocalMatrixAssembler(uxnedof, uxnedof, 0.0)
        kuyuy = LocalMatrixAssembler(uynedof, uynedof, 0.0)
        kuxuy = LocalMatrixAssembler(uxnedof, uynedof, 0.0)
        kuxp = LocalMatrixAssembler(uxnedof, pnedof, 0.0)
        kuyp = LocalMatrixAssembler(uynedof, pnedof, 0.0)
        for el in zip(elits...)
            uxel, uyel, pel = el
            init!(kuxux, eldofs(uxel), eldofs(uxel))
            init!(kuyuy, eldofs(uyel), eldofs(uyel))
            init!(kuxuy, eldofs(uxel), eldofs(uyel))
            init!(kuxp, eldofs(uxel), eldofs(pel))
            init!(kuyp, eldofs(uyel), eldofs(pel))
            for qp in zip(qpits...)
                uxqp, uyqp, pqp = qp
                Jac, J = jacjac(pel, pqp)
                JxW = J * weight(pqp)
                gradNp = bfungrad(pqp, Jac)
                gradNux = bfungrad(uxqp, Jac)
                gradNuy = bfungrad(uyqp, Jac)
                Np = bfun(pqp)
                for j in 1:uxnedof, i in 1:uxnedof
                    kuxux[i, j] += (mu * JxW) * (2 * gradNux[i][1] * gradNux[j][1] + gradNux[i][2] * gradNux[j][2])
                end
                for j in 1:uynedof, i in 1:uynedof
                    kuyuy[i, j] += (mu * JxW) * (gradNuy[i][1] * gradNuy[j][1] + 2 * gradNuy[i][2] * gradNuy[j][2])
                end
                for j in 1:uynedof, i in 1:uxnedof
                    kuxuy[i, j] += (mu * JxW) * (gradNux[i][1] * gradNuy[j][2])
                end
                for j in 1:pnedof, i in 1:uxnedof
                    kuxp[i, j] += (-JxW) * (gradNux[i][1] * Np[j])
                end
                for j in 1:pnedof, i in 1:uynedof
                    kuyp[i, j] += (-JxW) * (gradNuy[i][2] * Np[j])
                end
            end
            assemble!(ass, kuxux)
            assemble!(ass, kuxuy)
            assemble!(ass, transpose(kuxuy))
            assemble!(ass, kuyuy)
            assemble!(ass, kuxp)
            assemble!(ass, transpose(kuxp))
            assemble!(ass, kuyp)
            assemble!(ass, transpose(kuyp))
        end
        return ass
    end

    elits = (FEIterator(uxfesp), FEIterator(uyfesp), FEIterator(pfesp))
    qargs = (kind = :default, npts = 3,)
    qpits = (QPIterator(uxfesp, qargs), QPIterator(uyfesp, qargs), QPIterator(pfesp, qargs))
    ass = SysmatAssemblerSparse(0.0)
    start!(ass, tndof, tndof)
    integrateK!(ass, elits, qpits, mu)
    return finish!(ass)
end

function solve!(U, K, F, nu)
    KT = K * U
    U[1:nu] = K[1:nu, 1:nu] \ (F[1:nu] - KT[1:nu])
end

function test()
    mesh = genmesh()
    # Velocity spaces
    uxfesp = FESpace(Float64, mesh, FEH1_T3_BUBBLE(), 1)
    uyfesp = FESpace(Float64, mesh, FEH1_T3_BUBBLE(), 1)
    locs = geometry(mesh)
    inflate = A / N / 100
    # Part of the boundary that is immovable
    boxes = [[0.0 A 0.0 0.0], [0.0 0.0 0.0 A], [A A 0.0 A]]
    for box in boxes
        vl = vselect(locs; box = box, inflate = inflate)
        for i in vl
            setebc!(uxfesp, 0, i, 1, 0.0)
            setebc!(uyfesp, 0, i, 1, 0.0)
        end
    end
    # The lid
    uxbar = 1.0
    box = [0.0 A A A]
    vl = vselect(locs; box = box, inflate = inflate)
    for i in vl
        setebc!(uxfesp, 0, i, 1, uxbar)
        setebc!(uyfesp, 0, i, 1, 0.0)
    end
    # Pressure space
    pfesp = FESpace(Float64, mesh, FEH1_T3(), 1)
    setebc!(pfesp, 0, 1, 1, 0.0)
    # Number the degrees of freedom
    numberdofs!([uxfesp, uyfesp, pfesp])
    tndof = ndofs(uxfesp) + ndofs(uyfesp) + ndofs(pfesp)
    tnunk = nunknowns(uxfesp) + nunknowns(uyfesp) + nunknowns(pfesp)
    # Assemble the coefficient matrix
    K = assembleK(uxfesp, uyfesp, pfesp, tndof, mu)
    # p = spy(K, canvas = DotCanvas)
    # display(p)
    # Solve the system
    U = fill(0.0, tndof)
    gathersysvec!(U, [uxfesp, uyfesp, pfesp])
    F = fill(0.0, tndof)
    solve!(U, K, F, tnunk)
    scattersysvec!([uxfesp, uyfesp, pfesp], U)
    # Postprocessing
    makeattribute(pfesp, "p", 1)
    makeattribute(uxfesp, "ux", 1)
    makeattribute(uyfesp, "uy", 1)
    vtkwrite("m_stokes_driven_t3b-p", baseincrel(mesh), [(name = "p",), ])
    vtkwrite("m_stokes_driven_t3b-v", baseincrel(mesh), [(name = "ux",), (name = "uy",)])
    try rm("m_stokes_driven_t3b-p.vtu"); catch end
    try rm("m_stokes_driven_t3b-v.vtu"); catch end
    true
end

end
using .m_stokes_driven_t3b
m_stokes_driven_t3b.test()

module mstok1

using Test

"""
    th_p2_p1

The manufactured-solution colliding flow example from Elman et al 2014. The
Hood-Taylor formulation with quadratic triangles for the velocity and continuous
pressure on linear triangles.

The formulation is the one derived in Reddy, Introduction to the finite element
method, 1993. Page 486 ff.
"""
module th_p2_p1

using Test
using LinearAlgebra
using StaticArrays
using MeshCore: nrelations, nentities, ir_identity, attribute, VecAttrib
using MeshSteward: T6block, T6toT3
using MeshSteward: Mesh, attach!, baseincrel, boundary
using MeshSteward: vselect, geometry, summary, transform
using MeshSteward: vtkwrite
using Elfel.RefShapes: manifdim, manifdimv
using Elfel.FElements: FEH1_T6, FEH1_T3, refshape, Jacobian
using Elfel.FESpaces: FESpace, ndofs, setebc!, nunknowns, doftype
using Elfel.FESpaces: numberfreedofs!, numberdatadofs!, numberdofs!
using Elfel.FESpaces: scattersysvec!, makeattribute, gathersysvec!, edofcompnt
using Elfel.FESpaces: highestfreedofnum, highestdatadofnum
using Elfel.FEIterators: FEIterator, ndofsperel, elnodes, eldofs, eldofvals
using Elfel.FEIterators: jacjac, location
using Elfel.QPIterators: QPIterator, bfun, bfungrad, weight
using Elfel.Assemblers: SysmatAssemblerSparse, start!, finish!, assemble!
using Elfel.Assemblers: SysvecAssembler
using Elfel.LocalAssemblers: LocalMatrixAssembler, LocalVectorAssembler, init!
using UnicodePlots

mu = 1.0 # dynamic viscosity
A = 1.0 # half of the length of the side of the square
trueux = (x, y) -> 20 * x * y ^ 3
trueuy = (x, y) -> 5 * x ^ 4 - 5 * y ^ 4
truep = (x, y) -> 60 * x ^ 2 * y - 20 * y ^ 3

function genmesh(N)
    # Hood-Taylor pair of meshes is needed
    # This mesh will be for the velocities
    vmesh = Mesh()
    attach!(vmesh, T6block(2 * A, 2 * A, N, N), "velocity")
    ir = baseincrel(vmesh)
    transform(ir, x -> x .- A)
    # This mesh will be used for the pressures. Notice that it needs to be
    # "compatible" with the velocity mesh in the sense that they need to share
    # the nodes at the corners of the triangles.
    pmesh = Mesh()
    attach!(pmesh, T6toT3(baseincrel(vmesh, "velocity")), "pressure")
    return vmesh, pmesh
end

function assembleK(uxfesp, uyfesp, pfesp, tndof, mu)
    function integrateK!(ass, elits, qpits, mu)
        uxnedof, uynedof, pnedof = ndofsperel.(elits)
        kuxux = LocalMatrixAssembler(uxnedof, uxnedof, 0.0)
        kuyuy = LocalMatrixAssembler(uynedof, uynedof, 0.0)
        kuxuy = LocalMatrixAssembler(uxnedof, uynedof, 0.0)
        kuxp = LocalMatrixAssembler(uxnedof, pnedof, 0.0)
        kuyp = LocalMatrixAssembler(uynedof, pnedof, 0.0)
        for el in zip(elits...)
            uxel, uyel, pel = el
            init!(kuxux, eldofs(uxel), eldofs(uxel))
            init!(kuyuy, eldofs(uyel), eldofs(uyel))
            init!(kuxuy, eldofs(uxel), eldofs(uyel))
            init!(kuxp, eldofs(uxel), eldofs(pel))
            init!(kuyp, eldofs(uyel), eldofs(pel))
            for qp in zip(qpits...)
                uxqp, uyqp, pqp = qp
                Jac, J = jacjac(pel, pqp)
                JxW = J * weight(pqp)
                gradNp = bfungrad(pqp, Jac)
                gradNux = bfungrad(uxqp, Jac)
                gradNuy = bfungrad(uyqp, Jac)
                Np = bfun(pqp)
                for j in 1:uxnedof, i in 1:uxnedof
                    kuxux[i, j] += (mu * JxW) * (2 * gradNux[i][1] * gradNux[j][1] + gradNux[i][2] * gradNux[j][2])
                end
                for j in 1:uynedof, i in 1:uynedof
                    kuyuy[i, j] += (mu * JxW) * (gradNuy[i][1] * gradNuy[j][1] + 2 * gradNuy[i][2] * gradNuy[j][2])
                end
                for j in 1:uynedof, i in 1:uxnedof
                    kuxuy[i, j] += (mu * JxW) * (gradNux[i][1] * gradNuy[j][2])
                end
                for j in 1:pnedof, i in 1:uxnedof
                    kuxp[i, j] += (-JxW) * (gradNux[i][1] * Np[j])
                end
                for j in 1:pnedof, i in 1:uynedof
                    kuyp[i, j] += (-JxW) * (gradNuy[i][2] * Np[j])
                end
            end
            assemble!(ass, kuxux)
            assemble!(ass, kuxuy)
            assemble!(ass, transpose(kuxuy))
            assemble!(ass, kuyuy)
            assemble!(ass, kuxp)
            assemble!(ass, transpose(kuxp))
            assemble!(ass, kuyp)
            assemble!(ass, transpose(kuyp))
        end
        return ass
    end

    elits = (FEIterator(uxfesp), FEIterator(uyfesp), FEIterator(pfesp))
    qargs = (kind = :default, npts = 3,)
    qpits = (QPIterator(uxfesp, qargs), QPIterator(uyfesp, qargs), QPIterator(pfesp, qargs))
    ass = SysmatAssemblerSparse(0.0)
    start!(ass, tndof, tndof)
    integrateK!(ass, elits, qpits, mu)
    return finish!(ass)
end

function solve!(U, K, F, nu)
    KT = K * U
    U[1:nu] = K[1:nu, 1:nu] \ (F[1:nu] - KT[1:nu])
end

function evaluate_pressure_error(pfesp)
    function integrate!(elit, qpit, truep)
        pnedof = ndofsperel(elit)
        E = 0.0
        for el in elit
            dofvals = eldofvals(el)
            for qp in qpit
                Jac, J = jacjac(el, qp)
                JxW = J * weight(qp)
                Np = bfun(qp)
                pt = truep(location(el, qp)...)
                pa = 0.0
                for j in 1:pnedof
                    pa += (dofvals[j] * Np[j])
                end
                E += (JxW) * (pa - pt)^2
            end
        end
        return sqrt(E)
    end

    elit = FEIterator(pfesp)
    qargs = (kind = :default, npts = 3,)
    qpit = QPIterator(pfesp, qargs)
    return integrate!(elit, qpit, truep)
end

function evaluate_velocity_error(uxfesp, uyfesp)
    function integrate!(elits, qpits, trueux, trueuy)
        uxnedof, uynedof = ndofsperel.(elits)
        E = 0.0
        for el in zip(elits...)
            uxel, uyel = el
            uxdofvals = eldofvals(uxel)
            uydofvals = eldofvals(uyel)
            for qp in zip(qpits...)
                uxqp, uyqp = qp
                Jac, J = jacjac(uxel, uxqp)
                JxW = J * weight(uxqp)
                Np = bfun(uxqp)
                uxt = trueux(location(uxel, uxqp)...)
                uyt = trueuy(location(uyel, uyqp)...)
                uxa = 0.0
                uya = 0.0
                for j in 1:uxnedof
                    uxa += (uxdofvals[j] * Np[j])
                    uya += (uydofvals[j] * Np[j])
                end
                E += (JxW) * ((uxa - uxt)^2 + (uya - uyt)^2)
            end
        end
        return sqrt(E)
    end

    elits = (FEIterator(uxfesp), FEIterator(uyfesp),)
    qargs = (kind = :default, npts = 3,)
    qpits = (QPIterator(uxfesp, qargs), QPIterator(uyfesp, qargs),)
    return integrate!(elits, qpits, trueux, trueuy)
end

function run(N)
    vmesh, pmesh = genmesh(N)
    # Velocity spaces
    uxfesp = FESpace(Float64, vmesh, FEH1_T6(), 1)
    uyfesp = FESpace(Float64, vmesh, FEH1_T6(), 1)
    locs = geometry(vmesh)
    inflate = A / N / 100
    # The entire boundary
    boxes = [[-A A -A -A], [-A -A -A A], [A A -A A], [-A A A A]]
    for box in boxes
        vl = vselect(locs; box = box, inflate = inflate)
        for i in vl
            setebc!(uxfesp, 0, i, 1, trueux(locs[i]...))
            setebc!(uyfesp, 0, i, 1, trueuy(locs[i]...))
        end
    end
    # Pressure space
    pfesp = FESpace(Float64, pmesh, FEH1_T3(), 1)
    atcenter = vselect(geometry(pmesh); nearestto = [0.0, 0.0])
    setebc!(pfesp, 0, atcenter[1], 1, 0.0)
    # Number the degrees of freedom
    numberdofs!([uxfesp, uyfesp, pfesp])
    tndof = ndofs(uxfesp) + ndofs(uyfesp) + ndofs(pfesp)
    tnunk = nunknowns(uxfesp) + nunknowns(uyfesp) + nunknowns(pfesp)
    # Assemble the coefficient matrix
    K = assembleK(uxfesp, uyfesp, pfesp, tndof, mu)
    # p = spy(K, canvas = DotCanvas)
    # display(p)
    # Solve the system
    U = fill(0.0, tndof)
    gathersysvec!(U, [uxfesp, uyfesp, pfesp])
    F = fill(0.0, tndof)
    solve!(U, K, F, tnunk)
    scattersysvec!([uxfesp, uyfesp, pfesp], U)
    # Postprocessing
    makeattribute(pfesp, "p", 1)
    makeattribute(uxfesp, "ux", 1)
    makeattribute(uyfesp, "uy", 1)
    ep = evaluate_pressure_error(pfesp)
    ev = evaluate_velocity_error(uxfesp, uyfesp)
    # vtkwrite("th_p2_p1-p", baseincrel(pmesh), [(name = "p",), ])
    # vtkwrite("th_p2_p1-v", baseincrel(vmesh), [(name = "ux",), (name = "uy",)])
    # geom = geometry(pfesp.mesh)
    # ir = baseincrel(pfesp.mesh)
    # pt = VecAttrib([truep(geom[i]...) for i in 1:length(geom)])
    # ir.right.attributes["pt"] = pt
    # vtkwrite("th_p2_p1-pt", baseincrel(pmesh), [(name = "pt",), ])
    (ep, ev)
end

end

function test()
    ref = [(3.5171450671095306, 0.2968271617227661),
    (0.5999467323539439, 0.03781189670123018),
    (0.12350320261417459, 0.004741849976722882)]
    N = 4
    for loop in 1:3
        ep, ev = th_p2_p1.run(N)
        @test isapprox(vec([ep, ev]), vec([ref[loop]...]))
        N = N * 2
    end
end

end

using .mstok1
mstok1.test()

"""
    m_th_p2_p1_veclap_alt

The manufactured-solution colliding flow example from Elman et al 2014. The
Hood-Taylor formulation with quadratic triangles for the velocity and continuous
pressure on linear triangles.

This implementation is an alternative: for the velocity, a single finite element
space with multiple components (2) is used instead of multiple finite element
spaces.

The formulation is the one derived in Donea, Huerta, Introduction to the finite element
method, 1993. Page 486 ff, and Elman, et al., Finite elements and fast
iterative solvers, p. 132. In other words, it is the vector Laplacian version.
"""
module m_th_p2_p1_veclap_alt

using Test
using LinearAlgebra
using StaticArrays
using MeshCore: nrelations, nentities, ir_identity, attribute, VecAttrib
using MeshSteward: T6block, T6toT3
using MeshSteward: Mesh, attach!, baseincrel, boundary
using MeshSteward: vselect, geometry, summary, transform
using MeshSteward: vtkwrite
using Elfel.RefShapes: manifdim, manifdimv
using Elfel.FElements: FEH1_T6, FEH1_T3, refshape, Jacobian
using Elfel.FESpaces: FESpace, ndofs, setebc!, nunknowns, doftype
using Elfel.FESpaces: numberfreedofs!, numberdatadofs!, numberdofs!
using Elfel.FESpaces: scattersysvec!, makeattribute, gathersysvec!, edofcompnt
using Elfel.FESpaces: highestfreedofnum, highestdatadofnum
using Elfel.FEIterators: FEIterator, ndofsperel, elnodes, eldofs, eldofvals
using Elfel.FEIterators: jacjac, location
using Elfel.QPIterators: QPIterator, bfun, bfungrad, weight
using Elfel.Assemblers: SysmatAssemblerSparse, start!, finish!, assemble!
using Elfel.Assemblers: SysvecAssembler
using Elfel.LocalAssemblers: LocalMatrixAssembler, LocalVectorAssembler, init!
using UnicodePlots

mu = 1.0 # dynamic viscosity
A = 1.0 # half of the length of the side of the square
trueux = (x, y) -> 20 * x * y ^ 3
trueuy = (x, y) -> 5 * x ^ 4 - 5 * y ^ 4
truep = (x, y) -> 60 * x ^ 2 * y - 20 * y ^ 3

function genmesh(N)
    # Hood-Taylor pair of meshes is needed
    # This mesh will be for the velocities
    vmesh = Mesh()
    attach!(vmesh, T6block(2 * A, 2 * A, N, N), "velocity")
    ir = baseincrel(vmesh)
    transform(ir, x -> x .- A)
    # This mesh will be used for the pressures. Notice that it needs to be
    # "compatible" with the velocity mesh in the sense that they need to share
    # the nodes at the corners of the triangles.
    pmesh = Mesh()
    attach!(pmesh, T6toT3(baseincrel(vmesh, "velocity")), "pressure")
    return vmesh, pmesh
end

function assembleK(ufesp, pfesp, tndof, mu)
    function integrate!(ass, elits, qpits, mu)
        unedof, pnedof = ndofsperel.(elits)
        uedofcomp = edofcompnt(ufesp)
        kuu = LocalMatrixAssembler(unedof, unedof, 0.0)
        kup = LocalMatrixAssembler(unedof, pnedof, 0.0)
        for el in zip(elits...)
            uel, pel = el
            init!(kuu, eldofs(uel), eldofs(uel))
            init!(kup, eldofs(uel), eldofs(pel))
            for qp in zip(qpits...)
                uqp, pqp = qp
                Jac, J = jacjac(uel, uqp)
                JxW = J * weight(uqp)
                gradNu = bfungrad(uqp, Jac)
                Np = bfun(pqp)
                for j in 1:unedof, i in 1:unedof
                    if uedofcomp[i] == uedofcomp[j]
                        kuu[i, j] += (mu * JxW) * (dot(gradNu[i], gradNu[j]))
                    end
                end
                for j in 1:pnedof, i in 1:unedof
                    kup[i, j] += (-JxW * Np[j]) * gradNu[i][uedofcomp[i]]
                end
            end
            assemble!(ass, kuu)
            assemble!(ass, kup)
            assemble!(ass, transpose(kup))
        end
        return ass
    end

    elits = (FEIterator(ufesp), FEIterator(pfesp))
    qargs = (kind = :default, npts = 3,)
    qpits = (QPIterator(ufesp, qargs), QPIterator(pfesp, qargs))
    ass = SysmatAssemblerSparse(0.0)
    start!(ass, tndof, tndof)
    integrate!(ass, elits, qpits, mu)
    return finish!(ass)
end

function solve!(U, K, F, nu)
    KT = K * U
    U[1:nu] = K[1:nu, 1:nu] \ (F[1:nu] - KT[1:nu])
end

function evaluate_pressure_error(pfesp)
    function integrate!(elit, qpit, truep)
        pnedof = ndofsperel(elit)
        E = 0.0
        for el in elit
            dofvals = eldofvals(el)
            for qp in qpit
                Jac, J = jacjac(el, qp)
                JxW = J * weight(qp)
                Np = bfun(qp)
                pt = truep(location(el, qp)...)
                pa = 0.0
                for j in 1:pnedof
                    pa += (dofvals[j] * Np[j])
                end
                E += (JxW) * (pa - pt)^2
            end
        end
        return sqrt(E)
    end

    elit = FEIterator(pfesp)
    qargs = (kind = :default, npts = 3,)
    qpit = QPIterator(pfesp, qargs)
    return integrate!(elit, qpit, truep)
end

function evaluate_velocity_error(ufesp)
    function integrate!(elit, qpit, trueux, trueuy)
        unedof = ndofsperel(elit)
        uedofcomp = edofcompnt(ufesp)
        E = 0.0
        for el in elit
            udofvals = eldofvals(el)
            for qp in qpit
                Jac, J = jacjac(el, qp)
                JxW = J * weight(qp)
                Nu = bfun(qp)
                uxt = trueux(location(el, qp)...)
                uyt = trueuy(location(el, qp)...)
                uxa = 0.0
                uya = 0.0
                for j in 1:unedof
                    (uedofcomp[j] == 1) && (uxa += (udofvals[j] * Nu[j]))
                    (uedofcomp[j] == 2) && (uya += (udofvals[j] * Nu[j]))
                end
                E += (JxW) * ((uxa - uxt)^2 + (uya - uyt)^2)
            end
        end
        return sqrt(E)
    end

    elit = FEIterator(ufesp)
    qargs = (kind = :default, npts = 3,)
    qpit = QPIterator(ufesp, qargs)
    return integrate!(elit, qpit, trueux, trueuy)
end

function run(N)
    vmesh, pmesh = genmesh(N)
    # Velocity spaces
    ufesp = FESpace(Float64, vmesh, FEH1_T6(), 2)
    locs = geometry(vmesh)
    inflate = A / N / 100
    # The entire boundary
    boxes = [[-A A -A -A], [-A -A -A A], [A A -A A], [-A A A A]]
    for box in boxes
        vl = vselect(locs; box = box, inflate = inflate)
        for i in vl
            setebc!(ufesp, 0, i, 1, trueux(locs[i]...))
            setebc!(ufesp, 0, i, 2, trueuy(locs[i]...))
        end
    end
    # Pressure space
    pfesp = FESpace(Float64, pmesh, FEH1_T3(), 1)
    atcenter = vselect(geometry(pmesh); nearestto = [0.0, 0.0])
    setebc!(pfesp, 0, atcenter[1], 1, 0.0)
    # Number the degrees of freedom
    numberdofs!([ufesp, pfesp])
    tndof = ndofs(ufesp) + ndofs(pfesp)
    tnunk = nunknowns(ufesp) + nunknowns(pfesp)
    # Assemble the coefficient matrix
    K = assembleK(ufesp, pfesp, tndof, mu)
    # p = spy(K, canvas = DotCanvas)
    # display(p)
    # Solve the system
    U = fill(0.0, tndof)
    gathersysvec!(U, [ufesp, pfesp])
    F = fill(0.0, tndof)
    solve!(U, K, F, tnunk)
    scattersysvec!([ufesp, pfesp], U)
    # Postprocessing
    makeattribute(pfesp, "p", 1)
    makeattribute(ufesp, "ux", 1)
    makeattribute(ufesp, "uy", 2)
    ep = evaluate_pressure_error(pfesp)
    ev = evaluate_velocity_error(ufesp)
    # vtkwrite("m_th_p2_p1_veclap_alt-p", baseincrel(pmesh), [(name = "p",), ])
    # vtkwrite("m_th_p2_p1_veclap_alt-v", baseincrel(vmesh), [(name = "ux",), (name = "uy",)])
    return (ep, ev)
end

test() = begin
    ep, ev = m_th_p2_p1_veclap_alt.run(4)
    @test isapprox([ep, ev], [2.596076907594511, 0.3001331486426876])
end

end

using .m_th_p2_p1_veclap_alt
m_th_p2_p1_veclap_alt.test()

module m_th_p2_p1_veclap_alt_e

using Test
using LinearAlgebra
using StaticArrays
using MeshCore.Exports
using MeshSteward.Exports
using Elfel.Exports
using UnicodePlots

mu = 1.0 # dynamic viscosity
A = 1.0 # half of the length of the side of the square
trueux = (x, y) -> 20 * x * y ^ 3
trueuy = (x, y) -> 5 * x ^ 4 - 5 * y ^ 4
truep = (x, y) -> 60 * x ^ 2 * y - 20 * y ^ 3

function genmesh(N)
    # Hood-Taylor pair of meshes is needed
    # This mesh will be for the velocities
    vmesh = Mesh()
    attach!(vmesh, T6block(2 * A, 2 * A, N, N), "velocity")
    ir = baseincrel(vmesh)
    transform(ir, x -> x .- A)
    # This mesh will be used for the pressures. Notice that it needs to be
    # "compatible" with the velocity mesh in the sense that they need to share
    # the nodes at the corners of the triangles.
    pmesh = Mesh()
    attach!(pmesh, T6toT3(baseincrel(vmesh, "velocity")), "pressure")
    return vmesh, pmesh
end

function assembleK(ufesp, pfesp, tndof, mu)
    function integrate!(ass, elits, qpits, mu)
        unedof, pnedof = ndofsperel.(elits)
        uedofcomp = edofcompnt(ufesp)
        kuu = LocalMatrixAssembler(unedof, unedof, 0.0)
        kup = LocalMatrixAssembler(unedof, pnedof, 0.0)
        for el in zip(elits...)
            uel, pel = el
            init!(kuu, eldofs(uel), eldofs(uel))
            init!(kup, eldofs(uel), eldofs(pel))
            for qp in zip(qpits...)
                uqp, pqp = qp
                Jac, J = jacjac(uel, uqp)
                JxW = J * weight(uqp)
                gradNu = bfungrad(uqp, Jac)
                Np = bfun(pqp)
                for j in 1:unedof, i in 1:unedof
                    if uedofcomp[i] == uedofcomp[j]
                        kuu[i, j] += (mu * JxW) * (dot(gradNu[i], gradNu[j]))
                    end
                end
                for j in 1:pnedof, i in 1:unedof
                    kup[i, j] += (-JxW * Np[j]) * gradNu[i][uedofcomp[i]]
                end
            end
            assemble!(ass, kuu)
            assemble!(ass, kup)
            assemble!(ass, transpose(kup))
        end
        return ass
    end

    elits = (FEIterator(ufesp), FEIterator(pfesp))
    qargs = (kind = :default, npts = 3,)
    qpits = (QPIterator(ufesp, qargs), QPIterator(pfesp, qargs))
    ass = SysmatAssemblerSparse(0.0)
    start!(ass, tndof, tndof)
    integrate!(ass, elits, qpits, mu)
    return finish!(ass)
end

function solve!(U, K, F, nu)
    KT = K * U
    U[1:nu] = K[1:nu, 1:nu] \ (F[1:nu] - KT[1:nu])
end

function evaluate_pressure_error(pfesp)
    function integrate!(elit, qpit, truep)
        pnedof = ndofsperel(elit)
        E = 0.0
        for el in elit
            dofvals = eldofvals(el)
            for qp in qpit
                Jac, J = jacjac(el, qp)
                JxW = J * weight(qp)
                Np = bfun(qp)
                pt = truep(location(el, qp)...)
                pa = 0.0
                for j in 1:pnedof
                    pa += (dofvals[j] * Np[j])
                end
                E += (JxW) * (pa - pt)^2
            end
        end
        return sqrt(E)
    end

    elit = FEIterator(pfesp)
    qargs = (kind = :default, npts = 3,)
    qpit = QPIterator(pfesp, qargs)
    return integrate!(elit, qpit, truep)
end

function evaluate_velocity_error(ufesp)
    function integrate!(elit, qpit, trueux, trueuy)
        unedof = ndofsperel(elit)
        uedofcomp = edofcompnt(ufesp)
        E = 0.0
        for el in elit
            udofvals = eldofvals(el)
            for qp in qpit
                Jac, J = jacjac(el, qp)
                JxW = J * weight(qp)
                Nu = bfun(qp)
                uxt = trueux(location(el, qp)...)
                uyt = trueuy(location(el, qp)...)
                uxa = 0.0
                uya = 0.0
                for j in 1:unedof
                    (uedofcomp[j] == 1) && (uxa += (udofvals[j] * Nu[j]))
                    (uedofcomp[j] == 2) && (uya += (udofvals[j] * Nu[j]))
                end
                E += (JxW) * ((uxa - uxt)^2 + (uya - uyt)^2)
            end
        end
        return sqrt(E)
    end

    elit = FEIterator(ufesp)
    qargs = (kind = :default, npts = 3,)
    qpit = QPIterator(ufesp, qargs)
    return integrate!(elit, qpit, trueux, trueuy)
end

function run(N)
    vmesh, pmesh = genmesh(N)
    # Velocity spaces
    ufesp = FESpace(Float64, vmesh, FEH1_T6(), 2)
    locs = geometry(vmesh)
    inflate = A / N / 100
    # The entire boundary
    boxes = [[-A A -A -A], [-A -A -A A], [A A -A A], [-A A A A]]
    for box in boxes
        vl = vselect(locs; box = box, inflate = inflate)
        for i in vl
            setebc!(ufesp, 0, i, 1, trueux(locs[i]...))
            setebc!(ufesp, 0, i, 2, trueuy(locs[i]...))
        end
    end
    # Pressure space
    pfesp = FESpace(Float64, pmesh, FEH1_T3(), 1)
    atcenter = vselect(geometry(pmesh); nearestto = [0.0, 0.0])
    setebc!(pfesp, 0, atcenter[1], 1, 0.0)
    # Number the degrees of freedom
    numberdofs!([ufesp, pfesp])
    tndof = ndofs(ufesp) + ndofs(pfesp)
    tnunk = nunknowns(ufesp) + nunknowns(pfesp)
    # Assemble the coefficient matrix
    K = assembleK(ufesp, pfesp, tndof, mu)
    # p = spy(K, canvas = DotCanvas)
    # display(p)
    # Solve the system
    U = fill(0.0, tndof)
    gathersysvec!(U, [ufesp, pfesp])
    F = fill(0.0, tndof)
    solve!(U, K, F, tnunk)
    scattersysvec!([ufesp, pfesp], U)
    # Postprocessing
    makeattribute(pfesp, "p", 1)
    makeattribute(ufesp, "ux", 1)
    makeattribute(ufesp, "uy", 2)
    ep = evaluate_pressure_error(pfesp)
    ev = evaluate_velocity_error(ufesp)
    # vtkwrite("m_th_p2_p1_veclap_alt_e-p", baseincrel(pmesh), [(name = "p",), ])
    # vtkwrite("m_th_p2_p1_veclap_alt_e-v", baseincrel(vmesh), [(name = "ux",), (name = "uy",)])
    return (ep, ev)
end

test() = begin
    ep, ev = m_th_p2_p1_veclap_alt_e.run(4)
    @test isapprox([ep, ev], [2.596076907594511, 0.3001331486426876])
end

end

using .m_th_p2_p1_veclap_alt_e
m_th_p2_p1_veclap_alt_e.test()


"""
    Complexification{T <: Number} <: Plexification{T}

A complexification as a pair.
"""
struct Complexification{T <: Number} <: Plexification{T}
    l::T
    r::T
end

function (==)(x::Complexification, y::Complexification)
    x.l == y.l && x.r == y.r
end

function real(z::Complexification)
    real(z.l)
end

function zero(::Type{Complexification{T}}) where T <: Number
    Complexification{T}(zero(T), zero(T))
end

function zero(::Complexification{T}) where T <: Number
    zero(Complexification{T})
end

function one(::Type{Complexification{T}}) where T <: Number
    Complexification{T}(one(T), zero(T))
end

function one(::Complexification{T}) where T <: Number
    one(Complexification{T})
end

function (-)(z::Complexification)
    Complexification(-(z.l), -(z.r))
end

function (+)(x::Complexification{T}, y::Complexification{T}) where T <: Number
    Complexification{T}(x.l + y.l, x.r + y.r)
end

function (+)(a::T, z::Complexification{T}) where T <: Number
    Complexification{T}(a + z.l, z.r)
end

function (+)(z::Complexification{T}, a::T) where T <: Number
    Complexification{T}(z.l + a, z.r)
end

function (+)(a::Real, z::Complexification)
    Complexification(a + z.l, z.r)
end

function (+)(z::Complexification, a::Real)
    Complexification(z.l + a, z.r)
end

function (-)(x::Complexification{T}, y::Complexification{T}) where T <: Number
    Complexification{T}(x.l - y.l, x.r - y.r)
end

function (-)(a::T, z::Complexification{T}) where T <: Number
    Complexification{T}(a - z.l, -(z.r))
end

function (-)(z::Complexification{T}, a::T) where T <: Number
    Complexification{T}(z.l - a, z.r)
end

function (-)(a::Real, z::Complexification)
    Complexification(a - z.l, -(z.r))
end

function (-)(z::Complexification, a::Real)
    Complexification(z.l - a, z.r)
end

function (*)(x::Complexification{T}, y::Complexification{T}) where T <: Number
    Complexification{T}(
        (x.l * y.l) - (x.r * y.r),
        (x.l * y.r) + (x.r * y.l)
    )
end

function (*)(a::T, z::Complexification{T}) where T <: Number
    Complexification{T}(a * z.l, a * z.r)
end

function (*)(z::Complexification{T}, a::T) where T <: Number
    Complexification{T}(z.l * a , z.r * a)
end

function (*)(a::Real, z::Complexification)
    Complexification(a * z.l, a * z.r)
end

function (*)(z::Complexification, a::Real)
    Complexification(z.l * a, z.r * a)
end

function (/)(z::Complexification, a::Real)
    if iszero(a)
        error(ZeroDenominator)
    end

    Complexification(z.l / a, z.r / a)
end

function (\)(a::Real, z::Complexification)
    if iszero(a)
        error(ZeroDenominator)
    end

    Complexification(a \ z.l, a \ z.r)
end


type JointOrderStatistic{D,S<:MultivariateDistribution} <: AbstractJointOrderStatistic{D}
    sequence::S
    orders::Vector{Int64}
end

function JointOrderStatistic{D,S}(sequence::IIDRandomSequence{D,S},orders::Vector{Int64})
    return JointOrderStatistic{D,IIDRandomSequence{D,S}}(sequence,orders)
end

parentsequence(X::JointOrderStatistic) = X.sequence
orders(X::JointOrderStatistic) = X.orders

length(X::JointOrderStatistic) = length(orders(X))

function _rand!{T<:Real}(X::JointOrderStatistic, x::AbstractVector{T})
  y = sort(rand(parentsequence(X)))
  x = copy(y[orders(X)])
  return x
end

typealias IIDJointOrderStatistic{D<:ValueSupport,T<:UnivariateDistribution} JointOrderStatistic{D,IIDRandomSequence{D,T}}
typealias IIDContinuousJointOrderStatistic{T<:ContinuousUnivariateDistribution} IIDJointOrderStatistic{Continuous,T}
typealias IIDDiscreteJointOrderStatistic{T<:DiscreteUnivariateDistribution} IIDJointOrderStatistic{Discrete,T}

function _logpdf{T<:Real,S<:ContinuousUnivariateDistribution}(X::IIDContinuousJointOrderStatistic{S}, x::AbstractVector{T})
  N,ks = length(parentsequence(X)),orders(X)
  logps,Ps = logpdf(X.sequence.d,x),cdf(X.sequence.d,x)
  ΔPs,Δqs = diff([0.;Ps;1.]),diff([0;ks;ks[end]+1])-1
  return lfact(N)+sum(logps)+sum(Δqs.*log(ΔPs))- sum(lfact(Δqs))
end

function mean(X::IIDJointOrderStatistic)
  return map(mean,[ScalarOrderStatistic(parentsequence(X),order) for order in orders(X)])
end


# from https://github.com/FluxML/Zygote.jl/blob/master/test/gradcheck.jl

function ngradient(f, xs::AbstractArray...)
  grads = zero.(xs)
  for (x, Δ) in zip(xs, grads), i in 1:length(x)
    δ = sqrt(eps())
    tmp = x[i]
    x[i] = tmp - δ/2
    y1 = f(xs...)
    x[i] = tmp + δ/2
    y2 = f(xs...)
    x[i] = tmp
    Δ[i] = (y2-y1)/δ
  end
  return grads
end

function gradcheck(f, xs...)
    n_grads = ngradient(f, xs...)
    y_grads = Yota._grad(f, xs...)[2] |> collect
    all(isapprox.(n_grads, y_grads, rtol = 1e-5, atol = 1e-5))
end

# gradtest(f, xs::AbstractArray...) = gradcheck((xs...) -> sum(sin.(f(xs...))), xs...)
# gradtest(f, dims...) = gradtest(f, rand.(Float64, dims)...)


using GLAbstraction, GLPlot, Reactive

window = createdisplay(eyeposition=Vec3(2), lookat=Vec3(0.5))

##########################################################
# Volume

#So far just 1 dimensional color values are supported
N 		= 128
function func(x,y,z)
    R = sqrt(x^2 + y^2+z^2)
    sin(R)/R
end
volume 	= Float32[sin(x/15f0)+sin(y/15f0)+sin(z/15f0) for x=1:N, y=1:N, z=1:N]
max 	= maximum(volume)
min 	= minimum(volume)
volume 	= (volume .- min) ./ (max .- min)

#Filter keydown events
keypressed = window.inputs[:buttonspressed]

#Make some attributes intseractive
algorithm 	= foldl( (v0, v1) -> in('I', v1) ? 2f0 : in('M', v1) ? 1f0 : v0, 2f0, keypressed) # i for isosurface, m for MIP#
isovalue 	= foldl( (v0, v1) -> in(GLFW.KEY_UP, v1) ? (v0 + 0.01f0) : (in(GLFW.KEY_DOWN, v1) ? (v0 - 0.01f0) : v0), 0.5f0, keypressed)
stepsize 	= foldl( (v0, v1) -> in(GLFW.KEY_LEFT, v1) ? (v0 + 0.0001f0) : (in(GLFW.KEY_RIGHT, v1)  ? (v0 - 0.0001f0) : v0), 0.005f0, keypressed)
lift(println, isovalue)
glplot(volume, algorithm=algorithm, isovalue=isovalue, stepsize=stepsize, color=Vec3(1,0,0))

#glplot(imread("someexample.nrrd"), algorithm=algorithm, isovalue=isovalue, stepsize=stepsize, color=Vec3(1,0,0))

renderloop(window)


using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

compute(n::Int)::Int = (((4n + 3)n + 8)n - 9) ÷ 6

compute(5)

compute(1_001)

@benchmark compute(1_001)


type Graph
   backprop::Array{Function,1}
   doBackprop::Bool # backprop only needed during learning. Can turn off for prediction
   Graph() = new(Array(Function,0),true)
   Graph(backPropNeeded::Bool) = new(Array(Function,0),backPropNeeded)
end

function backprop!(g::Graph)
    for i = length(g.backprop):-1:1  g.backprop[i]() end
end

function rowpluck!(g::Graph, m::NNMatrix, ix::Int)
    # pluck a row of m and return it as a column vector
    out = NNMatrix(m.d, 1)
    out.w[:,1] = m.w[ix,:]'
    if g.doBackprop
        push!(g.backprop,
              function ()
                 m.dw[ix,:] += out.dw[:,1]'
              end )
    end
    return out
end

function concat!(g::Graph, ms::NNMatrix...)
    n = 0
    @inbounds for i in 1:length(ms)
        n += ms[i].n
    end
    out = NNMatrix(n, ms[1].d, zeros(n, ms[1].d), zeros(n, ms[1].d))
    n = 0
    @inbounds for m in ms
        for j = 1:m.d, i = 1:m.n
            out.w[i+n,j] = m.w[i,j]
        end
        n += m.n
    end
    if g.doBackprop
        push!(g.backprop,
              function ()
                  n = 0
                  @inbounds for m in ms
                      @inbounds for j = 1:m.d, i = 1:m.n
                          m.dw[i,j] += out.dw[i+n,j]
                      end
                      n += m.n
                  end
              end )
    end
    return out
end

function tanh!(g::Graph, m::NNMatrix)
    out = NNMatrix(m.n, m.d)
    out.w = tanh(m.w)
    if g.doBackprop
        push!(g.backprop,
              function ()
                  @inbounds for j = 1:m.d, i = 1:m.n
                      m.dw[i,j] += (1. - out.w[i,j]^2) * out.dw[i,j]
                  end
              end )
    end
    return out
end

function sigmoid!(g::Graph, m::NNMatrix)
    out = NNMatrix(m.n, m.d,
            [1.0 / (1.0 + exp(-m.w[i,j])) for i in 1:m.n, j in 1:m.d],
            zeros(m.n, m.d))
    if g.doBackprop
        push!(g.backprop,
              function ()
                  @inbounds for j = 1:m.d, i = 1:m.n
                      m.dw[i,j] +=  out.w[i,j] * (1. - out.w[i,j]) *  out.dw[i,j]
                  end
              end )
    end
    return out
end

function relu!(g::Graph, m::NNMatrix)
    out = NNMatrix(m.n, m.d)
    @inbounds for j = 1:m.d, i = 1:m.n
        out.w[i,j] = m.w[i,j] < 0. ? 0. : m.w[i,j]
    end
    if g.doBackprop
        push!(g.backprop,
          function ()
              @inbounds for j = 1:m.d, i = 1:m.n
                  m.dw[i,j] +=  m.w[i,j] < 0. ? 0. : out.dw[i,j]
              end
          end )
    end
    return out
end

function mul!(g::Graph, m1::NNMatrix, m2::NNMatrix)
    out = NNMatrix(m1.n, m2.d, m1.w * m2.w, zeros(m1.n, m2.d))
    if g.doBackprop
        push!(g.backprop,
            function ()
                @inbounds for i = 1:m1.n, j = 1:m2.d
                    b = out.dw[i,j]
                    for k = 1:m1.d
                        m1.dw[i,k] += m2.w[k,j] * b

                        m2.dw[k,j] += m1.w[i,k] * b
                    end
                end
            end )
    end
    return out
end

function mul!(g::Graph, m::NNMatrix, c::Float64)
    out = NNMatrix(m.n, m.d, m.w .* c, zeros(m.n, m.d))
    if g.doBackprop
        push!(g.backprop,
            function ()
                m.dw += out.dw .* c
            end )
    end
    return out
end

function add!(g::Graph, ms::NNMatrix...)
    out = NNMatrix(ms[1].n, ms[1].d, zeros(ms[1].n, ms[1].d), zeros(ms[1].n, ms[1].d))
    @inbounds for m in ms
        @inbounds for j in 1:m.d, i in 1:m.n
            out.w[i,j] += m.w[i,j]
        end
    end
    if g.doBackprop
        push!(g.backprop,
            function ()
                @inbounds for m in ms
                    @inbounds for j in 1:m.d, i in 1:m.n
                        m.dw[i,j] += out.dw[i,j]
                    end
                end
            end )
    end
    return out
end

function add!(g::Graph, m::NNMatrix, c::Float64)
    out = NNMatrix(m.n, m.d, m.w .+ c, zeros(m.n, m.d))
    if g.doBackprop
        push!(g.backprop,
            function ()
                m.dw += out.dw
            end )
    end
    return out
end

function eltmul!(g::Graph, m1::NNMatrix, m2::NNMatrix) # element-wise multiplication
    out = NNMatrix(m1.n, m2.d, m1.w .* m2.w, zeros(m1.n, m2.d))
    if g.doBackprop
        push!(g.backprop,
            function ()
                @inbounds for j in 1:m1.d, i in 1:m1.n
                  m1.dw[i,j] += m2.w[i,j] * out.dw[i,j]
                  m2.dw[i,j] += m1.w[i,j] * out.dw[i,j]
                end
              end )
        end
    return out
end


#%%
include("CellularBase.jl")
using .CellularBase
import .CellularBase: evolve!
using Plots
using Images
import StatsBase: sample
#%%
mutable struct WolframGrid <: AbstractGrid
	state::BitArray
	bc::BoundaryCondition
	neighborhood::AbstractNeighborhood
	possible_states::Array
	function WolframGrid(arr::BitArray)
		new(arr, Periodic, MooreNeighborhood(1, 1), [true, false])
	end
	function WolframGrid(size::Int)
		s = falses(size)
		s[size ÷ 2] = true
		WolframGrid(s)
	end
	function WolframGrid(trueratio::Float64, size::Int)
		ts = sample(1:size, ceil(Int, trueratio * size), replace=false)
		s = falses(size)
		s[ts] .= true
		WolframGrid(s)
	end
end

function (g::WolframGrid)(neighbors)
	sum = 0
	for (i, v) in enumerate(reverse(neighbors))
		sum += v * (2^(i - 1))
	end
	sum
end

function wr_table(rule::Int)
	Dict{Int,Int}([(i, rule >> i & 1) for i in 0:7])
end

function CellularBase.evolve!(grid::WolframGrid; kwargs...)
	tabular_evolution!(grid, wr_table(kwargs[:rule]))
end

function makeimg(k)
	k .|> (x -> x .== false) |> x->hcat(x...)' .|> Gray
end
#%%
k = simulate!(WolframGrid(100), 100, rule=26);
plot(makeimg(k))
#%%


#=
  Euler #18 in Julia.

  Problem 18:
  """
  By starting at the top of the triangle below and moving to adjacent
  numbers on the row below, the maximum total from top to bottom is 23.

  3
  7 4
  2 4 6
  8 5 9 3

  That is, 3 + 7 + 4 + 9 = 23.

  Find the maximum total from top to bottom of the triangle below:

   75
   95 64
   17 47 82
   18 35 87 10
   20 04 82 47 65
   19 01 23 75 03 34
   88 02 77 73 07 63 67
   99 65 04 28 06 16 70 92
   41 41 26 56 83 40 80 70 33
   41 48 72 33 47 32 37 16 94 29
   53 71 44 65 25 43 91 52 97 51 14
   70 11 33 28 77 73 17 78 39 68 17 57
   91 71 52 38 17 14 91 43 58 50 27 29 48
   63 66 04 68 89 53 67 30 73 16 69 87 40 31
   04 62 98 27 23 09 70 98 73 93 38 53 60 04 23

  NOTE: As there are only 16384 routes, it is possible to solve this problem
  by trying every route. However, Problem 67, is the same challenge with a
  triangle containing one-hundred rows; it cannot be solved by brute force,
  and requires a clever method! ;o)
  """

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

=#

include("Euler.jl")

# Recursive function for euler18()
function pp(row, col, sum, tri, max_val)
    if sum > max_val
        max_val = sum
    end
    row += 1
    if row <= length(tri)
        for i in 0:1
            max_val = pp(row, col + i, sum + tri[row][col + i], tri, max_val)
        end
    end
    return max_val
end

# 0.00015198s
function euler18a()
    triangle  =
         [[75],
         [95,64],
         [17,47,82],
         [18,35,87,10],
         [20, 4,82,47,65],
         [19, 1,23,75, 3,34],
         [88, 2,77,73, 7,63,67],
         [99,65, 4,28, 6,16,70,92],
         [41,41,26,56,83,40,80,70,33],
         [41,48,72,33,47,32,37,16,94,29],
         [53,71,44,65,25,43,91,52,97,51,14],
         [70,11,33,28,77,73,17,78,39,68,17,57],
         [91,71,52,38,17,14,91,43,58,50,27,29,48],
         [63,66, 4,68,89,53,67,30,73,16,69,87,40,31],
         [ 4,62,98,27,23, 9,70,98,73,93,38,53,60, 4,23]];
    max_val = 0
    max_val = pp(1, 1, triangle[1][1], triangle, max_val)

    return max_val
end

run_euler(euler18a)


using Distributions, StatsBase
using MixtureModels

function test_emstep(::Type{Beta}, N=10000)
    α = rand()
    β = rand()
    samples = rand(Beta(α, β), N)
    weights = uweights(N)

    pred = emstep(Beta(), samples, weights, ones(N))

    isapprox(pred.α, α, atol=0.05) && isapprox(pred.β, β, atol=0.05)
end

function test_emstep(::Type{Normal}, N=10000)
    μ = randn()
    σ = exp(randn())
    samples = rand(Normal(μ, σ), N)
    weights = uweights(N)

    pred = emstep(Normal(), samples, weights, ones(N))

    isapprox(pred.μ, μ, atol=0.05) && isapprox(pred.σ, σ, rtol=0.05)
end

function test_emstep(::Type{Exponential}, N=1000)
    θ = exp(randn())
    samples = rand(Exponential(θ), N)
    weights = uweights(N)

    pred = emstep(Exponential(), samples, weights, ones(N))

    isapprox(pred.θ, θ, rtol=0.05)
end

function test_emstep(::Type{Gamma}, N=10000)
    α = exp(randn())
    θ = exp(randn())
    samples = rand(Gamma(α, θ), N)
    weights = uweights(N)

    pred = emstep(Gamma(), samples, weights, ones(N))

    isapprox(pred.α, α, atol=0.05) && isapprox(pred.θ, θ, rtol=0.05)
end

function test_emstep(::Type{Erlang}, α, N=1000)
    θ = exp(randn())
    samples = rand(Erlang(α, θ), N)
    weights = uweights(N)

    pred = emstep(Erlang(α), samples, weights, ones(N))

    pred.α == α && isapprox(pred.θ, θ, rtol=0.05)
end

N_rep = 100
p_success = 0.8

for T in [ Beta, Normal, Exponential, Gamma ]
    @test count(test_emstep(T) for i in 1:N_rep) >= p_success * N_rep
end

@test count(test_emstep(Erlang, 5) for i in 1:N_rep) >= p_success * N_rep
@test count(test_emstep(Erlang, 20) for i in 1:N_rep) >= p_success * N_rep


# calculate the least squares line of the errors
function calc_line()
  dirs = Array(1:2)
  npts = length(dirs)
  println("maximum(dirs) = ", maximum(dirs))
  err_vals = zeros(maximum(dirs))
  h_vals = zeros(maximum(dirs))
  for i in dirs
    println("i = ", i)
    fname = "m$i/error_calc.dat"
#    fname2 = "m$i/counts_0.txt"
    val_i = readdlm(fname)
    err_vals[i] = val_i[1]

#    counts_i = readdlm(fname2)
#    h_vals[i] = 1/(sqrt(counts_i[3]/2))
    h_vals[i] = val_i[2]

  end

  h_vals = [0.5 0.2 0.1]

  min_idx = minimum(dirs)
  max_idx = maximum(dirs)
  err_vals = err_vals[min_idx:max_idx]
  h_vals = h_vals[min_idx:max_idx]
  npts = max_idx - min_idx + 1

  println("h values = ", h_vals)
  println("error values = ", err_vals)

  err_vals_log = zeros(npts)
  h_vals_log = zeros(npts)
  for i=1:npts
    err_vals_log[i] = log(err_vals[i])
    h_vals_log[i] = log(h_vals[i])
  end

  # now form system

  A = ones(npts, 2)
  A[:, 2] = h_vals_log

  # solve system
  AAt = A.'*A
  rhs = A.'*err_vals_log

  x = AAt\rhs

  println("x = ", x)

  for i=1:npts
    h_i = h_vals_log[i]
    err_exp = x[1] + h_i*x[2]
    err_act = err_vals_log[i]
    err_pct = 100*(err_exp - err_act)/err_act
    println("point $i percent error of best fit line = ", err_pct, "%")
  end

  # calculate slope from last two points
  slope2 = (err_vals_log[npts - 1] - err_vals_log[npts])/(h_vals_log[npts - 1] - h_vals_log[npts])

  println("2 point slope: ", slope2)
  println("least squares slope: ", x[2])

  # write data to file for plotting
  data_arr = hcat(h_vals, err_vals)
  writedlm("err_data.dat", data_arr)

  return x[2]  # return least squares slope
end


@turing_testset "is.jl" begin
    function reference(n)
        as = Vector{Float64}(undef, n)
        bs = Vector{Float64}(undef, n)
        logps = Vector{Float64}(undef, n)

        for i in 1:n
            as[i], bs[i], logps[i] = reference()
        end
        logevidence = logsumexp(logps) - log(n)

        return (as = as, bs = bs, logps = logps, logevidence = logevidence)
    end

    function reference()
        x = rand(Normal(4,5))
        y = rand(Normal(x,1))
        loglik = logpdf(Normal(x,2), 3) + logpdf(Normal(y,2), 1.5)
        return x, y, loglik
    end

    @model function normal()
        a ~ Normal(4,5)
        3 ~ Normal(a,2)
        b ~ Normal(a,1)
        1.5 ~ Normal(b,2)
        a, b
    end

    alg = IS()
    seed = 0
    n = 10

    model = normal()
    for i in 1:100
        Random.seed!(seed)
        ref = reference(n)

        Random.seed!(seed)
        chain = sample(model, alg, n)
        sampled = get(chain, [:a, :b, :lp])

        @test vec(sampled.a) == ref.as
        @test vec(sampled.b) == ref.bs
        @test vec(sampled.lp) == ref.logps
        @test chain.logevidence == ref.logevidence
    end

    @turing_testset "logevidence" begin
        Random.seed!(100)

        @model function test()
            a ~ Normal(0, 1)
            x ~ Bernoulli(1)
            b ~ Gamma(2, 3)
            1 ~ Bernoulli(x / 2)
            c ~ Beta()
            0 ~ Bernoulli(x / 2)
            x
        end

        chains = sample(test(), IS(), 10000)

        @test all(isone, chains[:x])
        @test chains.logevidence ≈ - 2 * log(2)
    end
end


pn = TimePetriNet(
    [:red_1,:yellow_1,:green_1, :red_2, :yellow_2, :green_2, :free],
    [:ry1,:yg1,:gr1, :ry2,:yg2,:gr2],
    sparse([
        1;0;0 ; 0;0;0 ; 1 ;; # traffic light 1 red->yellow
        0;1;0 ; 0;0;0 ; 0 ;; # traffic light 1 yellow->green
        0;0;1 ; 0;0;0 ; 0 ;; # traffic light 1 green->red
        0;0;0 ; 1;0;0 ; 1 ;; # traffic light 2 red->yellow
        0;0;0 ; 0;1;0 ; 0 ;; # traffic light 2 yellow->green
        0;0;0 ; 0;0;1 ; 0    # traffic light 2 green->red
    ]),
    sparse([
        -1; 1; 0 ;  0; 0; 0 ; -1 ;; # traffic light 1 red->yellow
         0;-1; 1 ;  0; 0; 0 ;  0 ;; # traffic light 1 yellow->green
         1; 0;-1 ;  0; 0; 0 ;  1 ;; # traffic light 1 green->red
         0; 0; 0 ; -1; 1; 0 ; -1 ;; # traffic light 2 red->yellow
         0; 0; 0 ;  0;-1; 1 ;  0 ;; # traffic light 2 yellow->green
         0; 0; 0 ;  1; 0;-1 ;  1    # traffic light 2 green->red
    ]),
    ones(6),
    ones(6),
    [5;0;0;0;1;0;0]
)

f,ax,p = graphplot(pn); f

#=
trange = 1:60*length(states)
record(f, "traffic_light.mp4", trange; framerate=30) do tt
    mm = states[1+div(tt-1,60)]
    animate_step!(p, pn, mm)
    update!(f.scene)
end
=#


using JuMP

type Recipe
	name::AbstractString
	time::Float64
	outs::Dict{AbstractString, Float64}
	ins::Dict{AbstractString, Float64}
	function Recipe(name, time, speed)
		recips[name] = new(name, time/speed, Dict{AbstractString, Float64}(), Dict{AbstractString, Float64}())
		recips[name]
	end
end

recips = Dict{AbstractString, Recipe}()

function production(r::Recipe, nfacs)
	@printf "%s: %d\n" r.name nfacs
	@printf "\tProduces\n"
	for (k,v) in r.outs
		@printf "\t\t%0.2f \"%s\" per min - %f per second\n" 60v * nfacs k v * nfacs
	end
	@printf "\tConsumes\n"
	for (k,v) in r.ins
		@printf "\t\t%0.2f \"%s\" per min - %f per second\n" 60v * nfacs k v * nfacs
	end
end

macro IN(iname, iqty)
	return quote
		recipe.ins[$iname] = $iqty / recipe.time
	end
end

macro OUT(iname, iqty)
	return quote
		recipe.outs[$iname] = $iqty / recipe.time
	end
end

@enum Rnum AOP HOC LOC SOL ROK

recips[AOP] = Recipe("Advanced Oil Processing", 5.0, 1.)
@IN AOP "Crude Oil" 200.
@IN AOP "Water" 100.
@OUT AOP "Heavy Oil" 10.
@OUT AOP "Light Oil" 45.
@OUT AOP "Petroleum Gas" 55.

recips[HOC] = Recipe("Heavy Oil Cracking", 3.0 , 1.25)
@IN HOC "Heavy Oil" 40.
@IN HOC "Water" 30.
@OUT HOC "Light Oil" 30.

recips[LOC] = Recipe("Light Oil Cracking", 5.0 , 1.25)
@IN LOC "Light Oil" 30.
@IN LOC "Water" 30.
@OUT LOC "Petroleum Gas" 20.

recips[SOL] = Recipe("Solid Fuel", 3., 1.25)
@IN SOL "Petroleum Gas" 20.
@OUT SOL "Solid Fuel" 1.

recips[ROK] = Recipe("Rocket Fuel", 30., 1.25)
@IN ROK "Solid Fuel" 10.
@OUT ROK "Rocket Fuel" 1.

m = Model()

@variable(m, AOPs >= 0, Int)
@variable(m, HOCs >= 0, Int)
@variable(m, LOCs >= 0, Int)
@variable(m, SOLs >= 0, Int)
@variable(m, ROKs >= 0, Int)

@constraint(m, AOPs * recips[AOP].ins["Crude Oil"] <= 2400)
@constraint(m, HOCs * recips[HOC].ins["Heavy Oil"] <= AOPs * recips[AOP].outs["Heavy Oil"])
@constraint(m, LOCs * recips[LOC].ins["Light Oil"] <= AOPs * recips[AOP].outs["Light Oil"] + HOCs * recips[HOC].outs["Light Oil"])
@constraint(m, SOLs * recips[SOL].ins["Petroleum Gas"] <= AOPs * recips[AOP].outs["Petroleum Gas"] + LOCs * recips[LOC].outs["Petroleum Gas"])
@constraint(m, ROKs * recips[ROK].ins["Solid Fuel"] <= SOLs * recips[SOL].outs["Solid Fuel"])

@objective(m, Max, ROKs)

solve(m)

production(recips[AOP], getvalue(AOPs))
production(recips[HOC], getvalue(HOCs))
production(recips[LOC], getvalue(LOCs))
production(recips[SOL], getvalue(SOLs))
production(recips[ROK], getvalue(ROKs))


abstract type AbstractFilter{T} end

struct KalmanFilter{T<:AbstractFloat} <: AbstractFilter{T}
    x::AbstractVector{T}
    F::AbstractMatrix{T}
    H::AbstractMatrix{T}
    B::AbstractMatrix{T}
    P::AbstractMatrix{T}
    Q::AbstractMatrix{T}
    R::AbstractMatrix{T}
    offset::AbstractVector{T}
    c::Base.RefValue{T}
    pdy::Base.RefValue{T}
end

struct MinimaxMMAE{T<:AbstractFloat} <: AbstractFilter{T}
    filterbank::AbstractVector{KalmanFilter{T}}
    optimizer::Function

    @doc raw"""
        mini = MinimaxMMAE(filterbank, γ, optimizer)
    Construct a Minimax Multiple Model Adaptive Estimator object.

    Create an estimator for the case where the dynamics belong to
    the finite set of linear systems
    equations
    ```math
    \begin{aligned}
    x_{t+1} & = F_ix_t + B_iu_t + w_t + \text{offset}_i\\
    y_t     & = H_ix_t + v_t,
    \end{aligned}
    ```
    and the associated objective
    ```math
    \min_{\hat y_N}\max_{x_0, \mathbf{w}^N, \mathbf{v}^N, i}\Bigg\{ |\hat y_N - H_ix_N|^2 - \gamma^2\left(|x_0 - \hat x_0|^2_{P_0^{-1}} + \sum_{t=0}^{N-1} |w_t|^2_{Q^{-1}} + |v_t|^2_{R^{-1}}\right)\Bigg\}.
    ```

    Filterbank is an array of KalmanFilter object associated with each $i$.
    """
    function MinimaxMMAE(
        filterbank::AbstractVector{KalmanFilter{T}},
        optimizer::Function) where T<: AbstractFloat
        new{T}(filterbank, optimizer)
    end
end

struct BayesianMMAE{T<:AbstractFloat} <: AbstractFilter{T}
    filterbank::AbstractVector{KalmanFilter{T}}

    @doc raw"""
        mmae = BayesianMMAE(filterbank)
    Construct a standard Multiple Model Adaptive Estimator object.

    Create an estimator for the case where the dynamics belong to
    the finite set of linear systems
    equations
    ```math
    \begin{aligned}
    x_{t+1} & = F_ix_t + B_iu_t + w_t + \text{offset}_i\\
    y_t     & = H_ix_t + v_t,
    \end{aligned}
    ```
    filterbank is an array of KalmanFilter objects associated with the
    dynamics and disturbance characteristics of each $i$.
    """
    function BayesianMMAE(filterbank::AbstractVector{KalmanFilter{T}}) where T<: AbstractFloat
        K = length(filterbank)
        for filter in filterbank
            filter.pdy[] = 1/K
        end
        new{T}(filterbank)
    end

end


struct ClusterBatchQs <: BatchPQs
    k::Int

    function ClusterBatchQs(; k::Int)::ClusterBatchQs
        (k < 1) && throw(ArgumentError("Invalid batch size k=$(k)."))
        return new(k)
    end
end

"""
For maximum diversity in batch selection.
Compute k-medoids clustering with k=batchsize, return cluster medoids
"""
function select_batch(qs::ClusterBatchQs, x::Array{T, 2}, labels::Dict{Symbol, Vector{Int}}, candidate_indices::Vector{Int})::Vector{Int} where T <: Real
    num_observations = length(candidate_indices)
    if num_observations <= qs.k
        return candidate_indices
    end

    candidates = x[:, candidate_indices]
    distances = Distances.pairwise(Distances.Euclidean(), candidates, dims=2)
    clustering = kmedoids(distances, qs.k)

    medoid_indices = clustering.medoids

    return candidate_indices[medoid_indices]
end


using TensorFlow

sess = TensorFlow.Session()

x = TensorFlow.constant(Float64[1,2])
y = TensorFlow.Variable(Float64[3,4])
z = TensorFlow.placeholder(Float64)

w = exp(x + z + -y)

run(sess, TensorFlow.global_variables_initializer())
res = run(sess, w, Dict(z=>Float64[1,2]))
Base.Test.@test res[1] ≈ exp(-1)


import FrankWolfe
using LinearAlgebra
using Random
using Test
using SparseArrays

n = Int(1e4)
k = 3000

s = 41
Random.seed!(s)

xpi = rand(n);
total = sum(xpi);
xp = xpi # ./ total;

f(x) = norm(x - xp)^2
function grad!(storage, x)
    @. storage = 2 * (x - xp)
end

lmo = FrankWolfe.KSparseLMO(100, 1.0)

x0 = FrankWolfe.compute_extreme_point(lmo, spzeros(size(xp)...))

x, v, primal, dual_gap, _ = FrankWolfe.blended_conditional_gradient(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Adaptive(),
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    L=2,
    verbose=true,
    trajectory=false,
    K=1.00,
    weight_purge_threshold=1e-10,
    epsilon=1e-9,
)

@test dual_gap ≤ 5e-4
@test f(x0) - f(x) ≥ 180

x0 = FrankWolfe.compute_extreme_point(lmo, spzeros(size(xp)...))

x, v, primal_cut, dual_gap, _ = FrankWolfe.blended_conditional_gradient(
    f,
    grad!,
    lmo,
    x0,
    max_iteration=k,
    line_search=FrankWolfe.Adaptive(),
    print_iter=k / 10,
    emphasis=FrankWolfe.memory,
    L=2,
    verbose=true,
    trajectory=false,
    K=1.00,
    weight_purge_threshold=1e-10,
    epsilon=1e-9,
    timeout=3.0,
)

@test primal ≤ primal_cut


"""
ERGM for directed networks with totan links and tortal reciprocated links as statistics
"""
Base.@kwdef struct  ErgmDirBin0Rec0 <: Ergm
    ergmTermsString::String="edges + mutual" # needs to be compatible with R ergm package
    nErgmPar = 2
end

ergm_term_string(x::ErgmDirBin0Rec0)  = "edges + mutual"

type_of_obs(model::ErgmDirBin0Rec0) =  Array{Float64, 1}

"""
given the 2 model's parameters compute the probabilities for each state of the diad
"""
function diadProbFromPars(Model::ErgmDirBin0Rec0, par)
    θ = par[1]
    η = par[2]
    eθ = exp(θ)
    e2θplusη = exp(2*θ + η)
    den = 1 + 2*eθ + e2θplusη
    p00 = 1/den
    p01 = eθ /den
    p10 = p01
    p11 = e2θplusη/den
    diadProbsVec = [p00, p01, p10, p11]
    return diadProbsVec
end

"""
given the vector of diad states probabilities sample one random matrix from the corresponding pdf
"""
function samplSingMatCan(Model::ErgmDirBin0Rec0, diadProbsVec::Array{<:Real,1}, N)
    out = zeros(Int8,N,N)
    #display((maximum(expMat),minimum(expMat)))
    for c = 1:N
        for r=1:c-1
            diadState = rand(Categorical(diadProbsVec))
            if diadState==2
                out[r,c] = 1
            elseif diadState==3
                out[c,r] = 1
            elseif diadState==4
                out[r,c] = 1
                out[c,r] = 1
            end
        end
    end
    return out
end

function sample_ergm(model::ErgmDirBin0Rec0, N, par_vec, nSample)

    θ_0, η_0 = par_vec

    diadProb = diadProbFromPars(model, [θ_0, η_0])

    A_vec =  reduce((x,y)->cat(x,y,dims=3), [samplSingMatCan(model, diadProb, N) for i=1:nSample])

    return A_vec
end

function stats_from_mat(Model::ErgmDirBin0Rec0, A ::Matrix{<:Real})
    L = sum(A)
    R = sum(A'.*A)/2
    N=size(A)[1]

    return [L, R, N]
end

function exp_val_stats(model::ErgmDirBin0Rec0, θ, η, N)
    x = exp(θ)
    x2y = exp(2*θ + η)
    z = 1 + 2*x + x2y
    Nlinks = n_pox_dir_links(N)
    α = (x + x2y)/z
    β = 0.5 * (x2y)/z
    L = Nlinks * α
    R = Nlinks * β
    return L, R
end

"""
return the ergm parameters that fix on average the input values for the ergm statistics:
    - L is  ∑_i>j A_ij + A_ji
    - R is  of ∑_i>j A_ij * A_ji /2
"""
function ergm_par_from_mean_vals(model::ErgmDirBin0Rec0, L, R, N)
    Nlinks = n_pox_dir_links(N)
    # average values per pair
    α = L/Nlinks
    β = R/Nlinks

    x = (2β-α)/(2α-2β-1)
    θ = log(x)

    y = 2(2*β^2  + β- 2*α*β )/((2β-α)^2)
    η =  log(y)

    return θ, η
end

function estimate(model::ErgmDirBin0Rec0, L, R, N)
    θ_est, η_est = ergm_par_from_mean_vals(model, L, R, N)
    vPar = [θ_est, η_est]
    return vPar
end

function estimate(Model::ErgmDirBin0Rec0, A::T where T <: Matrix)
    L, R, N = stats_from_mat(Model, A)
    return estimate(Model, L, R, N)
end

function logLikelihood(Model::ErgmDirBin0Rec0, L, R, N, par)
    θ, η = par
    z = 1 + 2*exp(θ) + exp(2*θ+η)
    return L * θ + R*η - (N*(N-1)/2)*log(z)
end

function logLikelihood(Model::ErgmDirBin0Rec0, A::Matrix, par)
    L, R, N = stats_from_mat(Model, A)
    return logLikelihood(Model, L, R, N, par)
end

estimate_sequence(model::ErgmDirBin0Rec0, obsT::Array{Array{T, 1},1} where T <: Real) = reduce(hcat, [estimate(model, obsT[t]... ) for t in 1:length(obsT)])

estimate_sequence(model::ErgmDirBin0Rec0, AT::Array{T, 3} where T <: Real) = reduce(hcat, [estimate(model, AT[:,:,t]) for t in 1:size(AT)[3]])


function number_of_trees(right,down)

   data= readdlm("input_day3", '\n')
   horizontal = div( (size(data)[1])*right, down )
   lline = length(data[1])

   #number of times to repet the pattern
   rep = div(horizontal,lline)+1

   #repeat the pattern the ammount needed
   for i in eachindex(data)
      data[i] = data[i]^rep
   end

   #horizontal index
   h_i = 1

   #number of trees
   n_t = 0

   for i in Iterators.countfrom(1,down)
      i > size(data)[1]-down && break
      line = data[i+down]
      h_i += right
      index = h_i

      if line[h_i] == '#'
        n_t += 1
      end
   end

   return (n_t)

end

using DelimitedFiles

r = [ 1, 3, 5, 7, 1 ]
d = [ 1, 1, 1, 1, 2 ]
slope = zeros(5)

product = 1

for i in eachindex(r)
   slope[i] = number_of_trees(r[i], d[i])
   global product *= slope[i]
   println(product)
end

println(product)


using ComponentArrays
using DifferentialEquations
using SimulationLogs
using Plots; gr()
using UnPack

## Simulation functions
# Simple car with velocity-square drag
function car!(D, vars, p, t; u=0.0)
    @unpack vel = vars
    @unpack c, m = p

    @log drag = c*vel^2

    D.pos = vel
    D.vel = (-drag*sign(vel) + u)/m
end

# Proportional-integral control
function PI!(D, vars, p, t; e)
    @unpack ∫e = vars
    @unpack kp, ki = p

    @log u = kp*e + ki*∫e

    D.∫e = e
    return u
end

# Car with cruise control
function feedback_car!(D, vars, p, t)
    @log r = p.ref(t)
    @log e = r - vars.car.vel

    u = PI!(D.control, vars.control, p.control, t; e=e)
    car!(D.car, vars.car, p.car, t; u)
end

## Inputs
# Parameters
p = (
    ref = t -> 10*(t>1),
    car = (
        m = 1000,
        c = 5,
    ),
    control = (
        kp = 800,
        ki = 40,
    )
)

# Initial conditions
ic = ComponentArray(
    car = (
        pos = 0.0,
        vel = 0.0,
    ),
    control = (
        ∫e = 0.0,
    ),
)

## Setup problem and solve
prob = ODEProblem(feedback_car!, ic, (0.0, 20.0), p)
sol = solve(prob)

## Plots
p1 = plot(sol, vars=2, title="Reference Tracking")
scope!(sol, :r)

p2 = scope(sol, [:u, :drag], title="Forces")

p3 = plot(sol, vars=3, title="Errors")
scope!(sol, :e)

plot(p1, p2, p3, layout=(3,1), size=(600,800))


using ProgressMeter

"""
    +(self, table[, N])

Computes with N (default 1,000) Monte Carlo simulations the odds of the hand
winning the showdown given the cards shown on the table

"""
function calculate_odds(self::Hand, table::Hand; N = 1_000)
    wins = 0
    for i = 1:N
        deck = fresh_deck - self - table
        opp, deck = deal(deck, 2)
        extra_cards, deck = deal(deck, 5 - length(table))
        if opp + table + extra_cards < self + table + extra_cards
            wins += 1
        end
    end
    return wins / N
end

"""
    get_preflop_odds([N])

Computes with N (default 1,000) Monte Carlo simulations the odds of the each
starting hand winning the showdown given no cards shown on the table yet.
This table is optionally used to speed up game calculations.

"""
function get_preflop_odds(; N = 1000)
    preflop_odds = zeros(52, 52)
    @showprogress for i = 1:52, j = 1:52
        if i != j
            preflop_odds[i, j] = calculate_odds(
                Hand([fresh_deck_cards[i], fresh_deck_cards[j]]),
                Hand([]), N = N)
        end
    end
    return preflop_odds
end


#=
 Algorithms for triple-word arithmetic
 Nicolas Fabiano, Jean-Michel Muller, Joris Picot
=#

function vecsum(x0::T,x1::T,x2::T) where {T}
   x0 === 0.0 && return (two_sum(x1, x2)..., 0.0)
   s1, e2 = two_sum(x1, x2)
   e0, e1 = two_sum(x0, s1)
   return (e0, e1, e2)
end

function vseb(a0::T, a1::T, a2::T, a3::T, a4::T, a5::T) where {T}
    a2 === 0.0 && return (a0, a1, a2)
	r0, e1 = two_sum(a0, a1)
	if e1 != 0.0
	  y0 = r0
	else
	  e1 = r0
	end
	y1, y2  = two_sum(e1, a2)
	return (y0, y1, y2)
end

function vseb(a0::T, a1::T, a2::T) where {T}
    a2 === 0.0 && return (a0, a1, a2)
	r0, e1 = two_sum(a0, a1)
	if e1 != 0.0
	  y0 = r0
	else
	  e1 = r0
	end
	y1, y2  = two_sum(e1, a2)
	return (y0, y1, y2)
end

function vecseb(x0::T, x1::T, x2::T) where {T}
   a0, a1, a2 = vecsum(x0, x1, x2)
   return vseb(a0, a1, a2)
end

function tripleword(a::T, b::T, c::T) where {T}
    d0, d1 = two_sum(a,b)
    # e0, e1, e2 = vecsum(d0, d1, c)
    s1, e2 = two_sum(d1, c)
    e0, e1 = two_sum(d0, s1)
    # s0, s1, s2 = vseb(e0, e1, e2)
    r0, e1 = two_sum(e0, e1)
	if e1 != 0.0
	  s0 = r0
	else
	  e1 = r0
	end
	s1, s2  = two_sum(e1, e2)
	return (s0, s1, s2)
end

using SetRounding

function Float64(x0::Float64, x1::Float64, x2::Float64)
    s, e = two_hilo_sum(x0, x1+x1)
    if e !== 0.0
        y = x0 + x1
    elseif x2 > 0.0
        setrounding(Float64, RoundUp)
        y = x0 + x1
        setrounding(Float64, RoundNearest)
    elseif x2 < 0.0
        setrounding(Float64, RoundDown)
        y = x0 + x1
        setrounding(Float64, RoundNearest)
    else
        y = x0 + x1
    end
    return y
end


using NeuralSimulation

# Hodgkin-Huxley channel gating functions
αn(V) = V == -55.0 ? 0.1 : (0.01*(V + 55.0))/(1.0 - exp(-(V + 55.0)/10.0))
βn(V) = 0.125 * exp(-(V + 65.0)/80.0)
αm(V) = V == -40.0 ? 1.0 : (0.1*(V + 40.0))/(1.0 - exp(-(V + 40.0)/10.0))
βm(V) = 4.0*exp(-(V + 65.0)/18.0)
αh(V) = 0.07*exp(-(V+65.0)/20.0)
βh(V) = 1.0/(1.0 + exp(-(V + 35.0)/10.0))

# Steady-state channel gating
m∞(V) = αm(V)/(αm(V) + βm(V))
h∞(V) = αh(V)/(αh(V) + βh(V))
n∞(V) = αn(V)/(αn(V) + βn(V))

# Updating gates (rewritten as explicit solution to backwards euler)
@inline function BEgates!(u, p, dt)
    V = view(u, :, 1)
    m = view(u, :, 2)
    h = view(u, :, 3)
    n = view(u, :, 4)
    @inbounds for i in 1:length(V)
        m[i] = (m[i] + αm(V[i])*dt)/(1 + (αm(V[i]) + βm(V[i]))*dt)
        h[i] = (h[i] + αh(V[i])*dt)/(1 + (αh(V[i]) + βh(V[i]))*dt)
        n[i] = (n[i] + αn(V[i])*dt)/(1 + (αn(V[i]) + βn(V[i]))*dt)
    end
end

# In-place update of the Hines coefficients + RHS
@inline function f_hines!(Hines::HinesMatrix{S, BackwardsEuler}, u, p, dt) where {S}
    V = view(u, :, 1)
    m = view(u, :, 2)
    h = view(u, :, 3)
    n = view(u, :, 4)
    a, rl, cm, gna, gk, gl, Ena, Ek, El, Iapp, dx = p
    @inbounds for k in 1:length(V)
        Hines.D[k] = 1 + (a*dt)/(rl*cm*dx^2) + (dt/cm)*(gna*m[k]^3*h[k] + gk*n[k]^4 + gl)
        V[k] +=  (dt/cm)*(gna*m[k]^3*h[k]*Ena + gk*n[k]^4*Ek + gl*El)
    end
    # current injected to first compartment; for testing only!
    V[1] += (Iapp*dt)/(pi*a*cm*dx)
    return nothing
end

##############################################################

dt = 0.025 # ms
Iapp = 1.0e-6 # mA (1 nA)
cable_length = 2000.0e-3 # 2000 µm length
N = 2000
dx = cable_length/N
radius = 2.0e-3 # µm in mm
area = 2*pi*radius*dx

rl = 1.0e6 # 1 kΩ in milliohms*mm
gl  = 0.003 #mS/mm^2
gna = 1.2 # mS/mm^2 (modified from 1.2 mS)
gk  = .36 # mS/mm^2
El  = -54.4 # mV
Ena = 50.0 # mV
Ek  = -77.0 # mV
cm  = 1.0e-5 # mF/mm^2

Vinit = -55.0 # mV
tspan = (0.0, 200.0)

p = [radius, rl, cm, gna, gk, gl, Ena, Ek, El, Iapp, dx]

cable = CableSection(radius, cable_length, rl, cm, N)

u0 = zeros(2000,4)
u0[:,1] .= Vinit
u0[:,2] .= m∞(Vinit)
u0[:,3] .= h∞(Vinit)
u0[:,4] .= n∞(Vinit)

fn = HHCableFunction{true}(f_hines!, BEgates!)
prob = HHCableProblem(fn, cable, u0, tspan, p)
@time sol = solve(prob, BackwardsEuler(), dt = 0.025)


function ei_global(g::SimpleGraph, groups::AbstractVector)
    L = ne(g)
    E = 0

    @inbounds for e in edges(g)
        groups[src(e)] == groups[dst(e)] && (E += 1)
    end

    I = L - E
    return (E - I) / (E + I)
end

function ei_group(g::SimpleGraph, groups::AbstractVector)

    ng = maximum(groups)
    E = zeros(ng)
    I = zeros(ng)

    @inbounds for v in vertices(g)
        @inbounds for nb in neighbors(g, v)
            if groups[v] == groups[nb]
                I[groups[v]] += 1
            else
                E[groups[v]] += 1
            end
        end
    end

    ei =  (E .- I ) ./ (E .+ I)

    return Dict(collect(1:ng) .=> ei)
end

function ei_indiv(g::SimpleGraph, groups::AbstractVector)

    I = zeros(nv(g))
    @inbounds for v in vertices(g)
        @inbounds for nb in neighbors(g, v)
            if groups[v] == groups[nb]
                I[v] += 1
            end
        end
    end

    deg = degree(g)
    E = deg .- I

    return  (E .- I) ./ deg
end