### Schur Eigenspace solvers

struct BaseSolver end
struct ArnoldiSolver end
struct KrylovSolver end

import ArnoldiMethod
function schurvectors(T, pi, n, israte, ::ArnoldiSolver)
	D = Diagonal(sqrt.(pi))
	T̃ = D * T * D^-1
	s, hist = ArnoldiMethod.partialschur(T̃; nev=n, which=israte ? ArnoldiMethod.LR() : ArnoldiMethod.LM())
	X̃ = collect(s.Q)
	X = D^-1 * X̃
	X ./= X[1,1]
	X, s.eigenvalues
end

import KrylovKit
function schurvectors(T, pi, n, israte, ::KrylovSolver)
	D = Diagonal(sqrt.(pi))
	T̃ = D * T * D^-1
	R, Q, v, info = KrylovKit.schursolve(T̃, pi, n, israte ? :LR : :LM, KrylovKit.Arnoldi())
	X̃ = reduce(hcat, Q)
	X = D^-1 * X̃
	X = X ./ X[1,1]
	X, v
end

function schurvectors(T, pi, n, israte, ::BaseSolver)
    Tw = Diagonal(sqrt.(pi))*T*Diagonal(1 ./ sqrt.(pi)) # rescale to keep markov property
    Sw = schur!(Tw)                       # returns orthonormal vecs by def
    Xw, λ = selclusters!(Sw, n, israte)
    X  = Diagonal(1 ./sqrt.(pi)) * Xw              # scale back
    X  = X[1,1]>0 ? X : -X
    X, λ
end

# select the schurvectors corresponding to the n abs-largest eigenvalues
# if reverse==true select highest abs value, otherwise select lowest (for rate matrices)
function selclusters!(S, n, ratematrix)
    ind = sortperm(abs.(S.values), rev=!ratematrix) # get indices for dominant eigenvalues
    select = zeros(Bool, size(ind))           # create selection vector
    select[ind[1:n]] .= true
    S = ordschur!(S, select)                  # reorder selected vectors to the left
    if !isapprox(S.T[n+1, n], 0)                       # check if we are cutting along a schur block
        @error("conjugated eigenvector missing")
        display(S.T)
    end
    S.vectors[:,1:n], S.values[1:n]       # select first n vectors
end


function build_F0_tr(
  nbset,  # tuple with the following elements.
    # nodes: array of Strings containing independent node names.
    # num_nodes: number of nodes.
    # branches: array of struct Branch.
    # num_branches: sum of branch phases.
    # num_dyns: number of dynamic branches.
  smset  # tuple with the following elements.
    # Ac: branch-node connectivity matrix.
    # Z: coefficient matrix of i.
    # Y: coefficient matrix of v.
    # B1: incidence matrix of alpha.
    # B2: incidence matrix of [ i v ].
  )  #=

builds the initial form of F which is the coefficient matrix of the sparse tableau equations for a transient simulation and returns it.

Author: Taku Noda
Started on Feb. 11, 2021
Update on June 29, 2021

Copyright 2021 Central Research Institute of Electric Power Industry

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

=#

( nodes, num_nodes, branches, num_branches, num_dyns ) = nbset
( Ac, Z, Y, B1, B2 ) = smset

N = num_nodes + 2*num_branches + num_dyns

F = spzeros( N, N )
range1 = 1:num_nodes
range2 = num_nodes .+ ( 1:num_branches )
range3 = num_branches .+ range2
range4 = ( num_nodes + 2*num_branches ) .+ ( 1:num_dyns )
range23 = num_nodes .+ ( 1:( 2*num_branches ) )
F[ range1, range2 ] = transpose( Ac )
F[ range2, range1 ] = Ac
F[ range2, range3 ] = sparse( -I, num_branches, num_branches )
F[ range3, range2 ] = Z
F[ range3, range3 ] = Y
F[ range3, range4 ] = B1
F[ range4, range23 ] = B2

return F

end


# # Matrix Formalism
#
# In this example we will consider the matrix formalism approach for a geometry of cylinders.

# LSP indexing solution                                                          #src
# https://github.com/julia-vscode/julia-vscode/issues/800#issuecomment-650085983 #src
if isdefined(@__MODULE__, :LanguageServer)                                       #src
    include("../src/SpinDoctor.jl")                                              #src
    using .SpinDoctor                                                            #src
end                                                                              #src

using SpinDoctor
using LinearAlgebra

if haskey(ENV, "GITHUB_ACTIONS")
    using CairoMakie
else
    using GLMakie
end

setup = CylinderSetup(;
    name = "Slice",
    ncell = 3,
    rmin = 2.0,
    rmax = 6.0,
    dmin = 0.2,
    dmax = 0.3,
    height = 1.0,
    bend = 0.0,
    twist = 0.0,
    ecs_shape = :convex_hull,
    ecs_ratio = 0.5,
)

# We also define coefficients for the different cell compartments `:in` (axon), `:out`
# (myelin), and `:ecs` (ECS).

coeffs = coefficients(
    setup;
    D = (; in = 0.002 * I(3), out = 0.002 * I(3), ecs = 0.002 * I(3)),
    T₂ = (; in = Inf, out = Inf, ecs = Inf),
    ρ = (; in = 1.0, out = 1.0, ecs = 1.0),
    κ = (; in_out = 1e-4, out_ecs = 1e-4, in = 0.0, out = 0.0, ecs = 0.0),
    γ = 2.67513e-4,
)

# We then proceed to build the geometry and finite element mesh.

mesh, = create_geometry(setup; recreate = true)
plot_mesh(mesh)

# The mesh looks good, so we may then proceed to assemble the biological model and the
# associated finite element matrices.

model = Model(; mesh, coeffs...);
matrices = assemble_matrices(model);

# We may also compute some useful quantities, including a scalar diffusion coefficient from
# the diffusion tensors.

volumes = get_cmpt_volumes(model.mesh)
D_avg = 1 / 3 * tr.(model.D)' * volumes / sum(volumes)

# The eigenfunctions of the diffusive part of the Bloch-Torrey operator forms a good basis
# for the finite element function space. The basis may be truncated at a certain level, thus
# reducing the number of degrees of freedom. We here opt for 400 eigenfunctions.

laplace = Laplace(; model, matrices, neig_max = 400)
lap_eig = solve(laplace)

# The resulting eigenvalues may be represented as length scales, describing the wavelength
# of the eigenfunctions.

length_scales = eig2length.(lap_eig.values, D_avg)

# We may also further truncate the eigenfunction basis, if we are satisfied skipping
# features below a threshold length scale of 3 micrometers.

length_scale = 3
λ_max = length2eig(length_scale, D_avg)
lap_eig = limit_lengthscale(lap_eig, λ_max)

# Each of the resulting eigenfunctions is represented in the same way as the initial
# magnetization field `ρ`.

ncompartment, nboundary = size(mesh.facets)
fig = Figure()
for i = 1:3, j = 1:4
    ieig = 6(i - 1) + j
    ϕ_cmpts = split_field(mesh, lap_eig.funcs[:, ieig])
    ax = Axis3(fig[i, j]; title = "n = $ieig, ℓ = $(length_scales[ieig])", aspect = :data)
    nboundary = size(mesh.facets, 2)
    scene = nothing
    first = true
    for icmpt = 1:ncompartment, iboundary = 1:nboundary
        facets = mesh.facets[icmpt, iboundary]
        points = mesh.points[icmpt]
        mesh!(ax, points', facets', color = ϕ_cmpts[icmpt], shading = false)
    end
end
fig

# We observe that the first functions have large features, while the higher-index functions
# have more rapidly varying features. We may now choose a gradient and compute the
# projection of magnetization field onto the truncated basis.

dir = [1.0, 0.0, 0.0]
profile = CosOGSE(5000.0, 5000.0, 2)
b = 1000
g = √(b / int_F²(profile)) / coeffs.γ
gradient = ScalarGradient(dir, profile, g)

# The matrix formalism problem is solved in the same way as the [`BTPDE`](@ref). The time
# profile is approximated on 500 points, since it is non-constant.

mf = MatrixFormalism(; model, matrices, lap_eig)
ξ = solve(mf, gradient; ninterval = 500)

# The resulting magnetization field may be plotted.

plot_field(model.mesh, ξ)


using LinearAlgebra

function find_extreme_points(points::Array{Vector{T}, 1},
                             zmin::Union{Missing, Vector{T}}=missing) where T <: Real
    M = length(points[1])

    @show(zmin)
    if ismissing(zmin)
        zmin = zeros(M)
    end
    @show zmin

    extremepts = [-1 for j = 1:M]
    for j = 1:M
        w = 0.00001 * ones(M)
        w[j] = 1.0
        minval = -Inf
        for (i, p) in enumerate(points)
            val = maximum((p-zmin) ./ w)
            if val > minval
                extremepts[j] = i
                minval = val
            end
        end
    end

    return extremepts
end

function find_hyperplane_params(points::Array{Vector{T}, 1}, b=1) where T <: Real
    M = length(points[1])

    interceptvals = zeros(M)

    Z = transpose(reduce(hcat, points))
    bv = b * ones(M)

    a = Z\bv
    return a
end

function find_axis_intercepts(points::Array{Vector{T}, 1}, b=1) where T <: Real
    return 1 ./ find_hyperplane_params(points, b)
end

function associate_points(
        canpoints::Array{Vector{T}, 1},
        refpoints::Array{Vector{T}, 1}) where T <: Real
    N = length(canpoints)
    R = length(refpoints)
    # for each candidate point, want to find the closest reference point,
    # measured as perpendicular distance to line through ref point and origin

    closest_ref_index = zeros(Int64, N)
    closest_ref_dist = zeros(Float64, N)
    dists = zeros(R)
    for n = 1:N
        s = canpoints[n]
        for r = 1:R
            w = refpoints[r]
            dists[r] = norm( s - (dot(w, s)/dot(w, w)) * w )
        end
        closest_ref_index[n] = argmin(dists)
        closest_ref_dist[n] = dists[closest_ref_index[n]]
    end
    return closest_ref_index, closest_ref_dist
end


using LinearAlgebra
#include("types.jl")

function isplanar(simplex::Simplex, fp::PointND, d1::Dim, d2::Dim)
  checkdims = [x for x in 1:length(fp) if x!=d1 && x!=d2]
  # need to check if the remaining dims are equal to the focus point
  all([all(simplex[:,d] .== fp[d]) for d in checkdims])
end

# TODO: ensure simplex and fp are compatible
function simplexPointIntersection(simplex::Simplex, fp::PointND, d1::Dim, d2::Dim)
  n = Dim(length(fp) + 1) # number of lambdas to check, dimensionality of the space + 1

  # Set up matrix for barycentric computation
  # Append a set of ones to the simplex and incorporate the focus point
  # to create a square matrix
  T = ones(Float64,(n,n)) # TODO: use square matrix type if exists
  T[1:n-1,1:n-1] = simplex'
  T[1:n-1,n] = fp

  if det(T) == 0 # extra focus point may be in the plane of the
                 # simplex or might be numerical problems
    if isplanar(simplex, fp, d1, d2)
      # if the simplex lies in the plane then just return the simplex as
      # a 2D projection of line segments
      nrow = size(simplex, 1)
      # all pairs of vertices in the simplex
      verts = filter(x -> x[1] != x[2], collect(Iterators.product(1:(nrow-1), 2:nrow)))
      v2d = [
        Intersect2D(simplex[pts[1], d1], simplex[pts[1], d2],
                    simplex[pts[2], d1], simplex[pts[2], d2])
        for pts = verts
      ]
      return v2d
    else
      T[1:n-1,n] = fp .+ (rand(n-1)*1e-9) # offset slightly to make matrix non-singluar
    end
  end

  # compute lambda as best we can (there will be 3 parts)
  rr = vcat(fp, 1.)
  rr[d1] = rr[d2] = 0
  rry = zeros(length(rr))
  rrx = zeros(length(rr))
  rrx[d1] = rry[d2] = 1.
  # We are trying to compute T^(-1) . [x-xn, y-yn,...,z-zn]
  Tlu = factorize(T)
  lc = Tlu \ rr # This way is more stable than inverting once
  lx = Tlu \ rrx
  ly = Tlu \ rry

  # find the d1 and d2 ranges that make each lambda 0
  # most indices are based on solving ax + by + c = 0
  # but keeping the other lambdas between 0 and 1

  # put y=mx+b into each other lambda formula and try and get a good range
  ccr(lx, ly, lc, n) # only consider the last lambda
end

## The common-cross-range functions
function ccr(lx::LambdaND, ly::LambdaND, lz::LambdaND, i::Dim)
  if lx[i] == 0 && ly[i] == 0
    return []
  elseif lx[i] == 0
    ccrX0(lx, ly, lz, i)
  elseif ly[i] == 0
    ccrY0(lx, ly, lz, i)
  else
    ccrXY(lx, ly, lz, i)
  end
end

function ccrX0(lx::LambdaND, ly::LambdaND, lc::LambdaND, i::Dim)
  # solve ly * y + c = 0
  y = -lc[i] / ly[i]
  # replace y in all other formulas and solve lx * x + ly * y + lc >=0 for x
  xs = lx[1:end .!= i]
  cs = ly[1:end .!= i] * y + lc[1:end .!= i]
  res = -cs ./ xs

  if any(xs .== 0) && any(cs[xs .== 0] .< 0) # no intersection
    return []
  else
    xsp = xs[map(isfinite,res)] # for filtering res later
    res = res[map(isfinite,res)]
    xrng = (maximum(res[xsp .>= 0]), minimum(res[xsp .< 0]))
    # some final bounds checking
    if xrng[2] - xrng[1] < -EPS
      return []
    end
    return [Intersect2D(xrng[1], y, xrng[2], y)]
  end
end

function ccrY0(lx::LambdaND, ly::LambdaND, lc::LambdaND, i::Dim)
  # solve ly * y + c = 0
  x = -lc[i] / lx[i]
  # replace x in all other formulas and solve lx * x + ly * y + lc >=0 for y
  ys = ly[1:end .!= i]
  cs = lx[1:end .!= i] * x + lc[1:end .!= i]
  res = -cs ./ ys

  if any(ys .== 0) && any(cs[ys .== 0] .< 0) # no intersection
    return []
  else
    ysp = ys[map(isfinite,res)] # for filtering res later
    res = res[map(isfinite,res)]
    yrng = (maximum(res[ysp .>= 0]), minimum(res[ysp .< 0]))
    # some final bounds checking
    if yrng[2] - yrng[1] < -EPS
      return []
    end
    return [Intersect2D(x, yrng[1], x, yrng[2])]
  end
end

# assumes lx[i] and ly[i] are non-zero
function ccrXY(lx::LambdaND, ly::LambdaND, lc::LambdaND, i::Dim)
  # solve lx * x + ly * y + c = 0 and substitute
  xs = lx[1:end .!= i] - ly[1:end .!= i] .* lx[i] ./ ly[i]
  cs = lc[1:end .!= i] - ly[1:end .!= i] .* lc[i] ./ ly[i]
  res = -cs ./ xs # solving for xs + cs >= 0

  if any(xs .== 0) && any(cs[xs .== 0] .< 0) # no intersection
    return []
  else
    xsp = xs[map(isfinite,res)] # for filtering res later
    res = res[map(isfinite,res)]
    xrng = (maximum(res[xsp .>= 0]), minimum(res[xsp .< 0]))
    yrng = (-lc[i] .- lx[i] .* xrng) ./ ly[i]
    # some final bounds checking
    if xrng[2] - xrng[1] < -EPS
      return []
    end
    return [Intersect2D(xrng[1], yrng[1], xrng[2], yrng[2])]
  end
end

