"""
    optblocklength(data, bi::BootInput)::Float64
    optblocklength(data ; kwargs...)::Float64

Provides an estimate of the optimal block-length to use with a dependent bootstrap.

For multivariate datasets, optimal block length is estimated for each column of data, and then
bi.fblocklengthcombine, which is a function that maps Vector{Float64} to Float64, is called
to reduce the multiple estimates to a single estimates. The default value for fblocklengthcombine
is median.

Block length methods currently implemented include: \n
     - Patton, Politis, White (2009) "Correction to Automatic Block Length Selection For the Dependent Bootstrap" \n

For all methods discussed above, bandwidth is estimated following Politis (2003) "Adaptive Bandwidth Choice", using the
flat-top kernel suggested in that paper.
"""
function optblocklength(x::AbstractVector{<:Number}, blm::BLPPW2009{P2003}, bootmethod::Tbm)::Float64 where {Tbm<:BootMethod}
	length(x) < 3 && error("You must have at least 3 observations to estimate block length")
    (M, xVar, covVec) = blocklength_ma_and_cor(x) #Bandwidth method currently forced to politis (2003)
	kernelCovVec = blocklength_kernel_cov(covVec, M)
	gHat = 0.0
	for k = 1:M
		gHat += 2 * k * kernelCovVec[k] #note, "2*" is because sum is from -M and M, but is symmetric about 0. Note, 0 term in sum vanishes since k=0 -> |k|=0
	end
    dHat = optblocklength_blppw2009_dhat(bootmethod, xVar, kernelCovVec)
    #Equation 9 and 14 from Politis, White (2004)
	blocklength = (2 * gHat^2 * length(x) / dHat)^(1/3)
	blocklength = min(blocklength, ceil(min(3*sqrt(length(x)), length(x) / 3))) #Enforce upper bound on block length suggested by Patton
	blocklength = max(blocklength, 1.0)
	return blocklength
end
(optblocklength_blppw2009_dhat(a::BootStationary, xvar::Float64, kernelcovvec::Vector{Float64})::Float64) = 2 * (xvar + 2*sum(kernelcovvec))^2 #note, in expression (1 + 2*sum(kernelCovVec)), "2*" is because sum is from -M to M, but is symmetric around 0. "1+" is the zero term of the sum which is simply equal to unity.
(optblocklength_blppw2009_dhat(a::T, xvar::Float64, kernelcovvec::Vector{Float64})::Float64) where {T<:Union{BootCircular,BootMoving,BootIID,BootNoOverlap}} = (4/3) * (xvar + 2*sum(kernelcovvec))^2 #note, in expression (1 + 2*sum(kernelCovVec)), "2*" is because sum is from -M to M, but is symmetric around 0. "1+" is the zero term of the sum which is simply equal to unity.

#Block length selection method of Paparoditis, Politis (2002) "The Tapered Block Bootstrap for General Statistics From Stationary Sequences"
# function optblocklength(x::AbstractVector{<:Number}, blm::BLPP2002{P2003,Tkf}, bootmethod::Tbm)::Float64 where {Tbm<:BootMethod,Tkf<:KernelFunctionMethod}
# 	length(x) < 3 && error("You must have at least 3 observations to estimate block length")
#     Tbm != BootTapered && println("WARNING: Optimal parameter values in pp2002 block length procedure are unknown for bootstrap method $(bootmethod).")
# 	(M, xVar, covVec) = blocklength_ma_and_cor(x)
# 	kernelCovVec = blocklength_kernel_cov(covVec, M)
# 	deltaUnknown = (xVar + 2*sum(kernelCovVec))^2 #Unknown parameter in Delta expression, section 4, Paparoditis, Politis (2002). (note, "1+" is for k=0 term in summation, "2*" is because summation is from -M to M, but is symmetric about 0)
# 	gammaUnknown = 0.0 #Unknown parameter in Gamma expression, start of section 4
# 	for k = 1:M
# 		gammaUnknown += 2 * k^2 * kernelCovVec[k] #note, "2*" is because sum is from -M and M, but is symmetric about 0. Note, 0 term in sum vanishes since k=0 -> k^2=0
# 	end
#     (gammaHat, deltaHat) = optblocklength_blpp2002_param(blm.kernelfunction, gammaUnknown, deltaUnknown)
#     #Equation 20 from Paraproditis, Politis (2002)
# 	blocklength = (4 * gammaHat^2 * length(x) / deltaHat)^(1/5)
# 	blocklength = min(blocklength, ceil(min(3*sqrt(length(x)), length(x) / 3))) #Enforce upper bound on block length suggested by Patton
# 	blocklength = max(blocklength, 1.0)
# 	return blocklength
# end
# (optblocklength_blpp2002_param(kf::KernelTrap, gu::Float64, du::Float64)::Float64) = (-5.45*gu, 1.099*du) #-5.45 = (1/2) * -10.9 = (1/2) * (w*w)''(0) / (w*w)(0) [OPTIMAL TRAP VALUES], see Paparoditis, Politis (2002), 1.099 = 2 * 0.5495 = 2 * int_{-1}^{1} ((w*w)^2(x) / (w*w)^2(0)) dx [OPTIMAL TRAP VALUES], see Paparoditis, Politis (2002)
# (optblocklength_blpp2002_param(kf::KernelSmooth, gu::Float64, du::Float64)::Float64) = (-5.175*gu, 1.1312*deltaUnknown) #-5.175 = (1/2) * -10.35 = (1/2) * (w*w)''(0) / (w*w)(0) [OPTIMAL SMOOTH VALUES], see Paparoditis, Politis (2002), 1.1312 = 2 * 0.5656 = 2 * int_{-1}^{1} ((w*w)^2(x) / (w*w)^2(0)) dx [OPTIMAL SMOOTH VALUES], see Paparoditis, Politis (2002)
# (optblocklength_blpp2002_param(kf::T, gu::Float64, du::Float64)::Float64) where {T} = error("Optimal parameters not known for input kernel function: $(kf)")
(optblocklength(x::AbstractVector{<:Number}, bi::BootInput{Tbm,BLDummy})::Float64) where {Tbm<:BootMethod} = error("Logic fail. It should not have been possible to call this method. Please file an issue with full stacktrace.")
(optblocklength(x::AbstractVector{<:Number}, bi::BootInput)::Float64) = optblocklength(x, bi.blocklengthmethod, bi.bootmethod)
#Multivariate dataset wrapper
(optblocklength(data, bl::Tbl, bm::Tbm, f::Tf)::Float64) where {Tbl<:BlockLengthMethod,Tbm<:BootMethod,Tf<:Function} = f([ optblocklength(local_get_var(data, i), bl, bm) for i = 1:num_var(data) ])
(optblocklength(data, bi::BootInput)::Float64) = bi.fblocklengthcombine([ optblocklength(local_get_var(data, i), bi) for i = 1:num_var(data) ])
#Keyword input wrapper
(optblocklength(data ; kwargs...)::Float64) = optblocklength(data, BootInput(data, kwargs...))

#These two functions are used by several of the block-length selection procedures
function blocklength_ma_and_cor(x::AbstractVector{T})::Tuple{Int, Float64, Vector{Float64}} where {T<:Number}
	(M, xVar, covVec) = bandwidth_politis_2003(x)
	if M > 0.5 * length(x)
		println("WARNING: Bandwidth in parameter estimation section of blocklength forced to half total number of observations. Data may contain excessive dependence.")
		M = Int(round(0.5 * length(x)))
	end
	M < 2 && (M = 2) #Even though M output of bandwidth function is always greater than 2, the above check for excessively large M can result in M being reduced below 2 (admittedly only in very unusual circumstances)
	length(covVec) < M && append!(covVec, autocov(x, length(covVec)+1:M)) #Get any additional autocovariances that we might need
	return(M, xVar, covVec)
end
function blocklength_kernel_cov(covVec::AbstractVector{Float64}, M::Int)::Vector{Float64}
	length(covVec) < M && error("Error in blocklength_kernel_cov likely caused by logic fail in blocklength_ma_and_cor. Please lodge a github issue, preferably with reproducible example.")
	kernelCovVec = Float64[ kernel_politis_2003_flat_top(k / M) * covVec[k] for k = 1:M ]
	return(kernelCovVec)
end

#Note, the following non-exported function is called in the ForecastEval package, so if you alter the function name you will
#need to adjust that package too.
"""
    bandwidth_politis_2003(x::AbstractVector{T})::Tuple{Int, Float64, Vector{Float64}} where {T<:Number}

Implements the methodology from Politis (2003) "Adaptive Bandwidth Choice" to obtain a data-driven bandwidth estimate.

Return tuple is, in order, the bandwidth estimate, the variance of x, and the autocorrelations used to get the bandwidth estimate.

Note, most users won't be interested in the second and third output, but sometimes this routine will be called by other
functions that need these terms, so they are returned to avoid duplicate computation.
"""
function bandwidth_politis_2003(x::AbstractVector{T})::Tuple{Int, Float64, Vector{Float64}} where {T<:Number}
    length(x) < 2 && error("Input data must have at least two observations")
    adjustmentTerm = 1.0 #This term is used by me for debugging. It serves no statistical purpose.
    politis_c = 2.0 #This value is again recommended in Politis (2003)
    K = max(5, Int(ceil(sqrt(log(10, length(x)))))) #This value is recommended in Politis (2003)
	mHat = 1
	corVec = Float64[]
	append!(corVec, autocor(x, 1:min(20, length(x)-1))) #I add autocorrelations in blocks as this is more efficient than doing it one at a time
	corBound = politis_c * sqrt(log(10, length(x)) / length(x)) #Note, use of log base 10 is deliberate and recommended in Politis (2003)
	KCounter = 0
	mHatFound = false
    for k = 1:length(x)-2
        #Add one to counter if bound is satisfied, otherwise reset counter to 0
        abs(corVec[k]) < corBound ? (KCounter += 1) : (KCounter = 0)
		if KCounter >= K #We found K autocorrelations in a row that satisfy the bound, so break.
			mHat = k - K + 1
			mHatFound = true
			break
		end
        #If we run out of autocorrelations to check, add another block of them to corVec
        k == length(corVec) && append!(corVec, autocor(x, length(corVec)+1:min(length(corVec)+20, length(x)-1)))
	end
    mHatFound == false && (mHat = length(x) - 1) #Bound mHat in the case where we fail to hit the break condition
	M = Int(ceil(adjustmentTerm * 2 * mHat)) #"2*" is a standard rule, see e.g. Politis (2003).
    M > length(x) - 1 && (M = length(x) - 1) #Apply upper bound to M
    M < 2 && (M = 2) #Apply lower bound to M
	xVar = ((length(x)-1) /  length(x)) * var(x) #Used to scale autocorrelations to autocovariances for return argument
	return (M, xVar, xVar * corVec)
end

"""
    kernel_politis_2003_flat_top(x::Float64)::Float64

Implements the flat-top kernel function discussed in Politis (2003) "Adaptive Bandwidth Choice"
"""
function kernel_politis_2003_flat_top(x::Float64)::Float64
    x_abs = abs(x)
    x_abs <= 0.5 && return 1.0
    x_abs <= 1.0 && return (2.0 * (1 - x_abs))
    return 0.0
end


# "Gaussian advection-diffusion" Spatial resolution convergence test

using PyPlot

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

using .ConvergenceTests
using .ConvergenceTests.OneDimensionalGaussianAdvectionDiffusion: run_test
using .ConvergenceTests.OneDimensionalUtils: plot_solutions!, plot_error_convergence!, unpack_errors

""" Run advection-diffusion test for all Nx in resolutions. """
function run_convergence_test(κ, U, resolutions...)

    # Determine save time-step
           Lx = 2.5
    stop_time = 0.25
            h = Lx / maximum(resolutions)
           Δt = min(0.1 * h / U, 0.01 * h^2 / κ)

    # Adjust time-step
    stop_iteration = round(Int, stop_time / Δt)
                Δt = stop_time / stop_iteration

    # Run the tests
    results = [run_test(Nx=Nx, Δt=Δt, stop_iteration=stop_iteration, U=U, κ=κ, width=0.1) for Nx in resolutions]

    return results
end

#####
##### Run test
#####

Nx = 2 .^ (6:8) # N = 64 through N = 256
diffusion_results = run_convergence_test(1e-1, 0, Nx...)
advection_results = run_convergence_test(1e-6, 3, Nx...)
advection_diffusion_results = run_convergence_test(1e-2, 1, Nx...)

#####
##### Plot solution and error profile
#####

all_results = (diffusion_results, advection_results, advection_diffusion_results)
names = ("diffusion only", "advection only", "advection-diffusion")
linestyles = ("-", "--", ":")
specialcolors = ("xkcd:black", "xkcd:indigo", "xkcd:wine red")

# Solutions
close("all")

fig, axs = subplots(nrows=2, figsize=(12, 6), sharex=true)

legends = plot_solutions!(axs, all_results, names, linestyles, specialcolors)

filepath = joinpath(@__DIR__, "figs", "gaussian_advection_diffusion_solutions.png")
savefig(filepath, dpi=480, bbox_extra_artists=legends, bbox_inches="tight")

# Error profile
fig, axs = subplots()

legend = plot_error_convergence!(axs, Nx, all_results, names)

filepath = joinpath(@__DIR__, "figs", "gaussian_advection_diffusion_error_convergence.png")
savefig(filepath, dpi=480, bbox_extra_artists=(legend,), bbox_inches="tight")

# Test rate of convergence
for (results, name) in zip(all_results, names)
    atol_L₁, atol_L∞ =
        name == "diffusion only" ? (0.05, 0.05) :
        name == "advection only" ? (0.10, 0.50) :
        name == "advection-diffusion" ? (0.05, 0.01) : nothing

    name = "1D Gaussian " * name
    @info "Testing rate of convergence for $name..."

    u_L₁, v_L₁, cx_L₁, cy_L₁, u_L∞, v_L∞, cx_L∞, cy_L∞  = unpack_errors(results)

    test_rate_of_convergence(u_L₁,  Nx, expected=-2.0, atol=atol_L₁, name=name*" u_L₁")
    test_rate_of_convergence(v_L₁,  Nx, expected=-2.0, atol=atol_L₁, name=name*" v_L₁")
    test_rate_of_convergence(cx_L₁, Nx, expected=-2.0, atol=atol_L₁, name=name*" cx_L₁")
    test_rate_of_convergence(cy_L₁, Nx, expected=-2.0, atol=atol_L₁, name=name*" cy_L₁")
    test_rate_of_convergence(u_L∞,  Nx, expected=-2.0, atol=atol_L∞, name=name*" u_L∞")
    test_rate_of_convergence(v_L∞,  Nx, expected=-2.0, atol=atol_L∞, name=name*" v_L∞")
    test_rate_of_convergence(cx_L∞, Nx, expected=-2.0, atol=atol_L∞, name=name*" cx_L∞")
    test_rate_of_convergence(cy_L∞, Nx, expected=-2.0, atol=atol_L∞, name=name*" cy_L∞")

    @test u_L₁ ≈ v_L₁ ≈ cx_L₁ ≈ cy_L₁
    @test u_L∞ ≈ v_L∞ ≈ cx_L∞ ≈ cy_L∞
