import AbstractAlgebra.Generic: FreeModuleElem

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

# Private functions (will eventually change these functions into internal/private functions)

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

@doc Markdown.doc"""
    projective_quadratic_triple(Q::MatrixElem{T}, L::MatrixElem{T}, c::T, k) -> Tuple

For 1 <= k <= length(L), this function returns a tuple, where the first item in the tuple is
a kxk (positive definite) symmetric matrix, the next entry is a column vector of length k
and the last entry is a rational number.

THEORY:
R, (x1,x2,x3) = PolynomialRing(QQ, ["x1", "x2", "x3"])
Given a Polynomial ring of n variables we proceed as follows:
Generate a list of projective quadratic triples of n-1, ..., 1 variables,
given an n-variabled quadratic triple QT=[Q,L,c], where Q is a nxn symmetric
matrix, L is a column vector of length n and c is a rational number.
NOTE: if k is an int < n then it returns the k-variabled
projective quadratic triple.
"""
function projective_quadratic_triple(Q::MatrixElem{T}, L::MatrixElem{T}, c::T, k) where T <: RingElem
  n = nrows(Q)
  r1 = Vector{T}(undef,n-1); m1 = Vector{T}(undef,n-1);
  Q1 = Vector{typeof(Q)}(undef,n-1); p1 = Vector{typeof(L)}(undef,n-1);
  L1 = Vector{typeof(L)}(undef,n-1);
  for i in 1:n-1
    r1[i] = Q[n+1-i, n+1-i]
    m1[i] = L[n+1-i,1:1][1]
    Q1[i] = Q[1:n-i, 1:n-i] #n-i symmetric matrix extracted from the matrix Q
    p1[i] = Q[1:n-i, n+1-i:n+1-i] #Q[ri:rj, ci:cj] returns enrties from a matrix in the respective rows r_ and columns c_
    L1[i] = L[1:n-i,1:1] #column vector of length n-i
  end
  #------------------------------------------------
  c1 = c; pQT = Vector{typeof(Q)}(undef,n-1); pLL = Vector{typeof(L)}(undef,n-1); cc1 = Vector{T}();
  for j in 1:length(Q1)
    pt1 = transpose(p1[j])
    pQT[j] = Q1[j] - r1[j]^-1 * (p1[j] * pt1)
    pLL[j] = L1[j] - ((m1[j] * r1[j]^-1) * p1[j])
    c1 = c1 - (m1[j]^2 * r1[j]^-1) #rational number
    push!(cc1, c1)
  end
  #------------------------------------------------
  if k < n
    return pQT[n-k], pLL[n-k], cc1[n-k]
  else
    return Q,L,c
  end
end