# TODO: force triangle to be a triangle in 3D space
function intersectTri(triangle, fp::PointND, d1::Dim, d2::Dim)
  # compute the normal to the slicing plane
  n = ones(3)
  n[d1] = n[d2] = 0.0

  # figure out the point projection
  i1 = dot(n, (fp - triangle[1,:])) / dot(n, (triangle[2,:] - triangle[1,:]))
  i2 = dot(n, (fp - triangle[1,:])) / dot(n, (triangle[3,:] - triangle[1,:]))
  i3 = dot(n, (fp - triangle[2,:])) / dot(n, (triangle[3,:] - triangle[2,:]))
  p1 = triangle[1,:] + i1 * (triangle[2,:] - triangle[1,:])
  p2 = triangle[1,:] + i2 * (triangle[3,:] - triangle[1,:])
  p3 = triangle[2,:] + i3 * (triangle[3,:] - triangle[2,:])

  pts = filter(x -> isfinite(x[1]), [p1, p2, p3])
  if length(pts) == 0
    return []
  else
    return [Intersect2D(pts[1][d1], pts[1][d2], pts[2][d1], pts[2][d2])]
  end
end


# dimensions
nq = 10
n = 2 * nq
m = 10
T = 20
nz = m * (T - 1) + n * (T - 1)
nd = n * (T - 1)

# problem data
Aa = [SMatrix{nq, nq}(Array(1.0 * Diagonal(0.1 * rand(nq) + ones(nq)))) for t = 1:T-1]
Ab = [SMatrix{nq, nq}(Array(1.0 * Diagonal(0.1 * rand(nq) + ones(nq)))) for t = 1:T-1]
Ac = [SMatrix{nq, nq}(Array(Diagonal(ones(nq)))) for t = 1:T-1]
Ba = [SMatrix{nq,m}(1.0 * rand(nq, m)) for t = 1:T-1]
Qa = [Diagonal(SVector{nq}(0.1 * rand(nq) + ones(nq))) for t = 1:T]
Qb = [Diagonal(SVector{nq}(0.1 * rand(nq) + ones(nq))) for t = 1:T]
Qv = [Diagonal(SVector{nq}(-1.0 * ones(nq) + 0.1 * rand(nq))) for t = 1:T]

# indices
u_idx = [collect((t - 1) * (m + n) .+ (1:m)) for t = 1:T-1]
x_idx = [collect((t - 1) * (m + n) + m .+ (1:n)) for t = 1:T-1]
qa_idx = [collect(x_idx[t][1:nq]) for t = 1:T-1]
qb_idx = [x_idx[t][nq .+ (1:nq)] for t = 1:T-1]
n_idx = [(t - 1) * n .+ (1:n) for t = 1:T-1]