end


function rising_cutoff(t)
    r = 0
    if t <= -1
        r = 0
    elseif t >= 1
        r = 1
    else
        r = sin(π / 4 * (1 + sin(π / 2 * t)))
    end
    return r
end

function standardize_fiedler(v)
    v1 = deepcopy(v)
    v1 ./= norm(v1, 2)
    if v1[1] < 0
        v1 = -v1
    end
    return round.(v1; digits = 14)
end

function Lrw_eigenvec(W; nev = 6)
    # N = size(W, 1)
    # D = Diagonal(sum(W, dims = 1)[:])
    # try
    #     val, vtmp = eigs(D - W, D,
    #                      nev = nev, sigma = eps(), v0 = ones(N) / sqrt(N))
    #     vtmp = vtmp[:, sortperm(val)[2:end]]
    #     vtmp ./= sqrt.(sum(vtmp.^2; dims = 1))
    #     vtmp *= Diagonal(1 .- (vtmp[1, :] .< 0) .* 2)
    # catch
    #
    # return round.(vtmp; digits = 14)
    deg = sum(W, dims = 1)[:]  # weighted degree vector
    if minimum(deg) <= 10^3 * eps()
        L = Diagonal(deg) - W
        𝛌, 𝚽 = eigen(Matrix(L))
        standardize_eigenvector_signs!(𝚽)
        return round.(𝚽[:, 2:nev]; digits = 14)
    end
    Lsym = Diagonal(deg.^(-1/2)) * (Diagonal(deg) - W) * Diagonal(deg.^(-1/2))
    𝛌sym, 𝚽sym = eigen(Matrix(Lsym))
    𝚽rw = Diagonal(deg.^(-1/2)) * 𝚽sym
    𝚽rw ./= sqrt.(sum(𝚽rw.^2; dims = 1))
    𝚽rw *= Diagonal(1 .- (𝚽rw[1, :] .< 0) .* 2)
    return round.(𝚽rw[:, 2:nev]; digits = 14)
end

function sortnodes_inmargin(active_region, Na, v, W, idx; sign = :positive)
    ind = sortperm(v[active_region]; rev = (sign == :positive))
    # if there is a tie, use the more dims eigenmaps and sort by lexical order
    if length(unique(v[ind])) < length(active_region)
        emb = Lrw_eigenvec(W[idx, idx]; nev = 3)'
        # for the second dimension coordinate we sort
        # the pos region and the neg region in the same order
        emb_tp = [Tuple(emb[:, i] .* [1, (sign == :positive) * 2 - 1]) for i in 1:length(idx)]
        ind = sortperm(emb_tp[active_region]; rev = (sign == :positive))
    end
    return ind[(end - Na + 1):end]
end