@doc Markdown.doc"""
    range_ellipsoid_dim1(Q::MatrixElem{T}, L::MatrixElem{T}, c::T) ->  StepRange{fmpz, fmpz}

Returns a finite set of integers for which the inhomogeneous quadratic function,
given a one variabled quadratic triple, is less than or equal to zero.
"""
function range_ellipsoid_dim1(Q::MatrixElem{T}, L::MatrixElem{T}, c::T) where T <: RingElem
  @assert nrows(Q) == 1
  @assert L[1]^2 - Q[1]*c >= 0
  #x1, x2 are roots of q s.t x1 <= x2
  #x1 = (-L[1] - sqrt(L[1]^2 - Q[1]*c)) / Q[1]
  #x2 = (-L[1] + sqrt(L[1]^2 - Q[1]*c)) / Q[1]
  # return the integers in the interval [x1,x2]
  sqrt_floor(a::fmpz) = isqrt(a)
  a = -L[1] // Q[1]; b = (L[1]^2 - Q[1]*c) // Q[1]^2;
  cc = lcm(denominator(a), denominator(b))
  aa = FlintZZ(a * cc); bb = FlintZZ(b * cc^2);
  l = round(fmpz, (aa - sqrt_floor(bb))//cc, RoundUp)
  r = round(fmpz, (aa + sqrt_floor(bb))//cc, RoundDown)
  return l:r
end

@doc Markdown.doc"""
    positive_quadratic_triple(a::fmpz, Q::MatrixElem{T}, L::MatrixElem{T}, c::T) -> Tuple

This function computes a n-1 variabled quadraric triple i*(a, QT), given an n>1
variabled projective quadratic triple QT and an integer a.
Return value: an n-1 variabled positive quadratic triple i*(a, QT)
"""
function positive_quadratic_triple(a::fmpz, Q::MatrixElem{T}, L::MatrixElem{T}, c::T) where T <: RingElem
  n = nrows(Q)
  @assert n > 1
  #------------------------------------------------
  Q1 = Q[2:n, 2:n] #n-1 symmetric matrix extracted from the matrix Q
  L1 = a * Q[2:n, 1:1] + L[2:n,1:1] #Q[ri:rj, ci:cj] -> enrties from a matrix in the respective rows r_ and columns c_, L[i:n,1:1] column vector of length n-i
  c1 = a^2 * Q[1, 1] + 2 * a * L[1] + c #rational number
  #------------------------------------------------
  return Q1, L1, c1 #n-1 variabled quadratic triple
end

@doc Markdown.doc"""
    positive_quadratic_triple2(aa::Vector{fmpz}, Q::MatrixElem{T},
                               L::MatrixElem{T}, c::T) -> Tuple

NOTE: aa is an array of length k.

This function computes an n-k variabled quadratic triple i*(aa, QT), given an n
variabled quadratic triple QT and a tuple aa = {a1, ..., ak}, where k < n.
Return value: an n-k variabled positive quadratic triple i*(aa, QT).
"""
function positive_quadratic_triple2(aa::Vector{fmpz}, Q::MatrixElem{T}, L::MatrixElem{T}, c::T) where T <: RingElem
  QT = Q, L, c
  k = length(aa)
  @req length(L) >= 2 "The input quadratic triple should be at least 2 variabled."
  @req k < length(L) "The first input parameter should have length one less than the third input parameter"
  for i in 1:k
    R = positive_quadratic_triple(aa[i], QT[1], QT[2], QT[3])
    QT = R[1], R[2], R[3]
  end
  return QT
end

@doc Markdown.doc"""
    list_positive_quadratic_triple(aa::Vector{fmpz}, Q::MatrixElem{T},
                                   L::MatrixElem{T}, c::T) -> Array

Computes a list (for v = 1) QT_m^1 == i*(aa, QT_m^0), for aa in E(QT_1^0)
"""
function list_positive_quadratic_triple(aa::Vector{fmpz}, Q::MatrixElem{T}, L::MatrixElem{T}, c::T) where T <: RingElem
  n = nrows(Q)
  v = length(aa)+1
  t = n-length(aa)
  ListI = Vector(undef, t)
  for m in v:n
    P = projective_quadratic_triple(Q, L, c, m) #QT_m^0
    ListI[m-v+1] = positive_quadratic_triple2(aa, P[1], P[2], P[3])
  end
  return ListI
end

@doc Markdown.doc"""
    list_positive_quadratic_triple2(b::fmpz, ListIv::Vector{Any}) -> Array

For a fixed v, this function computes a list of QT_m^v := i*(b, QT_m^{v-1}),
where b is in E(QT_v^{v-1}).
"""
function list_positive_quadratic_triple2(b::fmpz, ListIv::Vector{Any}) #ListIv is the list of quadratic triples QT_m^{v-1}
  L = Vector(undef, length(ListIv)-1);
  for i in 1:length(ListIv)-1
    L[i] = positive_quadratic_triple(b, ListIv[i+1][1], ListIv[i+1][2], ListIv[i+1][3])
  end
  return L
end

@doc Markdown.doc"""
    posQuadTriple_intVectors(QT::Vector{Any}, E::Vector{Any}) -> Tuple

INPUT:
  A list of quadratic triples QT a list of tuples E of the form {a_1, ..., a_k}
OUTPUT:
  A new list of quadratic triples in the one dimension higher, and an updated
  list E of the form {a_1, ..., a_k, a_{k+1}}
"""
function posQuadTriple_intVectors(QT::Vector{Any}, E::Vector{Any})
  # this function could be rewritten for better performance at some point
  QTT1 = [];
  QTT = Array{Array}(undef, length(QT))
  bbb = [];
  for j in 1:length(QT) #for a fixed v
    b = range_ellipsoid_dim1(QT[j][1][1], QT[j][1][2], QT[j][1][3]) #E(QT_v^{v-1}) = {b_1, ..., b_N}
    b1 = collect(b)
    QT1 = Array{Array}(undef, length(b1))
    for k in 1:length(b1)
      QT1[k] = list_positive_quadratic_triple2(b1[k], QT[j]) #QT_m^v for m = v+1, ..., n
      for kk in E
        vb = vcat(kk, b1[k])
        if !(vb in bbb)
          push!(bbb,vb) #bb is now a list of tuples [x_1, ..., x_{v-1},b[k]] that satisfiy a inhomogeneous quad function for a v variabled quad triple
        end
      end
    end
    QTT[j] = QT1
  end
  for j in 1:length(QT)
    append!(QTT1, QTT[j])
  end
  return QTT1, bbb
end

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

# Public functions

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

@doc Markdown.doc"""
    convert_type(G::MatrixElem{T}, K::MatrixElem{T}, d::T) -> Tuple{ZLat, MatrixElem{T}, T}
Where T is a concrete type, e.g. fmpz, fmpq, etc.
Converts a quadratic triple QT = [Q, K, d] to the input values required for closest vector problem (CVP).
"""
function convert_type(G::MatrixElem{T}, K::MatrixElem{T}, d::T) where T <: RingElem
  if G[1,1] > 0
    Q = G
  else
    Q = -G
  end
  vector = -inv(Q) * K
  upperbound = (-transpose(K) * inv(Q) * -K)[1] - d
  Lattice = Zlattice(gram = Q)
  return Lattice, vector, upperbound
end

@doc Markdown.doc"""
    convert_type(L::ZLat, v::MatrixElem{T}, c::T) -> Tuple{fmpq_mat, fmpq_mat, fmpq}

Where T is a concrete type, e.g. fmpz, fmpq, etc.
Converts the input values from closest vector enumeration (CVE) to the corresponding quadratic triple QT = [Q, K, d].
"""
function convert_type(L::ZLat, v::MatrixElem{T}, c::T) where T <: RingElem
  V = ambient_space(L)
  Q = gram_matrix(V)
  K = -Q*v
  v = vec([i for i in v])
  d = inner_product(V,v,v)-c
  return Q, K, d
end

@doc Markdown.doc"""
    closest_vectors(Q::MatrixElem{T}, L::MatrixElem{T}, c::T; equal::Bool=false, sorting::Bool=false)
                                                    -> Vector{Vector{fmpz}}

Return all the integer vectors `x` of length n such that the inhomogeneous
quadratic function `q_{QT}(x) := xQx + 2xL + c <= 0` corresponding to an n variabled
quadratic triple. If the optional argument ``equal = true``, it return
all vectors `x` such that `q_{QT}(x) = 0`. By default ``equal = false``.
If the argument ``sorting = true``, then we get a a list of sorted vectors.
The Default value for ``sorting`` is set to ``false``.
"""
function closest_vectors(G::MatrixElem{T}, L::MatrixElem{T}, c::T; equal::Bool=false, sorting::Bool=false) where T <: RingElem
  #1 < v <= n+1, a = [a_1, ..., a_{v-1}] int tuple &
  @req det(G) != 0 "The symmetric matrix is not definite."
  if G[1,1] > 0
    Q = G
  else
    Q = -G
  end
  n = nrows(Q)
  P = projective_quadratic_triple(Q, L, c, 1)
  a = range_ellipsoid_dim1(P[1], P[2], P[3])#E(QT_1^0)
  aa = collect(a)
  QT = Vector(undef, length(aa))
  for i in 1:length(aa)
    QT[i] = list_positive_quadratic_triple(fmpz[aa[i]], Q, L, c) #List  QT_m^1 i*(aa[i], QT_m^0) for m = 2,...,n. Here v = 1
  end
  QTT = []; bbb = []; EE = Array{Array{fmpz,1},1}();
  list_b1 = []
  for j in 1:length(QT) #v=2
    b = range_ellipsoid_dim1(QT[j][1][1], QT[j][1][2], QT[j][1][3])#E(QT_2^1) = {b_1, ..., b_N}
    b1 = collect(b)
    QT1 = Vector(undef, length(b1))
    for kk in aa
      for k in 1:length(b1)
        QT1[k] = list_positive_quadratic_triple2(b1[k], QT[j]) #QT_m^2 for m = 3, ..., n
        bb = [kk, b1[k]]
        if !(bb in bbb)
        push!(bbb,bb) #bbb is a list of tuples [a,b] that satisfiy a inhomogeneous quad function for a 2 variabled quad triple
        end
      end
    end
    QT[j] = QT1
  end
  for j in 1:length(QT)
    append!(QTT, QT[j])
  end
  QTT1 = QTT;
  E = bbb
  for v in 3:n
    QTT1, E = posQuadTriple_intVectors(QTT1, E)
  end
  if !equal
    for k in E
      t = length(k)
      if (transpose(matrix(QQ,t,1,k))*Q*matrix(QQ,t,1,k))[1] + (2*transpose(matrix(QQ,t,1,k))*L)[1] + c <= 0
        push!(EE, k)
      end
    end
    if !sorting
      return EE
    else
      return sort!(EE)
    end
  else
    for k in E
      t = length(k)
      if (transpose(matrix(QQ,t,1,k))*Q*matrix(QQ,t,1,k))[1] + (2*transpose(matrix(QQ,t,1,k))*L)[1] + c == 0
        push!(EE, k)
      end
    end
    if !sorting
      return EE
    else
      return sort!(EE)
    end
  end
end

@doc Markdown.doc"""
    closest_vectors(L:ZLat, v:MatrixElem{T}, upperbound::T; equal::Bool=false, sorting::Bool=false)
                                                    -> Vector{Vector{fmpz}}

Return all vectors `x` in `L` such that `b(v-x,v-x) <= c`, where `b` is the bilinear form on `L`.
If the optional argument ``equal = true`` then it return all vectors `x` in `L` such that `b(v-x,v-x) = c`.
By default ``equal = false``. If the argument ``sorting = true``, then we get a a list of sorted vectors.
The Default value for ``sorting`` is set to ``false``.
"""
function closest_vectors(L::ZLat, v::MatrixElem{T} , upperbound::T; equal::Bool=false, sorting::Bool=false) where T <: RingElem
  epsilon = QQ(1//10)   # some number > 0, not sure how it influences performance
  d = size(v)[1]
  if isdefinite(L) == false
    throw("Zlattice is indefinite.")
  end
  if rank(L) != d
    throw("Zlattice must have the same rank as the length of the vector in the second argument.")
  end
  g1 = gram_matrix(L)
  if g1[1,1] > 0
    G1 = g1
  else
    G1 = -g1
  end
  e = matrix(QQ, 1, 1, [upperbound//3+epsilon])
  G = diagonal_matrix(G1, e)
  B = diagonal_matrix(basis_matrix(L),matrix(QQ,1,1,[1]))
  for i in 1:d
    B[end,i] = -v[i]
  end
  N = Zlattice(B,gram=G)

  delta = QQ(4//3)*upperbound + epsilon
  sv = Hecke.short_vectors(N, delta)
  cv = Array{Array{fmpz,1},1}()
  for a in sv
    a = a[1]
    if a[end] == 0
    continue
    end
    if a[end] == -1
      a = -a
    end
    x = a[1:end-1]
    push!(cv, x)
  end
  V = ambient_space(L)
  v1 = Vector{T}()
  for i in v
    push!(v1, i)
  end
  cv2 = Array{Array{fmpz,1},1}()
  if !equal
    for x in cv
      t = x - v1
      dist = inner_product(V,t,t)
      if dist <= upperbound
          push!(cv2,x)
      end
    end
    if !sorting
      return cv2
    else
      return sort!(cv2)
    end
  else
    for x in cv
      t = x - v1
      dist = inner_product(V,t,t)
      if dist == upperbound
        push!(cv2,x)
      end
    end
    if !sorting
      return cv2
    else
      return sort!(cv2)
    end
  end
end


module straight_duct_examples
using FinEtools
using Gaston

function straight_duct_Q8_example()
    t0  =  time()

    rho = 1.21*phun("kg/m^3");# mass density
    c  = 343.0*phun("m/s");# sound speed
    bulk =  c^2*rho;
    omega =  54.5901*phun("rev/s")
    vn0 =  -1.0*phun("m/s")
    Lx = 10.0*phun("m");# length of the box, millimeters
    Ly = 1.0*phun("m"); # length of the box, millimeters
    n = 20;#number of elements along the length

    println("""

    Straight duct with anechoic termination.
    Example from Boundary element acoustics: Fundamentals and computer codes, TW Wu, page 44.
    Both real and imaginary components of the pressure should have amplitude of
    rho*c = $(rho*c).

    Two-dimensional model. Quadratic qquadrilateral mesh.
    """)

    fens,fes  =  Q8block(Lx,Ly,n,2); # Mesh
    bfes  =  meshboundary(fes)
    L0 = selectelem(fens,bfes,facing = true, direction = [-1.0 0.0])
    L10 = selectelem(fens,bfes,facing = true, direction = [+1.0 0.0])
    nLx = selectnode(fens,box = [0.0 Lx  0.0 0.0], inflate = Lx/1.0e5)

    geom  =  NodalField(fens.xyz)
    P  =  NodalField(zeros(ComplexF64,size(fens.xyz,1),1))

    numberdofs!(P)

    material = MatAcoustFluid(bulk,rho)
    femm  =  FEMMAcoust(IntegData(fes, GaussRule(2, 3)), material)

    S  =  acousticstiffness(femm, geom, P);
    C  =  acousticmass(femm, geom, P);

    E10femm  =  FEMMAcoustSurf(IntegData(subset(bfes,L10),GaussRule(1, 3)), material)
    D  =  acousticABC(E10femm, geom, P);

    E0femm  =  FEMMBase(IntegData(subset(bfes,L0), GaussRule(1,  3)))
    fi  =  ForceIntensity(-1.0im*omega*rho*vn0);
    F  =  distribloads(E0femm, geom, P, fi, 2);

    p = (-omega^2*S +omega*1.0im*D + C)\F
    scattersysvec!(P, p[:])

    println("Pressure amplitude bounds: ")
    println("  real $(minimum(real(P.values)))/$(maximum(real(P.values)))")
    println("  imag $(minimum(imag(P.values)))/$(maximum(imag(P.values)))")

    println("Total time elapsed  =  ",time() - t0,"s")

    File  =   "straight_duct_Q8.vtk"
    scalars = real(P.values);
    vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.Q8;  scalars = [("Pressure", scalars)])
    @async run(`"paraview.exe" $File`)

    # @pyimport matplotlib.pyplot as plt
    # plt.style[:use]("seaborn-whitegrid")
    # fig = plt.figure()
    # ax = plt.axes()
    # ix = sortperm(geom.values[nLx,1])
    # ax[:plot](geom.values[nLx,1][ix], real(P.values)[nLx][ix], marker=:o, color = :blue, label = "real")
    # ax[:plot](geom.values[nLx,1][ix], imag(P.values)[nLx][ix], marker=:d,  color = :red, label  =  "imag")
    # plt.legend()
    # plt.show()
    f = figure()
    set(axis="normal", plotstyle="linespoints", linewidth=2, pointsize = 2, color = "black", xlabel = "x", ylabel = "Pressure", grid="on", title = "")
    ix = sortperm(geom.values[nLx,1])
    plot(geom.values[nLx,1][ix], real(P.values)[nLx][ix], legend = "real", marker = "ecircle")
    plot!(geom.values[nLx,1][ix], imag(P.values)[nLx][ix], legend = "imag", color = "red", marker = "edmd")
    figure(f)

    true

end # straight_duct_Q8_example

function straight_duct_T3_example()
    t0  =  time()

    rho = 1.21*phun("kg/m^3");# mass density
    c  = 343.0*phun("m/s");# sound speed
    bulk =  c^2*rho;
    omega =  54.5901*phun("rev/s")
    vn0 =  -1.0*phun("m/s")
    Lx = 10.0*phun("m");# length of the box, millimeters
    Ly = 1.0*phun("m"); # length of the box, millimeters
    n = 20;#number of elements along the length

    println("""

    Straight duct with anechoic termination.
    Example from Boundary element acoustics: Fundamentals and computer codes, TW Wu, page 44.
    Both real and imaginary components of the pressure should have amplitude of
    rho*c = $(rho*c).

    Triangle mesh.
    """)

    fens,fes  =  T3block(Lx,Ly,n,2); # Mesh
    bfes  =  meshboundary(fes)
    L0 = selectelem(fens,bfes,facing = true, direction = [-1.0 0.0])
    L10 = selectelem(fens,bfes,facing = true, direction = [+1.0 0.0])
    nLx = selectnode(fens,box = [0.0 Lx  0.0 0.0], inflate = Lx/1.0e5)

    geom  =  NodalField(fens.xyz)
    P  =  NodalField(zeros(ComplexF64,size(fens.xyz,1),1))

    numberdofs!(P)

    material = MatAcoustFluid(bulk,rho)
    femm  =  FEMMAcoust(IntegData(fes, TriRule(1)), material)

    S  =  acousticstiffness(femm, geom, P);
    C  =  acousticmass(femm, geom, P);

    E10femm  =  FEMMAcoustSurf(IntegData(subset(bfes,L10),GaussRule(1, 2)), material)
    D  =  acousticABC(E10femm, geom, P);

    E0femm  =  FEMMBase(IntegData(subset(bfes,L0), GaussRule(1,  2)))
    fi  =  ForceIntensity(-1.0im*omega*rho*vn0);
    F  =  distribloads(E0femm, geom, P, fi, 2);

    p = (-omega^2*S +omega*1.0im*D + C)\F
    scattersysvec!(P, p[:])

    println("Pressure amplitude bounds: ")
    println("  real $(minimum(real(P.values)))/$(maximum(real(P.values)))")
    println("  imag $(minimum(imag(P.values)))/$(maximum(imag(P.values)))")

    println("Total time elapsed  =  ",time() - t0,"s")

    File  =   "straight_duct.vtk"
    scalars = real(P.values);
    vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.T3; scalars = [("Pressure", scalars)])
    @async run(`"paraview.exe" $File`)

    # @pyimport matplotlib.pyplot as plt
    # plt.style[:use]("seaborn-whitegrid")
    # fig = plt.figure()
    # ax = plt.axes()
    # ix = sortperm(geom.values[nLx,1])
    # ax[:plot](geom.values[nLx,1][ix], real(P.values)[nLx][ix], marker=:o, color = :blue, label = "real")
    # ax[:plot](geom.values[nLx,1][ix], imag(P.values)[nLx][ix], marker=:d,  color = :red, label  =  "imag")
    # plt.legend()
    # plt.show()
    # fig = figure()
    # figure(fig)
    set(axis="normal", plotstyle="linespoints", linewidth=2, pointsize = 2, color = "black", xlabel = "x", ylabel = "Pressure", grid="on", title = "")
    f = figure()
    ix = sortperm(geom.values[nLx,1])
    plot(geom.values[nLx,1][ix], real(P.values)[nLx][ix], legend = "real", marker = "ecircle")
    plot!(geom.values[nLx,1][ix], imag(P.values)[nLx][ix], legend = "imag", color = "red", marker = "edmd")
    figure(f)

    true

end # straight_duct_T3_example

function allrun()
    println("#####################################################")
    println("# straight_duct_Q8_example ")
    straight_duct_Q8_example()
    println("#####################################################")
    println("# straight_duct_T3_example ")
    straight_duct_T3_example()
end # function allrun

end # module straight_duct_examples


ENV["DYLD_LIBRARY_PATH"]="/usr/lib/"

using ConvDiffMIPDECO
# using PyPlot
using jInv.Mesh
using jInv.ForwardShare
using jInv.InverseSolve
# using jInvVis
using jInv.LinearSolvers
using MAT
using MUMPSjInv

# filename= "2DmodelLShaped.mat"
filename= joinpath(dirname(pathof(ConvDiffMIPDECO)),"..","examples","Peaks2D.mat")
file = matread(filename)
wtrue = vec(file["W"]) ## Sources on finest mesh
domain = file["domain"]
m      = file["m"]

M      = getRegularMesh(domain,m)
# viewImage2D(wtrue,M)

rec = rand(200,2) .*[domain[2] domain[4]]
# plot(rec[:,1],rec[:,2],".r")

# build linear interpolation matrix from nodes to receiver locations
x1,x2 = getNodalAxes(M)
P = interpmat(x1,x2, rec);

# build param that holds forward problem
pFor = getConvDiffFEMParam(M,Ainv=getMUMPSsolver(),P=P);

dtrue, pFor = getData(vec(wtrue),pFor);
utrue = pFor.Fields

file["utrue"]  = reshape(utrue,M.n[1]+1,M.n[2]+1);
file["dtrue"]  = dtrue
file["rec"]    = rec
file["sig"]    = pFor.sig
file["v"]      = pFor.v

matwrite(filename,file)

# viewImage2D(utrue,getPaddedMesh(M))
#plot(rec[:,1],rec[:,2],".r")
#
# title("true fields")


using Plots, DataFrames, Roots

import RecipesBase.plot

function plot(ht::pwr.htest; backend::Function = gr)

    methodlist = (
        "One-sample t-test power calculation",
        "Two-sample t-test power calculation",
        "Paired t-test power calculation",
        "T-test power calculation",
        "Difference of proportion power calculation for binomial distribution (same sample)",
        "Difference of proportion power calculation for binomial distribution (different sample)",
        "Balanced one-way analysis of variance power calculation",
        "Chi-square test power calculation",
        "Mean power calculation for normal distribution with known variance",
        "Proportion power calculation for binomial distribution (arcsine transformation)",
        "Approximate correlation power calculation (arctangh transformation)"
    )

    if backend in (gr, pyplot, plotly, plotlyjs) == false
        error(backend," is not supported")
    end

    if backend in (gr, pyplot)
        sepstr = "\n"
    else
        sepstr = "<br>"
    end

    ytitle = "Test Power = 1 - β"
    if backend == gr
        ytitle = "Test Power = 1 - \\beta"
    end

    d = ht.d

    if !(ht.title in methodlist)
        warn("The method `",ht.title,"` is not supported. Supported methods include: \n\t",join(methodlist,"\n\t"))
        exit(0)
    end

    breaks = 20

    # case: One-sample, Two-sample or Paired t test
    if ht.title in ("One-sample t-test power calculation",
        "Two-sample t-test power calculation",
        "Paired t-test power calculation")

        n = d["n"]
        n_upper = ceil(Int64,max(n*1.5, n+30))
        n_increment = ceil(Int64,(n_upper - 10)/breaks)
        sample_sizes = collect(10:n_increment:n_upper)

        # obtain power at each sample size
        power = [ss->powerTTest(n=ss,d=d["d"],alpha = d["alpha"],sampletype=d["sampletype"],alternative=d["alternative"]) for ss in sample_sizes]

        # create labels
        title = ht.title
        legend_string = string(
                    "tails = ", d["alternative"],sepstr,
                    "effect size d = ", d["d"],sepstr,
                    "alpha = ", d["alpha"])
        optimal_string = string(
                    "optimal sample size",sepstr,
                    "n = ", ceil(Int64,n),sepstr,d["note"])

    # case: Two-sample t test with n1 and n2
    elseif ht.title == "T-test power calculation"
        n = d["n1"] + d["n2"]
        n_upper = ceil(Int64,max(n*1.5, n+30)) # upper at least 30 above n
        n_rel = d["n1"] / n # relative sample size; will be kept constant in claculations

        # generate data
        n_increment = ceil(Int64,(n_upper - 10)/breaks)
        n_lower = ceil(Int64,fzero(x->(x*n_rel) - 2.1,2,d["n2"]))
        sample_sizes = collect(n_lower:n_increment:n_upper)

        function pwrt2test(ss)
            na = ceil(Int64,ss*n_rel)
            nb = ss - na
            if (na < 2 || nb < 2)
                return 0.0
            end
            return powerT2nTest(n1=na, n2=nb, d=d["d"], alpha = d["alpha"], alternative = d["alternative"])
        end

        power = [pwrt2test(ss) for ss in sample_sizes]

        # create labels
        title = ht.title
        legend_string = string(
                    "tails = ", d["alternative"], sepstr,
                    "effect size d = ", d["d"], sepstr,
                    "alpha = ", d["alpha"], sepstr,
                    "n1/n2 = ",round(n_rel,2))
        optimal_string = string("optimal sample size",sepstr,
                    "n = ",  d["n1"], " + ", d["n2"], " = ", n)

    # case: Difference of proportion (same sample size)
    elseif ht.title == "Difference of proportion power calculation for binomial distribution (same sample)"
        n = d["n"]
        n_upper = ceil(Int64, max(n*1.5, n+30)) # upper at least 30 above n
        n_increment = ceil(Int64,(n_upper - 10)/breaks)

        # generate data
        sample_sizes = collect(10:n_increment:n_upper)
        power = [ ss-> powerTwopTest(n=ss, h=d["h"], alpha = d["alpha"], alternative = d["alternative"]) for ss in sample_sizes]

        # create labels
        title = string("Difference of proportion power calculation",sepstr,"for binomial distribution (same sample)")
        legend_string = string(
                    "tails = ", d["alternative"], sepstr,
                    "effect size h = ", d["h"], sepstr,
                    "alpha = ", d["alpha"])
        optimal_string = string("optimal sample size", sepstr,
                    "n = ", ceil(Int64,n), sepstr,d["note"])

    # case: difference of proportion (different sample size)
    elseif ht.title == "Difference of proportion power calculation for binomial distribution (different sample)"
        n = d["n1"] + d["n2"]
        n_upper = ceil(Int64, max(n*1.5, n+30)) # upper at least 30 above n
        n_rel = d["n1"] / n # relative sample size; will be kept constant in claculations
        n_increment = ceil(Int64,(n_upper - 10)/breaks)
        n_lower = ceil(Int64,fzero(x->(x*n_rel) - 2.1,2,d["n2"]))

        # generate data
        sample_sizes = collect(n_lower:n_increment:n_upper)

        function pwr2p2ntest(ss)
            na = ceil(Int64,ss*n_rel)
            nb = ss - na
            if (na < 2 || nb < 2)
                return (0.0)
            end
            return powerTwop2nTest(n1=na, n2=nb, h=d["h"], alpha = d["alpha"], alternative = d["alternative"])
        end

        power = [ss->pwr2p2ntest(ss) for ss in sample_sizes]

        # create labels
        title = string("Difference of proportion power calculation",sepstr,"for binomial distribution (different sample)")
        legend_string = string(
                    "tails = ", d["alternative"],sepstr,
                    "effect size h = ", d["h"], sepstr,
                    "alpha = ", d["alpha"], sepstr,
                    "n1/n2 = ",round(n_rel,2))
        optimal_string = string("optimal sample size", sepstr,
                    "n = ",  d["n1"], " + ", d["n2"], " = ", n)

    # case: ANOVA
    elseif ht.title== "Balanced one-way analysis of variance power calculation"
        n = d["n"]
        n_upper = ceil(Int64, max(n*1.5, n+30)) # upper at least 30 above n
        n_increment = ceil(Int64,(n_upper - 10)/breaks)

        # generate data
        sample_sizes = collect(10:n_increment:n_upper)
        power = [ss->powerAnovaTest(n=ss, k=d["k"], f=d["f"], alpha = d["alpha"]) for ss in sample_sizes]

        # create labels
        title = ht.title
        legend_string = string(
                    "groups k = ", d["k"], sepstr,
                    "effect size f = ", d["f"], sepstr,
                    "alpha = ", d["alpha"])
        optimal_string = string("optimal sample size",sepstr,
                    "n = ", ceil(Int64,n),sepstr,d["note"])

    # case: Chi Squared
    elseif ht.title == "Chi-square test power calculation"
        n = d["N"]
        n_upper = ceil(Int64, max(n*1.5, n+30)) # upper at least 30 above n
        n_increment = ceil(Int64,(n_upper - 10)/breaks)

        # generate data
        sample_sizes = collect(10:n_increment:n_upper)
        power = [ss->powerChisqTest(N=ss, w=d["w"], alpha = d["alpha"], df=d["df"]) for ss in sample_sizes]

        # create labels
        title = ht.title
        legend_string = string(
                    "df = ", d["df"], sepstr,
                    "effect size w = ", d["w"], sepstr,
                    "alpha = ", d["alpha"])
        optimal_string = string("optimal sample size", sepstr,
                    "N = ", ceil(Int64,n),sepstr,d["note"])

    # case: Normal distribution
    elseif ht.title == "Mean power calculation for normal distribution with known variance"
        n = d["n"]
        n_upper = ceil(Int64, max(n*1.5, n+30)) # upper at least 30 above n
        n_increment = ceil(Int64,(n_upper - 10)/breaks)

        # generate data
        sample_sizes = collect(10:n_increment:n_upper)
        power = [ss->powerNormTest(n=ss, d=d["d"], alpha = d["alpha"], alternative = d["alternative"]) for ss in sample_sizes]

        # create labels
        title = string("Mean power calculation for normal distribution",sepstr,"with known variance")
        legend_string = string(
                    "tails = ", d["alternative"],sepstr,
                    "effect size d = ", d["d"],sepstr,
                    "alpha = ", d["alpha"])
        optimal_string = string("optimal sample size", sepstr,
                    "n = ", ceil(Int64,n))

    # case: proportion
    elseif ht.title == "Proportion power calculation for binomial distribution (arcsine transformation)"
        n = d["n"]
        n_upper = ceil(Int64, max(n*1.5, n+30)) # upper at least 30 above n
        n_increment = ceil(Int64,(n_upper - 10)/breaks)

        # generate data
        sample_sizes = collect(10:n_increment:n_upper)
        power = [ss->powerPTest(n=ss, h=d["h"], alpha = d["alpha"], alternative = d["alternative"]) for ss in sample_sizes]

        # create labels
        title = string("Proportion power calculation",sepstr,"for binomial distribution (arcsine transformation)")
        legend_string = string(
                    "tails = ", d["alternative"],sepstr,
                    "effect size h = ", d["h"],sepstr,
                    "alpha = ", d["alpha"])
        optimal_string = string("optimal sample size",sepstr,
                    "n = ", ceil(Int64,n))

    # case: correlation
    elseif ht.title == "Approximate correlation power calculation (arctangh transformation)"
        n = d["n"]
        n_upper = ceil(Int64, max(n*1.5, n+30)) # upper at least 30 above n
        n_increment = ceil(Int64,(n_upper - 10)/breaks)

        # generate data
        sample_sizes = collect(10:n_increment:n_upper)
        power = [ss->powerRTest(n=ss, r=d["r"], alpha = d["alpha"], alternative = d["alternative"]) for ss in sample_sizes]

        # create labels
        title = string("Approximate correlation power calculation",sepstr,"(arctangh transformation)")
        legend_string = string(
                    "tails = ", d["alternative"],sepstr,
                    "effect size r = ", d["r"],sepstr,
                    "alpha = ", d["alpha"])
        optimal_string = string("optimal sample size",sepstr,
                    "n = ", ceil(Int64,n))
    end

    # use DataFrame
    len = length(sample_sizes)
    if len < 2
        error("cannot display power for this test")
    end

    # select the backend
    backend()

    n_lower = sample_sizes[1]
    xlim_lower = sample_sizes[len] / 80

    # plot with title and x-axis and y-axis labels
    plot(sample_sizes,
        power,
        linecolor = :blue,
        marker=(2,.5,:circle,:blue),
        #left_margin = 8mm,
        top_margin = 8mm,
        right_margin = 4mm,
        title = title,
        yaxis = (ytitle,(-0.03,1.03),0.0:0.2:1.0),
        xaxis = ("Sample Size",(n_lower - xlim_lower,n_upper + xlim_lower)),
        legend=false,
        annotations=([((n_lower+xlim_lower),0.99,text(legend_string,9,:blue,:left,:top)),
            (sample_sizes[floor(Int64,len/2)],0.03,text(optimal_string,9,:red,:left,:bottom))]))
    vline!([n],line=(1,:dot,.8,:red))

end


import PolyLog

n = 10_000_000
z_min = -5
z_max = 5
cmpl_data = (z_max - z_min)*rand(ComplexF64, n) + z_min*(1.0 + 1.0im)*ones(n)
real_data = map(real, cmpl_data)

println("Benchmarking li2::Float64")

map(PolyLog.li2, real_data)       # trigger compilation
@time map(PolyLog.li2, real_data)
@time map(PolyLog.li2, real_data)
@time map(PolyLog.li2, real_data)

println("Benchmarking li2::ComplexF64")

map(PolyLog.li2, cmpl_data)       # trigger compilation
@time map(PolyLog.li2, cmpl_data)
@time map(PolyLog.li2, cmpl_data)
@time map(PolyLog.li2, cmpl_data)

println("Benchmarking li3::ComplexF64")

map(PolyLog.li3, cmpl_data)       # trigger compilation
@time map(PolyLog.li3, cmpl_data)
@time map(PolyLog.li3, cmpl_data)
@time map(PolyLog.li3, cmpl_data)

println("Benchmarking li4::ComplexF64")

map(PolyLog.li4, cmpl_data)       # trigger compilation
@time map(PolyLog.li4, cmpl_data)
@time map(PolyLog.li4, cmpl_data)
@time map(PolyLog.li4, cmpl_data)

println("Benchmarking li5::ComplexF64")

map(PolyLog.li5, cmpl_data)       # trigger compilation
@time map(PolyLog.li5, cmpl_data)
@time map(PolyLog.li5, cmpl_data)
@time map(PolyLog.li5, cmpl_data)

println("Benchmarking li6::ComplexF64")

map(PolyLog.li6, cmpl_data)       # trigger compilation
@time map(PolyLog.li6, cmpl_data)
@time map(PolyLog.li6, cmpl_data)
@time map(PolyLog.li6, cmpl_data)


@testset "Graph coloring" begin

@testset "49 US states + DC" begin
    com = CS.init()

    num_colors = 8
    washington = add_var!(com, 1, num_colors)
    montana = add_var!(com, 1, num_colors)
    maine = add_var!(com, 1, num_colors)
    north_dakota = add_var!(com, 1, num_colors)
    south_dakota = add_var!(com, 1, num_colors)
    wyoming = add_var!(com, 1, num_colors)
    wisconsin = add_var!(com, 1, num_colors)
    idaho = add_var!(com, 1, num_colors)
    vermont = add_var!(com, 1, num_colors)
    minnesota = add_var!(com, 1, num_colors)
    oregon = add_var!(com, 1, num_colors)
    new_hampshire = add_var!(com, 1, num_colors)
    iowa = add_var!(com, 1, num_colors)
    massachusetts = add_var!(com, 1, num_colors)
    nebraska = add_var!(com, 1, num_colors)
    new_york = add_var!(com, 1, num_colors)
    pennsylvania = add_var!(com, 1, num_colors)
    connecticut = add_var!(com, 1, num_colors)
    rhode_island = add_var!(com, 1, num_colors)
    new_jersey = add_var!(com, 1, num_colors)
    indiana = add_var!(com, 1, num_colors)
    nevada = add_var!(com, 1, num_colors)
    utah = add_var!(com, 1, num_colors)
    california = add_var!(com, 1, num_colors)
    ohio = add_var!(com, 1, num_colors)
    illinois = add_var!(com, 1, num_colors)
    washington_dc = add_var!(com, 1, num_colors)
    delaware = add_var!(com, 1, num_colors)
    west_virginia = add_var!(com, 1, num_colors)
    maryland = add_var!(com, 1, num_colors)
    colorado = add_var!(com, 1, num_colors)
    kentucky = add_var!(com, 1, num_colors)
    kansas = add_var!(com, 1, num_colors)
    virginia = add_var!(com, 1, num_colors)
    missouri = add_var!(com, 1, num_colors)
    arizona = add_var!(com, 1, num_colors)
    oklahoma = add_var!(com, 1, num_colors)
    north_carolina = add_var!(com, 1, num_colors)
    tennessee = add_var!(com, 1, num_colors)
    texas = add_var!(com, 1, num_colors)
    new_mexico = add_var!(com, 1, num_colors)
    alabama = add_var!(com, 1, num_colors)
    mississippi = add_var!(com, 1, num_colors)
    georgia = add_var!(com, 1, num_colors)
    south_carolina = add_var!(com, 1, num_colors)
    arkansas = add_var!(com, 1, num_colors)
    louisiana = add_var!(com, 1, num_colors)
    florida = add_var!(com, 1, num_colors)
    michigan = add_var!(com, 1, num_colors)

    states = [washington,montana,maine,north_dakota,south_dakota,wyoming,wisconsin,idaho,vermont,minnesota,oregon,new_hampshire,
    iowa,massachusetts,nebraska,new_york,pennsylvania,connecticut,rhode_island,new_jersey,indiana,nevada,utah,california,ohio,
    illinois,washington_dc,delaware,west_virginia,maryland,colorado,kentucky,kansas,virginia,missouri,arizona,oklahoma,north_carolina,
    tennessee,texas,new_mexico,alabama,mississippi,georgia,south_carolina,arkansas,louisiana,florida,michigan]

    add_constraint!(com, washington != oregon)
    add_constraint!(com, washington != idaho)
    add_constraint!(com, oregon != idaho)
    add_constraint!(com, oregon != nevada)
    add_constraint!(com, oregon != california)
    add_constraint!(com, california != nevada)
    add_constraint!(com, california != arizona)
    add_constraint!(com, nevada != idaho)
    add_constraint!(com, nevada != utah)
    add_constraint!(com, nevada != arizona)
    add_constraint!(com, idaho != montana)
    add_constraint!(com, idaho != wyoming)
    add_constraint!(com, idaho != utah)
    add_constraint!(com, utah != wyoming)
    add_constraint!(com, utah != colorado)
    add_constraint!(com, utah != new_mexico)
    add_constraint!(com, utah != arizona)
    add_constraint!(com, arizona != colorado)
    add_constraint!(com, arizona != new_mexico)
    add_constraint!(com, montana != north_dakota)
    add_constraint!(com, montana != south_dakota)
    add_constraint!(com, montana != wyoming)
    add_constraint!(com, wyoming != south_dakota)
    add_constraint!(com, wyoming != nebraska)
    add_constraint!(com, wyoming != colorado)
    add_constraint!(com, colorado != nebraska)
    add_constraint!(com, colorado != kansas)
    add_constraint!(com, colorado != oklahoma)
    add_constraint!(com, colorado != new_mexico)
    add_constraint!(com, new_mexico != oklahoma)
    add_constraint!(com, new_mexico != texas)
    add_constraint!(com, north_dakota != minnesota)
    add_constraint!(com, north_dakota != south_dakota)
    add_constraint!(com, south_dakota != minnesota)
    add_constraint!(com, south_dakota != iowa)
    add_constraint!(com, south_dakota != nebraska)
    add_constraint!(com, nebraska != iowa)
    add_constraint!(com, nebraska != missouri)
    add_constraint!(com, nebraska != kansas)
    add_constraint!(com, kansas != missouri)
    add_constraint!(com, kansas != oklahoma)
    add_constraint!(com, oklahoma != arkansas)
    add_constraint!(com, oklahoma != texas)
    add_constraint!(com, texas != arkansas)
    add_constraint!(com, texas != louisiana)
    add_constraint!(com, minnesota != wisconsin)
    add_constraint!(com, minnesota != iowa)
    add_constraint!(com, iowa != wisconsin)
    add_constraint!(com, iowa != illinois)
    add_constraint!(com, iowa != missouri)
    add_constraint!(com, missouri != illinois)
    add_constraint!(com, missouri != kentucky)
    add_constraint!(com, missouri != tennessee)
    add_constraint!(com, missouri != arkansas)
    add_constraint!(com, arkansas != tennessee)
    add_constraint!(com, arkansas != mississippi)
    add_constraint!(com, arkansas != louisiana)
    add_constraint!(com, louisiana != mississippi)
    add_constraint!(com, wisconsin != illinois)
    add_constraint!(com, wisconsin != michigan)
    add_constraint!(com, illinois != indiana)
    add_constraint!(com, illinois != kentucky)
    add_constraint!(com, kentucky != indiana)
    add_constraint!(com, kentucky != ohio)
    add_constraint!(com, kentucky != west_virginia)
    add_constraint!(com, kentucky != virginia)
    add_constraint!(com, kentucky != tennessee)
    add_constraint!(com, tennessee != virginia)
    add_constraint!(com, tennessee != north_carolina)
    add_constraint!(com, tennessee != georgia)
    add_constraint!(com, tennessee != alabama)
    add_constraint!(com, tennessee != mississippi)
    add_constraint!(com, mississippi != alabama)
    add_constraint!(com, michigan != indiana)
    add_constraint!(com, michigan != ohio)
    add_constraint!(com, indiana != ohio)
    add_constraint!(com, alabama != georgia)
    add_constraint!(com, alabama != florida)
    add_constraint!(com, ohio != pennsylvania)
    add_constraint!(com, ohio != west_virginia)
    add_constraint!(com, maine != new_hampshire)
    add_constraint!(com, new_hampshire != vermont)
    add_constraint!(com, new_hampshire != massachusetts)
    add_constraint!(com, vermont != massachusetts)
    add_constraint!(com, vermont != new_york)
    add_constraint!(com, massachusetts != new_york)
    add_constraint!(com, massachusetts != rhode_island)
    add_constraint!(com, massachusetts != connecticut)
    add_constraint!(com, connecticut != rhode_island)
    add_constraint!(com, new_york != connecticut)
    add_constraint!(com, new_york != pennsylvania)
    add_constraint!(com, new_york != new_jersey)
    add_constraint!(com, pennsylvania != new_jersey)
    add_constraint!(com, pennsylvania != delaware)
    add_constraint!(com, pennsylvania != maryland)
    add_constraint!(com, pennsylvania != west_virginia)
    add_constraint!(com, new_jersey != delaware)
    add_constraint!(com, maryland != washington_dc)
    add_constraint!(com, maryland != west_virginia)
    add_constraint!(com, maryland != virginia)
    add_constraint!(com, washington_dc != virginia)
    add_constraint!(com, west_virginia != virginia)
    add_constraint!(com, virginia != north_carolina)
    add_constraint!(com, north_carolina != south_carolina)
    add_constraint!(com, south_carolina != georgia)
    add_constraint!(com, georgia != florida)

    # :should be :Min or :Max
    @test_throws ErrorException set_objective!(com, :Minimize, CS.vars_max(states))

    set_objective!(com, :Min, CS.vars_max(states))

    status = solve!(com; keep_logs=true)
    CS.save_logs(com, "graph_color_optimize.json")
    rm("graph_color_optimize.json")

    @test status == :Solved
    @test com.best_sol == 4
    @test all([CS.isfixed(var) for var in states])
    @test maximum([value(var) for var in states]) == 4

end

@testset "49 US states + DC only 3 colors" begin
    com = CS.init()

    num_colors = 3
    washington = add_var!(com, 1, num_colors)
    montana = add_var!(com, 1, num_colors)
    maine = add_var!(com, 1, num_colors)
    north_dakota = add_var!(com, 1, num_colors)
    south_dakota = add_var!(com, 1, num_colors)
    wyoming = add_var!(com, 1, num_colors)
    wisconsin = add_var!(com, 1, num_colors)
    idaho = add_var!(com, 1, num_colors)
    vermont = add_var!(com, 1, num_colors)
    minnesota = add_var!(com, 1, num_colors)
    oregon = add_var!(com, 1, num_colors)
    new_hampshire = add_var!(com, 1, num_colors)
    iowa = add_var!(com, 1, num_colors)
    massachusetts = add_var!(com, 1, num_colors)
    nebraska = add_var!(com, 1, num_colors)
    new_york = add_var!(com, 1, num_colors)
    pennsylvania = add_var!(com, 1, num_colors)
    connecticut = add_var!(com, 1, num_colors)
    rhode_island = add_var!(com, 1, num_colors)
    new_jersey = add_var!(com, 1, num_colors)
    indiana = add_var!(com, 1, num_colors)
    nevada = add_var!(com, 1, num_colors)
    utah = add_var!(com, 1, num_colors)
    california = add_var!(com, 1, num_colors)
    ohio = add_var!(com, 1, num_colors)
    illinois = add_var!(com, 1, num_colors)
    washington_dc = add_var!(com, 1, num_colors)
    delaware = add_var!(com, 1, num_colors)
    west_virginia = add_var!(com, 1, num_colors)
    maryland = add_var!(com, 1, num_colors)
    colorado = add_var!(com, 1, num_colors)
    kentucky = add_var!(com, 1, num_colors)
    kansas = add_var!(com, 1, num_colors)
    virginia = add_var!(com, 1, num_colors)
    missouri = add_var!(com, 1, num_colors)
    arizona = add_var!(com, 1, num_colors)
    oklahoma = add_var!(com, 1, num_colors)
    north_carolina = add_var!(com, 1, num_colors)
    tennessee = add_var!(com, 1, num_colors)
    texas = add_var!(com, 1, num_colors)
    new_mexico = add_var!(com, 1, num_colors)
    alabama = add_var!(com, 1, num_colors)
    mississippi = add_var!(com, 1, num_colors)
    georgia = add_var!(com, 1, num_colors)
    south_carolina = add_var!(com, 1, num_colors)
    arkansas = add_var!(com, 1, num_colors)
    louisiana = add_var!(com, 1, num_colors)
    florida = add_var!(com, 1, num_colors)
    michigan = add_var!(com, 1, num_colors)

    states = [washington,montana,maine,north_dakota,south_dakota,wyoming,wisconsin,idaho,vermont,minnesota,oregon,new_hampshire,
    iowa,massachusetts,nebraska,new_york,pennsylvania,connecticut,rhode_island,new_jersey,indiana,nevada,utah,california,ohio,
    illinois,washington_dc,delaware,west_virginia,maryland,colorado,kentucky,kansas,virginia,missouri,arizona,oklahoma,north_carolina,
    tennessee,texas,new_mexico,alabama,mississippi,georgia,south_carolina,arkansas,louisiana,florida,michigan]

    add_constraint!(com, washington != oregon)
    add_constraint!(com, washington != idaho)
    add_constraint!(com, oregon != idaho)
    add_constraint!(com, oregon != nevada)
    add_constraint!(com, oregon != california)
    add_constraint!(com, california != nevada)
    add_constraint!(com, california != arizona)
    add_constraint!(com, nevada != idaho)
    add_constraint!(com, nevada != utah)
    add_constraint!(com, nevada != arizona)
    add_constraint!(com, idaho != montana)
    add_constraint!(com, idaho != wyoming)
    add_constraint!(com, idaho != utah)
    add_constraint!(com, utah != wyoming)
    add_constraint!(com, utah != colorado)
    add_constraint!(com, utah != new_mexico)
    add_constraint!(com, utah != arizona)
    add_constraint!(com, arizona != colorado)
    add_constraint!(com, arizona != new_mexico)
    add_constraint!(com, montana != north_dakota)
    add_constraint!(com, montana != south_dakota)
    add_constraint!(com, montana != wyoming)
    add_constraint!(com, wyoming != south_dakota)
    add_constraint!(com, wyoming != nebraska)
    add_constraint!(com, wyoming != colorado)
    add_constraint!(com, colorado != nebraska)
    add_constraint!(com, colorado != kansas)
    add_constraint!(com, colorado != oklahoma)
    add_constraint!(com, colorado != new_mexico)
    add_constraint!(com, new_mexico != oklahoma)
    add_constraint!(com, new_mexico != texas)
    add_constraint!(com, north_dakota != minnesota)
    add_constraint!(com, north_dakota != south_dakota)
    add_constraint!(com, south_dakota != minnesota)
    add_constraint!(com, south_dakota != iowa)
    add_constraint!(com, south_dakota != nebraska)
    add_constraint!(com, nebraska != iowa)
    add_constraint!(com, nebraska != missouri)
    add_constraint!(com, nebraska != kansas)
    add_constraint!(com, kansas != missouri)
    add_constraint!(com, kansas != oklahoma)
    add_constraint!(com, oklahoma != arkansas)
    add_constraint!(com, oklahoma != texas)
    add_constraint!(com, texas != arkansas)
    add_constraint!(com, texas != louisiana)
    add_constraint!(com, minnesota != wisconsin)
    add_constraint!(com, minnesota != iowa)
    add_constraint!(com, iowa != wisconsin)
    add_constraint!(com, iowa != illinois)
    add_constraint!(com, iowa != missouri)
    add_constraint!(com, missouri != illinois)
    add_constraint!(com, missouri != kentucky)
    add_constraint!(com, missouri != tennessee)
    add_constraint!(com, missouri != arkansas)
    add_constraint!(com, arkansas != tennessee)
    add_constraint!(com, arkansas != mississippi)
    add_constraint!(com, arkansas != louisiana)
    add_constraint!(com, louisiana != mississippi)
    add_constraint!(com, wisconsin != illinois)
    add_constraint!(com, wisconsin != michigan)
    add_constraint!(com, illinois != indiana)
    add_constraint!(com, illinois != kentucky)
    add_constraint!(com, kentucky != indiana)
    add_constraint!(com, kentucky != ohio)
    add_constraint!(com, kentucky != west_virginia)
    add_constraint!(com, kentucky != virginia)
    add_constraint!(com, kentucky != tennessee)
    add_constraint!(com, tennessee != virginia)
    add_constraint!(com, tennessee != north_carolina)
    add_constraint!(com, tennessee != georgia)
    add_constraint!(com, tennessee != alabama)
    add_constraint!(com, tennessee != mississippi)
    add_constraint!(com, mississippi != alabama)
    add_constraint!(com, michigan != indiana)
    add_constraint!(com, michigan != ohio)
    add_constraint!(com, indiana != ohio)
    add_constraint!(com, alabama != georgia)
    add_constraint!(com, alabama != florida)
    add_constraint!(com, ohio != pennsylvania)
    add_constraint!(com, ohio != west_virginia)
    add_constraint!(com, maine != new_hampshire)
    add_constraint!(com, new_hampshire != vermont)
    add_constraint!(com, new_hampshire != massachusetts)
    add_constraint!(com, vermont != massachusetts)
    add_constraint!(com, vermont != new_york)
    add_constraint!(com, massachusetts != new_york)
    add_constraint!(com, massachusetts != rhode_island)
    add_constraint!(com, massachusetts != connecticut)
    add_constraint!(com, connecticut != rhode_island)
    add_constraint!(com, new_york != connecticut)
    add_constraint!(com, new_york != pennsylvania)
    add_constraint!(com, new_york != new_jersey)
    add_constraint!(com, pennsylvania != new_jersey)
    add_constraint!(com, pennsylvania != delaware)
    add_constraint!(com, pennsylvania != maryland)
    add_constraint!(com, pennsylvania != west_virginia)
    add_constraint!(com, new_jersey != delaware)
    add_constraint!(com, maryland != washington_dc)
    add_constraint!(com, maryland != west_virginia)
    add_constraint!(com, maryland != virginia)
    add_constraint!(com, washington_dc != virginia)
    add_constraint!(com, west_virginia != virginia)
    add_constraint!(com, virginia != north_carolina)
    add_constraint!(com, north_carolina != south_carolina)
    add_constraint!(com, south_carolina != georgia)
    add_constraint!(com, georgia != florida)

    set_objective!(com, :Min, CS.vars_max(states))

    status = solve!(com)

    @test status == :Infeasible
end

end


@recipe function f(im::IntrinsicMetric{T}) where T
    2 ≤ ambientdim(im) ≤ 3 || error("Can't plot $(ambientdim(im))d points!")
    flat = ambientdim(im) == 2

    legend := false
    pts = points(im)
    g = adjgraph(im)

    @series begin
        seriestype := :scatter
        markersize --> 0.5
        label := :vertex
        if flat
            getindex.(pts, 1),
            getindex.(pts, 2)
        else
            getindex.(pts, 1),
            getindex.(pts, 2),
            getindex.(pts, 3)
        end
    end

    xs = T[]; ys = T[]; zs = T[]
    for e in edges(g)
        p1 = pts[src(e)]
        p2 = pts[dst(e)]
        append!(xs, [p1[1], p2[1], NaN])
        append!(ys, [p1[2], p2[2], NaN])
        !flat && append!(zs, [p1[3], p2[3], NaN])
    end

    @series begin
        markersize := :none
        seriestype := :path
        linewidth --> 1
        label := :edge
        if flat
            xs, ys
        else
            xs, ys, zs
        end
    end
end


import Base: ==, isless,
    +, -, *, ^, max, min, div, %, gcd, lcm,
    isqrt, isodd, iseven, one, isprime, factor, primes

# Remember that RN is typealiased to RomanNumeral

# Equality operators
==(n1::RN, n2::RN) = n1.val == n2.val

# Comparisons
<(n1::RN, n2::RN) = n1.val < n2.val
<=(n1::RN, n2::RN) = n1.val <= n2.val

## Arithmetic
# Multiple argument operators
for op in [:+, :-, :*, :^, :max, :min]
    @eval ($op)(ns::RN...) = $(op)(map(n -> n.val, ns)...) |> RN
end

# Two argument operators
for op in [:div, :%, :gcd, :lcm]
    @eval ($op)(n1::RN, n2::RN) = $(op)(n1.val, n2.val) |> RN
end

# One argument operators
for op in [:isqrt]
    @eval $(op)(num::RN) = $(op)(num.val) |> RN
end

# Integer properties
for op in [:isodd, :iseven, :isprime]
    @eval $(op)(num::RN) = $(op)(num.val)
end

# Who knew Romans did number theory
function Base.factor(num::RN)
    factors = Dict{RN,RN}()
    for (fac, mul) in factor(num.val)
        factors[RN(fac)] = RN(mul)
    end
    factors
end
Base.primes(num::RN) = map(RN, primes(num.val))


using JOLI
using Revise
using WAVEFORM
using LinearAlgebra

n = 100*[1;1;1];
d = 10.0*[1;1;1];
o = 0.0*[1;1;1];
v0 = 2000;
freq = 4.0;
λ = v0/freq;

t0 = 0.0;
f0 = 10.0;
unit = "m/s";
xsrc = [(n.*d)[1]/2];
ysrc = [(n.*d)[2]/2];
zsrc = [100.0];
xrec = range(0.0,stop=(n.*d)[1]::Float64,length=n[1]);
yrec = range(0.0,stop=(n.*d)[2]::Float64,length=n[2]);
zrec = [100.0];
freqs = [freq];
nsrc = length(xsrc)*length(ysrc)*length(zsrc);
nfreq = length(freqs)
model = Model{Int64,Float64}(n,d,o,t0,f0,unit,freqs,xsrc,ysrc,zsrc,xrec,yrec,zrec);

comp_n = n;
comp_d = d;
comp_o = o;
npml = convert(Int,λ/minimum(comp_d))*ones(Int64,(2,3));
npml = 30*ones(Int64,(2,3));
scheme = WAVEFORM.helm3d_operto27;
cut_pml = true;
implicit_matrix = true;
srcfreqmask = trues(nsrc,nfreq);
misfit = WAVEFORM.least_squares;
lsopts = LinSolveOpts(solver=:fgmres,maxinnerit=5,tol=1e-6);
lsopts.precond = :mlgmres;
opts = PDEopts{Int64,Float64}(scheme,comp_n,comp_d,comp_o,cut_pml,implicit_matrix,npml,misfit,srcfreqmask,lsopts);

nλ = maximum(n.*d/λ);
ppλ = λ/maximum(d);
v = v0*ones(Float64,n...);
v = vec(v);
opts.implicit_matrix = true;

(H,comp_grid,T,DT_adj,P) = helmholtz_system(v,model,freq,opts);

# ensure linear system is actually solved
nt = comp_grid.comp_n
q = zeros(eltype(H),tuple(nt...));
q[div.(nt,2)...] = 1.0
u = H\vec(q)
@test norm(H*u-vec(q)) < 1e-6*norm(q)
ut = H'\vec(q);
@test norm(H'*ut-vec(q)) < 1e-6*norm(q)

# test computed green's function ≈ analytical green's function in the non-pml region
(xt,yt,zt) = odn_to_grid(comp_grid.comp_o,comp_grid.comp_d,comp_grid.comp_n)
idx = findmax(abs.(q))[2]
qCartesian = CartesianIndices(q)
(ix,iy,iz) = qCartesian[idx].I
ω = (2*pi*freq)/v0;
R = [sqrt((x-xt[ix])^2 + (y-yt[ix])^2 + (z-zt[ix])^2) for x in xt, y in yt, z in zt];
G = [prod(model.d)*exp(im*ω*r)/(4*π*r) for r in R];
r = reshape(u,size(G))-G;
r[ix,iy,iz] = 0;
@test norm(comp_grid.comp_to_phys_grid*vec(r))/norm(comp_grid.comp_to_phys_grid*vec(u)) < 0.01

# test derivative behaviour
n = 50*[1;1;1];

xsrc = [(n.*d)[1]/2];
ysrc = [(n.*d)[2]/2];
zsrc = [100.0];
xrec = range(0.0,stop=(n.*d)[1]::Float64,length=n[1]);
yrec = range(0.0,stop=(n.*d)[2]::Float64,length=n[2]);
zrec = [100.0];
freqs = [freq];
nsrc = length(xsrc)*length(ysrc)*length(zsrc);
nfreq = length(freqs)
model = Model{Int64,Float64}(n,d,o,t0,f0,unit,freqs,xsrc,ysrc,zsrc,xrec,yrec,zrec);

comp_n = n;
comp_d = d;
comp_o = o;
npml = convert(Int,λ/minimum(comp_d))*ones(Int64,(2,3));
npml = 10*ones(Int64,(2,3));
scheme = WAVEFORM.helm3d_operto27;
cut_pml = true;
implicit_matrix = true;
srcfreqmask = trues(nsrc,nfreq);
misfit = WAVEFORM.least_squares;
lsopts = LinSolveOpts(solver=:fgmres,maxinnerit=5,tol=1e-10);
lsopts.precond = :mlgmres;
opts = PDEopts{Int64,Float64}(scheme,comp_n,comp_d,comp_o,cut_pml,implicit_matrix,npml,misfit,srcfreqmask,lsopts);

nλ = maximum(n.*d/λ);
ppλ = λ/maximum(d);
v = v0*ones(Float64,n...);
v = vec(v);
opts.implicit_matrix = true;

(H,comp_grid,T,DT_adj,P) = helmholtz_system(v,model,freq,opts);
nt = comp_grid.comp_n
q = zeros(eltype(H),tuple(nt...));
q[div.(nt,2)...] = 1.0;
q = vec(q);
u0 = H\q;
δv = randn(size(v));
δvx = comp_grid.phys_to_comp_grid*δv;
δu = H\(-T(u)*δvx);
step = 10.^(-5.0:1.0);
err0 = zeros(length(step))
err1 = zeros(length(step))
for i=1:length(step)
    h = step[i]
    H1 = helmholtz_system(v+δv*h,model,freq,opts)[1];
    u1 = H1\q;
    err0[i] = norm(u1-u0);
    err1[i] = norm(u1-u0-h*δu)
end

@test abs(median(diff(log10.(err0)))-1) < 0.2
@test abs(median(diff(log10.(err1)))-2) < 0.2


export ilan_benchmark

using IterativeSolvers
using LinearAlgebra
using Random
using Statistics

"""
    ilan_benchmark()
"""
ilan_benchmark(nep::NEP;params...)=ilan_benchmark(ComplexF64,nep;params...)
function ilan_benchmark(
    ::Type{T},
    nep::NEP;
    orthmethod::Type{T_orth}=DGKS,
    maxit=30,
    linsolvercreator=DefaultLinSolverCreator(),
    tol=eps(real(T))*10000,
    neigs=6,
    errmeasure::ErrmeasureType = DefaultErrmeasure,
    σ=zero(T),
    γ=one(T),
    v=randn(real(T),size(nep,1)),
    displaylevel=0,
    check_error_every=1,
    proj_solve=false,
    inner_solver_method=DefaultInnerSolver())where{T<:Number,T_orth<:IterativeSolvers.OrthogonalizationMethod}

    # Ensure types σ and v are of type T
    σ=T(σ)
    v=Array{T,1}(v)

    n = size(nep,1);
    m = maxit;

    # initialization
    V=zeros(T,n,m+1)
    Q=zeros(T,n,m+1)
    Qp=zeros(T,n,m+1)
    Qn=zeros(T,n,m+1)
    Z=zeros(T,n,m+1)
    H=zeros(T,m+1,m)
    HH=zeros(T,m+1,m)
    ω=zeros(T,m+1)
    a=Vector{T}(γ.^(0:m)); a[1]=zero(T); # TODO
    local M0inv::LinSolver = create_linsolver(linsolvercreator,nep,σ)
    err=ones(m,m);
    λ=zeros(T,m+1);

    # precompute the symmetrizer coefficients
    G=zeros(m+1,m+1);
    for i=1:m+1 G[i,1]=1/i end
    for j=1:m
        for i=1:m+1
            G[i,j+1]=(G[i,j]*j)/(i+j);
        end
    end

    # getting matrices and functions
    fv=get_fv(nep); p=length(fv);    Av=get_Av(nep)

    # precompute derivatives and FDH matrices
    SS=diagm(0 => σ*ones(T,2m+2)) + diagm(-1 => (1:2m+1))
    fD=zeros(T,2*m+2,p)
    for t=1:p fD[:,t]=fv[t](SS)[:,1] end
    FDH=Vector{Array{T,2}}(undef,p)
    for t=1:p FDH[t]=zeros(T,m+1,m+1)
        for i=1:m+1 for j=1:m+1
            FDH[t][i,j]=fD[i+j,t];
        end end
    end

    # setting initial step
    Q[:,1]=v/norm(v)
    ω[1]=Q[:,1]⋅compute_Mlincomb(nep,0,hcat(Q[:,1],Q[:,1]),[0,1]);
    V[:,1]=Q[:,1];

    k=1; conv_eig=0;
    local pnep::NEP;
    if (proj_solve)
        pnep=create_proj_NEP(nep);
    end

    while (k <= m) && (conv_eig<neigs)
        if (displaylevel>0) && ((rem(k,check_error_every)==0) || (k==m))
            println("Iteration:",k, " conveig:",conv_eig)
        end

        broadcast!(/,view(Qn,:,2:k+1),view(Q,:,1:k),(1:k)')
        Qn[:,1] = compute_Mlincomb!(nep,σ,Qn[:,1:k+1],a[1:k+1]);
        Qn[:,1] = -lin_solve(M0inv,Qn[:,1]);

        # B-multiplication
        Z=0*Z;
        for t=1:p
            Z[:,1:k+1]=Z[:,1:k+1]+Av[t]*Qn[:,1:k+1]*(G[1:k+1,1:k+1].*FDH[t][1:k+1,1:k+1]);
        end

        # orthogonalization (three terms recurrence)
        α=sum(sum(conj(Z).*Q,dims=1))
        if k>1 β=sum(sum(conj(Z).*Qp,dims=1)) end
        η=sum(sum(conj(Z).*Qn,dims=1))

        H[k,k]=α/ω[k]
        if k>1 H[k-1,k]=β/ω[k-1] end
        Qn=Qn-H[k,k]*Q;
        if k>1 Qn=Qn-H[k-1,k]*Qp end

        H[k+1,k]=norm(Qn);
        Qn=Qn/H[k+1,k];

        ω[k+1]=η-2*α*H[k,k]+ω[k]*H[k,k]^2;
        if k>1
            ω[k+1]=ω[k+1]-2*β*H[k-1,k]+ω[k-1]*H[k-1,k]^2;
        end
        ω[k+1]=ω[k+1]/H[k+1,k]^2;
        V[:,k+1]=Qn[:,1];

        orthogonalize_and_normalize!(view(V,:,1:k),view(V,:,k+1), view(HH,1:k,k), orthmethod)

        k=k+1;
        # shift the vectors
        Qp=Q;   Q=Qn; Qn=0*Qn;

    end
    k=k-1

    return V[:,1:k+1], H[1:k,1:k-1], ω[1:k], HH[1:k,1:k]
end


export ReinforcePolicy

using Flux: softmax
using LinearAlgebra: dot
using StatsBase

"""
    ReinforcePolicy(approximator, α, γ)
This is the very basic reinforce algorithm.
TODO: implement some other variants.
"""
Base.@kwdef struct ReinforcePolicy{A<:AbstractApproximator} <: AbstractPolicy
    approximator::A
    α::Float64
    γ::Float64
end

(π::ReinforcePolicy)(obs) =
    obs |> get_state |> π.approximator |> softmax |> x -> Weights(x, 1.0) |> sample

RLBase.get_prob(π::ReinforcePolicy, s) = s |> π.approximator |> softmax
RLBase.get_prob(π::ReinforcePolicy, s, a) = get_prob(π, s)[a]

# TODO: handle neural network q approximator
function RLBase.update!(π::ReinforcePolicy{<:LinearQApproximator}, t::AbstractTrajectory)
    if length(t) > 0 && get_trace(t, :terminal)[end]
        states, actions, rewards = get_trace(t, :state), get_trace(t, :action), get_trace(t, :reward)
        Q, α, γ = π.approximator, π.α, π.γ
        gains = discount_rewards(rewards, γ)
        γₜ = 1.0

        for (i, (s, a, g)) in enumerate(zip(states, actions, gains))
            # !!! we will multiply `Q.feature_func(s, a)` in the `LinearQApproximator` again!
            update!(
                Q,
                (s, a) => α * γₜ * g *
                          (Q.feature_func(s, a) .-
                           sum(x -> get_prob(π, s, x) .* Q.feature_func(s, x), Q.actions)),
            )
            γₜ *= γ
        end
    end
end


using LinearAlgebra

@testset "PointForce " begin

  @testset "Gaussian function" begin
    @test Gaussian!(1.0, 0.1) ≈ exp(-0.5*1.0^2*0.1^(-2))
    @test Gaussian(2.0)(0.5) ≈ Gaussian!(0.5, 2.0)
  end

  @testset "Spatial Gaussian distribution for NodeData" begin

    Re = 500; # Reynolds number
    xlim = (-1.0,5.0)
    ylim = (-2.0,2.0)
    ReΔ = 16
    Δx = ReΔ/Re
    Δt = min(0.25*Δx,0.5*Δx^2*Re);
    sys = Systems.NavierStokes(Re, Δx, xlim, ylim, Δt)

    w = Nodes(Dual,size(sys));
    wbuffer =  Nodes(Dual,size(sys));
    x0 = (1.0, 2.0)
    σx = 2.0

    SpatialGauss(wbuffer, w, x0, σx, sys)
    xx, yy = coordinates(w,dx=sys.grid.Δx,I0=Systems.origin(sys))
    @test wbuffer[10, 20] ≈ exp(-0.5*σx^(-2)*((xx[10]-x0[1])^2+(yy[20]-x0[2])^2))
  end

  @testset "Spatial Gaussian distribution for EdgeData" begin

    Re = 500; # Reynolds number
    xlim = (-1.0,5.0)
    ylim = (-2.0,2.0)
    ReΔ = 16
    Δx = ReΔ/Re
    Δt = min(0.25*Δx,0.5*Δx^2*Re);
    sys = Systems.NavierStokes(Re, Δx, xlim, ylim, Δt)

    q = Edges(Primal,size(sys));
    qbuffer =  Edges(Primal,size(sys));
    x0 = (1.0, 2.0)
    σx = 2.0

    SpatialGauss(qbuffer, q, x0, σx, sys)
    xx, xy, yx, yy = coordinates(q,dx=sys.grid.Δx,I0=Systems.origin(sys))
    @test qbuffer.u[10, 20] ≈ exp(-0.5*σx^(-2)*((xx[10]-x0[1])^2+(xy[20]-x0[2])^2))

    @test qbuffer.v[60, 41] ≈ exp(-0.5*σx^(-2)*((yx[60]-x0[1])^2+(yy[41]-x0[2])^2))

  end

@testset "PointForce(t) without spatial variation" begin
  Re = 500; # Reynolds number
  xlim = (-1.0,5.0)
  ylim = (-2.0,2.0)
  ReΔ = 16
  Δx = ReΔ/Re
  Δt = min(0.25*Δx,0.5*Δx^2*Re);
  sys = Systems.NavierStokes(Re, Δx, xlim, ylim, Δt)

  q = Edges(Primal,size(sys));
  qbuffer =  Edges(Primal,size(sys));
  xx, xy, yx, yy = coordinates(q,dx=sys.grid.Δx,I0=Systems.origin(sys))
  x0 = (xx[10], xy[30])
  f0 = (1.0, 2.0)
  σx = 2.0
  t0 = 2.0
  σt = 0.5
  pointf = PointForce(q, x0, f0, t0, σt, sys)

  # Test at the central location and time of the Gaussian distribution
  @test pointf(2.0).u[10, 30] ≈ (pointf.regop(pointf.ubuffer,deepcopy(pointf.fbuffer))).u[10, 30]

  # Test at the central time of the Gaussian distribution
  @test pointf(2.0).v.data  ≈  pointf.regop(pointf.ubuffer,deepcopy(pointf.fbuffer)).v.data

  @test pointf(2.5).u.data  ≈  pointf.regop(pointf.ubuffer,rmul!(deepcopy(pointf.fbuffer),
                              exp(-(2.5-pointf.t0)^2/pointf.σt^2))).u.data

  @test pointf(2.1).v.data  ≈  pointf.regop(pointf.ubuffer,rmul!(deepcopy(pointf.fbuffer),
                              exp(-(2.1-pointf.t0)^2/pointf.σt^2))).v.data

end

@testset "PointForce(t) with spatial variation" begin
        Re = 500; # Reynolds number
        xlim = (-1.0,5.0)
        ylim = (-2.0,2.0)
        ReΔ = 16
        Δx = ReΔ/Re
        Δt = min(0.25*Δx,0.5*Δx^2*Re);
        sys = Systems.NavierStokes(Re, Δx, xlim, ylim, Δt)

        q = Edges(Primal,size(sys));
        qbuffer =  Edges(Primal,size(sys));
        xx, xy, yx, yy = coordinates(q,dx=sys.grid.Δx,I0=Systems.origin(sys))
        x0 = (xx[10], xy[30])
        f0 = (1.0, 2.0)
        σx = 2.0
        t0 = 2.0
        σt = 0.5
        pointf = PointForce(q, x0, f0, t0, σt, sys, σx = σx)

        # Test at the central location and time of the Gaussian distribution
        @test pointf(2.0).u[10, 30] ≈  (pointf.regop(pointf.ubuffer,deepcopy(pointf.fbuffer))).u[10, 30]

        # Test at the central time of the Gaussian distribution
        @test pointf(2.0).v.data ≈  (pointf.xbuffer∘pointf.regop(pointf.ubuffer,deepcopy(pointf.fbuffer))).v.data

        @test pointf(2.5).u.data ≈  (pointf.xbuffer∘pointf.regop(pointf.ubuffer,rmul!(deepcopy(pointf.fbuffer),
                                    exp(-(2.5-pointf.t0)^2/pointf.σt^2)))).u.data

        @test pointf(2.1).v.data ≈  (pointf.xbuffer∘pointf.regop(pointf.ubuffer,rmul!(deepcopy(pointf.fbuffer),
                                    exp(-(2.1-pointf.t0)^2/pointf.σt^2)))).v.data

  end
end


""" Visualization of delta sets and dual complexes.

This module providers wrapper functions to plotting embedded delta sets using
Makie.jl, including wireframe, mesh, and scatter plots. Note that one of the
Makie backends (CairoMakie, GLMakie, or WGLMakie) must be imported in order for
plotting to work.
"""
module MeshGraphics

using Catlab.CategoricalAlgebra
using ..SimplicialSets
using ..MeshInterop
#import ...Makie: wireframe, wireframe!, mesh, mesh!, scatter, scatter!
#
#export wireframe, wireframe!, mesh, mesh!, scatter, scatter!

using GeometryBasics: Mesh
import ...Makie
import ...Makie: convert_arguments

""" This extends the "Mesh" plotting recipe for embedded deltasets by converting
an embedded deltaset into arguments that can be passed into Makie.mesh
"""

function convert_arguments(P::Union{Type{<:Makie.Wireframe},
                                    Type{<:Makie.Mesh},
                                    Type{<:Makie.Scatter}},
                           dset::HasDeltaSet)
  convert_arguments(P, Mesh(dset))
end

""" This extends the "LineSegments" plotting recipe for embedded deltasets by converting
an embedded deltaset into arguments that can be passed into Makie.linesegments
"""
function convert_arguments(P::Type{<:Makie.LineSegments}, dset::EmbeddedDeltaSet2D)
  edge_positions = zeros(ne(dset)*2,3)
  for e in edges(dset)
    edge_positions[2*e-1,:] = point(dset, src(dset, e))
    edge_positions[2*e,:] = point(dset, tgt(dset, e))
  end
  convert_arguments(P, edge_positions)
end

plottype(::EmbeddedDeltaSet2D) = Mesh

end


abstract type AbstractDie end

# Deterministic die
mutable struct DDie <: AbstractDie
  next :: Int
end

DDie() = DDie(0)

function next(d::DDie) :: Int
  d.next = d.next%100+1
  return d.next
end

############################################################################################
abstract type AbstractPlayer end

# Deterministic Player
mutable struct Player <: AbstractPlayer
  pos   :: Int
  score :: Int
end

Player(pos::Int) = Player(pos,0)
score(p::Player) = p.score

function play(p::Player,d::AbstractDie)
  step :: Int = next(d) + next(d) + next(d)
  p.pos = (p.pos + step - 1) % 10 + 1
  p.score += p.pos
end

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

function game(p1::Player,p2::Player,d::AbstractDie) :: Int
  turn = 1
  rolls = 0
  while true
    play(p1,d)
    rolls += 3
    (p1.score >= 1000) && (return p2.score * rolls)

    play(p2,d)
    rolls += 3
    (p2.score >= 1000) && (return p1.score * rolls)
    println("Turn ", turn, ": P1 -> (", p1.pos, ",", p1.score,")    P2 -> (", p2.pos, ",", p2.score,")")
    turn += 1
  end
end

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

# (Possible results => Probability) for 3 throws of a 3-face dice.
df = Dict{Int,Int}(
  3 => 1,
  4 => 3,
  5 => 6,
  6 => 7,
  7 => 6,
  8 => 3,
  9 => 1
)

# Play a turn for player 'pid' from position 'pos'.
function play(M:: Array{Int,4}, pos::CartesianIndex{4}, Mnew:: Array{Int,4}, pid::Int)
  (M[pos] == 0) && return
  if pid == 1 # Player 1
    for roll in 3:9 # For every possible result
      pnew = (pos[1] + roll -1) % 10 + 1
      Mnew[CartesianIndex(pnew, min(pos[2]+pnew,22), pos[3], pos[4])] += df[roll] * M[pos]
    end
  else # Player 2
    for roll in 3:9 # For every possible result
      pnew = (pos[3] + roll -1) % 10 + 1
      Mnew[CartesianIndex(pos[1], pos[2], pnew, min(pos[4]+pnew,22))] += df[roll] * M[pos]
    end
  end
end

function quantumGame(pos1::Int, pos2::Int) :: Int
  # All possible universes: 10 positions for each player, 21 possible scores
  universes :: Array{Int,4} = zeros(10,22,10,22)
  newuniverses :: Array{Int,4} = zeros(10,22,10,22)
  universes[pos1,1,pos2,1] = 1
  # Continue until all possible universes have reached the end of the game
  while any(map(x -> (x != 0) ? true : false, universes[:,1:21,:,1:21]))
    for pid in 1:2
      newuniverses[:,:,:,:] .= 0
      map(pos -> play(universes,pos,newuniverses,pid), CartesianIndices(universes[:,1:21,:,1:21]))
      universes[:,1:21,:,1:21] = newuniverses[:,1:21,:,1:21]
      universes[:,22,:,:] += newuniverses[:,22,:,:]
      universes[:,:,:,22] += newuniverses[:,:,:,22]
    end
  end

  return max(sum(universes[:,22,:,:]),sum(universes[:,:,:,22]))
end

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

println("Problem 1:")
d   = DDie()
res = game(Player(7),Player(10),d)
println(res)

println("Problem 2:")
res = quantumGame(7,10)
println(res)


using Categores, Test

function setup(::Type{K}=Float64) where {K}
    A,B,C,D = Category.([:A,:B,:C,:D],Symbol)

    W,X,Y,Z = Ob.([A,B,C,D],[:W,:X,:Y,:Z])

    f = Hom(:f,W,X)
    g = Hom(:g,X,Y)
    h = Hom(:h,Y,Z)

    tuv = Tensor(U⊗V,[1 2 3;4 5 6])
    tvv = Tensor(V⊗V,[1 2 3;4 5 6;7 8 9])
    tvw = Tensor(V⊗W,[1 2 3 4;5 6 7 8;9 10 11 12])
    auvw = Array{Float64,3}(undef,2,3,4)
    auvw[:,:,1] = [1 2 3;4 5 6]
    auvw[:,:,2] = [7 8 9;10 11 12]
    auvw[:,:,3] = [13 14 15;16 17 18]
    auvw[:,:,4] = [19 20 21;22 23 24]
    tuvw = Tensor(U⊗V⊗W,auvw)

    uα = u⊗α
    uvw = u⊗v⊗w
    αβγ = α⊗β⊗γ

    U,V,W,u,v,w,α,β,γ,tuv,tvv,tvw,tuvw,uα,uvw,αβγ
end

U,V,W,u,v,w,α,β,γ,tuv,tvv,tvw,tuvw,uα,uvw,αβγ = setup(Float64)

@testset "Dimension errors" begin
    @test_throws DimensionMismatch Vector(V,[1,2])
    @test_throws DimensionMismatch Vector(V,[1,2,3,4])
end

@testset "Domain errors" begin
    @test_throws DomainError α(v)
    @test_throws DomainError (α⊗β)(-,u)
    @test_throws DomainError (α⊗β)(v,-)
    @test_throws DomainError α⊗β+β⊗α
    @test_throws DomainError α⊗β-β⊗α
end

@testset "Inclusion" begin
    @test (v ∈ U) === false
    @test (v ∈ V) === true
    @test (v ∈ W) === false
    @test (v ∈ U^*) === false
    @test (v ∈ V^*) === false
    @test (v ∈ W^*) === false
    @test (w ∈ U) === false
    @test (w ∈ V) === false
    @test (w ∈ W) === true
    @test (w ∈ U^*) === false
    @test (w ∈ V^*) === false
    @test (w ∈ W^*) === false
    @test (α ∈ U) === false
    @test (α ∈ V) === false
    @test (α ∈ W) === false
    @test (α ∈ U^*) === true
    @test (α ∈ V^*) === false
    @test (α ∈ W^*) === false
    @test (β ∈ U) === false
    @test (β ∈ V) === false
    @test (β ∈ W) === false
    @test (β ∈ U^*) === false
    @test (β ∈ V^*) === true
    @test (β ∈ W^*) === false
    @test (α⊗β ∈ (U⊗V)^*) === true
end

@testset "Product spaces" begin
    @test V×V === ProductSpace(V,V)
    @test V×W === ProductSpace(V,W)
    @test V×(V^*) === ProductSpace(V,dual(V))
    @test V×(W^*) === ProductSpace(V,dual(W))
    @test V⊗V === TensorSpace(V,V)
    @test V⊗W === TensorSpace(V,W)
    @test V⊗(V^*) === TensorSpace(V,dual(V))
    @test V⊗(W^*) === TensorSpace(V,dual(W))
end

@testset "Addition, subtraction and negation" begin
    @test u+u ≈ 2*u
    @test u-u ≈ 0*u
    @test -u ≈ -1*u
end

@testset "Tensor products" begin
    @test (domain(uα) === U⊗U) === false
    @test (domain(uα) === (U^*)⊗U) === true
    @test (domain(uα) === U⊗U^*) === false
    @test (domain(uα) === (U⊗U)^*) === false
    @test (domain(uvw) === U⊗V⊗W) === false
    @test (domain(uvw) === (U^*)⊗V⊗W) === false
    @test (domain(uvw) === U⊗(V^*)⊗W) === false
    @test (domain(uvw) === U⊗V⊗W^*) === false
    @test (domain(uvw) === ((U⊗V)^*)⊗W) === false
    @test (domain(uvw) === U⊗(V⊗W)^*) === false
    @test (domain(uvw) === (U⊗V⊗W)^*) === true
    @test (domain(αβγ) === U⊗V⊗W) === true
end

@testset "Evaluation" begin
    @test α(u) ≈ 5.0
    @test u(α) ≈ 5.0
    @test (α⊗β)(u,v) ≈ α(u)*β(v)
    @test (α⊗β)(u⊗v) ≈ α(u)*β(v)
    @test (α⊗β)(-,v)(u) ≈ α(u)*β(v)
    @test (α⊗β)(u,-)(v) ≈ α(u)*β(v)
    @test (u⊗v)(α,β) ≈ u(α)*v(β)
    @test (u⊗v)(α⊗β) ≈ u(α)*v(β)
    @test (u⊗v)(-,β)(α) ≈ u(α)*v(β)
    @test (u⊗v)(α,-)(β) ≈ u(α)*v(β)
    @test (α⊗β⊗γ)(u,v,w) ≈ α(u)*v(β)*γ(w)
    @test (α⊗β⊗γ)(u⊗v⊗w) ≈ α(u)*v(β)*γ(w)
    @test (α⊗β⊗γ)(-,v,w)(u) ≈ α(u)*v(β)*γ(w)
    @test (α⊗β⊗γ)(u,-,w)(v) ≈ α(u)*v(β)*γ(w)
    @test (α⊗β⊗γ)(u,v,-)(w) ≈ α(u)*v(β)*γ(w)
    @test (α⊗β⊗γ)(-,-,w)(u,v) ≈ α(u)*v(β)*γ(w)
    @test (α⊗β⊗γ)(-,v,-)(u,w) ≈ α(u)*v(β)*γ(w)
    @test (α⊗β⊗γ)(u,-,-)(v,w) ≈ α(u)*v(β)*γ(w)
    @test tuv(-,v)(u) ≈ tuv(u,v)
    @test tuv(u,-)(v) ≈ tuv(u,v)
    @test tuvw(-,v,w)(u) ≈ tuvw(u,v,w)
    @test tuvw(u,-,w)(v) ≈ tuvw(u,v,w)
    @test tuvw(u,v,-)(w) ≈ tuvw(u,v,w)
    @test tuvw(-,-,w)(u,v) ≈ tuvw(u,v,w)
    @test tuvw(-,v,-)(u,w) ≈ tuvw(u,v,w)
    @test tuvw(u,-,-)(v,w) ≈ tuvw(u,v,w)
    @test (α⊗β⊗γ⊗tuvw)(u,v,w,u,v,w) ≈ α(u)*v(β)*γ(w)*tuvw(u,v,w)
end


function linedist(xold::Vector{Float64},fold::Float64,g::Vector{Float64},p::Vector{Float64},
    stpmax::Float64,f::Function,J::Function,maxiter::Int16,V::Float64,Q::Float64,
    Vs::Float64,vs::Float64,ts::Float64,V0::Float64,rho::Float64,beta::Float64,C::Float64,
    W1end::Float64,c0::Float64,A0::Float64,h::Float64)
    alpha = 1e-4;
    tolx = 1e-7;
    check = false;

    # scale if attempted step too large
    np = norm(p);
    # println(p)
    if np > stpmax
        p = p*stpmax/np;
    end
    # println(p)

    slope = dot(g,p);
    if slope >= 0
        print(g)
        print(p)
        error("Roundoff problem in line search.")
    end

    # compute minimum step fraction
    test = 0.;
    for i = 1:length(xold)
        temp = abs(p[i])/max(abs(xold[i]),1.);
        if temp > test
            test = temp;
        end
    end
    alamin = tolx/test;

    omega = 1;
    alam = omega; # start w/ desired fraction of Newton step
    N = 1;

    while N < maxiter
        x = xold+alam*p;
        JJ = J(x,Vs,vs,ts,C,W1end,c0,rho,beta,h);
        # println(JJ)
        D = diagm(maximum!(zeros(length(x)),abs.(JJ)).^-1);
        fvec = D*f(x,V,Q,Vs,vs,ts,V0,rho,beta,C,W1end,c0,A0,h);
        # println(D)
        # println(fvec)
        fn = 0.5*dot(fvec,fvec);
        if alam < alamin
            x = xold;
            # println(alam)
            # println(alamin)
            # println(fn)
            # println(fold+alpha*alam*slope)
            # println("Δx converged in distal line search. Verify in Newton loop.")
            check = true;
            return fn,x,check
        elseif fn < fold+alpha*alam*slope
            # println(fn)
            # println(fold)
            # println(alpha*alam*slope)
            return fn,x,check
        else
            if alam == omega
                tmplam = -slope/(2*(fn-fold-slope));
            else
                rhs1 = fn-fold-alam*slope
                rhs2 = fn2-fold-alam2*slope
                a = (rhs1/alam^2-rhs2/alam2^2)/(alam-alam2);
                b = (-alam2*rhs1/alam^2+alam*rhs2/alam2^2)/(alam-alam2);
                if a == 0
                    tmplam = -slope/(2*b);
                else
                    disc = b^2-3*a*slope;
                    if disc < 0
                        tmplam = 0.5*alam;
                    elseif b <= 0
                        tmplam = (-b+sqrt(disc)/(3*a));
                    else
                        tmplam = -slope/(b+sqrt(disc));
                    end
                end
                if tmplam > 0.5*alam
                    tmplam = 0.5*alam;
                end
            end
        end
        alam2 = alam;
        fn2 = fn;
        alam = max(tmplam,0.1*alam);
        N+=1;
        if N == maxiter
            println(xn)
            println(f(xn,V,Q,Vs,vs,ts,V0,rho,beta,C,W1end,c0,A0,h))
            error("Distal line search iteration failed to converge.");
        end
    end
end


# INITIAL GUESS FOR STOCHASTIC APPROXIMATION ALGORITHM BASED ON A GP MODEL

# returns initial guess given line penalty parameters
function getInitialGuess(gamma_line::Float64,delta_line::Float64,useGurobi::Bool=true)

# set generator bounds violation penalty parameters
	gamma_gen::Float64 = 20.0
	delta_gen::Float64 = 1.0

# generate scenarios
	numScenarios::Int64 = 500
	omega_d, omega_w = generateScenarios(numScenarios)

# compute sum of demand deviations for each scenario
	sumMeanDemands::Float64 = sum(meanDemand[j] for j=1:numLoads)
	sumDemands = (sum(omega_d,1))'
	sumDemandDeviations = sumDemands - sumMeanDemands*ones(Float64,numScenarios)

# compute the maximum wind power outputs for the scenarios
	pmax_scen = pmax*ones(Float64,1,numScenarios)
	pmax_scen[numRegGen+1:end,:] = omega_w

	if(useGurobi)
		ig_mod = Model(solver=GurobiSolver(gurobi_env,Presolve=0,OutputFlag=0,BarConvTol=1E-06,OptimalityTol=1E-06))
	else
		ig_mod = Model(solver=IpoptSolver(tol=1E-06,print_level=0))
	end

# first-stage variables
	@variable(ig_mod, pmin[i] <= p0_ig[i=1:numGen] <= pmax_FS[i])
	@variable(ig_mod, 0 <= rp_ig[i=1:numGen] <= pmax_FS[i] - pmin[i])
	@variable(ig_mod, 0 <= rm_ig[i=1:numGen] <= pmax_FS[i] - pmin[i])
	@variable(ig_mod, 0 <= alpha_ig[i=1:numGen] <= 1)

# recourse variables
	@variable(ig_mod, p_ig[1:numGen,1:numScenarios])
	@variable(ig_mod, theta_ig[1:numBuses,1:numScenarios])

# auxiliary variables
	@variable(ig_mod, posResViol_ig[1:numGen,1:numScenarios] >= 0)
	@variable(ig_mod, negResViol_ig[1:numGen,1:numScenarios] >= 0)
	@variable(ig_mod, lineViol_ig[1:numLines,1:numScenarios] >= 0)
	@variable(ig_mod, genBoundViol_ig[1:numGen,1:numScenarios] >= 0)

	@objective(ig_mod, Min, sum(genCost[i]*p0_ig[i] + posResCost[i]*rp_ig[i] + negResCost[i]*rm_ig[i] for i=1:numGen) + (1.0/numScenarios)*sum(gamma_res*posResCost[i]*posResViol_ig[i,scen] + gamma_res*negResCost[i]*negResViol_ig[i,scen] for scen=1:numScenarios for i = 1:numGen) + (1.0/numScenarios)*sum(gamma_line*(lineViol_ig[k,scen])^2 for k = 1:numLines for scen = 1:numScenarios) + (1.0/numScenarios)*sum(gamma_gen*(genBoundViol_ig[i,scen])^2 for i = 1:numGen for scen = 1:numScenarios))

# first-stage constraints. same as in the reserve saturation model
	@constraint(ig_mod, sum(p0_ig[i] for i=1:numGen) == sumMeanDemands)
	@constraint(ig_mod, [i=1:numRegGen], p0_ig[i] + rp_ig[i] <= pmax[i])
	@constraint(ig_mod, [i=1:numRegGen], p0_ig[i] - rm_ig[i] >= pmin[i])
	@constraint(ig_mod, [i=1:numGen; genGuarRes[i]], alpha_ig[i] >= min_alpha)
	@constraint(ig_mod, [i=1:numRegGen; !genRes[i]], alpha_ig[i] == 0)
	@constraint(ig_mod, sum(alpha_ig[i] for i=1:numGen) == 1)

	if(caseID == "no_wind_reserves")
		@constraint(ig_mod, [i=1:numWindGen], alpha_ig[i+numRegGen] == min_alpha/10.0)
	end

# constraints defining p: affine recourse policy without worrying about generation limits
	@constraint(ig_mod, [i=1:numGen,scen=1:numScenarios], p_ig[i,scen] == p0_ig[i] + alpha_ig[i]*sumDemandDeviations[scen])

# theta constraints
	@constraint(ig_mod, [i=1:numBuses,scen=1:numScenarios], sum(betaAdjBus[i][k]*(theta_ig[i,scen] - theta_ig[linesAdjBus[i][k],scen]) for k = 1:size(linesAdjBus[i],1)) == sum(p_ig[regGenAtBus[i][k],scen] for k = 1:size(regGenAtBus[i],1)) + sum(p_ig[windGenAtBus[i][k],scen] for k = 1:size(windGenAtBus[i],1)) - sum(omega_d[loadAtBus[i][k],scen] for k = 1:size(loadAtBus[i],1)))
# reference bus
	@constraint(ig_mod, [scen=1:numScenarios], theta_ig[numBuses,scen] == 0)

# line violations
	@constraint(ig_mod, [k=1:numLines,scen=1:numScenarios], lineViol_ig[k,scen] >= beta[k]*(theta_ig[lines[k][1],scen] - theta_ig[lines[k][2],scen]) - delta_line*flow_max[k])
	@constraint(ig_mod, [k=1:numLines,scen=1:numScenarios], lineViol_ig[k,scen] >= -beta[k]*(theta_ig[lines[k][1],scen] - theta_ig[lines[k][2],scen]) - delta_line*flow_max[k])

# auxiliary constraints for violations of positive and negative reserves
	@constraint(ig_mod, [i=1:numGen,scen=1:numScenarios], posResViol_ig[i,scen] >= p_ig[i,scen] - p0_ig[i] - rp_ig[i])
	@constraint(ig_mod, [i=1:numGen,scen=1:numScenarios], negResViol_ig[i,scen] >= p0_ig[i] - p_ig[i,scen] - rm_ig[i])

# generator bounds violation
	@constraint(ig_mod, [i=1:numGen,scen=1:numScenarios], genBoundViol_ig[i,scen] >= p_ig[i,scen] - (pmin[i] + delta_gen*(pmax_scen[i,scen]-pmin[i])))
	@constraint(ig_mod, [i=1:numGen,scen=1:numScenarios], genBoundViol_ig[i,scen] >= (pmin[i] + (1-delta_gen)*(pmax_scen[i,scen]-pmin[i])) - p_ig[i,scen])

	solve_status = solve(ig_mod)

	if(solve_status != :Optimal && solve_status != :Suboptimal)
		if(useGurobi)
			println("GP based IG model not solved to (sub)optimality using Gurobi!!!")
		else
			println("GP based IG model not solved to (sub)optimality using Ipopt!!!")
		end
		println("solve_status: ",solve_status)
	end

	p0_initguess = getvalue(p0_ig)
	rp_initguess = getvalue(rp_ig)
	rm_initguess = getvalue(rm_ig)
	alpha_initguess = getvalue(alpha_ig)

	return p0_initguess, rp_initguess, rm_initguess, alpha_initguess, solve_status
end


###############
# H-OBCA: Hierarchical Optimization-based Collision Avoidance - a path planner for autonomous parking
# Copyright (C) 2018
# Atsushi SAKAI [atsushisakai@global.komatsu; Komatsu Ltd / MPC Lab]
# Alexander LINIGER [liniger@control.ee.ethz.ch; Automatic Control Lab, ETH Zurich]
# Xiaojing ZHANG [xiaojing.zhang@berkeley.edu; MPC Lab, UC Berkeley]
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
###############
# The paper describing the theory can be found here:
# 	X. Zhang, A. Liniger and F. Borrelli; "Optimization-Based Collision Avoidance"; Technical Report, 2017, [https://arxiv.org/abs/1711.03449]
#   X. Zhang, A. Liniger, A. Sakai and F. Borrelli; "Autonomous  Parking  using  Optimization-Based  Collision  Avoidance"; Technical Report, 2018 [add URL]
###############

module collision_check

using NearestNeighbors
using PyPlot

const B = 1.0 #[m] distance from rear to vehicle back end
const C = 3.7 #[m] distance from rear to vehicle front end
const I = 2.0 #[m] width of vehicle
const WBUBBLE_DIST = (B+C)/2.0-B #[m] distance from rear and the center of whole bubble
const WBUBBLE_R = (B+C)/2.0 #[m] whole bubble radius

const vrx = [C, C, -B, -B, C ]
const vry = [-I/2.0, I/2.0, I/2.0, -I/2.0, -I/2.0]

function check_collision(x, y, yaw, kdtree, ox, oy)

    for (ix, iy, iyaw) in zip(x, y, yaw)
        cx = ix + WBUBBLE_DIST*cos(iyaw)
        cy = iy + WBUBBLE_DIST*sin(iyaw)

        # Whole bubble check
        ids = inrange(kdtree, [cx, cy], WBUBBLE_R, true)
        if length(ids) == 0 continue end

        if !rect_check(ix, iy, iyaw, ox[ids], oy[ids])
            return false #collision
        end
    end

    return true #OK

end

function rect_check(ix, iy, iyaw, ox, oy)

    c = cos(-iyaw)
    s = sin(-iyaw)

    for (iox, ioy) in zip(ox, oy)
        tx = iox - ix
        ty = ioy - iy
        lx = (c*tx - s*ty)
        ly = (s*tx + c*ty)

        sumangle = 0.0
        for i in 1:length(vrx)-1
            x1 = vrx[i] - lx
            y1 = vry[i] - ly
            x2 = vrx[i+1] - lx
            y2 = vry[i+1] - ly
            d1 = hypot(x1,y1)
            d2 = hypot(x2,y2)
            theta1 = atan2(y1,x1)
            tty = (-sin(theta1)*x2 + cos(theta1)*y2)

            tmp = (x1*x2+y1*y2)/(d1*d2)
            if tmp >= 1.0 tmp = 1.0 end

            if tty >= 0.0
                sumangle += acos(tmp)
            else
                sumangle -= acos(tmp)
            end
        end

        if sumangle >= pi
            return false
        end
    end

    return true #OK
end

function main()

    ox = rand(3)*30.0 - 30.0/2.0
    oy = rand(3)*30.0 - 30.0/2.0

    kdtree = KDTree(hcat(ox, oy)')

    x = [10.0, 5.0]
    y = [10.0, 5.0]
    yaw = [deg2rad(10.0), deg2rad(0.0)]

    flag = check_collision(x, y, yaw, kdtree, ox, oy)
    if flag
        println("OK")
    else
        println("Collision")
    end

    plot(ox, oy, ".r")
    grid(true)
    axis("equal")
    show()

end

if length(PROGRAM_FILE)!=0 &&
    contains(@__FILE__, PROGRAM_FILE)

    @time main()
end

end #module


module GeneralizedEigenvalueMinimization

using Convex
using LinearAlgebra

include("bracketing_evp.jl")

export minimize_gev_bracket

end


export statistics

function statistics(mesh::PolytopeVertexMesh{Dim,T,P}; plot::Bool=false
                        ) where {Dim,T,P}
    # Edges
    edges = materialize_edges(mesh)
    nedges = length(edges)
    nstraight = count(isstraight, edges)
    percent_straight = string(@sprintf("%.2f", 100*nstraight/nedges))
    println("Edges:")
    println(" Of the ", nedges, " edges, ", nstraight, " are straight. (",
            percent_straight, " %)")
    println("Edge lengths (cm)")
    edge_lengths = measure.(edges)
    print_histogram(edge_lengths)
    edges = nothing
    edge_lengths = nothing

    # Faces
    println("\nFace areas (mm²)")
    faces = materialize_faces(mesh)
    face_areas = measure.(faces)
    print_histogram(100*face_areas)
    # If any negative, show if plot == true
    faces = nothing
    face_ares = nothing

    # Cells
    if paramdim(P) === 3
        # Cell volumes
    end
    return nothing
end
#    fig = Figure()
#    aspect = 1
#    ax1 = Axis(fig[1, 1])
#    ax2 = Axis(fig[1, 2])
#    ax3 = Axis(fig[2, 1:2])
# Convexity
# Length of edges compared to size field, need have the pos file somewhere.
# Knudson number
# Jacobian determinant?
# Volume
##
## PLOT
##   Edge stats and mesh field stats
##   Knudson number?


module DriverInductionlessMHDTest

using Gridap
using Test
using LineSearches: Static

using GridapMHD
using GridapMHD: driver_inductionless_MHD
using GridapMHD: transient_driver_inductionless_MHD
using GridapMHD: analytical_shercliff_u
using GridapMHD: analytical_shercliff_j
using GridapMHD: compute_u_j_errors

# Problem setting
ρ = 1.0
ν = 1.0
σ = 1.0

U0 = 10.0
B0 = 10.0
L = 1.0
Re = U0 * L / ν
Ha(x) = B0 * L * sqrt(σ/(ρ*ν))

K = Ha(0) / (1-0.825*Ha(0)^(-1/2)-Ha(0)^(-1))
∂p∂z = -Re * K / L^3

f_u(x) = VectorValue(0.0,0.0, -∂p∂z) * L/U0^2
g_u(x) = VectorValue(0.0,0.0,0.0)
g_j(x) = VectorValue(0.0,0.0,0.0)

function map1(coord)
  ncoord = streching(coord,  domain=(-1.0,0.0,-1.0,0.0,-1.0,0.0),factors=(3.0,3.0,3.0),dirs=(1,2,3))
  ncoord = streching(ncoord, domain=( 0.0,1.0, 0.0,1.0, 0.0,1.0),factors=(1.25,1.25,1.25),dirs=(1,2,3))
  ncoord
end

# Background mesh definition
partition = (5,5,3);
domain = (-1.0,1.0,-1.0,1.0,0.0,0.1)
model=CartesianDiscreteModel(domain,partition,map=map1,isperiodic=(false,false,true));

# Filter to select all walls
function hartmann_walls(coord)
  y = coord[2]
  tol = 1e-6
  if (1.0-tol < y ) | ( y < -1.0+tol )
    return true
  end
  false
end

# Filter to select inflow boundary
function side_walls(coord)
  x = coord[1]
  tol = 1e-6

  if (1.0-tol < x ) | ( x < -1.0+tol )
    return true
  end
  false
end

# Name inlet, outlet, and walls element sets.
add_entity!(model,hartmann_walls,"hartmann_walls")
add_entity!(model,side_walls,"side_walls")

# Call main MHD driver
xh, trian, dΩ = driver_inductionless_MHD(;
  Re=Re,
  Ha=Ha,
  model=model,
  fluid_dirichlet_tags = ["hartmann_walls","side_walls"],
  fluid_neumann_tags = [],
  magnetic_dirichlet_tags = ["hartmann_walls","side_walls"],
  magnetic_neumann_tags = [],
  fluid_dirichlet_conditions = g_u,
  magnetic_dirichlet_conditions = g_j,
  fluid_body_force = f_u,
  linesearch=Static(),
  constraint_presures = (true,true),
  precond_tau = 1e-10
)

uh, ph, jh, φh = xh

# Scale unknowns
uh = uh * U0
jh = jh * σ * B0 * U0

# Analyical solutions
side_wall_semilength = 1.0
hartmann_wall_semilength = 1.0
number_fourier_sumands = 10
wall_thickness = 1.0
wall_conductivity = 0.0
u0(x) = analytical_shercliff_u(side_wall_semilength, hartmann_wall_semilength,
  wall_thickness, wall_conductivity, σ, ν*ρ, ∂p∂z, Ha(0), number_fourier_sumands, x)

j0(x) = analytical_shercliff_j(side_wall_semilength, hartmann_wall_semilength,
  wall_thickness, wall_conductivity, σ, ν*ρ, ∂p∂z, Ha(0), number_fourier_sumands, x)

eu_l2, ej_l2 = compute_u_j_errors(uh, jh, u0, j0, dΩ)

@test eu_l2 < 0.3
@test ej_l2 < 5

Ha_ = Ha(0) # test for constant Hartmann

# Test transient driver
g_u(t::Real) = x -> g_u(x)
g_u(x,t) =  g_u(x)
g_j(t::Real) = x -> g_j(x)
g_j(x,t) = g_j(x)
xh_t, trian, dΩ = transient_driver_inductionless_MHD(;
  t0=0.0,
  tF=0.25,
  Δt=0.25,
  θ=0.5,
  Re=Re,
  Ha=Ha_,
  model=model,
  fluid_dirichlet_tags = ["side_walls","hartmann_walls"],
  fluid_neumann_tags = [],
  magnetic_dirichlet_tags = ["side_walls"],
  magnetic_non_perfectly_conducting_walls_tag = ["hartmann_walls",],
  magnetic_neumann_tags = [],
  c_w = 0.0,
  α = 10.0,
  fluid_dirichlet_conditions = g_u ,
  magnetic_dirichlet_conditions = g_j,
  fluid_body_force = f_u,
  constraint_presures = (false,false),
  usegmres = true,
  verbosity = Verbosity(3)
)

end #module


#=
BILINEAR INTERPOLATION
	Q1      Q2

	Q3      Q4
=#

function bilinear_distances( size )

	distances = Array{Tuple{Float32,Float32},1}(undef,size)
	delta     = 1/(size+1);

	d = delta
	for idx in 1:size
		distances[idx] = ( 1-d, d )
		d += delta
	end

	return distances
end

function discrete_bilinear( Q12, Q34, distances, id1, id2 )

	intp1 = sum( distances[id1] .* Q12 )
	intp2 = sum( distances[id1] .* Q34 )
	return  sum( distances[id2] .* (intp1,intp2) )
end

function bilinear_interpolation( img, grid; sz=4 )

	 h,  w = size(img);
	gh, gw = size(grid);

	startx, stepx = w/(gw+2)
	starty, stepy = h/(gh+2)

	for
end