Q = [[Qa[t] Qv[t]; Qv[t]' Qb[t]] for t = 1:T]
R = [Diagonal(SVector{m}(0.1 * rand(m) + ones(m))) for t = 1:T-1]
A = [SMatrix{n,n}([zeros(nq, nq) I; Aa[t] Ab[t]]) for t = 1:T-1]
B = [SMatrix{n,m}([zeros(nq, m); Ba[t];]) for t = 1:T-1]
P = [SMatrix{n,n}([I zeros(nq, nq); zeros(nq, nq) Ac[t]]) for t = 1:T-1]

# direct problem
S = zeros(nz, nz)
C = zeros(nd, nz)

for t = 1:T-1
	S[u_idx[t], u_idx[t]] = R[t]
	S[x_idx[t], x_idx[t]] = Q[t+1]

	C[n_idx[t], u_idx[t]] = -B[t]
	C[n_idx[t], x_idx[t]] = P[t]
	t == 1 && continue
	C[n_idx[t], x_idx[t-1]] = -A[t]
end

Q̃ = [inv(Array(Q[t])) for t = 1:T]

Q̃a = [Q̃[t][1:nq, 1:nq] for t = 1:T]
Q̃b = [Q̃[t][nq .+ (1:nq), nq .+ (1:nq)] for t = 1:T]
Q̃v = [Q̃[t][1:nq, nq .+ (1:nq)] for t = 1:T]

R̃ = [inv(R[t]) for t = 1:T-1]
S̃ = zeros(nz, nz)

for t = 1:T-1
	S̃[u_idx[t], u_idx[t]] = R̃[t]
	S̃[x_idx[t], x_idx[t]] = Q̃[t+1]
end

norm(S̃ - inv(S))

# build kkt system
J = [S C'; C zeros(nd, nd)]
Js = sparse(J)
r = rand(nz + nd)
Δ = J \ r
Y = C * S̃ * C'
rlag = view(r, 1:nz)
rdyn = view(r, nz .+ (1:nd))
β = -rdyn + C * S̃ * rlag
Δν = Y \ β
Δz = S̃ * (rlag - C' * Δν)

# LU
solver_lu = lu_solver(Js)
Δlu = zeros(nz + nd)
@benchmark linear_solve!($solver_lu, $Δlu, $Js, $r)

solver_lu = lu_solver(J)
Δlu = zeros(nz + nd)
@benchmark linear_solve!($solver_lu, $Δlu, $J, $r)

# # QDLDL
# solver_ldl = ldl_solver(J)
# Δldl = zeros(nz + nd)
# @benchmark linear_solve!($solver_ldl, $Δldl, $J, $r)
# Δldl

norm(Δ - Δlu)
norm(Δ - [Δz; Δν])


using SpecialFunctions
using DistributedArrays
using LinearAlgebra

mutable struct BoundaryPrecomputedCalcs
  b::Boundary
  # Number of discretized steps (values go from 0:(2N-1))
  N::Int32
  k::Complex{Float64}
  j0NiMatrix::Array{Complex{Float64},2}
  j0NeMatrix::Array{Complex{Float64},2}
  j1NiMatrix::Array{Complex{Float64},2}
  j1NeMatrix::Array{Complex{Float64},2}
  h0NiMatrix::Array{Complex{Float64},2}
  h0NeMatrix::Array{Complex{Float64},2}
  h1NiMatrix::Array{Complex{Float64},2}
  h1NeMatrix::Array{Complex{Float64},2}
  absDiff::Array{Float64,2}
end

function absDiff(b::Boundary, t::Float64, τ::Float64)
  a = boundaryPoint(b, t)
  b = boundaryPoint(b, τ)
  norm(a-b)
end

function absDiff(b::Boundary, tτMatrix::Array{Array{Float64,1},2})
  map(pair -> norm(boundaryPoint(b, pair[1])-boundaryPoint(b, pair[2])), tτMatrix)
end

# @param l index (from 1:(2N-1) )
function γ(b::Boundary, N::Int64, l::Int64)
  norm(boundaryPointDeriv(b, l*π/N))+0im
end

# Quadrature weight factor
function R(j::Int64, N::Int64)
  cosAry = Array{Float64}(undef, N-1)
  for m in 1:N-1
    cosAry[m] = 1/m * cos(m*j*π/N)
  end
  -2π/N * sum(cosAry) - (-1)^j*π/N^2
end

################
j0(x::Complex{Float64}) = besselj(0, x)
j1(x::Complex{Float64}) = besselj(1, x)
h0(x::Complex{Float64}) = hankelh1(0, x)
# this is more accurate for some small values but it didn't seem to help me
# h0(x::Complex{Float64}) = j0(x)+bessely(0, x)*im
h1(x::Complex{Float64}) = hankelh1(1, x)
# h1(x::Complex{Float64}) = j1(x)+bessely(1, x)*im
nVals(b::Boundary, t::Float64) = (nInt(b), nExt(b))
xVals(b::Boundary, t::Float64, τ::Float64) = (convert(Array{Float64,1}, boundaryPoint(b, t)), convert(Array{Float64,1}, boundaryPoint(b, τ)))
log4Term(t::Float64, τ::Float64)::Float64 = log(4*sin((t-τ)/2)^2)

################
# H functions
function H1(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  if (t == τ)
    0.0+0im
  else
    # Note: Need to check that we can use the k values at t, not τ, this is a deviation from the Heider paper
    (ni, ne) = nVals(b, t)
    (xt, xτ) = xVals(b, t, τ)
    (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
    ad = bpc.absDiff[tIdx, τIdx]
    j1ni = bpc.j1NiMatrix[tIdx, τIdx]
    j1ne = bpc.j1NeMatrix[tIdx, τIdx]
    -1/(2π)*dot(normVector(b, τ), (xt-xτ)/ad)*(k*ne*j1ne-k*ni*j1ni)
  end
end

function H2(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  if (t == τ)
    0.0+0im
  else
    H(b, t, τ, k, bpc)-H1(b, t, τ, k, bpc)*log4Term(t, τ)
  end
end

function H(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  (ni, ne) = nVals(b, t)
  (xt, xτ) = xVals(b, t, τ)
  (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
  ad = bpc.absDiff[tIdx, τIdx]
  h1ni = bpc.h1NiMatrix[tIdx, τIdx]
  h1ne = bpc.h1NeMatrix[tIdx, τIdx]
  im/2*dot(normVector(b, τ), (xt-xτ)/ad)*(k*ne*h1ne-k*ni*h1ni)
end

H1star(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs) = H1(b, τ, t, k, bpc)*norm(boundaryPointDeriv(b, τ))
H2star(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs) = H2(b, τ, t, k, bpc)*norm(boundaryPointDeriv(b, τ))
Hstar(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs) = H(b, τ, t, k, bpc)*norm(boundaryPointDeriv(b, τ))

###################
# L functions
function L1(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  (ni, ne) = nVals(b, t)
  if (t == τ)
    ((k*ni)^2-(k*ne)^2) / (2π)
  else
    (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
    ad = bpc.absDiff[tIdx, τIdx]
    j0ni = bpc.j0NiMatrix[tIdx, τIdx]
    j0ne = bpc.j0NeMatrix[tIdx, τIdx]
    -1/(2π)*((k*ne)^2*j0ne-(k*ni)^2*j0ni)
  end
end

function L2(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  if (t == τ)
    (ni, ne) = nVals(b, t)
    (im/2 - Base.MathConstants.eulergamma/π - 1/π*log(norm(boundaryPointDeriv(b, t)) / 2)) * ((k*ne)^2-(k*ni)^2) - 1/π*((k*ne)^2*log(k*ne) - (k*ni)^2*log(k*ni))
  else
    L(b, t, τ, k, bpc)-L1(b, t, τ, k, bpc)*log4Term(t, τ)
  end
end

function L(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  (ni, ne) = nVals(b, t)
  (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
  ad = bpc.absDiff[tIdx, τIdx]
  h0ni = bpc.h0NiMatrix[tIdx, τIdx]
  h0ne = bpc.h0NeMatrix[tIdx, τIdx]
  im/2*((k*ne)^2*h0ne-(k*ni)^2*h0ni)
end

###################
# M functions
function M1(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  if (t == τ)
    0.0+0im
  else
    (ni, ne) = nVals(b, t)
    (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
    ad = bpc.absDiff[tIdx, τIdx]
    j0ni = bpc.j0NiMatrix[tIdx, τIdx]
    j0ne = bpc.j0NeMatrix[tIdx, τIdx]
    -1/(2π)*(j0ne-j0ni)
  end
end

function M2(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  if (t == τ)
    (ni, ne) = nVals(b, t)
    1/π*log(ni/ne)
  else
    M(b, t, τ, k, bpc) - M1(b, t, τ, k, bpc)*log4Term(t, τ)
  end
end

function M(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  (ni, ne) = nVals(b, t)
  (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
  ad = bpc.absDiff[tIdx, τIdx]
  h0ni = bpc.h0NiMatrix[tIdx, τIdx]
  h0ne = bpc.h0NeMatrix[tIdx, τIdx]
  im/2 * (h0ne-h0ni)
end

###################
# N functions
function N(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  (ni, ne) = nVals(b, t)
  xtP = boundaryPointDeriv(b, t)
  xτP = boundaryPointDeriv(b, τ)
  xpdot = dot(xtP, xτP)
  (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
  ad = bpc.absDiff[tIdx, τIdx]
  h0ni = bpc.h0NiMatrix[tIdx, τIdx]
  h0ne = bpc.h0NeMatrix[tIdx, τIdx]
  h1ni = bpc.h1NiMatrix[tIdx, τIdx]
  h1ne = bpc.h1NeMatrix[tIdx, τIdx]
  nbar = Nbar(b, t, τ, bpc)
  im/2*nbar*((k*ne)^2*h0ne-2*k*ne*h1ne/ad) +
    im/2*k*ne*xpdot/ad*h1ne - im/2*nbar*((k*ni)^2*h0ni-2*k*ni*h1ni/ad) -
    im/2*k*ni*xpdot/ad*h1ni
end

function Nbar(b::Boundary, t::Float64, τ::Float64, bpc::BoundaryPrecomputedCalcs)
  (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
  ad = bpc.absDiff[tIdx, τIdx]
  (xt, xτ) = xVals(b, t, τ)
  xtP = boundaryPointDeriv(b, t)
  xτP = boundaryPointDeriv(b, τ)

  xtm = (xt-xτ)
  dot(xtm, xtP)*dot(xtm, xτP)/ad^2
end

function N1(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  (ni, ne) = nVals(b, t)
  xtP = boundaryPointDeriv(b, t)

  if (t == τ)
    norm(xtP)^2*((k*ni)^2-(k*ne)^2)/(4π)
  else
    xτP = boundaryPointDeriv(b, τ)
    xpdot = dot(xtP, xτP)
    (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
    ad = bpc.absDiff[tIdx, τIdx]
    j0ni = bpc.j0NiMatrix[tIdx, τIdx]
    j0ne = bpc.j0NeMatrix[tIdx, τIdx]
    j1ni = bpc.j1NiMatrix[tIdx, τIdx]
    j1ne = bpc.j1NeMatrix[tIdx, τIdx]

    -1/(2π)*Nbar(b, t, τ, bpc)*((k*ne)^2*j0ne-2*k*ne*j1ne/ad) -
      k*ne*xpdot/(2π*ad)*j1ne +
      1/(2π)*Nbar(b, t, τ, bpc)*((k*ni)^2*j0ni-2*k*ni*j1ni/ad) +
      k*ni*xpdot/(2π*ad)*j1ni
  end
end

function N2(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  if (t == τ)
    (xt, xτ) = xVals(b, t, τ)
    (ni, ne) = nVals(b, t)
    xtP = boundaryPointDeriv(b, t)
    normXtP = norm(xtP)
    normXtPSq = norm(xtP)^2

    normXtPSq/(4π) * (
      ((k*ne)^2-(k*ni)^2)*(π*im-1-2*Base.MathConstants.eulergamma) -
        2*(k*ne)^2*log(k*ne*normXtP/2)+2*(k*ni)^2*log(k*ni*normXtP/2)
    )
  else
    N(b, t, τ, k, bpc)-N1(b, t, τ, k, bpc)*log4Term(t, τ)
  end
end

###################
# A matrix blocks
#
A11(b::Boundary, N::Int64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs) = @DArray [ -(R(abs(j-l), N)*H1(b, j*π/N, l*π/N, k, bpc) + π/N*H2(b, j*π/N, l*π/N, k, bpc)) for j in 0:(2N-1), l in 0:(2N-1) ]

function A12(b::Boundary, N::Int64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  @DArray [ -(R(abs(j-l), N)*M1(b, j*π/N, l*π/N, k, bpc)*γ(b, N, l) +
    π/N*M2(b, j*π/N, l*π/N, k, bpc)*γ(b,N,l))
    for j in 0:(2N-1), l in 0:(2N-1) ]
end

function A21entry(b::Boundary, N::Int64, k::Complex{Float64}, j::Int64, l::Int64, bpc::BoundaryPrecomputedCalcs)
  xjP = boundaryPointDeriv(b, j*π/N)
  xlP = boundaryPointDeriv(b, l*π/N)
  R(abs(j-l), N) * (L1(b, j*π/N, l*π/N, k, bpc)*dot(xjP, xlP) - N1(b, j*π/N, l*π/N, k, bpc)) +
    π/N*(L2(b, j*π/N, l*π/N, k, bpc)*dot(xjP, xlP) - N2(b, j*π/N, l*π/N, k, bpc))
end
A21(b::Boundary, N::Int64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs) = @DArray [ A21entry(b, N, k, j, l, bpc) for j in 0:(2N-1), l in 0:(2N-1) ]

A22(b::Boundary, N::Int64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs) = @DArray [ R(abs(j-l), N)*H1star(b, j*π/N, l*π/N, k, bpc) + π/N*H2star(b, j*π/N, l*π/N, k, bpc) for j in 0:(2N-1), l in 0:(2N-1) ]

###################
# Other matrix blocks
γMatrix(b::Boundary, N::Int64) = Diagonal([γ(b, N, l) for l in 0:(2N-1)])

###################
# Full matrix block
function A(b::Boundary, N::Int64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  a11 = A11(b, N, k, bpc)
  a12 = A12(b, N, k, bpc)
  a21 = A21(b, N, k, bpc)
  a22 = A22(b, N, k, bpc)
  2*[ Matrix(1.0I, 2N, 2N) zeros(2N,2N); zeros(2N,2N) γMatrix(b, N)] - [convertAndClose(a11) convertAndClose(a12); convertAndClose(a21) convertAndClose(a22)]
end

function A(b::Boundary, N::Int64, k::Complex{Float64})
  bpc = doPrecalculation(b, N, k)
  A(b, N, k, bpc)
end

###################
# A'(k) derivative functions

################
# H functions
function H1′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  if (t == τ)
    0.0+0im
  else
    # Note: Need to check that we can use the k values at t, not τ, this is a deviation from the Heider paper
    (ni, ne) = nVals(b, t)
    (xt, xτ) = xVals(b, t, τ)
    (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
    ad = bpc.absDiff[tIdx, τIdx]
    j0ni = bpc.j0NiMatrix[tIdx, τIdx]
    j0ne = bpc.j0NeMatrix[tIdx, τIdx]
    k*ad/(2π)*dot(normVector(b, τ), (xt-xτ)/ad)*(ni^2*j0ni-ne^2*j0ne)
  end
end

function H2′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  if (t == τ)
    0.0+0im
  else
    H′(b, t, τ, k, bpc)-H1′(b, t, τ, k, bpc)*log4Term(t, τ)
  end
end

function H′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  (ni, ne) = nVals(b, t)
  (xt, xτ) = xVals(b, t, τ)
  (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
  ad = bpc.absDiff[tIdx, τIdx]
  h0ni = bpc.h0NiMatrix[tIdx, τIdx]
  h0ne = bpc.h0NeMatrix[tIdx, τIdx]
  im*k*ad/2*dot(normVector(b, τ), (xt-xτ)/ad)*(ne^2*h0ne-ni^2*h0ni)
end

H1star′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs) = H1′(b, τ, t, k, bpc)*norm(boundaryPointDeriv(b, τ))
H2star′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs) = H2′(b, τ, t, k, bpc)*norm(boundaryPointDeriv(b, τ))
Hstar′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs) = H′(b, τ, t, k, bpc)*norm(boundaryPointDeriv(b, τ))

###################
# L functions
function L1′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  (ni, ne) = nVals(b, t)
  if (t == τ)
    k*(ni^2-ne^2) / π
  else
    (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
    ad = bpc.absDiff[tIdx, τIdx]
    j0ni = bpc.j0NiMatrix[tIdx, τIdx]
    j0ne = bpc.j0NeMatrix[tIdx, τIdx]
    j1ni = bpc.j1NiMatrix[tIdx, τIdx]
    j1ne = bpc.j1NeMatrix[tIdx, τIdx]
    1/(2π)*k*(k*ad*(ne^3*j1ne-ni^3*j1ni)-2*ne^2*j0ne+2*ni^2*j0ni)
  end
end

function L2′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  if (t == τ)
    (ni, ne) = nVals(b, t)
    k/π*(-2*ne^2*log(k*ne)+2*ni^2*log(k*ni)-(ne-ni)*(ne+ni)*(2*log(norm(boundaryPointDeriv(b, t)) / 2)-π*im+2*Base.MathConstants.eulergamma+1))
  else
    L′(b, t, τ, k, bpc)-L1′(b, t, τ, k, bpc)*log4Term(t, τ)
  end
end

function L′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  (ni, ne) = nVals(b, t)
  (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
  ad = bpc.absDiff[tIdx, τIdx]
  h0ni = bpc.h0NiMatrix[tIdx, τIdx]
  h0ne = bpc.h0NeMatrix[tIdx, τIdx]
  h1ni = bpc.h1NiMatrix[tIdx, τIdx]
  h1ne = bpc.h1NeMatrix[tIdx, τIdx]
  -k*im/2*(k*ad*(ne^3*h1ne-ni^3*h1ni)-2*ne^2*h0ne+2*ni^2*h0ni)
end

###################
# M functions
function M1′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  if (t == τ)
    0.0+0im
  else
    (ni, ne) = nVals(b, t)
    (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
    ad = bpc.absDiff[tIdx, τIdx]
    j1ni = bpc.j1NiMatrix[tIdx, τIdx]
    j1ne = bpc.j1NeMatrix[tIdx, τIdx]
    ad/(2π)*(ne*j1ne-ni*j1ni)
  end
end

function M2′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  if (t == τ)
    0.0+0im
  else
    M′(b, t, τ, k, bpc) - M1′(b, t, τ, k, bpc)*log4Term(t, τ)
  end
end

function M′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  (ni, ne) = nVals(b, t)
  (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
  ad = bpc.absDiff[tIdx, τIdx]
  h1ni = bpc.h1NiMatrix[tIdx, τIdx]
  h1ne = bpc.h1NeMatrix[tIdx, τIdx]
  -im*ad/2 * (ne*h1ne-ni*h1ni)
end

###################
# N functions
function N′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  (ni, ne) = nVals(b, t)
  xtP = boundaryPointDeriv(b, t)
  xτP = boundaryPointDeriv(b, τ)
  xpdot = dot(xtP, xτP)
  (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
  ad = bpc.absDiff[tIdx, τIdx]
  h0ni = bpc.h0NiMatrix[tIdx, τIdx]
  h0ne = bpc.h0NeMatrix[tIdx, τIdx]
  h1ni = bpc.h1NiMatrix[tIdx, τIdx]
  h1ne = bpc.h1NeMatrix[tIdx, τIdx]
  k*im/2 * (
    k*ad*Nbar(b, t, τ, bpc) * (ni^3*h1ni-ne^3*h1ne) +
    xpdot*(ne^2*h0ne-ni^2*h0ni)
  )
end

function N1′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  (ni, ne) = nVals(b, t)
  xtP = boundaryPointDeriv(b, t)
  xτP = boundaryPointDeriv(b, τ)

  if (t == τ)
    norm(xtP)^2*k*(ni^2-ne^2)/(2π)
  else
    xτP = boundaryPointDeriv(b, τ)
    xpdot = dot(xtP, xτP)
    (tIdx, τIdx) = tAndτIndices(bpc, t, τ)
    ad = bpc.absDiff[tIdx, τIdx]
    j0ni = bpc.j0NiMatrix[tIdx, τIdx]
    j0ne = bpc.j0NeMatrix[tIdx, τIdx]
    j1ni = bpc.j1NiMatrix[tIdx, τIdx]
    j1ne = bpc.j1NeMatrix[tIdx, τIdx]
    k/(2π)*(
      k*ad*Nbar(b, t, τ, bpc)*(ne^3*j1ne-ni^3*j1ni) +
      xpdot*(ni^2*j0ni-ne^2*j0ne)
    )
  end
end

function N2′(b::Boundary, t::Float64, τ::Float64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  if (t == τ)
    (ni, ne) = nVals(b, t)
    xtP = boundaryPointDeriv(b, t)
    normXtP = norm(xtP)
    normXtPSq = norm(xtP)^2

    k*normXtPSq/(2π) * (
      -2*ne^2*log(k/2*ne*normXtP) +
      2*ni^2*log(k/2*ni*normXtP) -
      (2+2*Base.MathConstants.eulergamma-π*im)*(ne-ni)*(ne+ni)
    )
  else
    N′(b, t, τ, k, bpc)-N1′(b, t, τ, k, bpc)*log4Term(t, τ)
  end
end

###################
# A matrix blocks′
#
A11′(b::Boundary, N::Int64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs) = @DArray [ -(R(abs(j-l), N)*H1′(b, j*π/N, l*π/N, k, bpc) + π/N*H2′(b, j*π/N, l*π/N, k, bpc)) for j in 0:(2N-1), l in 0:(2N-1) ]

function A12′(b::Boundary, N::Int64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  @DArray [ -(R(abs(j-l), N)*M1′(b, j*π/N, l*π/N, k, bpc)*γ(b, N, l) +
    π/N*M2′(b, j*π/N, l*π/N, k, bpc)*γ(b,N,l))
    for j in 0:(2N-1), l in 0:(2N-1) ]
end

function A21entry′(b::Boundary, N::Int64, k::Complex{Float64}, j::Int64, l::Int64, bpc::BoundaryPrecomputedCalcs)
  xjP = boundaryPointDeriv(b, j*π/N)
  xlP = boundaryPointDeriv(b, l*π/N)
  R(abs(j-l), N) * (L1′(b, j*π/N, l*π/N, k, bpc)*dot(xjP, xlP) - N1′(b, j*π/N, l*π/N, k, bpc)) +
    π/N*(L2′(b, j*π/N, l*π/N, k, bpc)*dot(xjP, xlP) - N2′(b, j*π/N, l*π/N, k, bpc))
end
A21′(b::Boundary, N::Int64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs) = @DArray [ A21entry′(b, N, k, j, l, bpc) for j in 0:(2N-1), l in 0:(2N-1) ]

A22′(b::Boundary, N::Int64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs) = @DArray [ R(abs(j-l), N)*H1star′(b, j*π/N, l*π/N, k, bpc) + π/N*H2star′(b, j*π/N, l*π/N, k, bpc) for j in 0:(2N-1), l in 0:(2N-1) ]

###################
# Full matrix block
function A′(b::Boundary, N::Int64, k::Complex{Float64}, bpc::BoundaryPrecomputedCalcs)
  a11′ = A11′(b, N, k, bpc)
  a12′ = A12′(b, N, k, bpc)
  a21′ = A21′(b, N, k, bpc)
  a22′ = A22′(b, N, k, bpc)
  - [convertAndClose(a11′) convertAndClose(a12′); convertAndClose(a21′) convertAndClose(a22′)]
end

function A′(b::Boundary, N::Int64, k::Complex{Float64})
  bpc = doPrecalculation(b, N, k)
  A′(b, N, k, bpc)
end

###################
# Utility functions
function tAndτIndices(bpc::BoundaryPrecomputedCalcs, t::Float64, τ::Float64)
  (Base.round(Int, t*bpc.N/π)+1, Base.round(Int, τ*bpc.N/π)+1)
end

function convertAndClose(a::DistributedArrays.DArray)
  tmp = Base.convert(Array, a)
  close(a)
  tmp
end

function doPrecalculation(b::Boundary, N::Int64, k::Complex{Float64})
  range = 0:(2N-1)
  nis = [nVals(b, l*π/N)[1] for l in range]
  nes = [nVals(b, l*π/N)[2] for l in range]

  tτMatrix = [[l*π/N,j*π/N] for l in range, j in range]
  absDiffM = absDiff(b, tτMatrix)

  localNiArgs = k*nis.*absDiffM
  localNeArgs = k*nes.*absDiffM
  niArgs = distribute(localNiArgs)
  neArgs = distribute(localNeArgs)
  j0NiM = besselj.(0, niArgs)
  j0NeM = besselj.(0, neArgs)
  j1NiM = besselj.(1, niArgs)
  j1NeM = besselj.(1, neArgs)

  # Set diagonal to NaN to prevent exception for hankel functions
  [localNiArgs[i, i] = NaN for i in 1:2N]
  [localNeArgs[i, i] = NaN for i in 1:2N]
  close(niArgs)
  close(neArgs)
  niArgs = distribute(localNiArgs)
  neArgs = distribute(localNeArgs)
  h0NiM = h0.(niArgs)
  h0NeM = h0.(neArgs)
  h1NiM = h1.(niArgs)
  h1NeM = h1.(neArgs)
  close(niArgs)
  close(neArgs)
  bpc = BoundaryPrecomputedCalcs(b, N, k, convertAndClose(j0NiM), convertAndClose(j0NeM), convertAndClose(j1NiM), convertAndClose(j1NeM),
      convertAndClose(h0NiM), convertAndClose(h0NeM), convertAndClose(h1NiM), convertAndClose(h1NeM), absDiffM)
  # @everywhere gc()
  bpc
end


@defcomp n2ocycle begin
    e_globalN2Oemissions=Parameter(index=[time],unit="Mtonne/year")
    e_0globalN2Oemissions=Parameter(unit="Mtonne/year", default=11.22281348065454)
    c_N2Oconcentration=Variable(index=[time],unit="ppbv")
    pic_preindustconcN2O=Parameter(unit="ppbv", default=270.)
    exc_excessconcN2O=Variable(unit="ppbv")
    c0_N2Oconcbaseyr=Parameter(unit="ppbv", default=324.)
    re_remainN2O=Variable(index=[time],unit="Mtonne")
    re_remainN2Obase=Variable(unit="Mtonne")
    nte_natN2Oemissions=Variable(index=[time],unit="Mtonne/year")
    air_N2Ofractioninatm=Parameter(unit="%", default=100.)
    tea_N2Oemissionstoatm=Variable(index=[time],unit="Mtonne/year")
    teay_N2Oemissionstoatm=Variable(index=[time],unit="Mtonne/t")
    y_year=Parameter(index=[time],unit="year")
    y_year_0=Parameter(unit="year")
    res_N2Oatmlifetime=Parameter(unit="year", default=114.)
    den_N2Odensity=Parameter(unit="Mtonne/ppbv", default=7.8)
    stim_N2Oemissionfeedback=Parameter(unit="Mtonne/degreeC", default=0.)
    rtl_g0_baselandtemp=Parameter(unit="degreeC", default=1.1683981941310047)
    rtl_g_landtemperature=Parameter(index=[time],unit="degreeC")

    function run_timestep(p, v, d, t)
        #note that Hope (2009) states that Equations 1-12 for methane also apply to N2O

        if is_first(t)
            #eq.3 from Hope (2006) - natural emissions feedback, using global temperatures calculated in ClimateTemperature component
            nte_0=p.stim_N2Oemissionfeedback*p.rtl_g0_baselandtemp
            v.nte_natN2Oemissions[t]=p.stim_N2Oemissionfeedback*p.rtl_g0_baselandtemp
            #eq.6 from Hope (2006) - emissions to atmosphere depend on the sum of natural and anthropogenic emissions
            v.tea_N2Oemissionstoatm[t]=(p.e_globalN2Oemissions[t]+v.nte_natN2Oemissions[t])*p.air_N2Ofractioninatm/100
            tea_0=(p.e_0globalN2Oemissions+nte_0)*p.air_N2Ofractioninatm/100
            v.teay_N2Oemissionstoatm[t]=(v.tea_N2Oemissionstoatm[t]+tea_0)*(p.y_year[t]-p.y_year_0)/2
            #adapted from eq.1 in Hope(2006) - calculate excess concentration in base year
            v.exc_excessconcN2O=p.c0_N2Oconcbaseyr-p.pic_preindustconcN2O
            #Eq. 2 from Hope (2006) - base-year remaining emissions
            v.re_remainN2Obase=v.exc_excessconcN2O*p.den_N2Odensity
            v.re_remainN2O[t]=v.re_remainN2Obase*exp(-(p.y_year[t]-p.y_year_0)/p.res_N2Oatmlifetime)+
                v.teay_N2Oemissionstoatm[t]*p.res_N2Oatmlifetime*(1-exp(-(p.y_year[t]-p.y_year_0)/p.res_N2Oatmlifetime))/(p.y_year[t]-p.y_year_0)
        else
            #eq.3 from Hope (2006) - natural emissions (carbon cycle) feedback, using global temperatures calculated in ClimateTemperature component
            #Here assume still using area-weighted average regional temperatures (i.e. land temperatures) for natural emissions feedback
            v.nte_natN2Oemissions[t]=p.stim_N2Oemissionfeedback*p.rtl_g_landtemperature[t-1]
            #eq.6 from Hope (2006) - emissions to atmosphere depend on the sum of natural and anthropogenic emissions
            v.tea_N2Oemissionstoatm[t]=(p.e_globalN2Oemissions[t]+v.nte_natN2Oemissions[t])*p.air_N2Ofractioninatm/100
            #eq.7 from Hope (2006) - average emissions to atm over time period
            v.teay_N2Oemissionstoatm[t]=(v.tea_N2Oemissionstoatm[t]+v.tea_N2Oemissionstoatm[t-1])*(p.y_year[t]-p.y_year[t-1])/2
            #eq.10 from Hope (2006) - remaining emissions in atmosphere
            v.re_remainN2O[t]=v.re_remainN2O[t-1]*exp(-(p.y_year[t]-p.y_year[t-1])/p.res_N2Oatmlifetime)+
                v.teay_N2Oemissionstoatm[t]*p.res_N2Oatmlifetime*(1-exp(-(p.y_year[t]-p.y_year[t-1])/p.res_N2Oatmlifetime))/(p.y_year[t]-p.y_year[t-1])
        end

        #eq.11 from Hope(2006) - N2O concentration
        v.c_N2Oconcentration[t]=p.pic_preindustconcN2O+v.exc_excessconcN2O*v.re_remainN2O[t]/v.re_remainN2Obase

    end
end


export grad, jacobian, jvp, j′vp, to_vec

"""
    grad(fdm, f, x::Vector{<:Number})

Approximate the gradient of `f` at `x` using `fdm`. Assumes that `f(x)` is scalar.
"""
function grad(fdm, f, x::Vector{T}) where T<:Number
    v, dx, tmp = fill(zero(T), size(x)), similar(x), similar(x)
    for n in eachindex(x)
        v[n] = one(T)
        dx[n] = fdm(function(ϵ)
                tmp .= x .+ ϵ .* v
                return f(tmp)
            end,
            zero(T),
        )
        v[n] = zero(T)
    end
    return dx
end

"""
    jacobian(fdm, f, x::Vector{<:Number}, D::Int)
    jacobian(fdm, f, x::Vector{<:Number})

Approximate the Jacobian of `f` at `x` using `fdm`. `f(x)` must be a length `D` vector. If
`D` is not provided, then `f(x)` is computed once to determine the output size.
"""
function jacobian(fdm, f, x::Vector{T}, D::Int) where {T<:Number}
    J = Matrix{T}(undef, D, length(x))
    for d in 1:D
        J[d, :] = grad(fdm, x->f(x)[d], x)
    end
    return J
end
jacobian(fdm, f, x::Vector{<:Number}) = jacobian(fdm, f, x, length(f(x)))

"""
    _jvp(fdm, f, x::Vector{<:Number}, ẋ::AbstractVector{<:Number})

Convenience function to compute `jacobian(f, x) * ẋ`.
"""
_jvp(fdm, f, x::Vector{<:Number}, ẋ::AV{<:Number}) = fdm(ε -> f(x .+ ε .* ẋ), zero(eltype(x)))

"""
    _j′vp(fdm, f, ȳ::AbstractVector{<:Number}, x::Vector{<:Number})

Convenience function to compute `transpose(jacobian(f, x)) * ȳ`.
"""
_j′vp(fdm, f, ȳ::AV{<:Number}, x::Vector{<:Number}) = transpose(jacobian(fdm, f, x, length(ȳ))) * ȳ

"""
    jvp(fdm, f, x, ẋ)

Compute a Jacobian-vector product with any types of arguments for which [`to_vec`](@ref)
is defined.
"""
function jvp(fdm, f, (x, ẋ)::Tuple{Any, Any})
    x_vec, vec_to_x = to_vec(x)
    _, vec_to_y = to_vec(f(x))
    return vec_to_y(_jvp(fdm, x_vec->to_vec(f(vec_to_x(x_vec)))[1], x_vec, to_vec(ẋ)[1]))
end
function jvp(fdm, f, xẋs::Tuple{Any, Any}...)
    x, ẋ = collect(zip(xẋs...))
    return jvp(fdm, xs->f(xs...), (x, ẋ))
end

"""
    j′vp(fdm, f, ȳ, x...)

Compute an adjoint with any types of arguments for which [`to_vec`](@ref) is defined.
"""
function j′vp(fdm, f, ȳ, x)
    x_vec, vec_to_x = to_vec(x)
    ȳ_vec, _ = to_vec(ȳ)
    return vec_to_x(_j′vp(fdm, x_vec->to_vec(f(vec_to_x(x_vec)))[1], ȳ_vec, x_vec))
end
j′vp(fdm, f, ȳ, xs...) = j′vp(fdm, xs->f(xs...), ȳ, xs)

"""
    to_vec(x) -> Tuple{<:AbstractVector, <:Function}

Transform `x` into a `Vector`, and return a closure which inverts the transformation.
"""
to_vec(x::Number) = ([x], first)

# Vectors
to_vec(x::Vector{<:Number}) = (x, identity)
function to_vec(x::Vector)
    x_vecs_and_backs = map(to_vec, x)
    x_vecs, backs = first.(x_vecs_and_backs), last.(x_vecs_and_backs)
    return vcat(x_vecs...), function(x_vec)
        sz = cumsum([map(length, x_vecs)...])
        return [backs[n](x_vec[sz[n]-length(x_vecs[n])+1:sz[n]]) for n in eachindex(x)]
    end
end

# Arrays
to_vec(x::Array{<:Number}) = vec(x), x_vec->reshape(x_vec, size(x))
function to_vec(x::Array)
    x_vec, back = to_vec(reshape(x, :))
    return x_vec, x_vec->reshape(back(x_vec), size(x))
end

# AbstractArrays
function to_vec(x::T) where {T<:LinearAlgebra.AbstractTriangular}
    x_vec, back = to_vec(Matrix(x))
    return x_vec, x_vec->T(reshape(back(x_vec), size(x)))
end
to_vec(x::Symmetric) = vec(Matrix(x)), x_vec->Symmetric(reshape(x_vec, size(x)))
to_vec(X::Diagonal) = vec(Matrix(X)), x_vec->Diagonal(reshape(x_vec, size(X)...))

function to_vec(X::Transpose)
    return vec(Matrix(X)), x_vec->Transpose(permutedims(reshape(x_vec, size(X))))
end
function to_vec(X::Adjoint)
    return vec(Matrix(X)), x_vec->Adjoint(conj!(permutedims(reshape(x_vec, size(X)))))
end

# Non-array data structures

function to_vec(x::Tuple)
    x_vecs, x_backs = zip(map(to_vec, x)...)
    sz = cumsum([map(length, x_vecs)...])
    return vcat(x_vecs...), function(v)
        return ntuple(n->x_backs[n](v[sz[n]-length(x_vecs[n])+1:sz[n]]), length(x))
    end
end

# Convert to a vector-of-vectors to make use of existing functionality.
function to_vec(d::Dict)
    d_vec_vec = [val for val in values(d)]
    d_vec, back = to_vec(d_vec_vec)
    return d_vec, function(v)
        v_vec_vec = back(v)
        return Dict([(key, v_vec_vec[n]) for (n, key) in enumerate(keys(d))])
    end
end


  # wrap function -  using periodic boundary conditions to resize i and j of cells
  function wrap_cell(dms,id,jd)
    i = dms
    j = dms

    if id < 1
      while id < 1
        id = id + i
      end
    elseif id > i
      while id > i
        id = id - i
      end
    end

    if jd < 1
      while jd < 1
        jd = jd + j
      end
    elseif jd > j
      while jd > j
        jd = jd - j
      end
    end
    return id, jd
  end

  function wrapcell(nc,i,j)
    if i < 1
      i = nc + i
    elseif i > nc
      i = i - nc
    end
    if j < 1
      j = nc + j
    elseif j > nc
     j = j - nc
    end
    return i,j
  end


### A Pluto.jl notebook ###
# v0.14.6

using Markdown
using InteractiveUtils

# ╔═╡ 660bbe15-74ae-4e61-9a7e-6b2bf1b1e1f5
begin
    import Pkg
    Pkg.activate(mktempdir())
    Pkg.add([
        Pkg.PackageSpec(name="Plots", version="1"),
        Pkg.PackageSpec(name="PlutoUI", version="0.7"),
	    Pkg.PackageSpec(name="DifferentialEquations", version="6"),
	    Pkg.PackageSpec(name="Parameters", version="0.12"),
        Pkg.PackageSpec(name="LabelledArrays", version="1"),
        Pkg.PackageSpec(name="Setfield", version="0.7"),
    ])
    using Plots, PlutoUI, LinearAlgebra, DifferentialEquations, Parameters, LabelledArrays, Setfield

	Plots.gr(fmt=:png, lw=2)

	PlutoUI.TableOfContents()
end

# ╔═╡ a65fcaa1-6d49-4121-b644-2e8945763fd2
md"""
# Figure 4.15, 4.16, and 4.17

Oscillatory network.
"""

# ╔═╡ 3fa4e3bb-2fa5-4f99-9e2d-e41936cd8e6d
"""
Model of oscillatory network from Figure 4.14. This code generates Figures
4.15, 4.16, and 4.17
"""
function model!(du, u, p, t)
    @unpack s1, s2 = u
    @unpack K0, K1, K2, N = p
    v0 = K0
    v1 = K1 * s1 * (1 + s2^N)
    v2 = K2 * s2
	du.s1 = v0 - v1
	du.s2 = v1 - v2
    return du
end

# ╔═╡ cf4047c0-bdec-11eb-0bcd-d958a6ba03d6
md"""
## Appendix
Running environment and some auxillary functions.

**You can use [this helper tool](https://fonsp.com/article-test-3/pkghelper.html) to generate these commands!**
"""

# ╔═╡ 7abc2f61-a1ed-4e61-bfa8-e557f944f615
begin
	# Convenience functions
    hill(x, k) = x / (x + k)
    hill(x, k, n) = hill(x^n, k^n)
    exprel(x) = ifelse(x≈zero(x), one(x), x / expm1(x))
end

# ╔═╡ c37327cf-a798-4820-8304-36966ac64e56
function figure0415(; param = (K0 = 8.0, K1 = 1.0, K2 = 5.0, N = 2),
	                  r = LinRange(0.0, 4.0, 20),
	                  tend = 8.0,
	                  figtitle="Fig 4.15")

	u0s = ( LVector(s1=1.5, s2=1.0), LVector(s1=0.0, s2=1.0),
        	LVector(s1=0.0, s2=3.0), LVector(s1=2.0, s2=0.0))
	sols = map(u0 -> solve(ODEProblem(model!, u0, tend, param)), u0s)

	# Fig 4.15 A
	p1 = plot(sols[1], xlabel="Time", ylabel="Concentration", title ="$figtitle (A)", xlims=(0.0, 8.0))

	# Fig 4.15 B: Vetor field
	function ∂F(x, y)
		u = LVector(s1=x, s2=y)
		dxdy = model!(similar(u), u, param, 0.0)
		return dxdy ./ (norm(dxdy)^0.5 * 20)
	end

	nullcline_s1(s2, p=param) = (p.K0 / p.K1) * hill(1, s2, p.N)
	nullcline_s2(s2, p=param) = (p.K2 * s2) / (p.K1 * (1 + s2^p.N))

	xx = [x for y in r, x in r]
	yy = [y for y in r, x in r]
	p2 = quiver(xx, yy, quiver=∂F, line=(:lightblue))

	for sol in sols
		plot!(p2, sol, vars=(1, 2), label=nothing)
	end

	rMin, rMax = r[1], r[end]

	plot!(p2, nullcline_s1, identity, rMin, rMax, label="Nullcline S1", line=(:dash, :red))
	plot!(p2, nullcline_s2, identity, rMin, rMax, label="Nullcline S2", line=(:dash, :blue))
	plot!(p2, title = "$figtitle (B)", xlabel="[S1]", ylabel="[S2]",
      xlims=(rMin, rMax), ylims=(rMin, rMax), aspect_ratio=:equal, size=(700, 700))

	return (p1, p2)
end

# ╔═╡ 497e9b54-40c1-493b-aabe-f6ad66120389
fig415a, fig415b = figure0415()

# ╔═╡ c72c49a8-c718-419b-aca6-31853282b150
fig415a

# ╔═╡ 934f6dc0-5c24-466e-b07c-879efbfd1f72
fig415b

# ╔═╡ 21027a60-2903-49d1-ba10-4c19f944fbb0
fig416a, fig416b = figure0415(param = (K0 = 8.0, K1 = 1.0, K2 = 5.0, N = 2.5), tend = 1000.0, figtitle="Fig 4.16")

# ╔═╡ 39d5557b-430b-4074-8944-94f21efcab7e
fig416a

# ╔═╡ 747493eb-9196-4c6b-b70c-42c6167c0b9f
fig416b

# ╔═╡ 6553c713-9663-4260-a416-5d0c8590ea97
let
	param = (K0 = 8.0, K1 = 1.0, K2 = 5.0, N = 2.5)
	sol417 = solve(ODEProblem(model!, LVector(s1=2.0, s2=1.5), 10.0, param))
	r = LinRange(0.0, 4.0, 20)
	xx = [x for y in r, x in r]
	yy = [y for y in r, x in r]

	"Vetor field"
	function ∂F(x, y)
		u = LVector(s1=x, s2=y)
		dxdy = model!(similar(u), u, param, 0.0)
		return dxdy ./ (norm(dxdy)^0.5 * 20)
	end

	nullcline_s1(s2, p=param) = (p.K0 / p.K1) * hill(1, s2, p.N)
	nullcline_s2(s2, p=param) = (p.K2 * s2) / (p.K1 * (1 + s2^p.N))

	xx = [x for y in r, x in r]
	yy = [y for y in r, x in r]
	p2 = quiver(xx, yy, quiver=∂F, line=(:lightblue))

	quiver(xx, yy, quiver=∂F, line=(:lightblue))
	plot!(sol417, vars=(1, 2), label=nothing, line=(:black), arrow=0.4)

	plot!(nullcline_s1, identity, 0.0, 4.0, label="Nullcline S1", line=(:dash, :red))
	plot!(nullcline_s2, identity, 0.0, 4.0, label="Nullcline S2", line=(:dash, :blue))
	plot!(title = "Fig. 4.17", xlabel="[S1]", ylabel="[S2]",
		  xlims=(1.0, 3.0), ylims=(1.0, 3.0), aspect_ratio=:equal, size=(700, 700))
end

# ╔═╡ Cell order:
# ╠═a65fcaa1-6d49-4121-b644-2e8945763fd2
# ╠═3fa4e3bb-2fa5-4f99-9e2d-e41936cd8e6d
# ╠═c37327cf-a798-4820-8304-36966ac64e56
# ╠═497e9b54-40c1-493b-aabe-f6ad66120389
# ╠═c72c49a8-c718-419b-aca6-31853282b150
# ╠═934f6dc0-5c24-466e-b07c-879efbfd1f72
# ╠═21027a60-2903-49d1-ba10-4c19f944fbb0
# ╠═39d5557b-430b-4074-8944-94f21efcab7e
# ╠═747493eb-9196-4c6b-b70c-42c6167c0b9f
# ╠═6553c713-9663-4260-a416-5d0c8590ea97
# ╠═cf4047c0-bdec-11eb-0bcd-d958a6ba03d6
# ╠═7abc2f61-a1ed-4e61-bfa8-e557f944f615
# ╠═660bbe15-74ae-4e61-9a7e-6b2bf1b1e1f5


# Toy linear ODE with 2 DOFs
# u_1_t - a * u_1 = 0
# u_2_t - b * u_1 - c * u_2 = 0

import GridapODEs.ODETools: ODEOperator
import GridapODEs.ODETools: AffineODEOperator
import GridapODEs.ODETools: ConstantODEOperator
import GridapODEs.ODETools: allocate_cache
import GridapODEs.ODETools: update_cache!
import GridapODEs.ODETools: allocate_residual
import GridapODEs.ODETools: jacobian!
import GridapODEs.ODETools: jacobian_t!
import GridapODEs.ODETools: jacobian_and_jacobian_t!
import GridapODEs.ODETools: allocate_jacobian
import GridapODEs.ODETools: residual!

struct ODEOperatorMock{T<:Real,C} <: ODEOperator{C}
  a::T
  b::T
  c::T
end

function residual!(r::AbstractVector,op::ODEOperatorMock,t::Real,u::AbstractVector,u_t::AbstractVector,ode_cache)
  r .= 0
  r[1] = u_t[1] - op.a * u[1]
  r[2] = u_t[2] - op.b * u[1] - op.c * u[2]
  r
end

function allocate_residual(op::ODEOperatorMock,u::AbstractVector,cache)
  zeros(2)
end

function jacobian!(J::AbstractMatrix,op::ODEOperatorMock,t::Real,u::AbstractVector,u_t::AbstractVector,ode_cache)
  J[1,1] += -op.a
  J[2,1] += -op.b
  J[2,2] += -op.c
  J
end

function jacobian_t!(J::AbstractMatrix,op::ODEOperatorMock,t::Real,u::AbstractVector,u_t::AbstractVector,du_t_u::Real,ode_cache)
  J[1,1] += 1.0*du_t_u
  J[2,2] += 1.0*du_t_u
  J
end

function jacobian_and_jacobian_t!(J::AbstractMatrix,op::ODEOperatorMock,t::Real,u::AbstractVector,u_t::AbstractVector,du_t_u::Real,ode_cache)
  jacobian!(J,op,t,u,u_t,ode_cache)
  jacobian_t!(J,op,t,u,u_t,du_t_u,ode_cache)
  J
end

function allocate_jacobian(op::ODEOperatorMock,u::AbstractVector,cache)
  zeros(2,2)
end

allocate_cache(op::ODEOperatorMock) = nothing
update_cache!(cache,op::ODEOperatorMock,t::Real) = nothing


using MAT
using MultipleViewGeometry, Base.Test
using MultipleViewGeometry.ModuleCostFunction
using MultipleViewGeometry.ModuleTypes
using BenchmarkTools, Compat
using StaticArrays

#"C:\Users\Spock\Desktop\tempdata"

file = matopen("test/debug/F.mat")
F = read(file,"F")
close(file)

file = matopen("test/debug/mPts1List.mat")
mpts1list = read(file,"mPts1List")
close(file)

file = matopen("test/debug/mPts2List.mat")
mpts2list = read(file,"mPts2List")
close(file)

file = matopen("test/debug/covs.mat")
covs = read(file,"covs")
close(file)

file = matopen("test/debug/Ct.mat")
Ct = read(file,"Ct")
close(file)

npts, dim  = size(mpts1list[1])

s²  = 1e-7
s = sqrt(s² )
Λ₁ =  [SMatrix{3,3}(s² *diagm([1.0,1.0,0.0])) for i = 1:npts]
Λ₂ =  [SMatrix{3,3}(s² *diagm([1.0,1.0,0.0])) for i = 1:npts]

ntrial = 10000
ℱ = zeros(9,ntrial)
for itrial = 1:ntrial
  mpts1 = mpts1list[itrial]
  mpts2 = mpts2list[itrial]
  𝒪 = [Point2DH(vcat(mpts1[i,:],[1])) for i =1:npts]
  𝒪ʹ= [Point2DH(vcat(mpts2[i,:],[1])) for i =1:npts]

  𝐅₀ = estimate(FundamentalMatrix(),DirectLinearTransform(), (𝒪, 𝒪ʹ))
  𝐅 = estimate(FundamentalMatrix(),
                           FundamentalNumericalScheme(reshape(𝐅₀,9,1), 5, 1e-10),
                                                             (Λ₁,Λ₂), (𝒪, 𝒪ʹ))
  𝐟 = reshape(𝐅,9,1)
  𝐟 = 𝐟 / norm(𝐟)
  ℱ[:,itrial] = 𝐟
end

norm(F[:,1])
norm(ℱ[:,1])

F[:,1]
ℱ[:,1]

index = 2
mpts1 = mpts1list[index]
mpts2 = mpts2list[index]
𝒪 = [Point2DH(vcat(mpts1[i,:],[1])) for i =1:npts]
𝒪ʹ= [Point2DH(vcat(mpts2[i,:],[1])) for i =1:npts]
T1 = hartley_transformation(𝒪)
T2 = hartley_transformation(𝒪ʹ)

𝐅₀  = estimate(FundamentalMatrix(), DirectLinearTransform(), (𝒪, 𝒪ʹ))
𝐅 = estimate(FundamentalMatrix(),
                        FundamentalNumericalScheme(reshape(𝐅₀,9,1), 5, 1e-10),
                                                          (Λ₁,Λ₂), (𝒪, 𝒪ʹ))
𝐟 = reshape(𝐅,9,1)
𝐟 = 𝐟 / norm(𝐟)

𝐅 = 𝐅 / norm(𝐅)

𝛍 = mean(ℱ,2)
𝛍 = 𝛍 / norm(𝛍)
d = length(𝛍)
𝐏 = eye(d) - norm(𝛍)^-2 * (𝛍*𝛍')
Cₜ = zeros((d,d))
for itrial = 1:ntrial
    𝐟 =  ℱ[:,itrial]
    Cₜ = Cₜ + 𝐏*(𝐟-𝛍) *(𝐟-𝛍)'*𝐏'
end
Cₜ = Cₜ / ntrial

#
# 𝛍 = mean(F,2)
# 𝛍 = 𝛍 / norm(𝛍)
# d = length(𝛍)
# 𝐏 = eye(d) - norm(𝛍)^-2 * (𝛍*𝛍')
# Cₜ = zeros((d,d))
# for itrial = 1:ntrial
#     𝐟 =  F[:,itrial]
#     Cₜ = Cₜ + 𝐏*(𝐟-𝛍) *(𝐟-𝛍)'*𝐏'
# end
# Cₜ = Cₜ / ntrial

# Q  = ℱ;
# for itrial = 1:ntrial
#     q =  Q[:,itrial]
#     Q[:,itrial] = Q[:,itrial] / sign(q[end-1])
# end
#
# 𝛍 = mean(Q,2)
# 𝛍 = 𝛍 / norm(𝛍)
# d = length(𝛍)
# 𝐏 = eye(d) - norm(𝛍)^-2 * (𝛍*𝛍')
# Cₜ = zeros((d,d))
# for itrial = 1:ntrial
#     𝐟 =  Q[:,itrial]
#     Cₜ = Cₜ + 𝐏*(𝐟-𝛍) *(𝐟-𝛍)'*𝐏'
# end
# Cₜ = Cₜ / ntrial

#
# A = Cₜ
#
#
# # Validate the covariance matrix of an estimate based on the AML cost function.
C = covariance_matrix(AML(),FundamentalMatrix(), reshape(𝐅,9,1), (Λ₁,Λ₂), (𝒪 , 𝒪ʹ))
C1 = covariance_matrix(AML(),FundamentalMatrix(),  𝐟, (Λ₁,Λ₂), (𝒪 , 𝒪ʹ))
C2 = covariance_matrix_debug(AML(),FundamentalMatrix(), reshape(𝐅,9,1), (Λ₁,Λ₂), (𝒪 , 𝒪ʹ))

C3 = covariance_matrix_debug(AML(),FundamentalMatrix(), 𝐟, (Λ₁,Λ₂), (𝒪 , 𝒪ʹ))

Cn = covariance_matrix_normalised(AML(),FundamentalMatrix(), 𝐟, (Λ₁,Λ₂), (𝒪 , 𝒪ʹ))

# (diag(C3)./diag(Ct) .-1)*100
#
# norm(diag(C3)-diag(Ct)) / norm(diag(Ct)) * 100
#
norm(C3-Ct) / norm(Ct) * 100
norm(C2-Ct) / norm(Ct) * 100
norm(Cn-Ct) / norm(Ct) * 100
norm(C1-Ct) / norm(Ct) * 100

# Cx = covariance_matrix(AML(),FundamentalMatrix(),  𝐟 , (Λ₁,Λ₂), (𝒪 , 𝒪ʹ))
#
#
# P = (1/norm(𝐟)) * (eye(9) - ((𝐟*𝐟') / norm(𝐟)^2) )
#
# P2 = eye(9) - norm(𝐟)^-2 * (𝐟*𝐟')
#
# cost(AML(),FundamentalMatrix(), 𝐟, (Λ₁,Λ₂), (𝒪 , 𝒪ʹ))

# A = Float64.([3 0 2; 2 0 -2; 0 1 1])
# size(A)
# B = similar(A)
# rows, cols = size(A)
# for r = 1:rows
#     for c =1:cols
#         B[r,c] = det(A)
#     end
# end

Q = reshape(𝐟,(3,3))


abstract type AbstractJacobiWeight{T} <: Weight{T} end

axes(::AbstractJacobiWeight{T}) where T = (Inclusion(ChebyshevInterval{T}()),)

==(w::AbstractJacobiWeight, v::AbstractJacobiWeight) = w.a == v.a && w.b == v.b

broadcasted(::LazyQuasiArrayStyle{1}, ::typeof(*), w::AbstractJacobiWeight, v::AbstractJacobiWeight) =
    JacobiWeight(w.a + v.a, w.b + v.b)

broadcasted(::LazyQuasiArrayStyle{1}, ::typeof(sqrt), w::AbstractJacobiWeight) =
    JacobiWeight(w.a/2, w.b/2)

broadcasted(::LazyQuasiArrayStyle{1}, ::typeof(Base.literal_pow), ::Base.RefValue{typeof(^)}, w::AbstractJacobiWeight, ::Base.RefValue{Val{k}}) where k =
    JacobiWeight(k * w.a, k * w.b)

struct JacobiWeight{T} <: AbstractJacobiWeight{T}
    a::T
    b::T
    JacobiWeight{T}(a, b) where T = new{T}(convert(T,a), convert(T,b))
end

JacobiWeight(a::V, b::T) where {T,V} = JacobiWeight{promote_type(T,V)}(a,b)
jacobiweight(a,b, d::AbstractInterval{T}) where T = JacobiWeight(a,b)[affine(d,ChebyshevInterval{T}())]

==(A::JacobiWeight, B::JacobiWeight) = A.b == B.b && A.a == B.a

function getindex(w::JacobiWeight, x::Number)
    x ∈ axes(w,1) || throw(BoundsError())
    (1-x)^w.a * (1+x)^w.b
end

sum(P::JacobiWeight) = jacobimoment(P.a, P.b)

struct LegendreWeight{T} <: AbstractJacobiWeight{T} end
LegendreWeight() = LegendreWeight{Float64}()
legendreweight(d::AbstractInterval{T}) where T = LegendreWeight{float(T)}()[affine(d,ChebyshevInterval{T}())]

function getindex(w::LegendreWeight{T}, x::Number) where T
    x ∈ axes(w,1) || throw(BoundsError())
    one(T)
end

getproperty(w::LegendreWeight{T}, ::Symbol) where T = zero(T)

sum(::LegendreWeight{T}) where T = 2one(T)

_weighted(::LegendreWeight, P) = P

broadcasted(::LazyQuasiArrayStyle{1}, ::typeof(*), ::LegendreWeight{T}, ::LegendreWeight{V}) where {T,V} =
    LegendreWeight{promote_type(T,V)}()

broadcasted(::LazyQuasiArrayStyle{1}, ::typeof(sqrt), w::LegendreWeight{T}) where T = w

broadcasted(::LazyQuasiArrayStyle{1}, ::typeof(Base.literal_pow), ::Base.RefValue{typeof(^)}, w::LegendreWeight, ::Base.RefValue{Val{k}}) where k = w

# support auto-basis determination

singularities(a::AbstractAffineQuasiVector) = singularities(a.x)
singularitiesbroadcast(_, L::LegendreWeight) = L # Assume we stay smooth
singularitiesbroadcast(::typeof(exp), L::LegendreWeight) = L
singularitiesbroadcast(::typeof(Base.literal_pow), ::typeof(^), L::LegendreWeight, ::Val) = L
for op in (:+, :-,:*)
    @eval begin
        singularitiesbroadcast(::typeof($op), ::LegendreWeight{T}, ::LegendreWeight{V}) where {T,V} = LegendreWeight{promote_type(T,V)}()
        singularitiesbroadcast(::typeof($op), L::LegendreWeight, ::NoSingularities) = L
        singularitiesbroadcast(::typeof($op), ::NoSingularities, L::LegendreWeight) = L
    end
end
singularitiesbroadcast(::typeof(/), ::NoSingularities, L::LegendreWeight) = L # can't find roots
singularitiesbroadcast(::typeof(*), x::AbstractJacobiWeight, y::AbstractJacobiWeight) = JacobiWeight(x.a+y.a, x.b+y.b)

_parent(::NoSingularities) = NoSingularities()
_parent(a) = parent(a)
_parentindices(a::NoSingularities, b...) = _parentindices(b...)
_parentindices(a, b...) = parentindices(a)
singularitiesbroadcast(F::Function, G::Function, V::SubQuasiArray, K) = singularitiesbroadcast(F, G, parent(V), K)[parentindices(V)...]
singularitiesbroadcast(F, V::Union{NoSingularities,SubQuasiArray}...) = singularitiesbroadcast(F, map(_parent,V)...)[_parentindices(V...)...]

singularitiesbroadcast(::typeof(*), ::LegendreWeight, b::JacobiWeight) = b
singularitiesbroadcast(::typeof(*), a::JacobiWeight, ::LegendreWeight) = a

abstract type AbstractJacobi{T} <: OrthogonalPolynomial{T} end

singularities(::AbstractJacobi{T}) where T = LegendreWeight{T}()
singularities(::Inclusion{T,<:AbstractInterval}) where T = LegendreWeight{T}()
singularities(d::Inclusion{T,<:Interval}) where T = LegendreWeight{T}()[affine(d,ChebyshevInterval{T}())]

struct Legendre{T} <: AbstractJacobi{T} end
Legendre() = Legendre{Float64}()

legendre() = Legendre()
legendre(d::AbstractInterval{T}) where T = Legendre{float(T)}()[affine(d,ChebyshevInterval{T}()), :]

==(::Legendre, ::Legendre) = true

OrthogonalPolynomial(w::LegendreWeight{T}) where {T} = Legendre{T}()
orthogonalityweight(::Legendre{T}) where T = LegendreWeight{T}()

function qr(P::Legendre)
    Q = Normalized(P)
    QuasiQR(Q, Diagonal(Q.scaling))
end

struct Jacobi{T} <: AbstractJacobi{T}
    a::T
    b::T
    Jacobi{T}(a, b) where T = new{T}(convert(T,a), convert(T,b))
end

Jacobi(a::V, b::T) where {T,V} = Jacobi{float(promote_type(T,V))}(a, b)

jacobi(a,b) = Jacobi(a,b)
jacobi(a,b, d::AbstractInterval{T}) where T = Jacobi(a,b)[affine(d,ChebyshevInterval{T}()), :]

Jacobi(P::Legendre{T}) where T = Jacobi(zero(T), zero(T))

OrthogonalPolynomial(w::JacobiWeight) = Jacobi(w.a, w.b)
orthogonalityweight(P::Jacobi) = JacobiWeight(P.a, P.b)

const WeightedJacobi{T} = WeightedBasis{T,<:JacobiWeight,<:Jacobi}

WeightedJacobi(a,b) = JacobiWeight(a,b) .* Jacobi(a,b)
WeightedJacobi{T}(a,b) where T = JacobiWeight{T}(a,b) .* Jacobi{T}(a,b)

axes(::AbstractJacobi{T}) where T = (Inclusion{T}(ChebyshevInterval{real(T)}()), OneTo(∞))
==(P::Jacobi, Q::Jacobi) = P.a == Q.a && P.b == Q.b
==(P::Legendre, Q::Jacobi) = Jacobi(P) == Q
==(P::Jacobi, Q::Legendre) = P == Jacobi(Q)
==(A::WeightedJacobi, B::WeightedJacobi) = A.args == B.args
==(A::WeightedJacobi, B::Jacobi{T}) where T = A == JacobiWeight(zero(T),zero(T)).*B
==(A::WeightedJacobi, B::Legendre) = A == Jacobi(B)
==(A::Jacobi{T}, B::WeightedJacobi) where T = JacobiWeight(zero(T),zero(T)).*A == B
==(A::Legendre, B::WeightedJacobi) = Jacobi(A) == B

###
# transforms
###

function grid(Pn::SubQuasiArray{T,2,<:AbstractJacobi,<:Tuple{<:Inclusion,<:AbstractUnitRange}}) where T
    kr,jr = parentindices(Pn)
    ChebyshevGrid{1,T}(maximum(jr))
end

function ldiv(::Legendre{V}, f::AbstractQuasiVector) where V
    T = ChebyshevT{V}()
    [cheb2leg(paddeddata(T \ f)); zeros(V,∞)]
end

function ldiv(P::Jacobi{V}, f::AbstractQuasiVector) where V
    T = ChebyshevT{V}()
    [cheb2jac(paddeddata(T \ f), P.a, P.b); zeros(V,∞)]
end

########
# Mass Matrix
#########

legendre_massmatrix(::Type{T}) where T = Diagonal(convert(T,2) ./ (2(0:∞) .+ 1))

function legendre_massmatrix(Ac, B)
    A = parent(Ac)
    P = Legendre{eltype(B)}()
    (P\A)'*legendre_massmatrix(eltype(P))*(P\B)
end

@simplify *(Ac::QuasiAdjoint{<:Any,<:Legendre}, B::Legendre) = legendre_massmatrix(Ac, B)
@simplify *(Ac::QuasiAdjoint{<:Any,<:AbstractJacobi}, B::AbstractJacobi) = legendre_massmatrix(Ac,B)

# 2^{a + b + 1} {\Gamma(n+a+1) \Gamma(n+b+1) \over (2n+a+b+1) \Gamma(n+a+b+1) n!}.

function jacobi_massmatrix(a, b)
    n = 0:∞
    Diagonal(2^(a+b+1) .* (exp.(loggamma.(n .+ (a+1)) .+ loggamma.(n .+ (b+1)) .- loggamma.(n .+ (a+b+1)) .- loggamma.(n .+ 1)) ./ (2n .+ (a+b+1))))
end

@simplify function *(wAc::QuasiAdjoint{<:Any,<:WeightedBasis{<:Any,<:AbstractJacobiWeight}}, wB::WeightedBasis{<:Any,<:AbstractJacobiWeight})
    w,A = arguments(parent(wAc))
    v,B = arguments(wB)
    A'*((w .* v) .* B)
end
@simplify function *(Ac::QuasiAdjoint{<:Any,<:AbstractJacobi}, wB::WeightedBasis{<:Any,<:AbstractJacobiWeight,<:AbstractJacobi})
    A = parent(Ac)
    w,B = arguments(wB)
    P = Jacobi(w.a, w.b)
    (P\A)' * jacobi_massmatrix(w.a, w.b) * (P \ B)
end

########
# Jacobi Matrix
########

jacobimatrix(::Legendre{T}) where T = _BandedMatrix(Vcat(((zero(T):∞)./(1:2:∞))', Zeros{T}(1,∞), ((one(T):∞)./(1:2:∞))'), ∞, 1,1)

# These return vectors A[k], B[k], C[k] are from DLMF. Cause of MikaelSlevinsky we need an extra entry in C ... for now.
function recurrencecoefficients(::Legendre{T}) where T
    n = zero(T):∞
    ((2n .+ 1) ./ (n .+ 1), Zeros{T}(∞), n ./ (n .+ 1))
end

function jacobimatrix(J::Jacobi)
    b,a = J.b,J.a
    n = 0:∞
    B = Vcat(2 / (a+b+2),  2 .* (n .+ 2) .* (n .+ (a+b+2)) ./ ((2n .+ (a+b+3)) .* (2n .+ (a+b+4))))
    A = Vcat((b-a) / (a+b+2), (b^2-a^2) ./ ((2n .+ (a+b+2)) .* (2n .+ (a+b+4))))
    C = 2 .* (n .+ a) .* (n .+ b) ./ ((2n .+ (a+b)) .* (2n .+ (a+b+1)))

    _BandedMatrix(Vcat(C',A',B'), ∞, 1,1)
end

function recurrencecoefficients(P::Jacobi)
    n = 0:∞
    ñ = 1:∞
    a,b = P.a,P.b
    A = Vcat((a+b+2)/2, (2ñ .+ (a+b+1)) .* (2ñ .+ (a+b+2)) ./ ((2*(ñ .+ 1)) .* (ñ .+ (a+b+1))))
    # n = 0 is special to avoid divide-by-zero
    B = Vcat((a-b)/2, (a^2 - b^2) * (2ñ .+ (a + b+1)) ./ ((2*(ñ .+ 1)) .* (ñ .+ (a+b+1)) .* (2ñ .+ (a+b))))
    C = ((n .+ a) .* (n .+ b) .* (2n .+ (a+b+2))) ./ (ñ .* (n .+ (a +b+1)) .* (2n .+ (a+b)))
    (A,B,C)
end

@simplify *(X::Identity, P::Legendre) = ApplyQuasiMatrix(*, P, P\(X*P))

##########
# Conversion
##########

\(A::Jacobi, B::Legendre) = A\Jacobi(B)
\(A::Legendre, B::Jacobi) = Jacobi(A)\B

function \(A::Jacobi, B::Jacobi)
    T = promote_type(eltype(A), eltype(B))
    a,b = B.a,B.b
    if A.a ≈ a && A.b ≈ b
        Eye{T}(∞)
    elseif isone(-a-b) && A.a == a && A.b == b+1
        _BandedMatrix(Vcat((((0:∞) .+ a)./((1:2:∞) .+ (a+b)))',
                            Vcat(1,((2:∞) .+ (a+b))./((3:2:∞) .+ (a+b)))'), ∞, 0,1)
    elseif isone(-a-b) && A.a == a+1 && A.b == b
        _BandedMatrix(Vcat((-((0:∞) .+ b)./((1:2:∞) .+ (a+b)))',
                            Vcat(1,((2:∞) .+ (a+b))./((3:2:∞) .+ (a+b)))'), ∞, 0,1)
    elseif A.a == a && A.b == b+1
        _BandedMatrix(Vcat((((0:∞) .+ a)./((1:2:∞) .+ (a+b)))',
                            (((1:∞) .+ (a+b))./((1:2:∞) .+ (a+b)))'), ∞, 0,1)
    elseif A.a == a+1 && A.b == b
        _BandedMatrix(Vcat((-((0:∞) .+ b)./((1:2:∞) .+ (a+b)))',
                            (((1:∞) .+ (a+b))./((1:2:∞) .+ (a+b)))'), ∞, 0,1)
    elseif A.a ≥ a+1
        J = Jacobi(a+1,b)
        (A \ J) * (J \ B)
    elseif A.b ≥ b+1
        J = Jacobi(a,b+1)
        (A \ J) * (J \ B)
    else
        error("not implemented for $A and $B")
    end
end

function \(A::Jacobi, w_B::WeightedJacobi)
    a,b = A.a,A.b
    (JacobiWeight(zero(a),zero(b)) .* A) \ w_B
end

function \(w_A::WeightedJacobi, B::Jacobi)
    a,b = B.a,B.b
    w_A \ (JacobiWeight(zero(a),zero(b)) .* B)
end

function broadcastbasis(::typeof(+), w_A::WeightedJacobi, w_B::WeightedJacobi)
    wA,A = w_A.args
    wB,B = w_B.args

    w = JacobiWeight(min(wA.a,wB.a), min(wA.b,wB.b))
    P = Jacobi(max(A.a,B.a + w.a - wB.a), max(A.b,B.b + w.b - wB.b))
    w .* P
end

function \(w_A::WeightedJacobi, w_B::WeightedJacobi)
    wA,A = w_A.args
    wB,B = w_B.args

    if wA == wB
        A \ B
    elseif B.a == A.a && B.b == A.b+1 && wB.b == wA.b+1 && wB.a == wA.a
        _BandedMatrix(Vcat((((2:2:∞) .+ 2A.b)./((2:2:∞) .+ (A.a+A.b)))', ((2:2:∞)./((2:2:∞) .+ (A.a+A.b)))'), ∞, 1,0)
    elseif B.a == A.a+1 && B.b == A.b && wB.b == wA.b && wB.a == wA.a+1
        _BandedMatrix(Vcat((((2:2:∞) .+ 2A.a)./((2:2:∞) .+ (A.a+A.b)))', -((2:2:∞)./((2:2:∞) .+ (A.a+A.b)))'), ∞, 1,0)
    elseif wB.a ≥ wA.a+1
        J = JacobiWeight(wB.a-1,wB.b) .* Jacobi(B.a-1,B.b)
        (w_A\J) * (J\w_B)
    elseif wB.b ≥ wA.b+1
        J = JacobiWeight(wB.a,wB.b-1) .* Jacobi(B.a,B.b-1)
        (w_A\J) * (J\w_B)
    else
        error("not implemented for $A and $wB")
    end
end

\(A::Legendre, wB::WeightedJacobi) = Jacobi(A) \ wB

##########
# Derivatives
##########

# Jacobi(a+1,b+1)\(D*Jacobi(a,b))
@simplify function *(D::Derivative{<:Any,<:AbstractInterval}, S::Jacobi)
    A = _BandedMatrix((((1:∞) .+ (S.a + S.b))/2)', ∞, -1,1)
    ApplyQuasiMatrix(*, Jacobi(S.a+1,S.b+1), A)
end

# Jacobi(a-1,b-1)\ (D*w*Jacobi(a,b))
@simplify function *(D::Derivative{<:Any,<:AbstractInterval}, WS::WeightedJacobi)
    w,S = WS.args
    a,b = S.a, S.b
    if w.a == 0 && w.b == 0
        D*S
    elseif iszero(w.a) && w.b == b #L_6
        A = _BandedMatrix((b:∞)', ∞, 0,0)
        ApplyQuasiMatrix(*, JacobiWeight(w.a,b-1) .* Jacobi(a+1,b-1), A)
    elseif iszero(w.b) && w.a == a #L_6^t
        A = _BandedMatrix(-(a:∞)', ∞, 0,0)
        ApplyQuasiMatrix(*, JacobiWeight(a-1,w.b) .* Jacobi(a-1,b+1), A)
    elseif w.a == a && w.b == b # L_1^t
        A = _BandedMatrix((-2*(1:∞))', ∞, 1,-1)
        ApplyQuasiMatrix(*, JacobiWeight(a-1,b-1) .* Jacobi(a-1, b-1), A)
    elseif iszero(w.a)
        W = (JacobiWeight(w.a, b-1) .* Jacobi(a+1, b-1)) \ (D * (JacobiWeight(w.a,b) .* S))
        J = Jacobi(a+1,b) # range Jacobi
        C1 = J \ Jacobi(a+1, b-1)
        C2 = J \ Jacobi(a,b)
        ApplyQuasiMatrix(*, JacobiWeight(w.a,w.b-1) .* J, (w.b-b) * C2 + C1 * W)
    elseif iszero(w.b)
        W = (JacobiWeight(a-1, w.b) .* Jacobi(a-1, b+1)) \ (D * (JacobiWeight(a,w.b) .* S))
        J = Jacobi(a,b+1) # range Jacobi
        C1 = J \ Jacobi(a-1, b+1)
        C2 = J \ Jacobi(a,b)
        ApplyQuasiMatrix(*, JacobiWeight(w.a-1,w.b) .* J, -(w.a-a) * C2 + C1 * W)
    else
        error("Not implemented")
    end
end

function \(L::Legendre, WS::WeightedBasis{Bool,JacobiWeight{Bool},Jacobi{Bool}})
    w,S = WS.args
    if w.b && w.a
        @assert S.b && S.a
        _BandedMatrix(Vcat(((2:2:∞)./(3:2:∞))', Zeros(1,∞), (-(2:2:∞)./(3:2:∞))'), ∞, 2,0)
    elseif w.b && !w.a
        @assert S.b && !S.a
        _BandedMatrix(Ones{eltype(L)}(2,∞), ∞, 1,0)
    elseif !w.b && w.a
        @assert !S.b && S.a
        _BandedMatrix(Vcat(Ones{eltype(L)}(1,∞),-Ones{eltype(L)}(1,∞)), ∞, 1,0)
    else
        error("Not implemented")
    end
end

@simplify function *(St::QuasiAdjoint{Bool,Jacobi{Bool}}, WS::WeightedBasis{Int,JacobiWeight{Int},Jacobi{Bool}})
    w = parent(W)
    (w.b == 2 && S.b && w.a == 2 && S.a && parent(St) == S) || throw(ArgumentError())
    W_sqrt = Diagonal(JacobiWeight(true,true))
    L = Legendre()
    A = PInv(L)*W_sqrt*S
    A'*(L'L)*A
end

###
# Splines
###

function \(A::Legendre, B::HeavisideSpline)
    @assert B.points == -1:2:1
    Vcat(1, Zeros(∞,1))
end

###
# sum
###

function _sum(P::Legendre{T}, dims) where T
    @assert dims == 1
    Hcat(convert(T, 2), Zeros{T}(1,∞))
end

_sum(p::SubQuasiArray{T,1,Legendre{T},<:Tuple{Inclusion,Int}}, ::Colon) where T = parentindices(p)[2] == 1 ? convert(T, 2) : zero(T)


function solve(filename)
    #load txt file
    values = readlines(filename*".txt")
    count0 = 0
    count1 = 0
    γ = ""
    ϵ = ""
    for i in 1:length(values[1])
        for j in 1:length(values)
            if values[j][i] == '0'
                count0 += 1
            else
                count1 += 1
            end
        end
        if count0 > count1
            γ = γ * "0"
            ϵ = ϵ * "1"
        else
            γ = γ * "1"
            ϵ = ϵ * "0"
        end
        count0 = 0
        count1 = 0
    end
    println(γ)
    println(parse(Int, γ, base = 2))
    γ = parse(Int, γ, base = 2)
    println(ϵ)
    println(parse(Int, ϵ, base = 2))
    ϵ = parse(Int, ϵ, base = 2)
    println("Result: "*string(γ * ϵ))
end

solve("smalltest")
solve("input")

function solve2(filename)
    #load txt file
    values = readlines(filename*".txt")
    count0 = 0
    count1 = 0
    count0_co2 = 0
    count1_co2 = 0
    currMostCommon = '0'
    currLeastCommon = '0'
    values_CO2 = copy(values)
    for i in 1:length(values[1])
        #println("i: "*string(i))
        #println("size of values: "*string(length(values)))
        #println("size of CO2vals: "*string(length(values_CO2)))
        if length(values) > 1 || length(values_CO2) > 1
            #compute most common bit
            for j in 1:length(values)
                if values[j][i] == '0'
                    count0 += 1
                else
                    count1 += 1
                end
            end
            for j in 1:length(values_CO2)
                if values_CO2[j][i] == '0'
                    count0_co2 += 1
                else
                    count1_co2 += 1
                end
            end

            #filter out all non-fitting entries
            if length(values) > 1
                if count0 > count1
                    currMostCommon = '0'
                else
                    currMostCommon = '1'
                end
                values = copy(filter!(x -> x[i] == currMostCommon, values))
                count0 = 0
                count1 = 0
                #println(values)
            end
            if length(values_CO2) > 1
                if count0_co2 > count1_co2
                    currLeastCommon = '1'
                else
                    currLeastCommon = '0'
                end
                values_CO2 = copy(filter!(x -> x[i] == currLeastCommon, values_CO2))
                count0_co2 = 0
                count1_co2 = 0
                #println(values_CO2)
            end
        else
            break
        end
    end
    println("Oxygen: "*string(parse(Int, values[1], base = 2)))
    println(values[1])
    println("CO2: "*string(parse(Int, values_CO2[1], base = 2)))
    println(values_CO2[1])
    println("Product: "*string(parse(Int, values[1], base = 2) * parse(Int, values_CO2[1], base = 2)))
end

solve2("smalltest")
solve2("input")
#solve2("addtest")


myParse(x) = parse.(Int, x)
heightmap = Matrix(hcat(myParse.(collect.(readlines("input.txt")))...)')

nRows, nCols = size(heightmap)

function surrounding(row, col)
    arr = [(row - 1, col), (row, col - 1), (row + 1, col), (row, col + 1)]
    return [heightmap[x, y] for (x, y) in arr if x in 1:nRows && y in 1:nCols]
end

function coordfromIndex(index)
    row = ((index - 1) % nRows) + 1
    col = ((index - 1) ÷ nRows) + 1
    return row, col
end

function indexfromCoord(row, col)
    return nRows * (col - 1) + row
end

function surrounding(index)
    row, col = coordfromIndex(index)
    return surrounding(row, col)
end

function surroundingIndices(index)
    row, col = coordfromIndex(index)
    arr = [(row - 1, col), (row, col - 1), (row + 1, col), (row, col + 1)]
    return [indexfromCoord(x, y) for (x, y) in arr if x in 1:nRows && y in 1:nCols]
end

#! Part 1

function checkLower(height, index)
    val = min(surrounding(index)...)
    return height < val
end

function sumP1()
    sum = 0

    for (i, height) in enumerate(heightmap)
        if checkLower(height, i)
            sum += height + 1
        end
    end

    return sum
end

sumP1()

#! Part 2

function getProductPt2()
    visited = [heightmap[i] == 9 ? true : false for i in 1:length(heightmap)]

    basinSizes = []

    function basinSearch(index, basin = [])
        if !visited[index]
            visited[index] = true
            push!(basin, index)
            for i in surroundingIndices(index)
                basinSearch(i, basin)
            end
        end
        return length(basin)
    end

    for (i, _) in enumerate(heightmap)
        size = basinSearch(i)
        if size != 0
            push!(basinSizes, size)
        end
    end

    sort!(basinSizes, rev=true)

    return prod(basinSizes[1:3])

end

getProductPt2()


@doc raw"""
    HardCoreGraph{T<:Integer} <: AbstractGraphPointProcess{T}

Concrete type representing a point process on the vertices of a `graph` ``=(V, E)`` parametrized by ``\beta \geq 0`` which characterizes the distribution on the [independent sets](https://en.wikipedia.org/wiki/Independent_set_(graph_theory)) of `graph`, where each vertex is present with marginal probability ``\frac{\beta}{1+\beta}``.

In other words, it can also be viewed as the product distribution ``\operatorname{Bernoulli}(\frac{\beta}{1+\beta})^{\otimes |V|}`` on the vertices of `graph` conditioned on forming an independent set,

```math
    \mathbb{P}\!\left[ \mathcal{X} = X \right]
    \propto
    \prod_{x\in X}
        \frac{\beta}{1+\beta}
     1_{X \text{forms an independent set}}.
```

**See also**

- Section 7.2 of [GuJeLi19](@cite),
- Example 4.1 of [MoKr20](@cite),
- [`PRS.HardCorePointProcess`](@ref), the spatial counterpart of [`PRS.HardCoreGraph`](@ref).

# Example

A realization from a ``5\times 5`` grid graph, with ``\beta = 0.5``.

![assets/hard_core_graph.png](assets/hard_core_graph.png)
"""
struct HardCoreGraph{T<:Integer} <: AbstractGraphPointProcess{T}
    "Graph"
    graph::LG.SimpleGraph{T}
    β::Float64
end

function Base.show(io::IO, pp::HardCoreGraph{T}) where {T}
    print(io, "HardCoreGraph{$T}\n- graph = $(pp.graph)\n- β = $(pp.β)")
end

"""
    HardCoreGraph(
        graph::LG.SimpleGraph{T},
        β::Real
    ) where {T<:Integer}

Construct a [`PRS.HardCoreGraph`](@ref).

```jldoctest; output = true
using PartialRejectionSampling
using LightGraphs; const LG = LightGraphs

g, β = LG.grid([5, 5]), 1
PRS.HardCoreGraph(g, β)

# output

HardCoreGraph{Int64}
- graph = {25, 40} undirected simple Int64 graph
- β = 1.0
```
"""
function HardCoreGraph(
    graph::LG.SimpleGraph{T},
    β::Real
) where {T<:Integer}
    @assert β >= 0
    return HardCoreGraph{T}(graph, β)
end

# Sampling

"""
    generate_sample(
        [rng=Random.AbstractRNG,]
        pp::HardCoreGraph{T}
    )::Vector{T} where {T}

Generate an exact sample from the [`PRS.SinkFreeGraph`](@ref).

Default sampler is [`PRS.generate_sample_prs`](@ref).
"""
function generate_sample(
    rng::Random.AbstractRNG,
    pp::HardCoreGraph{T}
)::Vector{T} where {T}
    return generate_sample_prs(rng, pp)
end

# generate_sample(pp::HardCoreGraph) = generate_sample(Random.default_rng(), pp)

## Partial Rejeciton Sampling (PRS)

"""
    generate_sample_prs(
        [rng=Random.AbstractRNG,]
        pp::HardCoreGraph{T}
    )::Vector{T} where {T}

Sample from [`PRS.HardCoreGraph`](@ref) using Partial Rejection Sampling (PRS), see Section 7.2 of [GuJeLi19](@cite)

**See also**

- Example 4.1 of [MoKr20](@ref).

# Example

An illustration of the procedure on a ``5\\times 5`` grid graph.

- red: variables involved in constraints violation (gray neighboring nodes)
- orange: variables to be resampled (red nodes and their neighborhood)

![assets/hard_core_graph.gif](assets/hard_core_graph.gif)
"""
function generate_sample_prs(
    rng::Random.AbstractRNG,
    pp::HardCoreGraph{T}
)::Vector{T} where {T}

    β_max = 1 / (2 * sqrt(exp(1)) * LG.Δ(pp.graph) - 1)
    if pp.β > β_max
        @warn "The arguments do not satisfy Theorem 35 of Guo, Jerrum and Liu (2019).\nPartial rejection sampling may not be efficient.\nGiven `graph`, consider choosing β ≤ $(β_max)"
    end

    proba = pp.β / (one(pp.β) + pp.β)

    adj = LG.adjacency_matrix(pp.graph)
    occupied = randsubseq(rng, LG.vertices(pp.graph), proba)
    while true
        # Check if occupied vertices form an independent set
        sub_graph = LG.SimpleGraph(adj[occupied, occupied])
        LG.ne(sub_graph) == 0 && break

        independent, resample = T[], Set{T}()
        for cc in LG.connected_components(sub_graph)
            if length(cc) == 1  # Identify current independent vertices
                append!(independent, occupied[cc])
            else  # Construct the resampling set of vertices
                union!(resample, occupied[cc])
                for v in cc
                    union!(resample, LG.neighbors(pp.graph, occupied[v]))
                end
            end
        end
        randsubseq!(rng, occupied, collect(resample), proba)
        append!(occupied, independent)
    end
    return occupied
end

# generate_sample_prs(pp::HardCoreGraph) = generate_sample_prs(Random.default_rng(), pp)


"""
    Domain

Abstract type for different domains.
"""
abstract type Domain end

"""
    ClosedIntervalDomain

Abstract type for interval domains.
"""
abstract type AbstractIntervalDomain <: Domain end

"""
    ClosedIntervalDomain(left, right)

Interval domain.
"""
struct ClosedIntervalDomain{T} <: AbstractIntervalDomain
    left::T
    right::T
end

"""
    PeriodicIntervalDomain(left, right)

Periodic interval domain.
"""
struct PeriodicIntervalDomain{T} <: AbstractIntervalDomain
    left::T
    right::T
end
PeriodicIntervalDomain(a, b) = PeriodicIntervalDomain(promote(a, b)...)


struct Satellite <: LieGroupModel
    J::Diagonal{Float64,SVector{3,Float64}}
end

Satellite() = Satellite(Diagonal(@SVector ones(3)))

RobotDynamics.control_dim(::Satellite) = 3
Base.position(::Satellite, x::SVector) = @SVector zeros(3)
RobotDynamics.orientation(::Satellite, x::SVector) = UnitQuaternion(x[4], x[5], x[6], x[7])

RobotDynamics.LieState(::Satellite) = RobotDynamics.LieState(UnitQuaternion{Float64}, (3,0))

function RobotDynamics.dynamics(model::Satellite, x::SVector, u::SVector)
    ω = @SVector [x[1], x[2], x[3]]
    q = normalize(@SVector [x[4], x[5], x[6], x[7]])
    J = model.J

    ωdot = J\(u - ω × (J*ω))
    qdot = 0.5*lmult(q)*hmat()*ω
    return [ωdot; qdot]
end

model = Satellite()
@test LieState(model) === RobotDynamics.QuatState(7, (4,))
@test state_dim(model) == 7
@test state_diff_size(model) == 6

x,u = rand(model)
s = LieState(model)
@test all(RobotDynamics.vec_states(s, x) .≈ (x[1:3],Float64[]))
@test all(RobotDynamics.rot_states(s, x) .≈ (UnitQuaternion(x[4:7]),))
@test all(RobotDynamics.vec_states(model, x) .≈ (x[1:3],Float64[]))
@test all(RobotDynamics.rot_states(model, x) .≈ (UnitQuaternion(x[4:7]),))

x2 = rand(s)
@test norm(x[4:7]) ≈ 1

dx = RobotDynamics.state_diff(model, x, x2)
@test length(dx) == 6
@test dx ≈ [x[1:3] - x2[1:3]; UnitQuaternion(x[4:7]) ⊖ UnitQuaternion(x2[4:7])]


type AtomicGrid{T,A<:AbstractMatrix,B<:AbstractMatrix,C<:AbstractGrid1d} <: AbstractGrid3d{T,2}
    c::A  # coords
    w::B  # weights
    r::C  # radial grid
    Ω::AngularGrid{T}  # angular grid
    R::Point3{T}  # origin of spherical coord system
    function AtomicGrid(c::AbstractMatrix{Point3{T}},w::AbstractMatrix{T},
            r::AbstractGrid1d{T,1},Ω::AngularGrid{T},R::Point3{T})
        @assert size(c) == size(w)
        new(c,w,r,Ω,R)
    end
end

AtomicGrid{T}(c::AbstractMatrix{Point3{T}},w::AbstractMatrix{T},
        r::AbstractGrid1d{T,1},Ω::AngularGrid{T},R::Point3{T}) =
    AtomicGrid{eltype(w),typeof(c),typeof(w),typeof(r)}(c,w,r,Ω,R)

function atomicgrid{T}(r::AbstractGrid1d{T,1}, Ω::AngularGrid{T}; R::Point3{Float64}=(0.,0.,0.))
    shp = (length(Ω),length(r))
    c = Array{Point3{Float64}}(shp)
    w = Array{Float64}(shp)
    for j in 1:length(r)
        for i in 1:length(Ω)
            c[i,j] = r[j]*Ω[i] + R
            w[i,j] = r[j]^2*r.w[j]*Ω.w[i]
        end
    end
    AtomicGrid(c,w,r,Ω,R)
end

# integrate_angular{T}(g::AtomicGrid{T}, f::AbstractMatrix{T}) = At_mul_B(f,g.Ω.w)

# integrate_angular{T}(g::AtomicGrid{T}, f::AbstractVector{T}) = reshape(f,size(g.w))'*g.Ω.w

# function integrate_angular{T}(g::AtomicGrid{T}, f::AbstractArray{T})
#     shp = (size(g.w,1),prod(size(f)[2:end]))
#     reshape(f,shp)'*g.Ω.w
# end

# type AtomicGrid{T,N,Tc<:AbstractArray,Tw<:AbstractArray,Tr<:AbstractGrid1d} <: AbstractGrid3d{T,N}
#     c::Tc  # coords
#     w::Tw  # weights
#     r::Tc  # radial grid
#     Ω::AngularGrid{T}  # angular grid
#     R::Point3{T}  # origin of spherical coord system
#     function AtomicGrid(c::AbstractArray{Point3{T},N},
#                         w::AbstractArray{T,N},
#                         r::AbstractGrid1d{T,1},
#                         Ω::AngularGrid{T},
#                         R::Point3{T})
#         @assert size(c) == size(w)
#         new(c,w,r,Ω,R)
#     end
# end

# function AtomicGrid{T,N}(c::AbstractArray{Point3{T},N},
#                          w::AbstractArray{T,N},
#                          r::AbstractGrid1d{T,1},
#                          Ω::AngularGrid{T},
#                          R::Point3{T})
#     AtomicGrid{eltype(w),N,typeof(c),typeof(w),typeof(r)}(c,w,r,Ω,R)
# end


using RPOMDPToolbox
using RPOMDPModels
using RobustValueIteration
using Base.Test
const RPBVI = RobustValueIteration

unc = 0.1
prob = SimpleBaby2RPOMDP(-5.0, -10.0, 0.9, unc)
b0 = DiscreteBelief(prob, [0.5, 0.5])
@test pdf(b0,:hungry) == 0.5
@test pdf(b0,:full) == 0.5

# testing updater initialization
bs = [[b, 1-b] for b in 0.0:0.05:1.0]
solver = RPBVISolver(beliefpoints = bs, max_iterations = 100)
policy = RPBVI.solve(solver, prob)
up = RobustUpdater(prob, policy.alphas)
isd = initial_state_distribution(prob)
b4 = initialize_belief(up, isd)
@test pdf(b4, :hungry) == pdf(isd, :hungry)
@test pdf(b4, :full) == pdf(isd, :full)

# testing update function; if we feed baby, it won't be hungry
a = :feed
o = :crying
b4p = update(up, b4, a, o)
@show b4p.b
@test pdf(b4p,:hungry) ≈ 0.0 atol = 1e-3
@test pdf(b4p,:full) ≈ 1.0 atol = 1e-3

a = :nothing
o = :crying
b4p = update(up, b4, a, o)
@show b4p.b

@test isapprox(pdf(b4p,:hungry), 0.470588, atol=1e-4)
@test isapprox(pdf(b4p,:full), 0.52941, atol=1e-4)


function koopman(g,prob,u0,p,args...;kwargs...)
    g(solve(remake(prob,u0=u0,p=p),args...;kwargs...))
end

function koopman_expectation(g,u0s,ps,prob,args...;maxiters=0,
                      batch = 0,
                      quadalg = HCubatureJL(),
                      ireltol = 1e-2, iabstol=1e-2,kwargs...)
  n = length(u0s)
  if batch == 0
    _f = function (x,p)
      u0 = x[1:n]
      p  = x[n+1:end]
      k = koopman(g,prob,u0,p,args...;kwargs...)
      w = prod(pdf(a,b) for (a,b) in zip(u0s,u0))*
          prod(pdf(a,b) for (a,b) in zip(ps,p))
      k*w
    end
  else
    _f = function (dx,x,p)
      trajectories = size(x,2)
      prob_func = (prob,i,repeat) -> remake(prob,u0=@view(x[1:n,i]),
                                                 p=@view(x[n+1:end,i]))
      output_func = function (sol,i)
        k = g(sol)
        u0= @view(x[1:n,i])
        p = @view(x[n+1:end,i])
        w = prod(pdf(a,b) for (a,b) in zip(u0s,u0))*
            prod(pdf(a,b) for (a,b) in zip(ps,p))
        k*w,false
      end

      ensembleprob = EnsembleProblem(prob,prob_func=prob_func,
                                     output_func = output_func)
      sol = solve(ensembleprob,args...;trajectories=trajectories,kwargs...)
      dx .= vec(sol.u)
      nothing
    end
  end
  xs = [u0s;ps]
  intprob = QuadratureProblem(_f,minimum.(xs),maximum.(xs),batch=batch)
  sol = solve(intprob,quadalg,reltol=ireltol,
              abstol=iabstol,maxiters = maxiters)
end

function montecarlo_expectation(g,u0s,ps,prob,args...;trajectories,kwargs...)
  prob_func = function (prob,i,repeat)
    remake(prob,u0=rand.(u0s),p=rand.(ps))
  end
  output_func = (sol,i) -> (g(sol),false)
  monte_prob = EnsembleProblem(prob;
                                 output_func = output_func,
                                 prob_func = prob_func)
  mean(solve(monte_prob,args...;trajectories=trajectories,kwargs...).u)
end


include("smote_exs.jl")
include("ub_smote.jl")
include("Boosting.jl")
include("simtools.jl")

# running single simulation
@time runsim(500, 20, 2.0, 0.1, seed = 111, ξ = 0.5)

# running multiple simulations
simresults = run_simulations(50, 500, 10, -2.5, 0.05, ξ = 0.6)
summarise_sims(simresults, 5)

n = 100
p = 10
# X = hcat(ones(n), randn(n, p-1))

Σ = [1.0 0.8 0.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0;
     0.8 1.0 0.6 0.2 0.0 0.0 0.0 0.0 0.0 0.0;
     0.4 0.6 1.0 0.6 0.4 0.0 0.0 0.0 0.0 0.0;
     0.0 0.2 0.6 1.0 0.8 0.0 0.0 0.0 0.0 0.0;
     0.0 0.0 0.4 0.8 1.0 0.0 0.0 0.0 0.0 0.0;
     0.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0;
     0.0 0.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0;
     0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0;
     0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0 0.0;
     0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0]

mvn = MvNormal(ones(p), Σ)

srand(round(Int, time()))

X = rand(mvn, n)'
cor(X)

ϵ = rand(Normal(2, 0.5), n)                                  # gives 10% minority class,
# ϵ = rand(Normal(-2.5, 0.5), n)                               # gives 5% minority class,
# ϵ = rand(Normal(4, 0.5), n)                                  # gives 15% minority class,

β = [-3, -4, 1, -5, -4, 0.0, 0.0, 0.0, 0.0, 0.0]
η = X*β + ϵ                                                    # linear predictor w/ error
pr = 1.0 ./ (1.0 + exp(-η))                                    # inv-logit

# simulate outcome variable
y = map(π -> rand(Binomial(1, π)), pr)
y[y .== 0] = -1

# y2, X2 = get_subset(y, X, 450, 0.10)
srand(111)
model, coeffs = build_adaboost_stumps(y, X, 50, 1.0);          # 50 boosting rounds, "sub-sample" 100%

# apply learned model
y_hat = apply_adaboost_stumps(model, coeffs, X)
mean(y .== y_hat)

mean(y)

# get the probability of each label
# apply_adaboost_stumps_proba(model, coeffs, [5.9, 3.0, 5.1, 1.9, 1.0, 2.1], [0, 1])

srand(111)
model, coeffs2 = build_adaboost_stumps(y, X, 50, 1.0, 0.2);

y_hat = apply_adaboost_stumps(model, coeffs2, X)
mean(y .== y_hat)
1 - mean(y)

# apply_adaboost_stumps_proba(model, coeffs, [5.9, 3.0, 5.1, 1.9, 1.0, 2.1], [0, 1])

find(y)

# run n-fold cross validation for boosted stumps, using 7 iterations and 3 folds
# accuracy = nfoldCV_stumps(y, X, 7, 3)

using DataFrames

ds = readtable("example_data.csv");
y = convert(Vector, ds[:, 1]);
X = convert(Array, ds[:, 2:end]);

using DecisionTree
using Distributions

n = 500
p = 4
# X = hcat(ones(n), randn(n, p-1))

Σ = [1.0 0.6 0.2 0.0;
          0.6 1.0 0.6 0.2;
          0.2 0.6 1.0 0.6;
          0.0 0.2 0.6 1.0]

mvn = MvNormal([1.0, -1.0, -3.0, 3.0], Σ)

X = rand(mvn, n)'
cor(X)

β = vcat([1.5, 1.1, -10.5, -100.5])
η = X*β                                                      # linear predictor
pr = 1.0 ./ (1.0 + exp(-η))                                  # inv-logit

# simulate outcome variable
y = map(π -> rand(Binomial(1, π)), pr)
mean(y)

srand(111)
model, coeffs = DecisionTree.build_adaboost_stumps(y, X, 50);

# apply learned model
y_hat = DecisionTree.apply_adaboost_stumps(model, coeffs, X)
mean(y .== y_hat)
mean(y)

# get the probability of each label
apply_adaboost_stumps_proba(model, coeffs, [5.9, 3.0, 5.1, 1.9], [0, 1])


using DoubleDouble

n = 1_000_000
X = 64*rand(n)
Y = 64*rand(n)

Zbig = Array(Float64,n)
Zdouble = Array(Float64,n)

function bigrem!(X,Y,Z)
    for i = 1:length(Z)
        Z[i] = float64(rem(big(X[i])*big(Y[i]),1.0))
    end
end

function doublerem!(X,Y,Z)
    for i = 1:length(Z)
        Z[i] = float64(rem(Single(X[i])*Single(Y[i]),1.0))
    end
end

bigrem!(X,Y,Zbig)
doublerem!(X,Y,Zdouble)

X = 64*rand(n)
Y = 64*rand(n)

@time bigrem!(X,Y,Zbig)
@time doublerem!(X,Y,Zdouble)
@assert Zbig == Zdouble


export TransposeOperator

type TransposeOperator{T<:Number,B<:Operator} <: Operator{T}
    op::B
end

TransposeOperator{T<:Number}(B::Operator{T})=TransposeOperator{T,typeof(B)}(B)

Base.convert{T}(::Type{Operator{T}},A::TransposeOperator)=TransposeOperator(convert(Operator{T},A.op))

domainspace(P::TransposeOperator)=rangespace(P.op)
rangespace(P::TransposeOperator)=domainspace(P.op)

domain(P::TransposeOperator)=domain(P.op)

bandinds(P::TransposeOperator)=-bandinds(P.op)[end],-bandinds(P.op)[1]

getindex(P::TransposeOperator,k::Integer,j::Integer) =
    P.op[j,k]

function Base.convert{T,TO<:TransposeOperator}(::Type{BandedMatrix},S::SubOperator{T,TO})
    kr,jr=parentindexes(S)
    BandedMatrix(view(parent(S).op,jr,kr)).'
end

Base.transpose(A::Operator)=TransposeOperator(A)
Base.ctranspose{T<:Real}(A::Operator{T})=TransposeOperator(A)


@testset "Static graphs" begin
    g = @inferred(CompleteDiGraph(5))
    @test nv(g) == 5 && ne(g) == 20
    g = @inferred(CompleteGraph(5))
    @test nv(g) == 5 && ne(g) == 10

    g = @inferred(CompleteBipartiteGraph(5, 8))
    @test nv(g) == 13 && ne(g) == 40

    g = @inferred(StarDiGraph(5))
    @test nv(g) == 5 && ne(g) == 4
    g = @inferred(StarGraph(5))
    @test nv(g) == 5 && ne(g) == 4
    g = @inferred(StarGraph(1))
    @test nv(g) == 1 && ne(g) == 0

    g = @inferred(PathDiGraph(5))
    @test nv(g) == 5 && ne(g) == 4
    g = @inferred(PathGraph(5))
    @test nv(g) == 5 && ne(g) == 4

    g = @inferred(CycleDiGraph(5))
    @test nv(g) == 5 && ne(g) == 5
    g = @inferred(CycleGraph(5))
    @test nv(g) == 5 && ne(g) == 5

    g = @inferred(WheelDiGraph(5))
    @test nv(g) == 5 && ne(g) == 8
    g = @inferred(WheelGraph(5))
    @test nv(g) == 5 && ne(g) == 8

    g = @inferred(Grid([3, 3, 4]))
    @test nv(g) == 3 * 3 * 4
    @test ne(g) == 75
    @test Δ(g) == 6
    @test δ(g) == 3

    g = @inferred(Grid([3, 3, 4], periodic=true))
    @test nv(g) == 3 * 3 * 4
    @test ne(g) == 108
    @test Δ(g) == 6
    @test δ(g) == 6

    g = @inferred(CliqueGraph(3, 5))
    @test nv(g) == 15 && ne(g) == 20
    @test g[1:3] == CompleteGraph(3)

    g = @inferred(crosspath(3, BinaryTree(2)))
    # f = Vector{Vector{Int}}[[2 3 4];
    # [1 5];
    # [1 6];
    # [1 5 6 7];
    # [2 4 8];
    # [3 4 9];
    # [4 8 9];
    # [5 7];
    # [6 7]
    # ]
    I = [1, 1, 1, 2, 2, 3, 3, 4, 4, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 7, 8, 8, 9, 9]
    J = [2, 3, 4, 1, 5, 1, 6, 1, 5, 6, 7, 2, 4, 8, 3, 4, 9, 4, 8, 9, 5, 7, 6, 7]
    V = ones(Int, length(I))
    Adj = sparse(I, J, V)
    @test Adj == sparse(g)

    g = @inferred(DoubleBinaryTree(3))
    # [[3, 2, 8]
    # [4, 1, 5]
    # [1, 6, 7]
    # [2]
    # [2]
    # [3]
    # [3]
    # [10, 9, 1]
    # [11, 8, 12]
    # [8, 13, 14]
    # [9]
    # [9]
    # [10]
    # [10]]
    I = [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 5, 6, 7, 8, 8, 8, 9, 9, 9, 10, 10, 10, 11, 12, 13, 14]
    J = [3, 2, 8, 4, 1, 5, 1, 6, 7, 2, 2, 3, 3, 10, 9, 1, 11, 8, 12, 8, 13, 14, 9, 9, 10, 10]
    V = ones(Int, length(I))
    Adj = sparse(I, J, V)
    @test Adj == sparse(g)

    rg3 = @inferred(RoachGraph(3))
    # [3]
    # [4]
    # [1, 5]
    # [2, 6]
    # [3, 7]
    # [4, 8]
    # [9, 8, 5]
    # [10, 7, 6]
    # [11, 10, 7]
    # [8, 9, 12]
    # [9, 12]
    # [10, 11]
    I = [1, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 7, 8, 8, 8, 9, 9, 9, 10, 10, 10, 11, 11, 12, 12]
    J = [3, 4, 1, 5, 2, 6, 3, 7, 4, 8, 9, 8, 5, 10, 7, 6, 11, 10, 7, 8, 9, 12, 9, 12, 10, 11]
    V = ones(Int, length(I))
    Adj = sparse(I, J, V)
    @test Adj == sparse(rg3)
end


using LASindex
using FileIO
using RegionTrees
using StaticArrays
using Test

workdir = dirname(@__FILE__)

# source: http://www.liblas.org/samples/ and lasindex
index = joinpath(workdir, "..", "sample", "Palm Beach Post Hurricane.lax")

# Header check
qt = load(index)

# Determine correct quadtree levels
@test LASindex.zlevels(0) == [0]
@test LASindex.zlevels(1) == [1]
@test LASindex.zlevels(4) == [4]
@test LASindex.zlevels(5) == [1,1]
@test LASindex.zlevels(20) == [4,4]
@test LASindex.zlevels(21) == [1,1,1]
@test LASindex.zlevels(22) == [2,1,1]
@test LASindex.zlevels(84) == [4,4,4]
@test LASindex.zlevels(85) == [1,1,1,1]

# Single integer indexing on RegionTrees cells
# based on Z order
root = Cell(SVector(0., 0), SVector(1., 1))
split!(root)
@test root[1] == root[1,1]
@test root[2] == root[2,1]
@test root[3] == root[1,2]
@test root[4] == root[2,2]
@test root[0] == root
@test_throws ErrorException root[5]

# Split cell based on given index level
# until cell with index is found, return it
root = Cell(SVector(0., 0), SVector(1., 1))
@test LASindex.quadtree!(root, 1).boundary.origin ≈ Cell(SVector(0., 0), SVector(0.5, 0.5)).boundary.origin
@test LASindex.quadtree!(root, 4).boundary.origin ≈ Cell(SVector(0.5, 0.5), SVector(0.5, 0.5)).boundary.origin
@test LASindex.quadtree!(root, 5).boundary.origin ≈ Cell(SVector(0., 0), SVector(0.25, 0.25)).boundary.origin
@test LASindex.quadtree!(root, 20).boundary.origin ≈ Cell(SVector(0.75, 0.75), SVector(0.25, 0.25)).boundary.origin
@test LASindex.quadtree!(root, 1).boundary.widths ≈ Cell(SVector(0., 0), SVector(0.5, 0.5)).boundary.widths
@test LASindex.quadtree!(root, 4).boundary.widths ≈ Cell(SVector(0.5, 0.5), SVector(0.5, 0.5)).boundary.widths
@test LASindex.quadtree!(root, 5).boundary.widths ≈ Cell(SVector(0., 0), SVector(0.25, 0.25)).boundary.widths
@test LASindex.quadtree!(root, 20).boundary.widths ≈ Cell(SVector(0.75, 0.75), SVector(0.25, 0.25)).boundary.widths

# Adding data to quadtree cell based on index
root = Cell(SVector(0., 0), SVector(1., 1), Vector{UnitRange{Int}}())
LASindex.quadtree!(root, 1, [2:5])
@test root[1,1].data == [2:5]

# Merge UnitRanges that do have overlap
r = Vector{UnitRange{Integer}}([1:5, 2:4, 12:13, 14:16])
@test LASindex.merge(r) == Vector{UnitRange{Integer}}([1:5, 12:16])
# and again to make sure it's not mutating
@test LASindex.merge(r) == Vector{UnitRange{Integer}}([1:5, 12:16])
r = Vector{UnitRange{Integer}}([1:5, 2:6, 12:13, 14:16, -1:2])
@test LASindex.merge(r) == Vector{UnitRange{Integer}}([-1:6, 12:16])
r = Vector{UnitRange{Integer}}()
@test LASindex.merge(r) == Vector{UnitRange{Integer}}()
r = Vector{UnitRange{Integer}}([1:5])
@test LASindex.merge(r) == Vector{UnitRange{Integer}}([1:5])

# Test spatial intersection
bbox = SVector(955000.0, 955002.0, 885000.0, 988800.0)  # small intersection
@test LASindex.intersect(qt, bbox) == Vector{UnitRange{Integer}}()
bbox = SVector(955000.0, 985002.0, 885000.0, 988800.0)  # fully overlaps
@test LASindex.intersect(qt, bbox) == Vector{UnitRange{Integer}}([1:1924631])


# %%
using Pkg
Pkg.activate("./GraphFusedElasticNet.jl/")
Pkg.instantiate()
using GraphFusedElasticNet
using LightGraphs
using DelimitedFiles
using JSON
using DataStructures

# %%
fname ="./processed_data/spatiotemporal_graph.csv"
edgeinfo = readdlm(fname, ',', Int, skipstart=1)
edgeinfo[1:10,:]

# %% use our utils from GFEN library for trail extraction
graph = graph_from_edgelist(edgeinfo[:,1:2], from_zero=true)
trails = find_trails(graph, ntrails=64)

# %% identify temporal edges
# create temporal edge dict
tmpdict = DefaultDict{Tuple{Int, Int}, Bool}(false)
cols = [edgeinfo[:,i] for i in 1:3]
for (k, (i, j, t)) in enumerate(zip(cols...))
    x, y = min(i, j), max(i, j)
    tmpdict[x, y] = t
end
# now tag along ptr
ptr = trails.ptr
istemporal = zeros(Int, length(ptr))
brks_ = OrderedSet(trails.brks)
for (k, (i, j)) in enumerate(zip(ptr[1:end-1], ptr[2:end]))
    if k + 1 ∉ brks_
        x, y = min(i, j), max(i, j)
        istemporal[k] = tmpdict[x, y]
    end
end

# %% save fields as json
data = Dict("ptr" => trails.ptr,
            "brks" => trails.brks,
            "wts" => trails.wts,
            "istemp" => istemporal,
            "num_nodes" => trails.num_nodes)
open("./processed_data/trails.json","w") do f
    JSON.print(f, data)
end