function find_pairinds(W; ϵ::Float64 = 0.2, idx = 1:size(W, 1), used_node = Set())
    v = partition_fiedler(W[idx, idx])[2]
    v = standardize_fiedler(v)
    vmax = norm(v, Inf)
    N = length(v)
    pos_active_region = (Int64)[]
    neg_active_region = (Int64)[]

    tol = 10^3 * eps()
    tol *= ((sum(v .>= tol) > sum(v .<= -tol)) * 2 - 1)

    for i in 1:N
        if idx[i] ∈ used_node
            continue
        end
        if tol <= v[i] < ϵ * vmax
            push!(pos_active_region, i)
        end
        if -ϵ * vmax < v[i] < tol
            push!(neg_active_region, i)
        end
    end
    Np = length(pos_active_region)
    Nn = length(neg_active_region)
    Na = min(Nn, Np)  # number of pair inds in action region

    # indp = sortperm(v[pos_active_region]; rev = true)[(end - Na + 1):end]
    # indn = sortperm(v[neg_active_region])[(end - Na + 1):end]
    indp = sortnodes_inmargin(pos_active_region, Na, v, W, idx; sign = :positive)
    indn = sortnodes_inmargin(neg_active_region, Na, v, W, idx; sign = :negative)
    pair_inds = vcat(idx[pos_active_region[indp]]',
                     idx[neg_active_region[indn]]')

    return pair_inds, v
end

function pair_inds_shadding(W, GP; ϵ = 0.3, J = 1)
    rs = GP.rs
    inds = GP.inds
    (N, jmax) = Base.size(inds)

    used_node = Set()
    shading = zeros(N)
    for j = 1:J
        regioncount = count(!iszero, rs[:, j]) - 1
        for r = 1:regioncount
            indr = rs[r, j]:(rs[r + 1, j] - 1)
            pair_inds = find_pairinds(W; ϵ = ϵ, idx = inds[indr, j], used_node = used_node)[1]
            for i in 1:size(pair_inds, 2)
                pv, nv = pair_inds[:, i]
                # if pv ∈ used_node || nv ∈ used_node
                #     continue
                # end
                if shading[pv] == 0
                    shading[pv] = J + 3 - j
                end
                if shading[nv] == 0
                    shading[nv] = -(J + 3 - j)
                end
            end
            union!(used_node, Set(pair_inds[:]))
        end
    end
    return shading
end

function keep_folding!(U, used_node, W, GP; ϵ = 0.3, j = 1)
    if j == 0
        return U
    end

    rs = GP.rs
    inds = GP.inds
    N = Base.size(inds, 1)

    regioncount = count(!iszero, rs[:, j]) - 1
    for r = 1:regioncount
        indr = rs[r, j]:(rs[r + 1, j] - 1)
        if length(indr) == 1
            continue
        end
        pair_inds, v = find_pairinds(W; ϵ = ϵ, idx = inds[indr, j], used_node = used_node)
        vmax = norm(v, Inf)
        for i in 1:size(pair_inds, 2)
            pv, nv = pair_inds[:, i]
            if pv ∈ used_node || nv ∈ used_node
                continue
            end
            # use the half distance between the 1D embeddings,
            # i.e., t = const⋅(v[pv] - v[nv] / 2), to compute the rising
            # cutoff function, which satisfy r(t)^2 + r(-t)^2 = 1
            t = (v[findfirst(inds[indr, j] .== pv)] - v[findfirst(inds[indr, j] .== nv)]) / (2 * ϵ * vmax)
            # t = v[findfirst(inds[indr, j] .== pv)] / (ϵ * vmax)
            U[pv, pv] = rising_cutoff(t)
            U[pv, nv] = rising_cutoff(-t)
            U[nv, pv] = -rising_cutoff(-t)
            U[nv, nv] = rising_cutoff(t)
        end
        union!(used_node, Set(pair_inds[:]))
    end
end

function unitary_folding_operator(W, GP; ϵ = 0.3, J = 1)
    rs = GP.rs
    inds = GP.inds
    (N, jmax) = Base.size(inds)

    U = Matrix{Float64}(I, N, N)
    used_node = Set()

    for j = 1:J
        keep_folding!(U, used_node, W, GP; ϵ = ϵ, j = j)
    end

    return U
end

"""
    lp_ngwp(𝚽::Matrix{Float64}, W_dual::SparseMatrixCSC{Float64, Int64},
            GP_dual::GraphPart; ϵ::Float64 = 0.3)

construct the lapped NGWP and GP.tag in place.

# Input Arguments
- `𝚽::Matrix{Float64}`: graph Laplacian eigenvectors
- `W_dual::SparseMatrixCSC{Float64, Int64}`: weight matrix of the dual graph
- `GP_dual::GraphPart`: GraphPart object of the dual graph

# Output Argument
- `wavelet_packet::Array{Float64,3}`: the lapped NGWP dictionary. The first
    index is for selecting wavelets at a fixed level; the second index is for
    selecting the level `j`; the third index is for selecting elements in the
    wavelet vector.

"""
function lp_ngwp(𝚽::Matrix{Float64}, W_dual::SparseMatrixCSC{Float64, Int64},
                 GP_dual::GraphPart; ϵ::Float64 = 0.3)
    rs = GP_dual.rs
    inds = GP_dual.inds
    (N, jmax) = Base.size(inds)

    GP_dual.tag = zeros(Int, N, jmax)
    GP_dual.tag[:, 1] = Vector{Int}(0:(N - 1))

    Uf = Matrix{Float64}(I, N, N)
    used_node = Set()

    wavelet_packet = zeros(N, jmax, N)
    wavelet_packet[:, 1, :] = Matrix{Float64}(I, N, N)
    for j = 2:jmax
        regioncount = count(!iszero, rs[:, j]) - 1
        # Uf = unitary_folding_operator(W_dual, GP_dual; ϵ = ϵ, J = j - 1)
        keep_folding!(Uf, used_node, W_dual, GP_dual; ϵ = ϵ, j = j - 1)
        for r = 1:regioncount
            indr = rs[r, j]:(rs[r + 1, j] - 1)
            GP_dual.tag[indr, j] = Vector{Int}(0:(length(indr) - 1))
            wavelet_packet[indr, j, :] = const_meyer_wavelets(𝚽, Uf; idx = inds[indr, j])'
        end
    end
    return wavelet_packet
end

function const_meyer_wavelets(𝚽, Uf; idx = 1:size(Uf, 1))
    N = size(𝚽, 1)
    # 1. nice interpretable way by smooth orthogonal projector
    # assemble smooth orthogonal projector custom to nodes `idx`
    P = Uf[idx, :]' * Uf[idx, idx]
    if diag(P) == χ(idx, N)
        B = 𝚽[:, idx]
    else
        # folding the eigenspace, i.e., 𝚽's column space
        Y = 𝚽 * P
        # find its column space's orthogonal basis
        B = svd(Y).U
    end
    # 2. alternative faster way by orthogonal unfolding operator
    # B = 𝚽 * Uf[:, idx]
    # perform varimax rotation to get the meyer_wavelets
    Wavelets = varimax(B)
    return Wavelets
end

"""
    lp_ngwp_analysis(G::GraphSig, 𝚽::Matrix{Float64}, W_dual::SparseMatrixCSC{Float64, Int64},
            GP_dual::GraphPart; ϵ::Float64 = 0.3)

perform the LP-NGWP transform of the graph signal(s) `G.f`.

# Input Arguments
- `G::GraphSig`: a GraphSig object
- `𝚽::Matrix{Float64}`: graph Laplacian eigenvectors
- `W_dual::SparseMatrixCSC{Float64, Int64}`: weight matrix of the dual graph
- `GP_dual::GraphPart`: GraphPart object of the dual graph

# Output Argument
- `wavelet_packet::Array{Float64,3}`: the lapped NGWP dictionary. The first
    index is for selecting wavelets at a fixed level; the second index is for
    selecting the level `j`; the third index is for selecting elements in the
    wavelet vector.

"""
function lp_ngwp_analysis(G::GraphSig, 𝚽::Matrix{Float64}, W_dual::SparseMatrixCSC{Float64, Int64},
                 GP_dual::GraphPart; ϵ::Float64 = 0.3)
    rs = GP_dual.rs
    inds = GP_dual.inds
    (N, jmax) = Base.size(inds)
    fcols = size(G.f, 2)

    GP_dual.tag = zeros(Int, N, jmax)
    GP_dual.tag[:, 1] = Vector{Int}(0:(N - 1))

    Uf = Matrix{Float64}(I, N, N)
    used_node = Set()

    dmatrix = zeros(N, jmax, fcols)
    dmatrix[:, 1, :] = G.f
    for j = 2:jmax
        regioncount = count(!iszero, rs[:, j]) - 1
        keep_folding!(Uf, used_node, W_dual, GP_dual; ϵ = ϵ, j = j - 1)
        for r = 1:regioncount
            indr = rs[r, j]:(rs[r + 1, j] - 1)
            GP_dual.tag[indr, j] = Vector{Int}(0:(length(indr) - 1))
            dmatrix[indr, j, :] = const_meyer_wavelets(𝚽, Uf; idx = inds[indr, j])' * G.f
        end
    end
    return dmatrix
end


module BoundingBoxes

export BoundingBox, overlaps

import GeometryTypes
import GeometryTypes: overlaps
import GeometryTypes.vertices
import GeometryTypes.intersects
import GeometryTypes.contains
import GeometryTypes.widths
import GeometryTypes.isinside

using GeometryTypes: Point, Vec, HyperSphere, HyperRectangle

import Base.convert
using StaticArrays: SVector
import NaNMath

"Axis-aligned bounding box. Constructor doesn't check all(minimum .<= maximum)"
struct BoundingBox{N,T} <: GeometryTypes.AbstractGeometry{N, T}
  minimum::SVector{N,T}
  maximum::SVector{N,T}
end

convert(::Type{HyperRectangle}, bb::BoundingBox) = HyperRectangle(Vec(minlimit(bb)), Vec(widths(bb)))
function Base.convert(::Type{HyperRectangle{N, T}}, bb::BoundingBox) where {N, T}
  HyperRectangle(Vec{N, T}(minlimit(bb)), Vec{N, T}(widths(bb)))
end

# splatting definition. TODO really should apply for anything whose size is not inferrable
BoundingBox(mi::Vector, ma::Vector) = BoundingBox(promote(SVector(mi...), SVector(ma...))...)
BoundingBox(mi, ma) = BoundingBox(promote(convert(SVector, mi), convert(SVector, ma))...)

minlimit(b::BoundingBox) = b.minimum
maxlimit(b::BoundingBox) = b.maximum
dimlimits(b::BoundingBox{N}) where {N} = tuple(((b.minimum[i], b.maximum[i]) for i = 1:N)...)

function contains(bb::BoundingBox{N, T}, p::Point{N, T}) where {N, T}
  return all(maxlimit(bb) .>= p .>= minlimit(bb))
end

"""symmetric overlap relation"""
function overlaps(b1::BoundingBox{N}, b2::BoundingBox{N}) where {N}
  # can't exit early with broadcasting, but branching might be more expensive

  #= # this is equivalent by DeMorgan's, except that for NaN behavior
  disjoint_dimensions = all(maximum(b1) .< minimum(b2)) .| all(maximum(b2) .< minimum(b1))
  return !(any(disjoint_dimensions))
  =#
  overlap_dimensions = (maxlimit(b1) .≥ minlimit(b2)) .& (maxlimit(b2) .≥ minlimit(b1))
  return all(overlap_dimensions)
end

"make bounding box from array of points"
function BoundingBox(points::AbstractVector) #{PT}) where {T, PT<:AbstractArray{T,1}}
  # errors if length(points) == 0
  #all(length.(points) .== length(points[1]))  # Hope it is elided for static arrays. Doesn't work on iterators
  #mi, ma = extrema(points) # need to reduce over min., max.
  if length(points) == 0
    N = length(eltype(points))
    T = eltype(eltype(points))
    nan = nan(T)
    BoundingBox(fill(nan, N), fill(nan, N))
  end
  mi = reduce((x,y)->min.(x,y), points) # TODO sad for performance
  ma = reduce((x,y)->max.(x,y), points)

  # works on empty collections
  #mi = reduce((x,y)->NaNMath.min.(x,y), nanvector, points) # TODO sad for performance
  #ma = reduce((x,y)->NaNMath.max.(x,y), nanvector, points)

  return BoundingBox(mi, ma)
end

import Base.union
function union(b1::BoundingBox{N,T1}, b2::BoundingBox{N,T2}) where {N, T1, T2}
  # NaN represents empty.
  BoundingBox( NaNMath.min.(minlimit(b1), minlimit(b2)), NaNMath.max.(maxlimit(b1), maxlimit(b2)) )
end

BoundingBox(point::Point) = BoundingBox(point, point)
BoundingBox(s::HyperSphere) = BoundingBox(s.center .- s.r, s.center .+ s.r)
# TODO use a lazy (geometric)-union type from another package
#BoundingBox(u::Union{UnionOfConvex, UnionOfGeometry}) = reduce(union, map(BoundingBox, u.pieces))

"""is ball inside box"""
@inline function contains(box::BoundingBox{N, T}, ball::HyperSphere{N, T}) where {N, T}
  r = ball.r
  c = ball.center
  for i = 1:N
    if !(
        box.maximum[i] > c[i] + r   &&
        box.minimum[i] < c[i] - r
        )
      return false
    end
  end
  return true
end

@inline function center(box::BoundingBox)
  widths(box) / 2 + box.minimum
end

@inline function widths(box::BoundingBox)
  box.maximum - box.minimum
end

@inline vertices(box::BoundingBox) = vertices(HyperRectangle(minlimit(box)..., widths(box)...))

end # module


for N in [Float64, Rational{Int}, Float32]
    # random line
    rand(Line2D)

    # construction
    a1 = N[0, 1]
    b1 = N(1)
    l1 = Line2D(a1, b1)
    l2 = Line2D(N[1, 0], N(1))
    l3 = Line2D(N[0, 1], N(2))
    l4 = Line2D(N[1, 1], N(0))
    l1_copy = Line2D(N[0, 1], N(1))

    # alternative construction from two points
    # line with positive slope
    L = Line2D([1.0, 0.0], [2.0, 0.5])
    @test [1.0, 0.0] ∈ L &&
          [2.0, 0.5] ∈ L &&
          [2.0, 1.5] ∉ L

    # vertical line
    L = Line2D([1.0, 0.0], [1.0, 2.5])
    @test [1.0, 0.0] ∈ L &&
          [1.0, 2.5] ∈ L &&
          [1.0, 10.0] ∈ L
          [10.0, 10.0] ∉ L

    # corner case: zero normal vector
    @test_throws AssertionError Line2D(N[0, 0], N(1))

    # dimension
    @test dim(l1) == 2

    # support vector
    σ(N[0, 1], l1)
    σ(N[1, 0], l2)
    σ(N[0, 1], l3)

    # boundedness
    @test !isbounded(l1)

    # is_polyhedral
    @test is_polyhedral(l1)

    # universality
    @test !isuniversal(l1)
    res, w = isuniversal(l1, true)
    @test !res && w ∉ l1

    # isempty
    @test !isempty(l1)

    # an_element and membership
    an_element(l1) ∈ l1
    an_element(l2) ∈ l2
    an_element(l3) ∈ l3
    an_element(l4) ∈ l4

    # constrained dimensions
    @test constrained_dimensions(l1) == [2]
    @test constrained_dimensions(l4) == [1, 2]

    # constraints_list
    @test ispermutation(constraints_list(l1),
                        [HalfSpace(a1, b1), HalfSpace(-a1, -b1)])

    # concrete intersection
    cap11 = intersection(l1, l1_copy)
    cap12 = intersection(l1, l2)
    cap13 = intersection(l1, l3)
    @test cap11 isa Line2D && cap11.a == l1.a && cap11.b == l1.b
    @test cap12 isa Singleton && element(cap12) == N[1, 1]
    @test cap13 isa EmptySet{N}

    # concrete linear map of a line
    L = Line2D(N[1, -1], N(0)) # x = y
    M = N[1 0; 0 0] # non-invertible matrix
    # projection is y = 0
    if test_suite_polyhedra
        lm = linear_map(M, L)
        if N == Float32 || N == Float64
            @test lm isa Line2D{Float64}
            @test lm.a ≈ N[0, -1] && lm.b ≈ N(0)

            # returned set is universal
            @test linear_map(N[1 1], L) == Universe{N}(1)
        elseif N == Rational{Int}
            @test lm isa Line2D{Rational{BigInt}}
            @test lm.a == N[0//1, -1//1] && lm.b == N(0//1)
        end
    end
    M = N[2 2; 0 1] # invertible matrix
    lm = linear_map(M, L)
    @test lm isa Line2D{N, Vector{N}}
    @test lm.a ≈ N[1/2, -2] && lm.b ≈ N(0)

    # projection
    L = Line2D(N[1, -1], N(0))  # x = y
    @test project(L, [1]) == project(L, [2]) == Universe{N}(1)
    @test project(L, [1, 2]) == L
    L = Line2D(N[2, 0], N(4))  # x = 2
    @test project(L, [1]) == Singleton(N[2])
    @test project(L, [2]) == Universe{N}(1)
    @test project(L, [1, 2]) == L

    # translation
    @test translate(l1, N[1, 2]) == Line2D(a1, N(3))

    # projecting a point onto a line
    L = Line2D(N[1, -1], N(0))  # x = y
    @test project(N[1, 0], L) ≈ N[1//2, 1//2]
end

# tests that only work with Float64 and Float32
for N in [Float64, Float32]
    # sampling
    L = Line2D(N[1, -1], N(0))  # x = y
    for x in sample(L, 10)
        # approximate membership test due to floating-point errors
        @test isapproxzero(dot(x, L.a) - L.b)
    end
end


#
# Data for the simulation
#

using Parameters

@with_kw struct Data

  eps :: Float64 = 5.0
  sig :: Float64 = 0.5

  cutoff :: Float64 = 2.0
  side :: Float64 = 100.

  N :: Int64 = 10_000

  T :: Float64 = 298.15
  R :: Float64 = 8.3145
  m :: Float64 = 1.0

  # Daata for computation
  eps4 :: Float64 = 4*eps
  sig6 :: Float64 = sig^6
  sig12 :: Float64 = sig^12
  nc  = trunc(Int64, side/cutoff)

end


### A Pluto.jl notebook ###
# v0.12.18

using Markdown
using InteractiveUtils

# ╔═╡ 1ad8f830-50e2-11eb-04ab-0b0ad17539a6
#I thought the Project.toml was supposed to do this? What's gone wrong here?
begin
	using Pkg
	Pkg.add("GaussianProcesses")
	Pkg.add("Random")
	Pkg.add("Plots")
#Pkg.status() #gets info for binder Project.toml (in console, not notebook)
end

# ╔═╡ c7b76b10-5059-11eb-1fd3-eff019f62e11
begin #dependencies: slow, possibly several minutes to precompile everything
	using GaussianProcesses
	using Random
	using Plots
end

# ╔═╡ a396df50-5058-11eb-2f43-a1622da4a533
md"Code shamelessly copied from [this tutorial](https://stor-i.github.io/GaussianProcesses.jl/latest/Regression/)
which is a walkthrough of [this library](https://github.com/STOR-i/GaussianProcesses.jl).

A lot of the appeal of going with Julia for this is actually [Stheno](https://github.com/willtebbutt/Stheno.jl), although python fans note there is also python version of this package! (The Julia version is supposed to enjoy a speed advantage and more importantly can interact with Soss and optim.jl, but on the other hand the python one is in python). I'm still too much of a beginner for the Stheno tools quite yet, but I think might be useful to try to head in that direction for possible future questions like 'can we infer things about the human kernel'
"

# ╔═╡ b42c1d00-50d8-11eb-1e8b-090170ad069f
md"## World setup"

# ╔═╡ d4f60080-50d5-11eb-35b1-fdff2d46911d
d=2 #number of dimensions

# ╔═╡ f08c59e0-50d8-11eb-0ddf-e3c4c3f848e9
n= 50 #number of observations

# ╔═╡ 54cea940-50d8-11eb-30b8-c56ad5646c56
x = 2π * rand(d, n); #Observations at these locations

# ╔═╡ 5b68b980-50d8-11eb-2054-7b3e4a525b2c
y = vec(sin.(x[1,:]).*sin.(x[2,:])) + 0.05*rand(n);#Values observed

# ╔═╡ 6bd48ff0-50d9-11eb-1819-9937373f4d69
md"## Hard thinking bit happens here: Set up a hypothesis space"

# ╔═╡ 5cf15550-50d8-11eb-13f2-57779f2b00b6
mZero = MeanZero(); # Zero mean function, not so controvertial

# ╔═╡ 65d618e0-50d8-11eb-3dff-f7e9a5f7244f
#choice of kernel on the other hand is doing a *lot* of the heavy lifting. These are the tutorial defaults.
kern = Matern(5/2,[0.0,0.0],0.0) + SE(0.0,0.0);    # Sum kernel with Matern 5/2 ARD kernel
                                               # with parameters [log(ℓ₁), log(ℓ₂)] = [0,0] and log(σ) = 0
                                               # and Squared Exponential Iso kernel with
                                               # parameters log(ℓ) = 0 and log(σ) = 0

# ╔═╡ f1b94bb0-50d9-11eb-2718-db3fac5608a2
md"## Crank the handle"

# ╔═╡ 734398e0-50d8-11eb-1f25-cb3fcf6730b9
gp = GP(x,y,mZero,kern,-2.0);          # Fit the GP!

# ╔═╡ 85f472c0-50d8-11eb-0294-f99d6ec955be
optimize!(gp);# Hyperparameter tuning

# ╔═╡ 9105da52-50d8-11eb-02e5-4d0ad0983b3a
plot(contour(gp) ,heatmap(gp); fmt=:png) #Woohoo! This thing is analogous to the line of best fit: not sure how to add a confidence ribbon in the 2d case, but could maybe draw a bunch of samples from the posterior and view them as an animation?

# ╔═╡ Cell order:
# ╟─a396df50-5058-11eb-2f43-a1622da4a533
# ╠═1ad8f830-50e2-11eb-04ab-0b0ad17539a6
# ╠═c7b76b10-5059-11eb-1fd3-eff019f62e11
# ╟─b42c1d00-50d8-11eb-1e8b-090170ad069f
# ╠═d4f60080-50d5-11eb-35b1-fdff2d46911d
# ╠═f08c59e0-50d8-11eb-0ddf-e3c4c3f848e9
# ╠═54cea940-50d8-11eb-30b8-c56ad5646c56
# ╠═5b68b980-50d8-11eb-2054-7b3e4a525b2c
# ╟─6bd48ff0-50d9-11eb-1819-9937373f4d69
# ╠═5cf15550-50d8-11eb-13f2-57779f2b00b6
# ╠═65d618e0-50d8-11eb-3dff-f7e9a5f7244f
# ╟─f1b94bb0-50d9-11eb-2718-db3fac5608a2
# ╠═734398e0-50d8-11eb-1f25-cb3fcf6730b9
# ╠═85f472c0-50d8-11eb-0294-f99d6ec955be
# ╠═9105da52-50d8-11eb-02e5-4d0ad0983b3a


using Test
using Distributions
using Random

using EnsembleKalmanProcesses.DataStorage

@testset "DataStorage" begin
    seed=2021
    Random.seed!(seed)

    parameter_samples = rand(MvNormal(2,0.1),10) #10 samples of 4D params
    data_samples = rand(MvNormal(12,2),10) #10 samples of 12D data
    data_samples_short = data_samples[:,1:end-1]

    new_parameter_samples = rand(MvNormal(2,0.1),10) #10 samples of 4D params
    new_data_samples = rand(MvNormal(12,2),10) #10 samples of 12D data
    new_data_samples_short = new_data_samples[:,1:end-1]

    #test DataContainer
    parameter_samples_T = permutedims(parameter_samples, (2,1))
    idata = DataContainer(parameter_samples)
    idata_T = DataContainer(parameter_samples_T, data_are_columns=false)

    @test get_data(idata) == parameter_samples
    @test get_data(idata_T) == parameter_samples
    @test size(idata) == size(parameter_samples)

    #test Paired Container
    idata = DataContainer(parameter_samples)
    odata = DataContainer(data_samples)
    odata_short = DataContainer(data_samples_short)

    @test_throws DimensionMismatch PairedDataContainer(parameter_samples, data_samples_short)
    iopairs = PairedDataContainer(parameter_samples,data_samples)

    @test_throws DimensionMismatch PairedDataContainer(idata,odata_short)
    iopairs2 = PairedDataContainer(idata,odata)

    @test get_data(iopairs) == get_data(iopairs2)
    @test get_inputs(iopairs) == parameter_samples
    @test get_outputs(iopairs) == data_samples

end


## Calibrate camera
    # Initial guess at reference camera

    # Richard Shenton email:
    # Camera height: 2.065m
    #
    # Camera offset from centre line: 0.6m
    # Pitch: 19.16 degrees down
    # Yaw: 7.28 degrees to the left
    # Roll: Assumed to be zero

    # The camera frame rate is 25fps for this data. Regarding train details, there are some details here https://en.wikipedia.org/wiki/British_Rail_Class_357 .
    #
    # It looks to me that the carriage with the cab (seems to be called DMOS) is 20.75(L) x 2.8 (W) x 3.78(H) and weighs 40.7 t.

using TrainView

using CSV, DataFrames
using Statistics, LinearAlgebra

file = "../data/rail_detection_results_20210127.csv";
df = CSV.read(file, DataFrame);

Lu1 = df." Left_Rail_Bottom_X";
Lv1 = df." Left_Rail_Bottom_Y";
Lu2 = df." Left_Rail_Top_X";
Lv2 = df." Left_Rail_Top_Y";

Ru1 = df." Right_Rail_Bottom_X";
Rv1 = df." Right_Rail_Bottom_Y";
Ru2 = df." Right_Rail_Top_X";
Rv2 = df." Right_Rail_Top_Y";

# max_v = max(maximum(Lv1),maximum(Rv1))
max_v = 720
vs = [Lv1, Lv2, Rv1, Rv2];
for v in vs
    v[1:end] = v .- (max_v / 2.0)
end

# max_u = max(maximum(Lu2),maximum(Ru2))
max_u = 1280
us = [Lu1, Lu2, Ru1, Ru2];
for v in us
    v[1:end] = v .- (max_u / 2.0)
end

Lu = LinRange(mean(Lu1),mean(Lu2),20);
Lv = LinRange(mean(Lv1),mean(Lv2),20);
Ru = LinRange(mean(Ru1),mean(Ru2),20);
Rv = LinRange(mean(Rv1),mean(Rv2),20);

left_uvs =  [[Lu[i],Lv[i]] for i in eachindex(Lu)];
right_uvs = [[Ru[i],Rv[i]] for i in eachindex(Ru)];

    trackprop = TrackProperties(track_gauge = 1.435 + 0.065)

    cameraposition_reference = [0.0,0.6,-2.165];
    ψθφ_ref = [0.0,-19.1,-7.28] .* (pi/180.0)

    camera_reference = VideoCamera(cameraposition_reference;
        focal_length = 5.8e-3,
        # focal_length = 5.8e-3,
        pixelspermeter = 1 / 5.5e-6,
        ψθφ = ψθφ_ref
    )

    uLs = left_track_image_u(camera_reference, trackprop, Lv);
    uRs = right_track_image_u(camera_reference, trackprop, Rv);

    distortion = rail_uvs_to_distortion(left_uvs, right_uvs, camera_reference, trackprop;
        choose_distortions = [:Y,:Z,:θ,:φ],
        # choose_distortions = [:Y,:Z,:f,:θ,:φ,:ψ,:α,:β],
        iterations = 6)

    camera = VideoCamera(camera_reference,distortion)

    uLs2 = left_track_image_u(camera, trackprop, Lv);
    uRs2 = right_track_image_u(camera, trackprop, Rv);

    using Plots
    gr(yflip = true)

    plot(Lu,Lv, lab = "")
    plot!(Ru,Rv, lab = "")

    plot!(uLs2,Lv,linestyle=:dash, lab = "")
    plot!(uRs2,Rv,linestyle=:dash, lab = "")

## Analise journey
    choose_distortions = [:Y,:Z,:θ,:ψ];
    # choose_distortions = [:Y,:Z,:θ,:φ,:ψ,:α,:β]

    distortions = map(eachindex(Lu1)) do j
        Lu = LinRange(Lu1[j],Lu2[j],20);
        Lv = LinRange(Lv1[j],Lv2[j],20);
        Ru = LinRange(Ru1[j],Ru2[j],20);
        Rv = LinRange(Rv1[j],Rv2[j],20);

        left_uvs =  [[Lu[i],Lv[i]] for i in eachindex(Lu)];
        right_uvs = [[Ru[i],Rv[i]] for i in eachindex(Ru)];

        distortion = rail_uvs_to_distortion(left_uvs, right_uvs, camera, trackprop;
            choose_distortions = choose_distortions,
            iterations = 1
        );
        return [distortion[d] for d in choose_distortions]
    end

    data = transpose(hcat(distortions...))

    dict = Dict(choose_distortions[j] => data[:,j] for j in eachindex(choose_distortions))

    df = DataFrame(dict);

    moving_average(vs,n) = [sum(@view vs[i:(i+n-1)])/n for i in 1:(length(vs)-(n-1))]

    ## Plotting using the @df macro specifying colum names as symbol:
    using StatsPlots
    n = 15;
    n = 25;
    fps = 25;
    time = (1:size(data,1)) ./ fps
    df.time = time;

    @df df plot(moving_average(:time,n), [moving_average(:Y,n),moving_average(:Z,n)], labels = ["Horizontal" "Vertical"], xlab = "seconds")

    @df df plot(:time,[:Y,:Z], labels = ["Horizontal" "Vertical"], xlab = "seconds")

    @df df plot(moving_average(:time,n), [moving_average(:θ,n) * 180/π, moving_average(:ψ,n) * 180/π], labels = ["Pitch" "Roll"], xlab = "seconds", ylab = "Degree")


using PyPlot
using CSV
using DataFrames
using SparseArrays
using Statistics
import StatsBase: sample
using LinearAlgebra
using GPUMatFac

showln(x) = (show(x); println())

################################
# LOAD DATA
data = CSV.read("mnist_matrix.tsv", DataFrame; header=false, delim="\t");
#data = convert(Matrix, data);
data = Matrix(data);
full_X = data[:,1:end-1];
full_labels = data[:,end];
classes = sort(unique(full_labels));

################################
# TRAIN/TEST SPLIT
M = size(data,1)

train_idx = sort(sample(1:M, Int(0.9*M), replace=false))
train_idx_set = Set(train_idx)
test_idx = Int[i for i=1:M if !in(i, train_idx_set)]

train_X = full_X[train_idx,:]
train_labels = full_labels[train_idx]
showln("TRAIN SET: ")
showln(size(train_X))

test_X = full_X[test_idx,:]
test_labels = full_labels[test_idx]
showln("TEST SET: ")
showln(size(test_X))

full_X = vcat(train_X, test_X)
full_labels = vcat(train_labels, test_labels)
showln("FULL DATA: ")
showln(size(full_X))

################################
# SOME USEFUL FUNCTIONS
unflatten_mnist(vec) = transpose(reshape(vec, (28,28)))

function coord_to_idx(x, y, H)
    return (x-1)*H + y
end

function idx_to_pixel_coord(idx, H)
    return (div((idx-1),H), ((idx-1) % H))
end

################################
# BUILD THE PIXEL GRAPH
function build_pixel_graph(H, W)
    pixel_graph = [Int64[] for idx=1:(H*W)]

    for x=1:W
        for y=1:H
            u_idx = coord_to_idx(x,y, H)

            # Left
            if x > 1
                v_idx = coord_to_idx(x-1, y, H)
                push!(pixel_graph[u_idx], v_idx)
                push!(pixel_graph[v_idx], u_idx)
            end
            # Up
            if y > 1
                v_idx = coord_to_idx(x, y-1, H)
                push!(pixel_graph[u_idx], v_idx)
                push!(pixel_graph[v_idx], u_idx)
            end

        end
    end

    return pixel_graph
end

function plot_graph(g, idx_to_coord_func)
    for (u, neighbors) in enumerate(g)
        u_xy = idx_to_coord_func(u)
        for v in neighbors
            if v < u
                v_xy = idx_to_coord_func(v)
                plot([u_xy[1], v_xy[1]], [u_xy[2], v_xy[2]])
            end
        end
    end
end

pixel_graph = build_pixel_graph(28,28);

mnist_idx_to_coord_func(idx) = idx_to_pixel_coord(idx,28)

####################################
# BUILD THE TRAIN INSTANCE GRAPH
function build_instance_graph(labels)

    classes = sort(unique(labels))
    class_to_idx = Dict(cls => length(labels)+i for (i, cls) in enumerate(classes))
    n_classes = length(classes)
    n_bk = n_classes #+ 1
                                      # Leaf nodes;  # class nodes;  # root node
    instance_graph = [Int64[] for i=1:(length(labels) + n_bk) ];

    for (leaf_idx, label) in enumerate(labels)
        cls_idx = class_to_idx[label]
        push!(instance_graph[leaf_idx], cls_idx)
        push!(instance_graph[cls_idx], leaf_idx)
    end

    #root_idx = length(instance_graph)
    #for cls in classes
    #    cls_idx = class_to_idx[cls]
    #    push!(instance_graph[cls_idx], root_idx)
    #    push!(instance_graph[root_idx], cls_idx)
    #end

    return instance_graph
end

####################################
# BUILD THE TRAIN INSTANCE MATRIX
train_instance_graph = build_instance_graph(train_labels);
#aug_X = vcat( train_X, fill(NaN, length(classes)+1, 28*28) );
aug_X = vcat( train_X, fill(NaN, length(classes), 28*28) );

showln("AUGMENTED X: ")
showln(size(aug_X))

showln("BUILT PIXEL AND INSTANCE GRAPHS")

function graph_to_spmat(graph; epsilon=0.001, standardize=false)

    I = Int64[]
    J = Int64[]
    V = Float64[]
    N = length(graph)

    diag_entries = fill(epsilon, N)

    for (u, neighbors) in enumerate(graph)
        for v in neighbors
            push!(I, u)
            push!(J, v)
            push!(V, -1.0)
        end
        diag_entries[u] += length(neighbors)
    end
    for (i, v) in enumerate(diag_entries)
        push!(I, i)
        push!(J, i)
        push!(V, v)
    end

    mat = sparse(I, J, V, N, N)

    if standardize
        standardizer = sparse(1:N, 1:N, 1.0./sqrt.(diag_entries), N, N)
        mat = standardizer * mat * standardizer
    end

    return mat
end

pixel_spmat = graph_to_spmat(pixel_graph; standardize=true);
train_inst_spmat = graph_to_spmat(train_instance_graph; standardize=true);

#M_train = length(train_labels)
#train_inst_spmat = sparse(I, M_train, M_train)

####################################
# BUILD THE TEST SET INSTANCE MATRIX
function build_test_graph(train_graph, train_labels, test_labels)

    M_train = length(train_labels)
    M_train_aug = length(train_graph)
    M_test = length(test_labels)
    M_full = M_train_aug + M_test

    classes = sort(unique(train_labels))
    n_classes = length(classes)
    n_bk = n_classes + 1
    @assert n_bk == M_train_aug - M_train

    label_to_idx = Dict(l => (M_train + i) for (i,l) in enumerate(classes))

    test_graph = copy(train_graph)

    for i=1:M_test
        test_idx = M_train_aug + i
        parent_idx = label_to_idx[test_labels[i]]

        push!(test_graph, [parent_idx])
        push!(test_graph[parent_idx], test_idx)
    end

    return test_graph
end

#test_instance_graph = build_test_graph(train_instance_graph, train_labels, test_labels)
#test_spmat = graph_to_spmat(test_instance_graph)
#println("TEST SPARSE MATRIX:")
#showln(test_spmat)

showln("BUILT PIXEL AND INSTANCE MATRICES")

############################################
# Make a matrix of instance covariates -- just
# a one-hot encoding of the label
train_covariates = zeros(length(train_labels) + 11, 10)
for i=1:length(train_labels)
    train_covariates[i,Integer(train_labels[i])+1] = 1.0
end
for i=1:10
    train_covariates[length(train_labels)+i,i] = 1.0
end
#train_covariates[end,:] .= 0.1

############################################
# CONSTRUCT THE MODEL
k = 10

feature_spmat_vec = [copy(pixel_spmat) for i=1:k-1]
train_instance_spmat_vec = [copy(train_inst_spmat) for i=1:k]
#test_instance_spmat_vec = [copy(test_spmat) for i=1:k]

losses = [LogisticLoss(1.0) for i=1:28*28]
model = MatFacModel(train_instance_spmat_vec, feature_spmat_vec, losses;
                    #instance_covariate_coeff_reg=pixel_spmat,
                    feature_offset_reg=pixel_spmat, K=k);

showln("INITIALIZED MODEL")

############################################
# TRAIN MODEL
showln("ABOUT TO FIT")
#fit!(model, aug_X; instance_covariates=train_covariates,
#     method="nesterov", inst_reg_weight=0.0, feat_reg_weight=1.0,
#     max_iter=2000, loss_iter=10, lr=0.5, momentum=0.1,
#     rel_tol=1e-9, a_0_tau=1.0, b_0_tau=1.0)
fit!(model, aug_X;
     #instance_covariates=train_covariates,
     inst_reg_weight=160.0, feat_reg_weight=10.0,
     max_iter=1000, loss_iter=10, lr=0.1,
     rel_tol=1e-9, a_0_tau=1.0, b_0_tau=1.0)
#fit!(model, aug_X; inst_reg_weight=0.1, feat_reg_weight=0.1, max_iter=1000, lr=0.5, rel_tol=1e-9)
#fit!(model, aug_X; method="line_search", inst_reg_weight=0.1, feat_reg_weight=0.1, alpha=0.1, line_search_max_iter=5, grow=1.5, shrink=0.5, c1=1e-4, c2=0.5, max_iter=1000, rel_tol=1e-6)

println("FACTOR IMPORTANCES:")
println(model.factor_importances)

#####################
# SAVE MODEL
save_hdf("mnist_model.hdf", model)

#####################
# RELOAD MODEL
model = load_hdf("mnist_model.hdf")

train_labels = convert(Vector{Int64}, train_labels)
test_labels = convert(Vector{Int64}, test_labels)
colors = collect(keys(PyPlot.colorsm.TABLEAU_COLORS))

function embedding_scatter(X, labels; F=nothing)
    projected = nothing
    if F == nothing
        F = LinearAlgebra.svd(X)
        projected = F.Vt
    else
        projected = diagm(1.0./F.S) * (transpose(F.U)* X)
    end
    M = length(labels)
    for lab in sort(unique(labels))
        lab_idx = findall(labels .== lab)
        scatter3D(projected[1, lab_idx], projected[2,lab_idx], projected[3,lab_idx], color=colors[lab+1], label=string(lab))
    end
    legend()

    return F
end

embedding_F = embedding_scatter(model.X[1:(end-1),:], train_labels)
savefig("embedding_scatter.png", dpi=200)

for i=1:k
    matshow(unflatten_mnist(model.Y[i,:]))
    colorbar()
    savefig(string("embedding_basis_",i,".png"), dpi=200)
end

matshow(unflatten_mnist(model.feature_offset))
colorbar()
savefig(string("embedding_offset.png"), dpi=200)

for i=1:size(model.instance_covariate_coeff,1)
    matshow(unflatten_mnist(model.instance_covariate_coeff[i,:]))
    colorbar()
    savefig(string("covariate_coeff_",i,".png"), dpi=200)
end

######################################
# Use the model to impute data
Z = impute_values(model)
println(size(Z))

#####################################
# TRANSFORM THE HELDOUT DATA

new_factor = 0.001*randn(k, length(test_labels))
new_factor[k,:] .= 1.0

showln("NEW FACTOR:")
showln(size(new_factor))

# TODO UPDATE THE TRANSFORM METHOD
#transformed = GPUMatFac.transform(model, test_X; inst_reg_weight=0.1, max_iter=1000, loss_iter=1, lr=0.1, momentum=0.8, rel_tol=1e-9, X_new=new_factor) #, new_inst_reg_mats=test_instance_spmat_vec)

#println("TRANSFORMED:")
#println(size(transformed))

#embedding_scatter(transformed[1:(end-1),:], test_labels; F=embedding_F)
#savefig("transformed_scatter.png", dpi=200)


@testset "Maximal Independent Set" begin

    g3 = StarGraph(5)
    for g in testgraphs(g3)
        z = @inferred(independent_set(g, MaximalIndependentSet()))
        @test (length(z)== 1 || length(z)== 4)
    end

    g4 = CompleteGraph(5)
    for g in testgraphs(g4)
        z = @inferred(independent_set(g, MaximalIndependentSet()))
        @test length(z)== 1 #Exactly one vertex
    end

    g5 = PathGraph(5)
    for g in testgraphs(g5)
        z = @inferred(independent_set(g, MaximalIndependentSet()))
        sort!(z)
        @test (z == [2, 4] || z == [2, 5] || z == [1, 3, 5] || z == [1, 4])
    end
end


using Revise
using ADCME
using ADCMEKit
using NNFEM
using PyPlot
using ProgressMeter
using Printf
using MAT
using Clustering
using JLD2

close("all")

dat = matread("data/dippingfault_viscosity_inversion.mat")
σ_ = dat["sigma"]
d_ = dat["d"]
η_ = dat["eta"]
d0 = dat["d0"]
v0 = dat["v0"]
a0 = dat["a0"]
ϵ0 = dat["epsilon0"]
σ0 = dat["sigma0"]
NT = dat["NT"]
Δt = dat["dt"]

@load "data/dippingfault_viscosity_inversion.jld2" domain
#################### Get surface observation ####################

## get surface slip
slip_idx = findall(domain.nodes[:,2] .≈ 0.0)
coords = domain.nodes[slip_idx,1]
ii = sortperm(coords)
sorted_slip_idx = slip_idx[ii]
coords = coords[ii]

########################### Inversion ###########################
# create Variable eta
gnodes = getGaussPoints(domain)
using Random; Random.seed!(0)
kr = kmeans(gnodes', 20)
A = kr.assignments

figure()
for i = 1:20
  scatter(gnodes[A .== i,1], gnodes[A .== i,2])
end
gca().invert_yaxis()
autoscale(enable=true, axis="both", tight=true)
title("Patch division for inversion")
ylabel("y")
xlabel("x")
savefig("figures/dipslip-patch.png")

vs = Variable(1.5*ones(20))
ETA = constant(zeros(getNGauss(domain)))
for i = 1:20
  mask = zeros(getNGauss(domain))
  mask[A.==i] .= 1.0
  global ETA += vs[i] * mask
end
ETA *= 1e10
η = ETA

μ = zeros(getNGauss(domain))
λ = zeros(getNGauss(domain))
k = 0
for i = 1:domain.neles
  e = domain.elements[i]
  for mat in e.mat
    global k += 1
    μ[k] = mat.μ
    λ[k] = mat.λ
  end
end

globaldata = GlobalData(missing, missing, missing, missing, domain.neqs, nothing, nothing, nothing)
ts = GeneralizedAlphaSolverTime(Δt, NT)
ubd, abd = compute_boundary_info(domain, globaldata, ts)
Fext = compute_external_force(domain, globaldata, ts)

d, v, a, σ, ϵ = ViscoelasticitySolver(
  globaldata, domain, d0, v0, a0, σ0, ϵ0, Δt, NT, μ, λ, η, Fext, ubd, abd
)

dat = matread("data/dippingfault_viscosity_inversion.mat")
d_ = dat["d"]
y_id = [sorted_slip_idx; sorted_slip_idx .+ domain.nnodes]
loss = sum((d[:, y_id]-d_[:, y_id])^2)
# loss = sum((d - d_)^2)
sess = Session(); init(sess)

figure()
cb = (vs, iter, loss)->begin
    η_est = vs[1]
    if mod(iter, 10) == 0
        clf()
        visualize_scalar_on_scoped_body(η_est, zeros(domain.nnodes*2), domain, vmin=minimum(η_), vmax=maximum(η_))
        title("Iter = $iter")
        autoscale(enable=true, axis="both", tight=true)
        savefig("figures_tmp/dippingfault_viscosity_inversion_$(lpad(iter,5,"0")).png",bbox_size="tight")
    end
    printstyled("[#iter $iter] loss = $loss\n", color = :green)
end

@show run(sess, loss)
BFGS!(sess, loss, 200, callback=cb, vars=[η[1:9:end]])

figure()
η_est = run(sess, η[1:9:end])
visualize_scalar_on_scoped_body(η_est, zeros(domain.nnodes*2), domain, vmin=minimum(η_), vmax=maximum(η_))
autoscale(enable=true, axis="both", tight=true)
title("Inverted viscosity result")
savefig("figures/dipslip-inv-linear-visco-model.png")

figure()
η_ref = η_[1:9:end]
visualize_scalar_on_scoped_body(η_ref, zeros(domain.nnodes*2), domain)
autoscale(enable=true, axis="both", tight=true)
title(L"True viscosity ($\eta$) model")
savefig("figures/dipslip-linear-visco-model.png")


@testset "test_discrete" begin

C_111 = ss([-5], [2], [3], [0])
C_212 = ss([-5 -3; 2 -9], [1; 2], [1 0; 0 1], [0; 0])
C_221 = ss([-5 -3; 2 -9], [1 0; 0 2], [1 0], [0 0])
C_222_d = ss([-5 -3; 2 -9], [1 0; 0 2], [1 0; 0 1], [1 0; 0 1])

@test c2d(ss(4*[1 0; 0 1]), 0.5, :zoh) == (ss(4*[1 0; 0 1], 0.5), zeros(0, 2))
@test c2d(ss(4*[1 0; 0 1]), 0.5, :foh) == (ss(4*[1 0; 0 1], 0.5), zeros(0, 2))
@test_c2d(c2d(C_111, 0.01, :zoh),
    ss([0.951229424500714], [0.019508230199714396], [3], [0], 0.01), [1 0], true)
@test_c2d(c2d(C_111, 0.01, :foh),
    ss([0.951229424500714], [0.01902855227625244], [3], [0.029506188017136226], 0.01),
    [1 -0.009835396005712075], true)
@test_c2d(c2d(C_212, 0.01, :zoh),
    ss([0.9509478368863918 -0.027970882212682433; 0.018647254808454958 0.9136533272694819],
    [0.009466805409666932; 0.019219966830212765], [1 0; 0 1], [0; 0], 0.01),
    [1 0 0; 0 1 0], true)
@test_c2d(c2d(C_212, 0.01, :foh),
    ss([0.9509478368863921 -0.027970882212682433; 0.018647254808454954 0.913653327269482],
    [0.008957940478201584; 0.018468989584974498], [1 0; 0 1], [0.004820885889482196;
    0.009738343195298675], 0.01), [1 0 -0.004820885889482196; 0 1 -0.009738343195298675], true)
@test_c2d(c2d(C_221, 0.01, :zoh),
    ss([0.9509478368863918 -0.027970882212682433; 0.018647254808454958 0.9136533272694819],
    [0.009753161420545834 -0.0002863560108789034; 9.54520036263011e-5 0.019124514826586465],
    [1.0 0.0], [0.0 0.0], 0.01), [1 0 0 0; 0 1 0 0], true)
@test_c2d(c2d(C_221, 0.01, :foh),
    ss([0.9509478368863921 -0.027970882212682433; 0.018647254808454954 0.913653327269482],
    [0.009511049106772921 -0.0005531086285713394; 0.00018436954285711309 0.018284620042117387],
    [1 0], [0.004917457305816479 -9.657141633428213e-5], 0.01),
    [1 0 -0.004917457305816479 9.657141633428213e-5;
    0 1 -3.219047211142736e-5 -0.009706152723187249], true)
@test_c2d(c2d(C_222_d, 0.01, :zoh),
    ss([0.9509478368863918 -0.027970882212682433; 0.018647254808454958 0.9136533272694819],
    [0.009753161420545834 -0.0002863560108789034; 9.54520036263011e-5 0.019124514826586465],
    [1 0; 0 1], [1 0; 0 1], 0.01), [1 0 0 0; 0 1 0 0], true)
@test_c2d(c2d(C_222_d, 0.01, :foh),
    ss([0.9509478368863921 -0.027970882212682433; 0.018647254808454954 0.913653327269482],
    [0.009511049106772921 -0.0005531086285713394; 0.00018436954285711309 0.018284620042117387],
    [1 0; 0 1], [1.0049174573058164 -9.657141633428213e-5; 3.219047211142736e-5 1.0097061527231872], 0.01),
    [1 0 -0.004917457305816479 9.657141633428213e-5; 0 1 -3.219047211142736e-5 -0.009706152723187249], true)

# Test some false
@test_c2d(c2d(2C_111, 0.01, :zoh),  # Factor 2
    ss([0.951229424500714], [0.019508230199714396], [3], [0], 0.01), [1 0], false)
@test_c2d(c2d(C_111, 0.01, :foh),   #Wrong C
    ss([0.951229424500714], [0.01902855227625244], [3*3], [0.029506188017136226], 0.01),
    [1 -0.009835396005712075], false)
@test_c2d(c2d(C_212, 0.1, :zoh),    # Wrong Ts
    ss([0.9509478368863918 -0.027970882212682433; 0.018647254808454958 0.9136533272694819],
    [0.009466805409666932; 0.019219966830212765], [1 0; 0 1], [0; 0], 0.01),
    [1 0 0; 0 1 0], false)

# Test c2d for transfer functions
G = tf([1, 1], [1, 3, 1])
Gd = c2d(G, 0.2)
@test Gd ≈ tf([0, 0.165883310712090, -0.135903621603238], [1.0, -1.518831946985175, 0.548811636094027], 0.2) rtol=1e-14

# Issue #391
@test c2d(tf(C_111), 0.01, :zoh) ≈ tf(c2d(C_111, 0.01, :zoh)[1])
@test c2d(tf(C_111), 0.01, :foh) ≈ tf(c2d(C_111, 0.01, :foh)[1])

# c2d on a zpk model should arguably return a zpk model
@test_broken typeof(c2d(zpk(G), 1)) <: TransferFunction{<:ControlSystems.SisoZpk}

# ERRORS
@test_throws ErrorException c2d(ss([1], [2], [3], [4], 0.01), 0.01)   # Already discrete
@test_throws ErrorException c2d(ss([1], [2], [3], [4], -1), 0.01)     # Already discrete

# d2c
@static if VERSION > v"1.4" # log(matrix) is buggy on previous versions, should be fixed in 1.4 and back-ported to 1.0.6
    @test d2c(c2d(C_111, 0.01)[1]) ≈ C_111
    @test d2c(c2d(C_212, 0.01)[1]) ≈ C_212
    @test d2c(c2d(C_221, 0.01)[1]) ≈ C_221
    @test d2c(c2d(C_222_d, 0.01)[1]) ≈ C_222_d
    @test d2c(Gd) ≈ G

    sys = ss([0 1; 0 0], [0;1], [1 0], 0)
    sysd = c2d(sys, 1)[1]
    @test d2c(sysd) ≈ sys
end

end


# Define the theoretic model
################### Assemble matrix and vector #####################
# Weak form of the Helmholtz equation :
# Hz: a(p,u,v)=Λ⋅∇v⋅ξ(p)Λ⋅∇u-k²μv⋅u
# Material distribution
#ξ0(x,ϵ1,ϵ2) = x[2]>0 ? 1/ϵ1 : 1/ϵ2
#ξd(p,ϵmin,ϵmax)= 1/(ϵmin + (ϵmax-ϵmin)*p) - 1/ϵmin # in the design region
#a0(u,v,ϵ1,ϵ2,μ,σs,k,LHp,LHn,dpml) = ((x->Λ(x,σs,k,LHp,LHn,dpml))⋅∇(v))⊙((x->ξ0(x,ϵ1,ϵ2))*((x->Λ(x,σs,k,LHp,LHn,dpml))⋅∇(u))) - k^2*μ*(v*u)
#a_d(u,v,ph,k,ϵ1,ϵd) = ∇(v)⊙(((p -> ξd(p, ϵ1, ϵd))∘ph)*∇(u))
# Ez: a(p,u,v)=1/μ*Λ⋅∇v⋅Λ⋅∇u-k²ξ(p)uv
ξd(p,ϵmin,ϵmax)= ϵmin+(ϵmax-ϵmin)*p # in the design region
a0(u,v,ph,ϵ1,ϵ2,μ,σs,k,LHp,LHn,Hd,dpml) = ((x->Λ(x,σs,Hd,k,LHp,LHn,dpml))⋅∇(v))⊙((1/μ)*((x->Λ(x,σs,Hd,k,LHp,LHn,dpml))⋅∇(u))) - k^2*(((p -> ξd(p, ϵ1, ϵ2))∘ph)*(v*u))
#a_d(u,v,ph,k,ϵ1,ϵd) = -k^2*((p -> ξd(p, ϵ1, ϵd))∘ph)*(v*u)
# Source term (Gaussian point source approximation at center)

#b(v,x0,δ) = v*(x->GaussianD(x,x0,δ))

function MatrixA(ph,ϵ1,ϵ2,μ,σs,k,LHp,LHn,dpml,Hd,U,V,dΩ)
    # Assemble the matrix
    return assemble_matrix(U,V) do u,v
        ∫( a0(u,v,ph,ϵ1,ϵ2,μ,σs,k,LHp,LHn,Hd,dpml) )dΩ
    end
end

function MatrixB(x0,δ,Amp,V,dΩ,dΓ)
    l_temp(v) = ∫( Amp*v*x->GaussianY(x,x0,δ) )*dΩ+∫( 0*v )*dΓ
    #op = AffineFEOperator((u,v)->(∫(u*v)*dΩ),l_temp,U,V)
    #op = FESource(l_temp,U,V)
    assemble_vector(l_temp,V)
end


function time_update(filter::SquareRootFilter, x, k)
    F = jacobian_dynamics_x(filter, x, k)
    v = get_process_noise(filter, k) # Zero mean RV
    G_v = jacobian_dynamics_v(filter, x, k)

    Q_sqrt = chol(v)
    P_sqrt = chol(x)

    x_pred = predict_x(filter, x, k)

    L = lq(hcat(F*P_sqrt, G_v*Q_sqrt)).L
    n_x = size(F, 1)
    P_pred_sqrt = L[1:n_x, 1:n_x]

    return MvNormal(x_pred, P_pred_sqrt)
end
function measurement_update(filter::SquareRootFilter, y, x_pred, k)

    H = get_H(filter, k)
    P_pred_sqrt = chol(x_pred)
    R_sqrt = chol(y)

    y_pred = H*mean(x_pred) # Include other input and shit
    e = mean(y) - y_pred # Innovation

    n_x = size(P_pred_sqrt, 1)
    n_y = size(R_sqrt, 1)
    z = zeros(n_x, n_y)
    M = vcat(
        hcat(R_sqrt, H*P_pred_sqrt),
        hcat(z,  P_pred_sqrt)
    )
    L = lq(M).L

    S_sqrt = L[1:n_y, 1:n_y]
    P_sqrt = L[1+n_y:n_y + n_x, 1+n_y : n_y + n_x]
    KS_sqrt = L[1 + n_y:n_x + ny, 1:n_y]

    K = KS_sqrt/S_sqrt # Should be fast since triangular

    x_filtered = mean(x_pred) + K*e

    return MvNormal(x_filtered, P_sqrt), MvNormal(e, S_sqrt)
end


using Printf
using CSV, DataFrames
using Arpack
using PolynomialBasis
using ImplicitDomainQuadrature
using Revise
using CutCell
include("../../test/useful_routines.jl")

function displacement(alpha, x)
    u1 = alpha * x[2] * sin(pi * x[1])
    u2 = alpha * (x[1]^3 + cos(pi * x[2]))
    return [u1, u2]
end

function body_force(lambda, mu, alpha, x)
    b1 = alpha * (lambda + 2mu) * pi^2 * x[2] * sin(pi * x[1])
    b2 =
        -alpha * (6mu * x[1] + (lambda + mu) * pi * cos(pi * x[1])) +
        alpha * (lambda + 2mu) * pi^2 * cos(pi * x[2])
    return [b1, b2]
end

function onboundary(x, L, W)
    return x[2] ≈ 0.0 || x[1] ≈ L || x[2] ≈ W || x[1] ≈ 0.0
end

function solve_elasticity(x0, normal, nelmts, polyorder, numqp, penaltyfactor; eta = 1.0)
    L = 1.0
    W = 1.0
    lambda, mu = 1.0, 2.0
    dx = 1.0 / nelmts
    penalty = penaltyfactor / dx * (lambda + mu)
    alpha = 0.1
    stiffness = CutCell.HookeStiffness(lambda, mu, lambda, mu)

    basis = TensorProductBasis(2, polyorder)
    mesh = CutCell.Mesh([0.0, 0.0], [L, W], [nelmts, nelmts], basis)
    levelset = InterpolatingPolynomial(1, basis)
    levelsetcoeffs =
        CutCell.levelset_coefficients(x -> plane_distance_function(x, normal, x0), mesh)

    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    cellquads = CutCell.CellQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    interfacequads = CutCell.InterfaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    facequads = CutCell.FaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)

    bilinearforms = CutCell.BilinearForms(basis, cellquads, stiffness, cutmesh)
    interfacecondition =
        CutCell.InterfaceCondition(basis, interfacequads, stiffness, cutmesh, penalty)

    displacementbc = CutCell.DisplacementCondition(
        x -> displacement(alpha, x),
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> onboundary(x, L, W),
        penalty,
    )

    sysmatrix = CutCell.SystemMatrix()
    sysrhs = CutCell.SystemRHS()

    CutCell.assemble_bilinear_form!(sysmatrix, bilinearforms, cutmesh)
    CutCell.assemble_interface_condition!(sysmatrix, interfacecondition, cutmesh, eta = eta)
    CutCell.assemble_body_force_linear_form!(
        sysrhs,
        x -> body_force(lambda, mu, alpha, x),
        basis,
        cellquads,
        cutmesh,
    )
    CutCell.assemble_penalty_displacement_bc!(sysmatrix, sysrhs, displacementbc, cutmesh)

    matrix = CutCell.make_sparse(sysmatrix, cutmesh)
    rhs = CutCell.rhs(sysrhs, cutmesh)

    sol = matrix \ rhs
    disp = reshape(sol, 2, :)

    err = mesh_L2_error(disp, x -> displacement(alpha, x), basis, cellquads, cutmesh)
    normalizer =
        integral_norm_on_cut_mesh(x -> displacement(alpha, x), cellquads, cutmesh, 2)
    return err ./ normalizer
end

function condition_number(x0, normal, nelmts, polyorder, numqp, penaltyfactor; eta = 1.0)
    L = 1.0
    W = 1.0
    lambda, mu = 1.0, 2.0
    dx = 1.0 / nelmts
    penalty = penaltyfactor / dx * (lambda + mu)
    alpha = 0.1
    stiffness = CutCell.HookeStiffness(lambda, mu, lambda, mu)

    basis = TensorProductBasis(2, polyorder)
    mesh = CutCell.Mesh([0.0, 0.0], [L, W], [nelmts, nelmts], basis)
    levelset = InterpolatingPolynomial(1, basis)
    levelsetcoeffs =
        CutCell.levelset_coefficients(x -> plane_distance_function(x, normal, x0), mesh)

    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    cellquads = CutCell.CellQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    interfacequads = CutCell.InterfaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    facequads = CutCell.FaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)

    bilinearforms = CutCell.BilinearForms(basis, cellquads, stiffness, cutmesh)
    interfacecondition =
        CutCell.InterfaceCondition(basis, interfacequads, stiffness, cutmesh, penalty)

    displacementbc = CutCell.DisplacementCondition(
        x -> displacement(alpha, x),
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> onboundary(x, L, W),
        penalty,
    )

    sysmatrix = CutCell.SystemMatrix()
    sysrhs = CutCell.SystemRHS()

    CutCell.assemble_bilinear_form!(sysmatrix, bilinearforms, cutmesh)
    CutCell.assemble_interface_condition!(sysmatrix, interfacecondition, cutmesh, eta = eta)

    matrix = CutCell.make_sparse(sysmatrix, cutmesh)

    lambda1, phi1 = eigs(matrix,nev=1,which=:SM)
    lambda2, phi2 = eigs(matrix,nev=1,which=:LM)

    k = abs(real(lambda2[1])/real(lambda1[1]))

    return k
end

function write_convergence_data_to_file(x0, theta, nelmts, polyorder, numqp, penaltyfactor)
    normal = normal_from_angle(theta)
    errors =
        [solve_elasticity(x0, normal, ne, polyorder, numqp, penaltyfactor) for ne in nelmts]

    dx = 1.0 ./ nelmts
    u1err = [er[1] for er in errors]
    u1rate = [
        (log(u1err[i]) - log(u1err[i-1])) / (log(dx[i]) - log(dx[i-1]))
        for i = 2:length(nelmts)
    ]
    pushfirst!(u1rate, 0.0)

    u2err = [er[2] for er in errors]
    u2rate = [
        (log(u2err[i]) - log(u2err[i-1])) / (log(dx[i]) - log(dx[i-1]))
        for i = 2:length(nelmts)
    ]
    pushfirst!(u2rate, 0.0)

    df = DataFrame(
        [dx, u1err, u1rate, u2err, u2rate],
        ["Element Size", "U1 Error", "U1 Rate", "U2 Error", "U2 Rate"],
    )

    penaltystr = @sprintf "%1.1f" penaltyfactor

    foldername = "examples/cut-convergence/theta-" *
    string(round(Int, theta)) *
    "-penalty-" * penaltystr

    if !ispath(foldername)
        mkpath(foldername)
    end

    filename = foldername *
        "/polyorder_" *
        string(polyorder) *
        "_.csv"
    CSV.write(filename, df)
end

x0 = [0.7, 0.0]
theta = 30
normal = normal_from_angle(theta)
powers = 1:7
nelmts = [2^p + 1 for p in powers]
polyorder = 1
numqp = required_quadrature_order(polyorder)+2
penaltyfactor = 1e1
eta = 1.0

k = [condition_number(x0,normal,ne,polyorder,numqp,penaltyfactor) for ne in nelmts]
# errors = [solve_elasticity(x0,normal,ne,polyorder,numqp,penaltyfactor,eta=eta) for ne in nelmts]

# dx = 1.0 ./ nelmts
# u1err = [er[1] for er in errors]

# rate = diff(log.(u1err)) ./ diff(log.(dx))

# write_convergence_data_to_file(x0, theta, nelmts, polyorder, numqp, penaltyfactor)


using Flux
using Base: @get!
using MacroTools: @forward

const ϵ = 1e-8

# TODO: should use weak refs

"""
    Descent(η)

Classic gradient descent optimiser with learning rate `η`.
For each parameter `p` and its gradient `δp`, this runs `p -= η*δp`.
"""
mutable struct Descent
  eta::Float64
end

Descent() = Descent(0.1)

function apply!(o::Descent, x, Δ)
  Δ .*= o.eta
end

"""
    Momentum(params, η = 0.01; ρ = 0.9)

Gradient descent with learning rate `η` and momentum `ρ`.
"""
mutable struct Momentum
  eta::Float64
  rho::Float64
  velocity::IdDict
end

Momentum(η = 0.01, ρ = 0.9) = Momentum(η, ρ, IdDict())

function apply!(o::Momentum, x, Δ)
  η, ρ = o.eta, o.rho
  v = get!(o.velocity, x, zero(x))::typeof(data(x))
  @. v = ρ * v - η * Δ
  @. Δ = -v
end

"""
    Nesterov(eta, ρ = 0.9)

Gradient descent with learning rate  `η` and Nesterov momentum `ρ`.
"""
mutable struct Nesterov
  eta::Float64
  rho::Float64
  velocity::IdDict
end

Nesterov(η = 0.001, ρ = 0.9) = Nesterov(η, ρ, IdDict())

function apply!(o::Nesterov, x, Δ)
  η, ρ = o.eta, o.rho
  v = get!(o.velocity, x, zero(x))::typeof(data(x))
  d = @. ρ^2 * v - (1+ρ) * η * Δ
  @. v = ρ*v - η*Δ
  @. Δ = -d
end

"""
    RMSProp(η = 0.001, ρ = 0.9)

[RMSProp](https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf)
optimiser. Parameters other than learning rate don't need tuning. Often a good
choice for recurrent networks.
"""
mutable struct RMSProp
  eta::Float64
  rho::Float64
  acc::IdDict
end

RMSProp(η = 0.001, ρ = 0.9) = RMSProp(η, ρ, IdDict())

function apply!(o::RMSProp, x, Δ)
  η, ρ = o.eta, o.rho
  acc = get!(o.acc, x, zero(x))::typeof(data(x))
  @. acc = ρ * acc + (1 - ρ) * Δ^2
  @. Δ *= η / (√acc + ϵ)
end

"""
    ADAM(η = 0.001, β = (0.9, 0.999))

[ADAM](https://arxiv.org/abs/1412.6980v8) optimiser.
"""
mutable struct ADAM
  eta::Float64
  beta::Tuple{Float64,Float64}
  state::IdDict
end

ADAM(η = 0.001, β = (0.9, 0.999)) = ADAM(η, β, IdDict())

function apply!(o::ADAM, x, Δ)
  η, β = o.eta, o.beta
  mt, vt, βp = get!(o.state, x, (zero(x), zero(x), β))
  @. mt = β[1] * mt + (1 - β[1]) * Δ
  @. vt = β[2] * vt + (1 - β[2]) * Δ^2
  @. Δ =  mt / (1 - βp[1]) / (√(vt / (1 - βp[2])) + ϵ) * η
  o.state[x] = (mt, vt, βp .* β)
  return Δ
end

"""
    AdaMax(params, η = 0.001; β1 = 0.9, β2 = 0.999, ϵ = 1e-08)

[AdaMax](https://arxiv.org/abs/1412.6980v9) optimiser. Variant of ADAM based on
the ∞-norm.
"""
mutable struct AdaMax
  eta::Float64
  beta::Tuple{Float64,Float64}
  state::IdDict
end

AdaMax(η = 0.001, β = (0.9, 0.999)) = AdaMax(η, β, IdDict())

function apply!(o::AdaMax, x, Δ)
  η, β = o.eta, o.beta
  mt, ut, βp = get!(o.state, x, (zero(x), zero(x), β))
  @. mt = β[1] * mt + (1 - β[1]) * Δ
  @. ut = max(β[2] * ut, abs(Δ))
  @. Δ = (η/(1 - βp[1])) * mt/(ut + ϵ)
  o.state[x] = (mt, ut, βp .* β)
  return Δ
end

"""
    ADAGrad(η = 0.1; ϵ = 1e-8)

[ADAGrad](http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf) optimiser.
Parameters don't need tuning.
"""
mutable struct ADAGrad
  eta::Float64
  acc::IdDict
end

ADAGrad(η = 0.1) = ADAGrad(η, IdDict())

function apply!(o::ADAGrad, x, Δ)
  η = o.eta
  acc = get!(o.acc, x, fill(ϵ, size(x)))::typeof(data(x))
  @. acc += Δ^2
  @. Δ *= η / (√acc + ϵ)
end

"""
    ADADelta(ρ = 0.9, ϵ = 1e-8)

[ADADelta](https://arxiv.org/abs/1212.5701) optimiser. Parameters don't need
tuning.
"""
mutable struct ADADelta
  rho::Float64
  state::IdDict
end

ADADelta(ρ = 0.9) = ADADelta(ρ, IdDict())

function apply!(o::ADADelta, x, Δ)
  ρ = o.rho
  acc, Δacc = get!(o.state, x, (zero(x), zero(x)))
  @. acc = ρ * acc + (1 - ρ) * Δ^2
  @. Δ *= √Δacc/ (√acc + ϵ)
  @. Δacc = ρ * Δacc + (1 - ρ) * Δ^2
  return Δ
end

"""
    AMSGrad(η = 0.001, β = (0.9, 0.999))

[AMSGrad](https://openreview.net/forum?id=ryQu7f-RZ) optimiser. Parameters don't need
tuning.
"""
mutable struct AMSGrad
  eta::Float64
  beta::Tuple{Float64, Float64}
  state::IdDict
end

AMSGrad(η = 0.001, β = (0.9, 0.999)) = AMSGrad(η, β, IdDict())

function apply!(o::AMSGrad, x, Δ)
  η, β = o.eta, o.beta
  mt, vt, v̂t = get!(o.state, x, (fill(ϵ, size(x)), fill(ϵ, size(x)), fill(ϵ, size(x))))
  @. mt = β[1] * mt + (1 - β[1]) * Δ
  @. vt = β[2] * vt + (1 - β[2]) * Δ ^ 2
  @. v̂t = max.(v̂t, vt)
  @. Δ = η * mt / (√v̂t + ϵ)
end

"""
    NADAM(η = 0.001, β = (0.9, 0.999))

[NADAM](http://cs229.stanford.edu/proj2015/054_report.pdf) optimiser. Parameters don't need
tuning.
"""
mutable struct NADAM
  eta::Float64
  beta::Tuple{Float64, Float64}
  state::IdDict
end

NADAM(η = 0.001, β = (0.9, 0.999)) = NADAM(η, β, IdDict())

function apply!(o::NADAM, x, Δ)
  η, β = o.eta, o.beta
  β1p, β2p = o.beta
  mt, vt = get!(o.state, x, (zero(x), zero(x)))
  @. mt = β[1] * mt + (1 - β[1]) * Δ
  @. vt = β[2] * vt + (1 - β[2]) * Δ^2
  @. Δ = (β[1] * mt / (1 - β[1] * β1p) + (1 - β[1]) * Δ / (1 - β1p)) / (√(vt * β[2] / (1 - β2p)) + ϵ) * η
  o.state[x] = (mt, vt, (β1p * β[1], β2p * β[2]))
  return Δ
end

"""
    ADAMW((η = 0.001, β = (0.9, 0.999), decay = 0)

[ADAMW](https://arxiv.org/abs/1711.05101) fixing weight decay regularization in Adam.
"""
ADAMW(η = 0.001, β = (0.9, 0.999), decay = 0) =
  Optimiser(ADAM(η, β), WeightDecay(decay))

# Compose optimizers

"""
    Optimiser(a, b, c...)

Combine several optimisers into one; each optimiser produces a modified gradient
that will be fed into the next, and this is finally applied to the parameter as
usual.
"""
mutable struct Optimiser
  os::Vector{Any}
end

Optimiser(o...) = Optimiser(Any[o...])

@forward Optimiser.os Base.getindex, Base.first, Base.last, Base.lastindex, Base.push!, Base.setindex!
@forward Optimiser.os Base.iterate

Base.getindex(c::Optimiser, i::AbstractArray) = Optimiser(c.os[i]...)

function apply!(o::Optimiser, x, Δ)
  for opt in o.os
    Δ = apply!(opt, x, Δ)
  end
  return Δ
end

"""
`InvDecay(γ)`

Apply inverse time decay to an optimiser
```julia
  Optimiser(InvDecay(..), Opt(..))
```
"""
mutable struct InvDecay
  gamma::Float64
  state::IdDict
end

InvDecay(γ = 0.001) = InvDecay(γ, IdDict())

function apply!(o::InvDecay, x, Δ)
  γ = o.gamma
  n = get!(o.state, x, 1)
  Δ .*= 1 / (1 + γ * n)
  o.state[x] = n + 1
  return Δ
end

"""
`ExpDecay(eta, decay, decay_step, clip)`

Schedule the learning rate `eta` by `decay` every `decay_step` till a minimum of `clip`.

To apply exponential decay to an optimiser:
```julia
  Optimiser(ExpDecay(..), Opt(..))
```
"""
mutable struct ExpDecay
  eta::Float64
  decay::Float64
  step::Int64
  clip::Float64
  current::IdDict
end

ExpDecay(opt = 0.001, decay = 0.1, decay_step = 1000, clip = 1e-4) = ExpDecay(opt, decay, decay_step, clip, IdDict())

function apply!(o::ExpDecay, x, Δ)
  η, s, decay = o.eta, o.step, o.decay
  n = o.current[x] = get(o.current, x, 0) + 1
  if o.current[x]%s == 0 && count(x -> x%s == 0, values(o.current)) == 1
    η = max(η * decay^(s / n), o.clip)
    o.eta = η
  end
  @. Δ *= η
end

"""
`WeightDecay(wd)`

Decay the weight parameter by `wd`
"""
mutable struct WeightDecay
  wd::Real
end

WeightDecay() = WeightDecay(0)

function apply!(o::WeightDecay, x, Δ)
  wd = o.wd
  @. Δ += wd * data(x)
end


R, x = RationalFunctionField(QQ, "x")
L = Localization(R, degree)

@testset "DegreeLocalization" begin

    @testset "Constructor" begin

      @test parent_type(KInftyElem{fmpq}) == KInftyRing{fmpq}
      @test elem_type(L) == KInftyElem{fmpq}
      @test function_field(L) == R

    end

    @testset "Canonicalisation" begin

      for i in 1:300
        a = rand(L, 0:10, -10:10)
        c = canonical_unit(a)
        @test (iszero(c) && !isunit(c)) || isunit(c)
      end
    end

    @testset "Valuation" begin

      @test valuation(L(x//(x^2 + 1))) == 1
      @test valuation(L(5//5)) == 0
      @test valuation(L((x^2 + 1)//(3//5*x^3))) == 1
      @test valuation(L((2//5)//(x^3 + 3x + 1))) == 3

      for i in 1:300
        a = rand(L, 0:10, -10:10)
        while iszero(a)
          a = rand(L, 0:10, -10:10)
        end
        @test valuation(a) >= 0
      end

    end

    @testset "Parent object call overloading" begin

      @test parent(L(3)) == L
      @test parent(L(1//x)) == L
      @test parent(L(fmpz(3))) == L
      @test parent(L(fmpq(3, 2))) == L
      @test parent(L()) == L
      @test L(6//3) == L(2)
      @test L(R(5)) == L(5)
      @test_throws ErrorException L(x)
    end

    @testset "GCDX" begin
      for i in 1:550
        a = rand(L, 0:10, -10:10)
        b = rand(L, 0:10, -10:10)
        (g, u, v) = gcdx(a, b)
        @test g == u*a + v*b
      end
    end

    @testset "GCD" begin
      for i in 1:550
        a = rand(L, 0:10, -10:10)
        b = rand(L, 0:10, -10:10)
        g = gcd(a, b)
        @test divides(a, g)[1] && divides(b, g)[1]
      end
    end

    @testset "Exact division" begin

      @test divides(L(1//x), L(1//x)) == (true, L(1))
      @test_throws DivideError divexact(L(1//(x + 1)), L())

      for i in 1:300
        a = rand(L, 0:10, -10:10)
        b = rand(L, 0:10, -10:10)
        d = divides(a, b)
        if d[1]
          @test a == d[2] * b
        end
      end
    end

    @testset "Inversion" begin

      @test inv(L((x + 1)//(x + 2))) == L((x + 2)//(x + 1))
      @test inv(L(23)) == L(1//23)
      @test_throws DivideError inv(L())
    end

    @testset "Binary operators" begin

      @test L((-1//3)//(x^2 - 3*x)) + L(-1//(x - 1//2)) == L((-x^2 + 8//3*x + 1//6)//(x^3 - 7//2*x^2 + 3//2*x))
      @test L((-1//3)//(x^2 - 3*x)) - L(-1//(x - 1//2)) == L((x^2 - 10//3*x + 1//6)//(x^3 - 7//2*x^2 + 3//2*x))
      @test L((-1//3)//(x - 3))*L(-1//(x - 1//2)) == (1//3)//(x^2 - 7//2*x + 3//2)

      @test L(18//2) + L(2//1) == L(11)
      @test L(18//3) - L(1) == L(5)
      @test L(32) * L(4) == L(128)
    end

    @testset "Basic manipulation" begin

      @test iszero(L()) == true
      @test iszero(L(0)) == true
      @test isone(L(1//x)) == false
      @test isunit(L((x + 1)//(x + 2))) == true

      @test iszero(L(4//3)) == false
      @test isone(L(13//13)) == true
      @test isunit(L()) == false
    end
end


# TODO: add a type for kinetic energy

struct Hamiltonian{M<:AbstractMetric, Tlogπ, T∂logπ∂θ}
    metric::M
    logπ::Tlogπ
    ∂logπ∂θ::T∂logπ∂θ
end

# Create a `Hamiltonian` with a new `M⁻¹`
(h::Hamiltonian)(M⁻¹) = Hamiltonian(h.metric(M⁻¹), h.logπ, h.∂logπ∂θ)

∂H∂θ(h::Hamiltonian, θ::AbstractVector) = -h.∂logπ∂θ(θ)

∂H∂r(h::Hamiltonian{<:UnitEuclideanMetric}, r::AbstractVector) = copy(r)
∂H∂r(h::Hamiltonian{<:DiagEuclideanMetric}, r::AbstractVector) = h.metric.M⁻¹ .* r
∂H∂r(h::Hamiltonian{<:DenseEuclideanMetric}, r::AbstractVector) = h.metric.M⁻¹ * r

function hamiltonian_energy(h::Hamiltonian, θ::AbstractVector, r::AbstractVector)
    K = kinetic_energy(h, r, θ)
    if isnan(K)
        K = Inf
        @warn "Kinetic energy is `NaN` and is set to `Inf`."
    end
    V = potential_energy(h, θ)
    if isnan(V)
        V = Inf
        @warn "Potential energy is `NaN` and is set to `Inf`."
    end
    return K + V
end

potential_energy(h::Hamiltonian, θ::AbstractVector) = -h.logπ(θ)

# Kinetic energy
# NOTE: the general form of K depends on both θ and r
kinetic_energy(h::Hamiltonian{<:UnitEuclideanMetric}, r, θ) = sum(abs2, r) / 2
function kinetic_energy(h::Hamiltonian{<:DiagEuclideanMetric}, r, θ)
    return sum(abs2(r[i]) * h.metric.M⁻¹[i] for i in 1:length(r)) / 2
end
function kinetic_energy(h::Hamiltonian{<:DenseEuclideanMetric}, r, θ)
    mul!(h.metric._temp, h.metric.M⁻¹, r)
    return dot(r, h.metric._temp) / 2
end

# Momentum sampler
function rand_momentum(rng::AbstractRNG, h::Hamiltonian{<:UnitEuclideanMetric})
    return randn(rng, h.metric.dim)
end
function rand_momentum(rng::AbstractRNG, h::Hamiltonian{<:DiagEuclideanMetric})
    r = randn(rng, h.metric.dim)
    r ./= h.metric.sqrtM⁻¹
    return r
end
function rand_momentum(rng::AbstractRNG, h::Hamiltonian{<:DenseEuclideanMetric})
    r = randn(rng, h.metric.dim)
    ldiv!(h.metric.cholM⁻¹, r)
    return r
end

rand_momentum(h::Hamiltonian) = rand_momentum(GLOBAL_RNG, h)


using ArrayInterface, Test
using Base: setindex
import ArrayInterface: has_sparsestruct, findstructralnz, fast_scalar_indexing, lu_instance
@test ArrayInterface.ismutable(rand(3))

using StaticArrays
@test ArrayInterface.ismutable(@SVector [1,2,3]) == false
@test ArrayInterface.ismutable(@MVector [1,2,3]) == true
@test ArrayInterface.ismutable(1:10) == false
@test ArrayInterface.ismutable((0.1,1.0)) == false
@test isone(ArrayInterface.known_first(typeof(StaticArrays.SOneTo(7))))
@test ArrayInterface.known_last(typeof(StaticArrays.SOneTo(7))) == 7

using LinearAlgebra, SparseArrays

D=Diagonal([1,2,3,4])
@test has_sparsestruct(D)
rowind,colind=findstructralnz(D)
@test [D[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,2,3,4]
@test length(rowind)==4
@test length(rowind)==length(colind)

Bu = Bidiagonal([1,2,3,4], [7,8,9], :U)
@test has_sparsestruct(Bu)
rowind,colind=findstructralnz(Bu)
@test [Bu[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,7,2,8,3,9,4]
Bl = Bidiagonal([1,2,3,4], [7,8,9], :L)
@test has_sparsestruct(Bl)
rowind,colind=findstructralnz(Bl)
@test [Bl[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,7,2,8,3,9,4]

Tri=Tridiagonal([1,2,3],[1,2,3,4],[4,5,6])
@test has_sparsestruct(Tri)
rowind,colind=findstructralnz(Tri)
@test [Tri[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,2,3,4,4,5,6,1,2,3]
STri=SymTridiagonal([1,2,3,4],[5,6,7])
@test has_sparsestruct(STri)
rowind,colind=findstructralnz(STri)
@test [STri[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,2,3,4,5,6,7,5,6,7]

Sp=sparse([1,2,3],[1,2,3],[1,2,3])
@test has_sparsestruct(Sp)
rowind,colind=findstructralnz(Sp)
@test [Tri[rowind[i],colind[i]] for i in 1:length(rowind)]==[1,2,3]

@test ArrayInterface.ismutable(spzeros(1, 1))
@test ArrayInterface.ismutable(spzeros(1))

@test !fast_scalar_indexing(qr(rand(10, 10)).Q)
@test !fast_scalar_indexing(qr(rand(10, 10), Val(true)).Q)
@test !fast_scalar_indexing(lq(rand(10, 10)).Q)

using BandedMatrices

B=BandedMatrix(Ones(5,5), (-1,2))
B[band(1)].=[1,2,3,4]
B[band(2)].=[5,6,7]
@test has_sparsestruct(B)
rowind,colind=findstructralnz(B)
@test [B[rowind[i],colind[i]] for i in 1:length(rowind)]==[5,6,7,1,2,3,4]
B=BandedMatrix(Ones(4,6), (-1,2))
B[band(1)].=[1,2,3,4]
B[band(2)].=[5,6,7,8]
rowind,colind=findstructralnz(B)
@test [B[rowind[i],colind[i]] for i in 1:length(rowind)]==[5,6,7,8,1,2,3,4]

using BlockBandedMatrices
BB=BlockBandedMatrix(Ones(10,10),[1,2,3,4],[4,3,2,1],(1,0))
BB[Block(1,1)].=[1 2 3 4]
BB[Block(2,1)].=[5 6 7 8;9 10 11 12]
rowind,colind=findstructralnz(BB)
@test [BB[rowind[i],colind[i]] for i in 1:length(rowind)]==
    [1,5,9,2,6,10,3,7,11,4,8,12,
    1,1,1,1,1,1,1,1,1,1,1,1,1,1,
    1,1,1,1,1,1,1,1,1,1,1,1,1,1,
    1,1,1,1,1]

dense=collect(Ones(8,8))
for i in 1:8
    dense[:,i].=[1,2,3,4,5,6,7,8]
end
BBB=BandedBlockBandedMatrix(dense, [4, 4] ,[4, 4], (1, 1), (1, 1))
rowind,colind=findstructralnz(BBB)
@test [BBB[rowind[i],colind[i]] for i in 1:length(rowind)]==
    [1,2,3,1,2,3,4,2,3,4,5,6,7,5,6,7,8,6,7,8,
     1,2,3,1,2,3,4,2,3,4,5,6,7,5,6,7,8,6,7,8]

@testset "setindex" begin
    @testset "$(typeof(x))" for x in [
        zeros(3),
        falses(3),
        spzeros(3),
    ]
        y = setindex(x, true, 1)
        @test iszero(x)  # x is not mutated
        @test y[1] == true
        @test iszero(x[CartesianIndices(size(x)) .== [CartesianIndex(1)]])

        y2 = setindex(x, one.(x), :)
        @test iszero(x)
        @test all(isone, y2)
    end

    @testset "$(typeof(x))" for x in [
        zeros(3, 3),
        falses(3, 3),
        spzeros(3, 3),
    ]
        y = setindex(x, true, 1, 1)
        @test iszero(x)  # x is not mutated
        @test y[1, 1] == true
        @test iszero(x[CartesianIndices(size(x)) .== [CartesianIndex(1, 1)]])

        y2 = setindex(x, one.(x), :, :)
        @test iszero(x)
        @test all(isone, y2)
    end

    @testset "$(typeof(x))" for x in [
        zeros(3, 3, 3),
        falses(3, 3, 3),
    ]
        y = setindex(x, true, 1, 1, 1)
        @test iszero(x)  # x is not mutated
        @test y[1, 1, 1] == true
        @test iszero(x[CartesianIndices(size(x)) .== [CartesianIndex(1, 1, 1)]])

        y2 = setindex(x, one.(x), :, :, :)
        @test iszero(x)
        @test all(isone, y2)
    end
end

using SuiteSparse
@testset "lu_instance" begin
  for A in [
    randn(5, 5),
    @SMatrix(randn(5, 5)),
    @MMatrix(randn(5, 5)),
    sprand(50, 50, 0.5)
  ]
    @test lu_instance(A) isa typeof(lu(A))
  end
  @test lu_instance(1) === 1
end

using Random
using ArrayInterface: issingular
@testset "issingular" begin
    for T in [Float64, ComplexF64]
        R = randn(MersenneTwister(2), T, 5, 5)
        S = Symmetric(R)
        L = UpperTriangular(R)
        U = LowerTriangular(R)
        @test all(!issingular, [R, S, L, U, U'])
        R[:, 2] .= 0
        @test all(issingular, [R, L, U, U'])
        @test !issingular(S)
        R[2, :] .= 0
        @test issingular(S)
        @test all(!issingular, [UnitLowerTriangular(R), UnitUpperTriangular(R), UnitUpperTriangular(R)'])
    end
end

using ArrayInterface: zeromatrix
@test zeromatrix(rand(4,4,4)) == zeros(4*4*4,4*4*4)

using ArrayInterface: parent_type
@testset "Parent Type" begin
    x = ones(4, 4)
    @test parent_type(view(x, 1:2, 1:2)) <: typeof(x)
    @test parent_type(reshape(x, 2, :)) <: typeof(x)
    @test parent_type(transpose(x)) <: typeof(x)
    @test parent_type(Symmetric(x)) <: typeof(x)
    @test parent_type(UpperTriangular(x)) <: typeof(x)
end

@testset "Range Interface" begin
    @test isnothing(ArrayInterface.known_first(typeof(1:4)))
    @test isone(ArrayInterface.known_first(typeof(Base.OneTo(4))))

    @test isnothing(ArrayInterface.known_last(typeof(1:4)))

    @test isnothing(ArrayInterface.known_step(typeof(1:0.2:4)))
    @test isone(ArrayInterface.known_step(typeof(1:4)))
end


"""
    dsep(g::AbstractGraph, u, v, s; verbose = false)

Check  whether `u` and `v` are d-separated given set `s`.
Algorithm: unrolled https://arxiv.org/abs/1304.1505

"""
function dsep(g::AbstractGraph, u::Integer, v::Integer, S; verbose = false)
    T = eltype(g)
    in_seen = falses(nv(g)) # nodes reached earlier backwards
    out_seen = falses(nv(g)) # nodes reached earlier forwards
    descendant = falses(nv(g)) # descendant in s
    blocked = falses(nv(g))

    for ve in S
        in_seen[ve] = true
        blocked[ve] = true
    end

    (in_seen[u] || in_seen[v]) && throw(ArgumentError("S should not contain u or v"))

    u == v && throw(ArgumentError("u == v"))

    next = Vector{T}()

    # mark vertices with descendants in S
    next = Vector{T}()
    for w in S
        push!(next, w)
        descendant[w] = true
    end

    while !isempty(next)
        for w in inneighbors(g, popfirst!(next))
            if !descendant[w]
                push!(next, w) # push onto queue
                descendant[w] = true
            end
        end
    end

    verbose && println(descendant)

    in_next = Vector{T}()
    out_next = Vector{T}()

    push!(in_next, u) # treat u as vertex reached backwards
    in_seen[u] = true

    while true
        sin = isempty(in_next)
        sout = isempty(out_next)
        sin && sout && return true # no vertices in the queue

        if !sin # some vertex reach backwards in the queue
            src = popfirst!(in_next)
            for w in outneighbors(g, src) # possible collider at destination
                if !out_seen[w] && (!blocked[w] || descendant[w])
                    verbose && println("<- $src -> $w")
                    w == v && return false
                    push!(out_next, w)
                    out_seen[w] = true
                end
            end
            for w in inneighbors(g, src)
                if !in_seen[w]
                    verbose && println("<- $src <- $w")
                    w == v && return false
                    push!(in_next, w)
                    in_seen[w] = true
                end
            end
        end
        if !sout # some vertex reach forwards in the queue
            src = popfirst!(out_next)
            for w in outneighbors(g, src) # possible collider at destination
                if !out_seen[w] && !blocked[src] && (!blocked[w] || descendant[w])
                    verbose && println("-> $src -> $w")
                    w == v && return false
                    push!(out_next, w)
                    out_seen[w] = true
                end
            end
            for w in inneighbors(g, src) # collider at source
                if !in_seen[w] && descendant[src] # shielded collider
                    verbose && println("-> $src <- $w")
                    w == v && return false
                    push!(out_next, w)
                    in_seen[w] = true
                end
            end
        end

    end
end


using Compat
using Compat.Test # Base.Test on 0.6, and Test on 0.7
using Compat.LinearAlgebra # dot

using Einsum

## Test that vars in Main aren't overwritten by einsum
let
    i = -1
    y = randn(10)
    @einsum x[i] := y[i]
    @test i == -1
end

## Test that B is overwritten by := operator
let
    B = randn(10, 10)
    A = randn(5, 10)
    @einsum B[i, j] := A[i, j] # this should run without a problem
    @test size(B) == size(A)
end

## CP decomposition test case ##
let

    # preallocated test case
    A = zeros(5, 6, 7)
    B = similar(A)
    C = similar(A)

    X = randn(5, 2)
    Y = randn(6, 2)
    Z = randn(7, 2)

    @einsum A[i, j, k] = X[i, r] * Y[j, r] * Z[k, r]
    @einsimd B[i, j, k] = X[i, r] * Y[j, r] * Z[k, r]
    @vielsum C[i, j, k] = X[i, r] * Y[j, r] * Z[k, r]

    for i = 1:5
        for j = 1:6
            for k = 1:7
                s = 0.0
                for r = 1:2
                    s += X[i, r] * Y[j, r] * Z[k, r]
                end
                @test isapprox(A[i, j, k], s)
                @test isapprox(B[i, j, k], s)
                @test isapprox(C[i, j, k], s)
            end
        end
    end

    # without preallocation
    @einsum A2[i, j, k] := X[i, r] * Y[j, r] * Z[k, r]
    @test isapprox(A, A2)

end

# Interesting test case, can throw an error that
# local vars are declared twice.
let
    A = zeros(5, 6, 7)
    X = randn(5, 2)
    Y = randn(6, 2)
    Z = randn(7, 2)
    if true
        @einsum A[i, j, k] = X[i, r] * Y[j, r] * Z[k, r]
    else
        @einsum A[i, j, k] = X[i, r] * Y[j, r] * Z[k, r]
    end
end

# From #21: local `T` does not interfer with internal T
let
    function test(x::Vector{T}, y::Vector{T}) where T
        @einsum z := x[i] * y[i]
        return z
    end
    @test_nowarn test(rand(3), rand(3))
end
# From #20: local `s` does not interfere with internal s
let
    x = rand(2, 3)
    @test_nowarn @einsum y[i] := x[i, s]
end

# At one point this threw an error because the lhs
# had no indices/arguments
let
    x = randn(10)
    y = randn(10)
    @einsum k := x[i] * y[i]
    @test isapprox(k, dot(x, y))
end

# Elementwise multiplication (this should create nested loops with no
# no summation.)
let
    x = randn(10)
    y = randn(10)
    @einsum k[i] := x[i] * y[i]
    @einsimd k2[i] := x[i] * y[i]
    @vielsum k3[i] := x[i] * y[i]
    @test isapprox(k, x .* y)
    @test isapprox(k2, x .* y)
    @test isapprox(k3, x .* y)
end

# Transpose a block matrix
let
    z = [rand(2, 2) for i = 1:2, j = 1:2]
    @einsum t[i, j] := transpose(z[j, i])
    @test isapprox(z[1, 1],  t[1, 1]')
    @test isapprox(z[2, 2],  t[2, 2]')
    @test isapprox(z[1, 2],  t[2, 1]')
    @test isapprox(z[2, 1],  t[1, 2]')
end

# Mapping functions
let
    A = randn(10, 10)
    @einsum B[i, j] := exp(A[i, j])
    @test isapprox(exp.(A), B)
end

# Example from numpy
let
    A = reshape(collect(1:25), 5, 5)
    @einsum B[i] := A[i, i]
    @test all(B .== [1, 7, 13, 19, 25])
end

# TODO: consider adding support for this:
# @einsum A[i, j] = A[i,j] + 50

## Test in-place operations ##
let
    A = randn(5, 6, 7)
    B = randn(5, 6, 7)
    A1 = copy(A)
    B1 = copy(B)

    X = randn(5, 2)
    Y = randn(6, 2)
    Z = randn(7, 2)

    @einsum A[i, j, k] += X[i, r] * Y[j, r] * Z[k, r]
    @einsimd B[i, j, k] += X[i, r] * Y[j, r] * Z[k, r]

    for i = 1:5
        for j = 1:6
            for k = 1:7
                s = 0.0
                for r = 1:2
                    s += X[i, r] * Y[j, r] * Z[k, r]
                end
                @test isapprox(A[i, j, k], A1[i, j, k] + s)
                @test isapprox(B[i, j, k], B1[i, j, k] + s)
            end
        end
    end

    x = randn(10)
    y = randn(10)
    k0 = randn()
    k = k0
    @einsum k += x[i] * y[i]
    @test isapprox(k, k0 + dot(x, y))

    # test multiplication

    A1[:] = A[:]
    B1[:] = B[:]

    @einsum A[i, j, k]  *= X[i, r] * Y[j, r] * Z[k, r]
    @einsimd B[i, j, k] *= X[i, r] * Y[j, r] * Z[k, r]

    for i = 1:5
        for j = 1:6
            for k = 1:7
                s = 0.0
                for r = 1:2
                    s += X[i, r] * Y[j, r] * Z[k, r]
                end
                @test isapprox(A[i, j, k], A1[i, j, k] * s)
                @test isapprox(B[i, j, k], B1[i, j, k] * s)
            end
        end
    end

    x = randn(10)
    y = randn(10)
    k0 = randn()
    k = k0
    @einsum k *= x[i] * y[i]
    @test isapprox(k, k0 * dot(x, y))
end

# Test offsets
let
    X = randn(10)

    # without preallocation
    @einsum A[i] := X[i + 5]
    @test size(A) == (5,)
    @test all(A .== X[6:end])

    # with preallocation
    B = zeros(10)
    @einsum B[i] = X[i + 5]
    @test size(B) == (10,)
    @test all(B[1:5] .== X[6:end])
end

# Test symbolic offsets
let
    offset = 5
    X = randn(10)

    # without preallocation
    # @einsum A[i] := X[i + :offset] # error on 1.0
    # @test size(A) == (5,)
    # @test all(A .== X[6:end])

    # with preallocation
    B = zeros(10)
    # @einsum B[i] = X[i + :offset] # error on 1.0
    # @test size(B) == (10,)
    # @test all(B[1:5] .== X[6:end])
end

# Test adding/subtracting constants
let
    k = 5
    X = randn(10)

    # without preallocation
    @einsum A[i] := X[i] + k
    @einsum B[i] := X[i] - k
    @test isapprox(A, X .+ k)
    @test isapprox(B, X .- k)

    # with preallocation
    C, D = zeros(10), zeros(10)
    @einsum C[i] = X[i] + k
    @einsum D[i] = X[i] - k
    @test isapprox(C, X .+ k)
    @test isapprox(D, X .- k)
end

# Test multiplying/dividing constants
let
    k = 5
    X = randn(10)

    # without preallocation
    @einsum A[i] := X[i] * k
    @einsum B[i] := X[i] / k
    @test isapprox(A, X .* k)
    @test isapprox(B, X ./ k)

    # with preallocation
    C, D = zeros(10), zeros(10)
    @einsum C[i] = X[i] * k
    @einsum D[i] = X[i] / k
    @test isapprox(C, X .* k)
    @test isapprox(D, X ./ k)
end

# Test indexing with a constant
let
    A = randn(10, 2)
    j = 2
    # @einsum B[i] := A[i, :j] # error on 1.0
    # @test all(B .== A[:, j])
    @einsum C[i] := A[i, 1]
    @test all(C .== A[:, 1])

    D = zeros(10, 3)
    # @einsum D[i, 1] = A[i, :j]
    # @test isapprox(D[:, 1], A[:, j])
    # @einsum D[i, :j] = A[i, :j]
    # @test isapprox(D[:, j], A[:, j])
end

# Better type inference on allocating arrays
let
    B1 = ones(Int, 5)
    B2 = ones(Float32, 5)
    B3 = ones(5)
    C = randn(5)
    @einsum A1[i, j] := B1[i] * C[j]
    @einsum A2[i, j] := B2[i] * C[j]
    @einsum A3[i, j] := B3[i] * C[j]

    @test eltype(A1) == Float64
    @test eltype(A2) == Float64
    @test eltype(A3) == Float64
    @test isapprox(A1, A3)
    @test isapprox(A2, A3)
end


using jPCA
using Base.Test

# main
X = rand(100,100)
Y = fit(JPCA, X, 10)

@test size(Y.M,1) == size(Y.M,2)

@test_throws ErrorException fit(JPCA, rand(5,5), 6)
@test_throws ErrorException fit(JPCA, rand(5,10))

# skew sym
import jPCA.reshape_skew
# create a random skew sym matrix
Z = Z = rand(100,100)
Z .= 0.5(Z - Z')
@test Z == reshape_skew(reshape_skew(Z))

@test_throws ErrorException reshape_skew(Z[:,1:99])
Z[1,1] += 0.1
@test_throws ErrorException reshape_skew(Z)