### RandomReal

@mkapprule RandomReal

@doap RandomReal() = return rand()

### Random

@sjdoc Random """
    Random(Integer)

return 0 or 1 with equal probability.

    Random(Integer, [lo,hi])

return a random integer between `lo` and `hi`, inclusive.

    Random(Real)

return a random floating-point number between 0 and 1.

    Random(Complex)

return a random complex floating-point number with parts between 0 and 1.
"""

@mkapprule Random

@doap Random() = rand()

@doap function Random(sym::SJSym)
    if sym == :Integer
        rand(0:1)
    elseif sym == :Real
        rand()
    elseif sym == :Complex
        complex(rand(),rand())
    else
        mx
    end
end

@doap function Random(sym::SJSym, lst::ListT)
    if sym == :Integer
        rand(UnitRange(margs(lst)...))
    end
end

### RandomInteger

@sjdoc RandomInteger """
    Random([nmin,nmax])

return a random integer between `nmin` and `nmax`.

    RandomInteger(nmax)

return a random integer between 0 and `nmax`.

    RandomInteger()

return 0 or 1 with equal probablility.

    RandomInteger(range, n)

return an array of `n` random integers
"""

@mkapprule RandomInteger

@doap RandomInteger() = rand(0:1)
@doap RandomInteger(n::Integer) = rand(0:n)

@doap RandomInteger(x::Mxpr{:List}) = _randintrange(margs(x)...)

_randintrange(n::Integer,m::Integer) = rand(n:m)

@doap function RandomInteger(x::Mxpr{:List},n::Integer)
    lims = x[1]:x[2]
    nargs = newargs(n)
    for i=1:n  nargs[i] = rand(lims) end
    mxpr(:List,nargs)
end


module ExtrusionPolytopesTests

using Test
using InplaceArrays.Helpers
using InplaceArrays.TensorValues
using InplaceArrays.Arrays
using InplaceArrays.Fields
using InplaceArrays.ReferenceFEs

@test num_faces(VERTEX) == 1
@test num_dims(VERTEX) == 0

p = Polytope(HEX_AXIS, HEX_AXIS)
test_polytope(p,optional=true)

@test get_faces(p,0,0) == [[1], [2], [3], [4]]
@test get_faces(p,1,0) == [[1, 2], [3, 4], [1, 3], [2, 4]]
@test get_faces(p,2,0) == [[1, 2, 3, 4]]
@test get_faces(p,0,1) == [[1, 3], [1, 4], [2, 3], [2, 4]]
@test get_faces(p,1,1) == [[1], [2], [3], [4]]
@test get_faces(p,2,1) == [[1, 2, 3, 4]]
@test get_faces(p,0,2) == [[1], [1], [1], [1]]
@test get_faces(p,1,2) == [[1], [1], [1], [1]]
@test get_faces(p,2,2) == [[1]]
@test get_facedims(p) == [0,0,0,0,1,1,1,1,2]
@test get_offsets(p) == [0,4,8]
@test num_facets(p) == 4
@test num_edges(p) == 4

r = Point{2,Float64}[(1, 0), (1, 0), (0, 1), (0, 1)]
@test get_edge_tangents(p) == r

r = [
  [1, 2, 3, 4], [1, 3, 2, 4], [2, 1, 4, 3], [2, 4, 1, 3],
  [3, 1, 4, 2], [3, 4, 1, 2], [4, 2, 3, 1], [4, 3, 2, 1]]
@test get_vertex_permutations(p) == r

r = Point{2,Float64}[(0, -1), (0, 1), (-1, 0), (1, 0)]
@test get_facet_normals(p) == r

@test num_faces(p) == 9
@test num_dims(p) == 2

p = Polytope(TET_AXIS, TET_AXIS, TET_AXIS)
test_polytope(p,optional=true)

@test num_faces(p) == 15
@test num_dims(p) == 3

x = Point{3,Float64}[(0, 0, 0), (1, 0, 0), (0, 1, 0), (0, 0, 1)]
@test get_vertex_coordinates(p) == x

p = SEGMENT
test_polytope(p,optional=true)
@test get_vertex_coordinates(p) == VectorValue{1,Float64}[(0),(1)]
@test get_edge_tangents(p) == VectorValue{1,Float64}[(1)]
@test get_facet_normals(p) == VectorValue{1,Float64}[(-1),(1)]

p = VERTEX
test_polytope(p,optional=true)
@test get_vertex_coordinates(p) == VectorValue{0,Float64}[()]
@test get_edge_tangents(p) == VectorValue{0,Float64}[]
@test get_facet_normals(p) == VectorValue{0,Float64}[]

test_polytope(TRI,optional=true)
test_polytope(QUAD,optional=true)
test_polytope(TET,optional=true)
test_polytope(HEX,optional=true)

@test num_facets(SEGMENT) == 2
@test num_facets(TRI) == 3
@test num_facets(QUAD) == 4
@test num_facets(TET) == 4
@test num_facets(HEX) == 6

@test num_vertices(PYRAMID) == 5
# TODO There is a problem with PYRAMID
#@test num_edges(PYRAMID) == 8
#@test num_facets(PYRAMID) == 5

end # module


"""
Computes the maximum multiroute flow (for any real number of routes)
between the source and target vertexes in a flow graph using the
[Extended Multiroute Flow algorithm](http://dx.doi.org/10.1016/j.disopt.2016.05.002).
If a number of routes is given, returns the value of the multiroute flow as
well as the final flow matrix, along with a multiroute cut if
Boykov-Kolmogorov max-flow algorithm is used as a subroutine.
Otherwise, it returns the vector of breaking points of the parametric
multiroute flow function.
Use a default capacity of 1 when the capacity matrix isn\'t specified.
Requires arguments:
- flow_graph::ADiGraph                    # the input graph
- source::Int                            # the source vertex
- target::Int                            # the target vertex
- capacity_matrix::AbstractMatrix{T}   # edge flow capacities
- flow_algorithm::AbstractFlowAlgorithm  # keyword argument for algorithm
- routes::Int                            # keyword argument for routes
"""
function emrf(
  flow_graph::ADiGraph,                   # the input graph
  source::Int,                           # the source vertex
  target::Int,                           # the target vertex
  capacity_matrix::AbstractMatrix{T},  # edge flow capacities
  flow_algorithm::AbstractFlowAlgorithm, # keyword argument for algorithm
  routes::R = 0
  ) where {T<:AbstractFloat, R<:Real}
  breakingpoints = breakingPoints(flow_graph, source, target, capacity_matrix)
  if routes > zero(R)
    x, f = intersection(breakingpoints, routes)
    return maximum_flow(flow_graph, source, target, capacity_matrix,
                        algorithm = flow_algorithm, restriction = x)
  end
  return breakingpoints
end

"""
Output a set of (point,slope) that compose the restricted max-flow function.
One point by possible slope is enough (hence O(λ×max_flow) complexity).
Requires arguments:
- flow_graph::ADiGraph                    # the input graph
- source::Int                            # the source vertex
- target::Int                            # the target vertex
- capacity_matrix::AbstractMatrix{T}   # edge flow capacities
"""
function auxiliaryPoints(
  flow_graph::ADiGraph,                   # the input graph
  source::Int,                           # the source vertex
  target::Int,                           # the target vertex
  capacity_matrix::AbstractMatrix{T}   # edge flow capacities
  ) where T<:AbstractFloat
  # Problem descriptors
  λ = maximum_flow(flow_graph, source, target)[1] # Connectivity
  n = nv(flow_graph) # number of nodes
  r1, r2 = minmaxCapacity(capacity_matrix) # restriction left (1) and right (2)
  auxpoints = fill((0., 0.), λ + 1)

  # Initialisation of left side (1)
  f1, F1, cut1 = maximum_flow(flow_graph, source, target, capacity_matrix,
                 algorithm = BoykovKolmogorovAlgorithm(), restriction = r1)
  s1 = slope(flow_graph, capacity_matrix, cut1, r1) # left slope
  auxpoints[λ + 1 - s1] = (r1, f1) # Add left initial auxiliary point

  # Initialisation of right side (2)
  f2, F2, cut2 = maximum_flow(flow_graph, source, target, capacity_matrix,
                 algorithm = BoykovKolmogorovAlgorithm(), restriction = r2)
  s2 = slope(flow_graph, capacity_matrix, cut2, r2) # right slope
  auxpoints[λ + 1 - s2] = (r2, f2) # Add right initial auxiliary point

  # Loop if the slopes are distinct by at least 2
  if s1 > s2 + 1
    queue = [((f1, s1, r1), (f2, s2, r2))]

    while !isempty(queue)
      # Computes an intersection (middle) with a new associated slope
      (f1, s1, r1), (f2, s2, r2) = pop!(queue)
      r, expectedflow = intersection(r1, f1, s1, r2, f2, s2)
      f, F, cut = maximum_flow(flow_graph, source, target, capacity_matrix,
                  algorithm = BoykovKolmogorovAlgorithm(), restriction = r)
      s = slope(flow_graph, capacity_matrix, max.(cut, 1), r) # current slope
      auxpoints[λ + 1 - s] = (r, f)
      # If the flow at the intersection (middle) is as expected
      if expectedflow ≉ f # approximatively not equal (enforced by floating precision)
        # if the slope difference between (middle) and left is at least 2
        # push (left),(middle)
        if s1 > s + 1 && (r2, f2) ≉ (r, f)
          q = (f1, s1, r1), (f, s, r)
          push!(queue, q)
        end
        # if the slope difference between (middle) and right is at least 2
        # push (middle),(right)
        if s > s2 + 1 && (r1, f1) ≉ (r, f)
          q = (f, s, r), (f2, s2, r2)
          push!(queue, q)
        end
      end
    end
  end
  return auxpoints
end

"""
Calculates the breaking of the restricted max-flow from a set of auxiliary points.
Requires arguments:
- flow_graph::ADiGraph                    # the input graph
- source::Int                            # the source vertex
- target::Int                            # the target vertex
- capacity_matrix::AbstractMatrix{T}   # edge flow capacities
"""
function breakingPoints(
  flow_graph::ADiGraph,                   # the input graph
  source::Int,                           # the source vertex
  target::Int,                           # the target vertex
  capacity_matrix::AbstractMatrix{T}   # edge flow capacities
  ) where T<:AbstractFloat
  auxpoints = auxiliaryPoints(flow_graph, source, target, capacity_matrix)
  λ = length(auxpoints) - 1
  left_index = 1
  breakingpoints = Vector{Tuple{T, T, Int}}()

  for (id, point) in enumerate(auxpoints)
    if id == 1
      push!(breakingpoints, (0., 0., λ))
    else
      pleft = breakingpoints[left_index]
      if point[1] != 0
        x, y = intersection(pleft[1], pleft[2], pleft[3],
                            point[1], point[2], λ + 1 - id)
        push!(breakingpoints,(x, y, λ + 1 - id))
        left_index += 1
      end
    end
  end
  return breakingpoints
end

"""
Function to get the nonzero min and max function of a Matrix.

Note: this is more efficient than maximum() / minimum() / extrema()
since we have to ignore zero values.since we have to ignore zero values.

Requires argument:
- capacity_matrix::AbstractMatrix{T}   # edge flow capacities
"""
function minmaxCapacity(
  capacity_matrix::AbstractMatrix{T}    # edge flow capacities
  ) where T<:AbstractFloat
# Function to get the nonzero min and max function of a Matrix
# note: this is more efficient than maximum() / minimum() / extrema()
#       since we have to ignore zero values.
  cmin, cmax = typemax(T), typemin(T)
  for c in capacity_matrix
    if c > zero(T)
      cmin = min(cmin, c)
    end
    cmax = max(cmax, c)
  end
  return cmin, cmax
end

"""
Function to get the slope of the restricted flow. The slope is initialized at 0
and is incremented for each non saturated edge in the restricted min-cut.
Requires argument:
  flow_graph::ADiGraph,                   # the input graph
  capacity_matrix::AbstractMatrix{T},  # edge flow capacities
  cut::Vector{Int},                      # cut information for vertices
  restriction::T                         # value of the restriction
"""
function slope(
  flow_graph::ADiGraph,                   # the input graph
  capacity_matrix::AbstractMatrix{T},  # edge flow capacities
  cut::Vector{Int},                      # cut information for vertices
  restriction::T                         # value of the restriction
  ) where T<:AbstractFloat
  # Function to get the slope of the restricted flow
  slope = 0
  for e in edges(flow_graph)
      if cut[dst(e)] == 2 > cut[src(e)] &&
            capacity_matrix[src(e), dst(e)] > restriction
        slope += 1
    end
  end
  return slope
end

"""
Computes the intersection between:
1) Two lines
2) A set of segments and a linear function of slope k passing by the origin.
Requires argument:
1) - x1, y1, a1, x2, y2, a2::T<:AbstractFloat # Coordinates/slopes
2) - points::Vector{Tuple{T, T, Int}}         # vector of points with T<:AbstractFloat
   - k::R<:Real                             # number of routes (slope of the line)
"""
function intersection(
  x1::T,          # x coordinate of point 1
  y1::T,          # y coordinate of point 1
  a1::Int,        # slope passing by point 1
  x2::T,          # x coordinate of point 2
  y2::T,          # y coordinate of point 2
  a2::R           # slope passing by point 2
  ) where {T<:AbstractFloat, R<:Real}

  (a1 == a2) && return -1., -1. # result will be ignored in other intersection method
  b1 = y1 - a1 * x1
  b2 = y2 - a2 * x2
  x = (b2 - b1) / (a1 - a2)
  y = a1 * x + b1
  return x, y
end

# Compute the intersection between a set of segment and a line of slope k passing by the origin
function intersection(
  points::Vector{Tuple{T, T, Int}},  # vector of breaking points
  k::R                               # number of routes (slope of the line)
  ) where {T<:AbstractFloat, R<:Real}
  λ = points[1][1] # Connectivity

  # Loop over the segments (pair of breaking points)
  for (id, p) in enumerate(points[1:(end - 1)])
    if id == 1
      (k ≈ λ) && return points[2]
    else
      x, y = intersection(p[1], p[2], p[3], 0., 0., k)
      (p[1] ≤ x ≤ points[id + 1][1]) && return x, y
    end
  end
  p = points[end]
  return intersection(p[1], p[2], p[3], 0., 0., k)
end

"""
Redefinition of ≈ (isapprox) for a pair of points
Requires argument:
  a::Tuple{T, T},         # Point A with floating coordinates
  b::Tuple{T, T}          # Point B with floating coordinates
"""
function ≈(
  a::Tuple{T, T},         # Point A with floating coordinates
  b::Tuple{T, T}          # Point B with floating coordinates
  ) where T<:AbstractFloat
  return a[1] ≈ b[1] && a[2] ≈ b[2]
end


# compare FD to a traditional rank-k approximation
# Danny Perry (dperry@cs.utah.edu)
# Oct 2014
#

using PyPlot
using Debug
using RDatasets

using Sketch

X = []
Labels = []

#name = "iris"
name = "baseball"

if name == "iris"
	iris = dataset("datasets","iris")

	X = [ convert(Array{Float64,1}, iris[1]) convert(Array{Float64,1}, iris[2]) convert(Array{Float64,1}, iris[3]) convert(Array{Float64,1}, iris[4]) ]
	n = size(X,1)
	println("X: ", size(X))
	Labels = ASCIIString[]
	for i=1:n
		push!(Labels, iris[5][i])
	end
else

	baseball = dataset("plyr","baseball")

	numerical_cols = [1,4,8,9,10,11,12,13,17]
	X = zeros(size(baseball,1),length(numerical_cols))

	for i=1:length(numerical_cols)
		println(i,": ",numerical_cols[i])
		X[:,i] = convert(Array{Float64,1}, baseball[numerical_cols[i]])
	end

end

n = size(X,1)
X_centered = X .- (ones(n,1)*mean(X,1))

@debug function compare(X)

	n = size(X,1)
	d = size(X,2)

	errs = Float64[]
	minerrs = Float64[]
	maxerrs = Float64[]

	figure()
	for l=2:20
		B = FD(X,l)
		BB = B'*B
		U,S,V = svd(BB)

		projerr = sum((X .- (X*V*V')).^2,2)
		err = mean(projerr)
		minerr = minimum(projerr)
		maxerr = maximum(projerr)

		push!(errs,err)
		push!(minerrs,minerr)
		push!(maxerrs,maxerr)

		clf()
		plot(1:length(errs),errs,"r-")
		hold(:on)
		plot(1:length(errs),minerrs,"r--")
		plot(1:length(errs),maxerrs,"r--")
		title("FD projection error")
		xlabel("size of sketch")
		ylabel("proj err squared")
		sleep(0.0001)
	end

	@bp
end

compare(X_centered)


export samplecdf,normalizedcumsum!

##bisection inverse

bisectioninv(f::Fun{S,T},x::Real;opts...) where {S,T} = first(bisectioninv(f,[x];opts...))

# gives a generalization of midpoint for when `a` or `b` is infinite
function genmidpoint(a::T,b::T) where T
    if isinf(a) && isinf(b)
        zero(T)
    elseif isinf(a)
        b - 100
    elseif isinf(b)
        a + 100
    else
        (a+b)/2
    end
end

function bisectioninv(f::Fun{S,T},x::Float64;numits::Int=47) where {S,T}
    d=domain(f)
    a = minimum(d);b = maximum(d)

    for k=1:numits  #TODO: decide 47
        m= genmidpoint(a,b)
        val = f(m)

            (val<= x) ? (a = m) : (b = m)
    end
    .5*(a+b)
end

bisectioninv(f::Fun{S,T},x::AbstractVector;opts...) where {S,T} = Float64[bisectioninv(f,xx;opts...) for xx in x]

## Clenshaw bisection

function chebbisectioninv(c::AbstractVector{Float64},x::Float64;numits::Int=47)
    a = -1.;b = 1.

    C=Chebyshev()
    for k=1:numits  #TODO: decide 47
        m=.5*(a+b)
        val = clenshaw(C,c,m)

            (val<= x) ? (a = m) : (b = m)
    end
    .5*(a+b)
end

chebbisectioninv(c::AbstractVector{Float64},xl::AbstractVector{Float64}) =
    (n=length(xl);chebbisectioninv(c,xl,ClenshawPlan(Float64,Chebyshev(),length(c),n)))
function chebbisectioninv(c::AbstractVector{Float64},xl::AbstractVector{Float64},plan::ClenshawPlan{Chebyshev{D,R},Float64}) where {D<:Domain,R}
    n = length(xl)
    a = -fill(1.0,n)
    b = fill(1.0,n)

    for k=1:47  #TODO: decide 47
        m=.5*(a+b);
        vals = clenshaw(c,m,plan)

        for j = 1:n
            (vals[j] <= xl[j]) ? (a[j] = m[j]) : (b[j] = m[j])
        end
    end
    m=.5*(a+b)
end

#here, xl is vector w/ length == #cols of c
chebbisectioninv(c::AbstractMatrix{Float64},xl::AbstractVector{Float64}) =
    (n=length(xl);chebbisectioninv(c,xl,ClenshawPlan(Float64,Chebyshev(),size(c,1),n)))
function chebbisectioninv(c::AbstractMatrix{Float64},xl::AbstractVector{Float64},plan::ClenshawPlan{Chebyshev{D,R},Float64}) where {D<:Domain,R}
    @assert size(c)[2] == length(xl)

    n = length(xl)
    a = -fill(1.0,n)
    b = fill(1.0,n)

    for k=1:47  #TODO: decide 47
        m=.5*(a+b);
        vals = clenshaw(c,m,plan)

        for j = 1:n
            (vals[j] <= xl[j]) ? (a[j] = m[j]) : (b[j] = m[j])
        end
    end
    m=.5*(a+b)
end

for TYP in (:Vector,:Float64)
    @eval begin
        bisectioninv(cf::Fun{SP,Float64},x::$TYP;opts...) where {SP<:Chebyshev} =
            fromcanonical.(Ref(space(cf)),chebbisectioninv(coefficients(cf),x;opts...))
#        bisectioninv{SP<:LineSpace}(cf::Fun{SP,Float64},x::$TYP;opts...)=fromcanonical(cf,chebbisectioninv(coefficients(cf),x;opts...))
    end
end

##normalized cumsum

function normalizedcumsum(f::Fun)
    cf = cumsum(f)
    cf = cf/last(cf)

    cf
end

function subtract_zeroatleft!(f::AbstractVector{Float64})
    for k=2:length(f)
        @inbounds f[1] += (-1.)^k.*f[k]
    end

    f
end

function subtract_zeroatleft!(f::AbstractMatrix{Float64})
    for k=2:size(f)[1],j=1:size(f)[2]
        @inbounds f[1,j] += (-1.)^k.*f[k,j]
    end

    f
end

function multiply_oneatright!(f::AbstractVector{Float64})
    val=0.0
    for k=1:length(f)
        val += f[k]
    end

    val = 1/val

    for k=1:length(f)
        @inbounds f[k] *= val
    end

    f
end

function multiply_oneatright!(f::AbstractMatrix{Float64})

    for j=1:size(f)[2]
        val=0.0
        for k=1:size(f)[1]
            val += f[k,j]
        end

        val = 1/val

        for k=1:size(f)[1]
            @inbounds f[k,j] *= val
        end
    end

    f
end

function chebnormalizedcumsum!(f)
    ultraconversion!(f)
    ultraint!(f)
    subtract_zeroatleft!(f)
    multiply_oneatright!(f)
end

# For RandomMatrices compatibility
normalizedcumsum!(f::AbstractVector{Float64})=chebnormalizedcumsum!(f)

## Sampling

sample(f::Fun,n::Integer) = samplecdf(normalizedcumsum(f),n)

samplecdf(cf::Fun,n::Integer) = bisectioninv(cf,rand(n))

sample(f::Fun) = sample(f,1)[1]
samplecdf(f::Fun) = samplecdf(f,1)[1]
samplecdf(v::AbstractVector) = chebbisectioninv(v,rand())

##2D sample

sample(f::Fun{TS},k::Integer) where {TS<:AbstractProductSpace}  =sample(ProductFun(f),k)

function sample(f::LowRankFun,n::Integer)
    rx=sample(sum(f,2),n)
    fA=evaluate.(f,rx,:)
    ry=map(sample,fA)
    [rx ry]
end

function sample(f::LowRankFun{C,C,TensorSpace{Tuple{C,C},DD,RR},Float64},n::Integer) where {C<:Chebyshev,DD<:Domain2d,RR<:Real}
    ry=sample(sum(f,1),n)
    fA=evaluate.(f.A,ry')
    CB=coefficientmatrix(f.B)
    AB=CB*fA
    chebnormalizedcumsum!(AB)
    rx=chebbisectioninv(AB,rand(n))
    [fromcanonical.(Ref(domain(f,1)),rx) ry]
end

sample(f::ProductFun,n)=sample(LowRankFun(f),n)
sample(f::MultivariateFun)=sample(f,1)[1,:]

## Special spaces

# Rays may be schwartz at right endpoint so we project
function sample(f::Fun{JacobiWeight{SS,DD,RR,TT},Float64}, n::Integer) where {SS,DD<:Ray,RR,TT}
    if space(f).α == 0
        samplecdf(normalizedcumsum(f),n)
    else
        sample(default_Fun(f,JacobiWeight(1,space(f).α,domain(f))),n)
    end
end


#===================================================================================================
NSGA-II

[1] Deb, Kalyanmoy, et al. "A fast and elitist multiobjective genetic algorithm: NSGA-II."
IEEE transactions on evolutionary computation 6.2 (2002): 182-197.
===================================================================================================#

"""
    nondominated_sort(fitness::AbstractMatrix{<:Real}, senses::AbstractVector{<:Integer})::Vector{Vector{Int}}

Fast nondominated sorting on page 3 of [1]. The fronts are returned, where the first one is the Pareto
optimal front, and the i-th front contains the indices of individuals in that front.
"""
function nondominated_sort(fitness::AbstractMatrix{<:Real}, senses::AbstractVector{<:Integer})::Vector{Vector{Int}}
    n = size(fitness, 2) # number of individuals
    d_counter = zeros(Int, n)
    dominated = [Int[] for _ = 1:n]
    fronts = Vector{Int}[]
    @inbounds for ip = 1:n, iq = ip+1:n
        p = view(fitness, :, ip)
        q = view(fitness, :, iq)
        if dominate(p, q, senses)
            push!(dominated[ip], iq)
            d_counter[iq] += 1
        elseif dominate(q, p, senses)
            push!(dominated[iq], ip)
            d_counter[ip] += 1
        end
    end
    # the first front
    push!(fronts, Int[])
    f1 = fronts[1]
    @inbounds for i = 1:n
        if d_counter[i] == 0
            push!(f1, i)
        end
    end
    n_in_fronts = length(f1)
    # the other fronts
    current = 1
    while true
        new_front = Int[]
        sizehint!(new_front, n - n_in_fronts)
        for ip in fronts[current]
            for iq in dominated[ip]
                d_counter[iq] -= 1
                if d_counter[iq] == 0
                    push!(new_front, iq)
                end
            end
        end
        if isempty(new_front)
            break
        else
            push!(fronts, new_front)
            n_in_fronts += length(new_front)
        end
        current += 1
    end
    return fronts
end

"""
    get_nondomination_rank(fronts::AbstractVector{<:AbstractVector{<:Integer}})::Vector{Int}

Given the sorted front, determine which front each individual resides (starting from 1).
The value at the i-th position of the returned array is the rank of the i-th individual in the
population.
"""
function get_nondomination_rank(fronts::AbstractVector{<:AbstractVector{<:Integer}})::Vector{Int}
    n = sum(length, fronts)
    ranks = zeros(Int, n)
    for (r, front) in enumerate(fronts)
        ranks[front] .= r
    end
    return ranks
end

"""
    assign_crowding_distance(fitness::AbstractMatrix{<:Real},
        fronts::AbstractVector{<:AbstractVector{<:Integer}})::Vector{Real}

Compute the crowding distance for each individual in the whole population characterized by its
`fitness` and Pareto front in  `fronts`. Refer to the algorithm on the bottom of page 4 of Ref. [1].

A vector of crowding distances `cds` is returned, in which `cds[i]` denotes the crowding distance of
the individual `i`.
"""
function assign_crowding_distance(fitness::AbstractMatrix{<:Real},
    fronts::AbstractVector{<:AbstractVector{<:Integer}})::Vector{Float64}
    m, n = size(fitness)
    cds = zeros(n)  # crowding distance
    fronts = deepcopy(fronts)
    # handle each front separately
    for front in fronts # front contains the index of an individual in the population
        if length(front) <= 2
            cds[front] .= Inf
            continue
        end
        for io = 1:m  # each objective
            sort!(front; by = idx -> fitness[io, idx])
            scale = fitness[io, front[end]] - fitness[io, front[1]]
            # all individuals in this front have the identical `io`-th objective value, ignore it
            if scale > 0
                cds[front[1]] = Inf
                cds[front[end]] = Inf
                for i = 2:length(front)-1
                    cds[front[i]] += (fitness[io, front[i+1]] - fitness[io, front[i-1]]) / scale
                end
            end
        end
    end
    return cds
end

"""
   crowded_compare(i_rank, i_dist, j_rank, j_dist)::Bool

Crowed-comparison operator. Return -1 if we prefer the individual `i`; otherwise, 1 means
the individual `j` is preferable. If both individuals `i` and `j` are equally good, 0 is returned.
"""
function crowded_compare(i_rank, i_dist, j_rank, j_dist)::Int
    if i_rank < j_rank
        return -1
    elseif i_rank > j_rank
        return 1
    end
    # now i and j belong to the same Pareto front: prefer the less crowded one
    if i_dist > j_dist
        return -1
    elseif i_dist < j_dist
        return 1
    end
    return 0
end


include("../../module/CustomStencil.jl")
include("../../misc/misc.jl")
include("../../misc/cpu_stencils.jl")
## Star Stencil definition with radius = 2
coefs = fill(0.06, 5,5,5)
coefs[2:4, 2:4, 2:4] .= 0.125
coefs[3,3,3] = 0.25
st_def = CreateStencilDefinition(coefs)
st_inst = NewStencilInstance(st_def, m_step=false)  #don't combine time steps

## Input Data size Definition
radius = 2
nx = 5
ny = 5
nz = 5
dx = 1<<(nx)
dy = 1<<(ny)
dz = 1<<(nz)

data = CreateData(dx,dy,dz)

t_steps = 5
## GPU Implementation
gpu_out = ApplyStencil(st_inst, data, t_steps)

## CPU result
cpu_out = cpu_dense_stencil(data, radius, coefs, t_steps)

## Check
mean_er = sum(abs.(cpu_out .- gpu_out))./(reduce(*, size(gpu_out)))
@assert(mean_er < 0.001)

## Visualize Error
using Plots

shift = 0
Int(round(dz/2))+shift
heatmap(gpu_out[:,:,Int(round(dz/2))+shift]-cpu_out[:,:,Int(round(dz/2))+shift])
heatmap(cpu_out[:,:,Int(round(dz/2))+shift])
heatmap(gpu_out[:,:,Int(round(dz/2))+shift])


using Dice
using Dice: num_flips, num_nodes, to_dice_ir

function reachable(adjacency::Matrix, src::Int, dest::Int)
    n = size(adjacency, 1)
    reachable = Vector(undef, n)
    reachable .= false
    reachable[src] = true
    for i=1:n
        for node=1:n
            reachable[node] |= reduce(|, reachable .& adjacency[:, node])
        end
    end
    reachable[dest]
end

n = 5

# run on sampled deterministic graph
adjacency_sampled = rand(Bool, n, n)
r = reachable(adjacency_sampled, 1, n)
println("Sampled graph reachability: ", r)

# run on random graph
code = @dice begin
    adjacency_rand = [flip(0.5) for i=1:n, j=1:n]
    reachable(adjacency_rand, 1, n)
end

# BDD analysis
bdd = compile(code)
num_flips(bdd), num_nodes(bdd)
@assert num_flips(bdd) == n*n-3n+3
infer(code, :bdd)

# IR analysis
to_dice_ir(code)

# # TODO fix loopy let identifiers
# has_dice_binary() && rundice(code)
# has_dice_binary() && infer(code, :ocaml)


"""
operations
"""

function Base.:+(x::AbstractNumberSymbolic{Tv}) where Tv
    SymbolicExpression{Tv}(x.params, :+, [x], [])
end

function Base.:-(x::AbstractNumberSymbolic{Tv}) where Tv
    SymbolicExpression{Tv}(x.params, :-, [x], [])
end

for op in [:+, :-, :*, :/, :^]
    @eval function Base.$op(x::AbstractNumberSymbolic{Tx}, y::AbstractNumberSymbolic{Ty}) where {Tx<:Number,Ty<:Number}
        s = union(x.params, y.params)
        Tv = promote_type(Tx,Ty)
        SymbolicExpression{promote_type(Tx,Ty)}(s, $(Expr(:quote, op)), [x, y], [])
    end

    @eval function Base.$op(x::Tx, y::AbstractNumberSymbolic{Ty}) where {Tx<:Number,Ty<:Number}
        Tv = promote_type(Tx,Ty)
        SymbolicExpression{promote_type(Tx,Ty)}(y.params, $(Expr(:quote, op)), [x, y], [])
    end

    @eval function Base.$op(x::AbstractNumberSymbolic{Tx}, y::Ty) where {Tx<:Number,Ty<:Number}
        Tv = promote_type(Tx,Ty)
        SymbolicExpression{promote_type(Tx,Ty)}(x.params, $(Expr(:quote, op)), [x, y], [])
    end
end

function Base.exp(x::AbstractNumberSymbolic{Tv}) where Tv
    SymbolicExpression{Tv}(x.params, :exp, [x], [])
end

function Base.log(x::AbstractNumberSymbolic{Tv}) where Tv
    SymbolicExpression{Tv}(x.params, :log, [x], [])
end

function Base.sqrt(x::AbstractNumberSymbolic{Tv}) where Tv
    SymbolicExpression{Tv}(x.params, :sqrt, [x], [])
end

###

function Base.sum(x::Vector{<:AbstractNumberSymbolic{T}}) where {T<:Number}
    sum(convert(AbstractVectorSymbolic{T}, x))
end

function Base.sum(x::AbstractVectorSymbolic{T}) where {T<:Number}
    SymbolicExpression{T}(x.params, :sum, [x], [])
end

function dot(x::Vector{<:AbstractNumberSymbolic{Tx}}, y::Vector{<:AbstractNumberSymbolic{Ty}}) where {Tx<:Number,Ty<:Number}
    dot(convert(AbstractVectorSymbolic{Tx}, x), convert(AbstractVectorSymbolic{Ty}, y))
end

function dot(x::AbstractVectorSymbolic{Tx}, y::Vector{<:AbstractNumberSymbolic{Ty}}) where {Tx<:Number,Ty<:Number}
    dot(x, convert(AbstractVectorSymbolic{Ty}, y))
end

function dot(x::Vector{<:AbstractNumberSymbolic{Tx}}, y::AbstractVectorSymbolic{Ty}) where {Tx<:Number,Ty<:Number}
    dot(convert(AbstractVectorSymbolic{Tx}, x), y)
end

function dot(x::AbstractVectorSymbolic{Tx}, y::AbstractVectorSymbolic{Ty}) where {Tx<:Number,Ty<:Number}
    Tv = promote_type(Tx,Ty)
    s = union(x.params, y.params)
    SymbolicExpression{Tv}(s, :dot, AbstractVectorSymbolic{Tv}[x, y], [])
end

for op in [:+, :-]
    @eval function Base.$op(x::AbstractVectorSymbolic{Tx}, y::AbstractVectorSymbolic{Ty}) where {Tx<:Number,Ty<:Number}
        s = union(x.params, y.params)
        Tv = promote_type(Tx,Ty)
        SymbolicVectorExpression{promote_type(Tx,Ty)}(s, $(Expr(:quote, op)), [x, y], x.dim, [])
    end

    @eval function Base.$op(x::Tx, y::AbstractVectorSymbolic{Ty}) where {Tx<:Number,Ty<:Number}
        Tv = promote_type(Tx,Ty)
        SymbolicVectorExpression{promote_type(Tx,Ty)}(y.params, $(Expr(:quote, op)), [x, y], x.dim, [])
    end

    @eval function Base.$op(x::AbstractVectorSymbolic{Tx}, y::Ty) where {Tx<:Number,Ty<:Number}
        Tv = promote_type(Tx,Ty)
        SymbolicVectorExpression{promote_type(Tx,Ty)}(x.params, $(Expr(:quote, op)), [x, y], x.dim, [])
    end
end


module _DispEn2D
export DispEn2D
using Clustering: kmeans, assignments
using Statistics: std, mean
using StatsFuns: normcdf
    """
        Disp2D, RDE = DispEn2D(Mat)

    Returns the bidimensional dispersion entropy estimate (`Disp2D`) and reverse
    bidimensional dispersion entropy (`RDE`) estimated for the data matrix (`Mat`)
    using the default parameters:  time delay = 1, symbols = 3, logarithm = natural,
    data transform = normalised cumulative density function (`'ncdf'`), logarithm = natural,
    template matrix size = [floor(H/10) floor(W/10)], (where H and W represent
    the height (rows) and width (columns) of the data matrix `Mat`) \n
    ** The minimum number of rows and columns of Mat must be > 10.**

        Disp2D, RDE = DispEn2D(Mat::AbstractArray{T,2} where T<:Real;
                            m::Union{Int,Tuple{Int,Int}}=floor.(Int, size(Mat)./10), tau::Int=1,
                                c::Int=3, Typex::String="ncdf", Logx::Real=exp(1), Norm::Bool=false, Lock::Bool=true)

    Returns the bidimensional dispersion entropy (`Disp2D`) and reverse
    bidimensional distribution entropy (`RDE`) estimate for the data matrix (`Mat`)
    using the specified 'keyword' arguments:

    # Arguments:
    `m`     - Template submatrix dimensions, an integer scalar (i.e. the same
              height and width) or a two-element tuple of integers
              [height, width] with a minimum value > 1.
              [default: [floor(H/10) floor(W/10)]] \n
    `tau`   - Time Delay, a positive integer       [default: 1]  \n
    `c`     - Number of symbols, an integer > 1
    `Typex` - Type of symbolic mapping transform, one of the following:
              {`linear`, `kmeans`, `ncdf`, `equal`}
              See the `EntropyHub Guide` for more info on these transforms.
    `Logx`  - Logarithm base, a positive scalar        [default: natural]\n
              ** enter 0 for natural logarithm.**\n
    `Norm`  - Normalisation of `Disp2D` value, a boolean:
                - [false]   no normalisation - default
                - [true]    normalises w.r.t number of possible dispersion patterns.
    `Lock`  - By default, DispEn2D only permits matrices with a maximum
              size of 128 x 128 to prevent memory errors when storing data on RAM.
              e.g. For Mat = [200 x 200], m = 3, and tau = 1, DispEn2D
              creates a vector of 753049836 elements. To enable matrices
              greater than [128 x 128] elements, set `Lock` to false.
              [default: 'true']
              `WARNING: unlocking the permitted matrix size may cause your Julia
               IDE to crash.`

    # See also `DispEn`, `DistEn2D`, `SampEn2D`, `FuzzEn2D`, `MSEn`

    # References:
        [1] Hamed Azami, et al.,
            "Two-dimensional dispersion entropy: An information-theoretic
            method for irregularity analysis of images."
            Signal Processing: Image Communication,
            75 (2019): 178-187.

    """
    function DispEn2D(Mat::AbstractArray{T,2} where T<:Real; m::Union{Int,Tuple{Int,Int}}=floor.(Int, size(Mat)./10), tau::Int=1,
                        c::Int=3, Typex::String="ncdf", Logx::Real=exp(1), Norm::Bool=false, Lock::Bool=true)

    Logx == 0  ? Logx = exp(1) : nothing

    NL, NW = size(Mat)
    ((NL > 128 || NW > 128) && Lock) ?
        error("To prevent memory errors, matrix width & length must have <= 128 elements.
            To estimatDispEn2D  for the current matrix ($NL,$NW) change Lock to 'false'.
            Caution: unlocking the safe matrix size may cause the Julia IDE to crash.") :
            nothing

    length(m)==1 ?  (mL = m; mW = m) : (mL = m[1];  mW = m[2])
    (NL > 10 && NW > 10) ? nothing :
    error("Number of rows and columns in Mat must be > 10")
    (minimum(m) > 1) ? nothing : error("m:   must be an integer > 1, or 2-element integer tuple w/ values > 1")
    (tau > 0) ? nothing : error("tau:   must be an integer > 0")
    (c > 1) ? nothing :   error("c:     must be an integer > 1")
    (lowercase(Typex) in ["linear", "kmeans", "ncdf", "equal"]) ? nothing :
        error("Typex:    must be one of the following strings - 'linear','kmeans','ncdf','equal'")
    (Logx>0) ? nothing :  error("Logx:  must be a positive number > 0")

    if lowercase(Typex) == "linear"
        Edges = range(minimum(Mat),maximum(Mat),length=c+1)
        Zi = map(x -> sum(Edges[1:c].<=x), Mat)

    elseif lowercase(Typex) == "kmeans"
        Temp = kmeans(transpose(Mat[:]), c; maxiter=200)
        Zx = assignments(Temp)
        Clux = Temp.centers
        xx = sortperm(Clux[:]);   Zi = zeros(Int,length(Mat))
        for k = 1:c
            Zi[Zx.==xx[k]] .= k;
        end
        Zi = reshape(Zi,size(Mat))

    elseif lowercase(Typex) == "ncdf"
        Zx = normcdf.(mean(Mat),std(Mat,corrected=false),Mat);
        Zi = map(x -> sum(range(0,1,length=c+1)[1:c].<=x), Zx);

    elseif lowercase(Typex) == "equal"
        ix = sortperm(Mat[:]);
        xx = Int.(round.(range(0,length(Mat),length=c+1)))
        Zi = zeros(Int, length(Mat))
        for k = 1:c
            Zi[ix[xx[k]+1:xx[k+1]]] .= k
        end
        Zi = reshape(Zi,size(Mat))
    end

    NL = NL - (mL-1)*tau
    NW = NW - (mW-1)*tau
    X = zeros(Int,NL*NW,mL*mW)
    p = 0
    for k = 1:NL
        for n = 1:NW
            p += 1
            X[p,:] = Zi[k:tau:(mL-1)*tau+k,n:tau:(mW-1)*tau+n][:]
        end
    end

    p != NL*NW ? @warn("Potential error with submatrix division.") : nothing

    T = unique(X,dims=1)
    Nx = size(T)[1]
    Counter = zeros(Nx)
    for n = 1:Nx
        Counter[n] = sum(all(X .- transpose(T[n,:]) .==0, dims=2))
    end

    Ppi = Counter[Counter.!= 0]/size(X)[1]

    big(c)^(mL*mW) > 10^16 ? error("RDE cannot be estimated with c = $c and
    a submatrix of size $mL x $mW. Required floating point precision exceeds 10^16.
    Consider reducing the template submatrix size (m) or the number of symbols (c).") : nothing

    #RDE = sum(Ppi.^2) - (1/Nx)
    RDE = sum((Ppi .- (1/(c^(mL*mW)))).^2);

    if round(sum(Ppi),digits=4) != 1
        @warn("Potential Error calculating probabilities")
    end

    Disp2D = -sum(Ppi.*log.(Logx, Ppi))
    if Norm
        Disp2D = Disp2D/log(Logx, c^(mL*mW))
        RDE = RDE./(1 - (1/(c^(mL*mW))))
    end

    return Disp2D, RDE
    end

end

"""
Copyright 2021 Matthew W. Flood, EntropyHub

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

For Terms of Use see https://github.com/MattWillFlood/EntropyHub
"""


function setup_update_matrices()
    options = Dict(:verbose => false,
                   :eps_abs => 1e-08,
                   :eps_rel => 1e-08,
                   :polish => false,
                   :check_termination => 1)

    seed!(1)

    n = 5
    m = 8
    p = 0.7
    Pt = sprandn(n, n, p)
    Pt_new = copy(Pt)
    P = Pt * Pt' + sparse(I, n, n)
    #  PtI = findall(!iszero, P)
    #  (Pti, Ptj) = (getindex.(PtI, 1), getindex.(PtI, 2))
    Ptx = copy(Pt.nzval)

    Pt_newx = Ptx + 0.1 * randn(length(Ptx))
    # Pt_new = sparse(Pi, Pj, Pt_newx)
    P_new = Pt_new * Pt_new' + sparse(I, n, n)
    q = randn(n)
    A = sprandn(m, n, p)

    (Ai, Aj, _) = findnz(A)
    Ax = copy(A.nzval)
    A_newx = Ax + randn(length(Ax))
    A_new = sparse(Ai, Aj, A_newx)

    # A_new = copy(A)
    # A_new.nzval += randn(length(A_new.nzval))
    l = zeros(m)
    u = 30 .+ randn(m)

    problem = Dict()
    problem[:P] = P
    problem[:P_new] = P_new
    problem[:q] = q
    problem[:A] = A
    problem[:A_new] = A_new
    problem[:l] = l
    problem[:u] = u
    problem[:m] = m
    problem[:n] = n
    return problem, options
end

tol = 1e-5

# This unit test relies on a precomputed reference solution for a randomly generated problem (done in Julia v1.0).
# However, in newer versions of Julia the random number stream changed leading to a different problem being generated.
if VERSION < v"1.1"

        @testset "update_matrices" begin

            @testset "solve" begin
                problem, options = setup_update_matrices()

                (n, m, P, q, A, l, u) = (problem[:n], problem[:m], problem[:P], problem[:q], problem[:A], problem[:l], problem[:u])

                model = OSQP.Model()
                OSQP.setup!(model; P=P, q=q, A=A, l=l, u=u, options...)
                results = OSQP.solve!(model)

                # # Solve with Gurobi
                # using Gurobi
                # env = Gurobi.Env()
                # setparam!(env, "OutputFlag", 1)
                # model = Gurobi.Model(env, "model")
                # add_cvars!(model, q); update_model!(model)
                # add_constrs!(model, A, repmat(['<'], m), u); update_model!(model)
                # add_constrs!(model, A, repmat(['>'], m), l); update_model!(model)
                # add_qpterms!(model, P); update_model!(model)
                # optimize(model)
                # x_test = get_solution(model)
                # obj_test= get_objval(model)
                # y_test = -Gurobi.get_dblattrarray(model, "Pi", 1, 2 * m)
                # y_test = y_test[m+1:end] + y_test[1:m]
                        #
                # println("x_test = $(x_test)")
                # println("y_test = $(y_test)")
                # println("obj_test = $(obj_test)")

                x_test = [-0.324865, 0.598681, -0.066646, -0.00653471, -0.0736556]
                y_test = [-2.72542e-10, -1.927e-12, -0.547374, -2.25907e-10, -0.645086, -4.09874e-9, -1.4083e-11, -1.26234e-11]
                obj_test = -0.6736977821770016

                @test isapprox(results.x, x_test, atol=tol)
                @test isapprox(results.y, y_test, atol=tol)
                @test isapprox(results.info.obj_val, obj_test, atol=tol)
                @test results.info.status == :Solved

            end

                @testset "update_P" begin

                problem, options = setup_update_matrices()

                (n, m, P, q, A, l, u) = (problem[:n], problem[:m],
                             problem[:P], problem[:q], problem[:A], problem[:l], problem[:u])
                P_new = problem[:P_new]

                model = OSQP.Model()
                OSQP.setup!(model; P=P, q=q, A=A, l=l, u=u, options...)

                # Update matrix
                Pnew_triu = triu(P_new)
                Pnew_triu_idx = collect(1 : length(Pnew_triu.nzval))
                OSQP.update!(model, Px=Pnew_triu.nzval, Px_idx=Pnew_triu_idx)
                results = OSQP.solve!(model)

                # # Solve with Gurobi
                # using Gurobi
                # env = Gurobi.Env()
                # setparam!(env, "OutputFlag", 1)
                # model = Gurobi.Model(env, "model")
                # add_cvars!(model, q); update_model!(model)
                # add_constrs!(model, A, repmat(['<'], m), u); update_model!(model)
                # add_constrs!(model, A, repmat(['>'], m), l); update_model!(model)
                # add_qpterms!(model, P_new); update_model!(model)
                # optimize(model)
                # x_test = get_solution(model)
                # obj_test= get_objval(model)
                # y_test = -Gurobi.get_dblattrarray(model, "Pi", 1, 2 * m)
                # y_test = y_test[m+1:end] + y_test[1:m]
                        #
                # println("x_test = $(x_test)")
                # println("y_test = $(y_test)")
                # println("obj_test = $(obj_test)")

                x_test = [-0.324865, 0.598681, -0.066646, -0.00653471, -0.0736556]
                y_test = [-2.72542e-10, -1.927e-12, -0.547374, -2.25907e-10, -0.645086, -4.09874e-9, -1.4083e-11, -1.26234e-11]
                obj_test = -0.6736977821770016

                @test isapprox(results.x, x_test, atol=tol)
                @test isapprox(results.y, y_test, atol=tol)
                @test isapprox(results.info.obj_val, obj_test, atol=tol)
                @test results.info.status == :Solved
            end

                @testset "update_P_allind" begin

                problem, options = setup_update_matrices()

                (n, m, P, q, A, l, u) = (problem[:n], problem[:m],
                             problem[:P], problem[:q], problem[:A], problem[:l], problem[:u])
                P_new = problem[:P_new]

                model = OSQP.Model()
                OSQP.setup!(model; P=P, q=q, A=A, l=l, u=u, options...)

                # Update matrix
                Pnew_triu = triu(P_new)
                OSQP.update!(model, Px=Pnew_triu.nzval)
                results = OSQP.solve!(model)

                # # Solve with Gurobi
                # using Gurobi
                # env = Gurobi.Env()
                # setparam!(env, "OutputFlag", 1)
                # model = Gurobi.Model(env, "model")
                # add_cvars!(model, q); update_model!(model)
                # add_constrs!(model, A, repmat(['<'], m), u); update_model!(model)
                # add_constrs!(model, A, repmat(['>'], m), l); update_model!(model)
                # add_qpterms!(model, P_new); update_model!(model)
                # optimize(model)
                # x_test = get_solution(model)
                # obj_test= get_objval(model)
                # y_test = -Gurobi.get_dblattrarray(model, "Pi", 1, 2 * m)
                # y_test = y_test[m+1:end] + y_test[1:m]
                        #
                # println("x_test = $(x_test)")
                # println("y_test = $(y_test)")
                # println("obj_test = $(obj_test)")

                x_test = [-0.324865, 0.598681, -0.066646, -0.00653471, -0.0736556]
                y_test = [-2.72542e-10, -1.927e-12, -0.547374, -2.25907e-10, -0.645086, -4.09874e-9, -1.4083e-11, -1.26234e-11]
                obj_test = -0.6736977821770016

                @test isapprox(results.x, x_test, atol=tol)
                @test isapprox(results.y, y_test, atol=tol)
                @test isapprox(results.info.obj_val, obj_test, atol=tol)
                @test results.info.status == :Solved
            end

                @testset "update_A" begin

                problem, options = setup_update_matrices()

                (n, m, P, q, A, l, u) = (problem[:n], problem[:m],
                             problem[:P], problem[:q], problem[:A], problem[:l], problem[:u])
                A_new = problem[:A_new]

                model = OSQP.Model()
                OSQP.setup!(model; P=P, q=q, A=A, l=l, u=u, options...)

                # Update matrix
                A_new_idx = collect(1 : length(A_new.nzval))
                OSQP.update!(model, Ax=A_new.nzval, Ax_idx=A_new_idx)
                results = OSQP.solve!(model)

                # # Solve with Gurobi
                # using Gurobi
                # env = Gurobi.Env()
                # setparam!(env, "OutputFlag", 1)
                # model = Gurobi.Model(env, "model")
                # add_cvars!(model, q); update_model!(model)
                # add_constrs!(model, A_new, repmat(['<'], m), u); update_model!(model)
                # add_constrs!(model, A_new, repmat(['>'], m), l); update_model!(model)
                # add_qpterms!(model, P); update_model!(model)
                # optimize(model)
                # x_test = get_solution(model)
                # obj_test= get_objval(model)
                # y_test = -Gurobi.get_dblattrarray(model, "Pi", 1, 2 * m)
                # y_test = y_test[m+1:end] + y_test[1:m]
                        #
                # println("x_test = $(x_test)")
                # println("y_test = $(y_test)")
                # println("obj_test = $(obj_test)")
                        #
                x_test = [-0.0318085, 0.067069, -0.0242966, -0.0593736, -0.0321274]
                y_test = [-2.16989e-10, -1.01848, -0.516013, -0.0227263, -3.19721e-10, -1.04482, -3.4596e-10, -4.51608e-10]
                obj_test = -0.02187014865840703

                @test isapprox(results.x, x_test, atol=tol)
                @test isapprox(results.y, y_test, atol=tol)
                @test isapprox(results.info.obj_val, obj_test, atol=tol)
                @test results.info.status == :Solved
            end

            @testset "update_A_allind" begin

                problem, options = setup_update_matrices()

                (n, m, P, q, A, l, u) = (problem[:n], problem[:m],
                             problem[:P], problem[:q], problem[:A], problem[:l], problem[:u])
                A_new = problem[:A_new]

                model = OSQP.Model()
                OSQP.setup!(model; P=P, q=q, A=A, l=l, u=u, options...)

                # Update matrix
                OSQP.update!(model, Ax=A_new.nzval)
                results = OSQP.solve!(model)

                # # Solve with Gurobi
                # using Gurobi
                # env = Gurobi.Env()
                # setparam!(env, "OutputFlag", 1)
                # model = Gurobi.Model(env, "model")
                # add_cvars!(model, q); update_model!(model)
                # add_constrs!(model, A_new, repmat(['<'], m), u); update_model!(model)
                # add_constrs!(model, A_new, repmat(['>'], m), l); update_model!(model)
                # add_qpterms!(model, P); update_model!(model)
                # optimize(model)
                # x_test = get_solution(model)
                # obj_test= get_objval(model)
                # y_test = -Gurobi.get_dblattrarray(model, "Pi", 1, 2 * m)
                # y_test = y_test[m+1:end] + y_test[1:m]
                        #
                # println("x_test = $(x_test)")
                # println("y_test = $(y_test)")
                # println("obj_test = $(obj_test)")
                        #
                x_test = [-0.0318085, 0.067069, -0.0242966, -0.0593736, -0.0321274]
                y_test = [-2.16989e-10, -1.01848, -0.516013, -0.0227263, -3.19721e-10, -1.04482, -3.4596e-10, -4.51608e-10]
                obj_test = -0.02187014865840703

                @test isapprox(results.x, x_test, atol=tol)
                @test isapprox(results.y, y_test, atol=tol)
                @test isapprox(results.info.obj_val, obj_test, atol=tol)
                @test results.info.status == :Solved
            end

                @testset "update_P_A_indP_indA" begin

                problem, options = setup_update_matrices()

                (n, m, P, q, A, l, u) = (problem[:n], problem[:m],
                             problem[:P], problem[:q], problem[:A], problem[:l], problem[:u])
                P_new = problem[:P_new]
                A_new = problem[:A_new]

                model = OSQP.Model()
                OSQP.setup!(model; P=P, q=q, A=A, l=l, u=u, options...)

                # Update matrices P and A
                P_new_triu = triu(P_new)
                P_new_triu_idx = collect(1 : length(P_new_triu.nzval))
                A_new_idx = collect(1 : length(A_new.nzval))

                OSQP.update!(model, Px=P_new_triu.nzval, Px_idx=P_new_triu_idx,
                         Ax=A_new.nzval, Ax_idx=A_new_idx)
                results = OSQP.solve!(model)

                # # Solve with Gurobi
                # using Gurobi
                # env = Gurobi.Env()
                # setparam!(env, "OutputFlag", 1)
                # model = Gurobi.Model(env, "model")
                # add_cvars!(model, q); update_model!(model)
                # add_constrs!(model, A_new, repmat(['<'], m), u); update_model!(model)
                # add_constrs!(model, A_new, repmat(['>'], m), l); update_model!(model)
                # add_qpterms!(model, P_new); update_model!(model)
                # optimize(model)
                # x_test = get_solution(model)
                # obj_test= get_objval(model)
                # y_test = -Gurobi.get_dblattrarray(model, "Pi", 1, 2 * m)
                # y_test = y_test[m+1:end] + y_test[1:m]
                        #
                # println("x_test = $(x_test)")
                # println("y_test = $(y_test)")
                # println("obj_test = $(obj_test)")
                        #
                x_test = [-0.0318085, 0.067069, -0.0242966, -0.0593736, -0.0321274]
                y_test = [-2.16989e-10, -1.01848, -0.516013, -0.0227263, -3.19721e-10, -1.04482, -3.4596e-10, -4.51608e-10]
                obj_test = -0.02187014865840703

                @test isapprox(results.x, x_test, atol=tol)
                @test isapprox(results.y, y_test, atol=tol)
                @test isapprox(results.info.obj_val, obj_test, atol=tol)
                @test results.info.status == :Solved
            end

                @testset "update_P_A_indP" begin

                problem, options = setup_update_matrices()

                (n, m, P, q, A, l, u) = (problem[:n], problem[:m],
                             problem[:P], problem[:q], problem[:A], problem[:l], problem[:u])
                P_new = problem[:P_new]
                A_new = problem[:A_new]

                model = OSQP.Model()
                OSQP.setup!(model; P=P, q=q, A=A, l=l, u=u, options...)

                # Update matrices P and A
                P_new_triu = triu(P_new)
                P_new_triu_idx = collect(1 : length(P_new_triu.nzval))

                OSQP.update!(model, Px=P_new_triu.nzval, Px_idx=P_new_triu_idx,
                         Ax=A_new.nzval)
                results = OSQP.solve!(model)

                # # Solve with Gurobi
                # using Gurobi
                # env = Gurobi.Env()
                # setparam!(env, "OutputFlag", 1)
                # model = Gurobi.Model(env, "model")
                # add_cvars!(model, q); update_model!(model)
                # add_constrs!(model, A_new, repmat(['<'], m), u); update_model!(model)
                # add_constrs!(model, A_new, repmat(['>'], m), l); update_model!(model)
                # add_qpterms!(model, P_new); update_model!(model)
                # optimize(model)
                # x_test = get_solution(model)
                # obj_test= get_objval(model)
                # y_test = -Gurobi.get_dblattrarray(model, "Pi", 1, 2 * m)
                # y_test = y_test[m+1:end] + y_test[1:m]
                        #
                # println("x_test = $(x_test)")
                # println("y_test = $(y_test)")
                # println("obj_test = $(obj_test)")
                        #
                x_test = [-0.0318085, 0.067069, -0.0242966, -0.0593736, -0.0321274]
                y_test = [-2.16989e-10, -1.01848, -0.516013, -0.0227263, -3.19721e-10, -1.04482, -3.4596e-10, -4.51608e-10]
                obj_test = -0.02187014865840703

                @test isapprox(results.x, x_test, atol=tol)
                @test isapprox(results.y, y_test, atol=tol)
                @test isapprox(results.info.obj_val, obj_test, atol=tol)
                @test results.info.status == :Solved
            end

                @testset "update_P_A_indA" begin

                problem, options = setup_update_matrices()

                (n, m, P, q, A, l, u) = (problem[:n], problem[:m],
                             problem[:P], problem[:q], problem[:A], problem[:l], problem[:u])
                P_new = problem[:P_new]
                A_new = problem[:A_new]

                model = OSQP.Model()
                OSQP.setup!(model; P=P, q=q, A=A, l=l, u=u, options...)

                # Update matrices P and A
                P_new_triu = triu(P_new)
                A_new_idx = collect(1 : length(A_new.nzval))

                OSQP.update!(model, Px=P_new_triu.nzval,
                         Ax=A_new.nzval, Ax_idx=A_new_idx)
                results = OSQP.solve!(model)

                # # Solve with Gurobi
                # using Gurobi
                # env = Gurobi.Env()
                # setparam!(env, "OutputFlag", 1)
                # model = Gurobi.Model(env, "model")
                # add_cvars!(model, q); update_model!(model)
                # add_constrs!(model, A_new, repmat(['<'], m), u); update_model!(model)
                # add_constrs!(model, A_new, repmat(['>'], m), l); update_model!(model)
                # add_qpterms!(model, P_new); update_model!(model)
                # optimize(model)
                # x_test = get_solution(model)
                # obj_test= get_objval(model)
                # y_test = -Gurobi.get_dblattrarray(model, "Pi", 1, 2 * m)
                # y_test = y_test[m+1:end] + y_test[1:m]
                        #
                # println("x_test = $(x_test)")
                # println("y_test = $(y_test)")
                # println("obj_test = $(obj_test)")
                        #
                x_test = [-0.0318085, 0.067069, -0.0242966, -0.0593736, -0.0321274]
                y_test = [-2.16989e-10, -1.01848, -0.516013, -0.0227263, -3.19721e-10, -1.04482, -3.4596e-10, -4.51608e-10]
                obj_test = -0.02187014865840703

                @test isapprox(results.x, x_test, atol=tol)
                @test isapprox(results.y, y_test, atol=tol)
                @test isapprox(results.info.obj_val, obj_test, atol=tol)
                @test results.info.status == :Solved
            end

                @testset "update_P_A_allind" begin

                problem, options = setup_update_matrices()

                (n, m, P, q, A, l, u) = (problem[:n], problem[:m],
                             problem[:P], problem[:q], problem[:A], problem[:l], problem[:u])
                P_new = problem[:P_new]
                A_new = problem[:A_new]

                model = OSQP.Model()
                OSQP.setup!(model; P=P, q=q, A=A, l=l, u=u, options...)

                # Update matrices P and A
                P_new_triu = triu(P_new)

                OSQP.update!(model, Px=P_new_triu.nzval, Ax=A_new.nzval)
                results = OSQP.solve!(model)

                # # Solve with Gurobi
                # using Gurobi
                # env = Gurobi.Env()
                # setparam!(env, "OutputFlag", 1)
                # model = Gurobi.Model(env, "model")
                # add_cvars!(model, q); update_model!(model)
                # add_constrs!(model, A_new, repmat(['<'], m), u); update_model!(model)
                # add_constrs!(model, A_new, repmat(['>'], m), l); update_model!(model)
                # add_qpterms!(model, P_new); update_model!(model)
                # optimize(model)
                # x_test = get_solution(model)
                # obj_test= get_objval(model)
                # y_test = -Gurobi.get_dblattrarray(model, "Pi", 1, 2 * m)
                # y_test = y_test[m+1:end] + y_test[1:m]
                        #
                # println("x_test = $(x_test)")
                # println("y_test = $(y_test)")
                # println("obj_test = $(obj_test)")
                        #
                x_test = [-0.0318085, 0.067069, -0.0242966, -0.0593736, -0.0321274]
                y_test = [-2.16989e-10, -1.01848, -0.516013, -0.0227263, -3.19721e-10, -1.04482, -3.4596e-10, -4.51608e-10]
                obj_test = -0.02187014865840703

                @test isapprox(results.x, x_test, atol=tol)
                @test isapprox(results.y, y_test, atol=tol)
                @test isapprox(results.info.obj_val, obj_test, atol=tol)
                @test results.info.status == :Solved
            end

        end
end


function softmax(x;k=1.)
    thresh = 500.
    g = map(x) do xi
        if k*xi < -thresh
            (log(1. + exp(k*xi)))/k
        elseif k*xi > thresh
            (k*xi + log(1. + exp(-k*xi)))/k
        else
            log(exp(k*xi) + 1.)/k
        end
    end
    g
end

function gsoftmax(x;k=1.)
    thresh = 500.
    g = map(x) do xi
        if k*xi > thresh
            1. / (1. + exp(-k*xi))
        else
            exp(k*xi) / (1. + exp(k*xi))
        end
    end
    Matrix(Diagonal(g))
end

function L(x,λ,f,h,c)
    f(x) - dot(λ,h(x)) + .5*c*dot(h(x),h(x))
end

function auglag(fun, num_eqs, num_ineqs, x0, options)
    num_fosteps = get(options, "num_fosteps", 3)
    num_sosteps = get(options, "num_sosteps", 2)
    c = get(options, "c", 1.)
    c_fos = get(options, "c_fos", 10.)
    c_sos = get(options, "c_sos", 10.)
    ls_method = get(options, "ls_method", :pinv)

    num_x = length(x0)
    rtol = eps(1.) * (num_eqs + num_ineqs)

    x = vcat(copy(x0), zeros(num_ineqs))
    λ = zeros(num_eqs + num_ineqs)

    for i = 1:num_fosteps
        # TODO inplace
        obj, eqs, ineqs, gobj, geqs, gineqs, Hobj = fun(x[1:num_x])

        ∇f = vcat(gobj, zeros(eltype(gobj), num_ineqs))
        Hf = zeros(eltype(Hobj), num_x+num_ineqs, num_x+num_ineqs)
        Hf[1:num_x, 1:num_x] .= Hobj
        hx = vcat(eqs, ineqs + softmax(x[num_x+1:end]))
        ∇h = hcat(vcat(geqs, gineqs), vcat(zeros(eltype(x),num_eqs,num_ineqs),gsoftmax(x[num_x+1:end])))

        gL = ∇f - ∇h'*λ + c*∇h'*hx
        HL = Hf + c*∇h'*∇h

        δx = (HL + (sum(HL.^2)+rtol)*I) \ (-gL)
        δλ = -c * hx

        x += δx
        λ += δλ

        c *= c_fos
    end

    for i = 1:num_sosteps
        # TODO inplace
        obj, eqs, ineqs, gobj, geqs, gineqs, Hobj = fun(x[1:num_x])

        ∇f = vcat(gobj, zeros(eltype(gobj),num_ineqs))
        Hf = zeros(eltype(Hobj), num_x+num_ineqs, num_x+num_ineqs)
        Hf[1:num_x, 1:num_x] .= Hobj

        hx = vcat(eqs, ineqs + softmax(x[num_x+1:end]))
        ∇h = hcat(vcat(geqs, gineqs), vcat(zeros(eltype(x),num_eqs,num_ineqs),gsoftmax(x[num_x+1:end])))

        gL = ∇f - ∇h'*λ + c*∇h'*hx
        HL = Hf + c*∇h'*∇h

        A = vcat(hcat(HL,∇h'),hcat(∇h,zeros(eltype(∇h),num_eqs+num_ineqs,num_eqs+num_ineqs)))
        b = -vcat(gL,hx)

        # for numerical stability
        tol = 1e-12
        A[abs.(A) .<= tol] .= 0.
        b[abs.(b) .<= tol] .= 0.

        if (ls_method == :pinv)
            Apinv = pinv(A)
            δxλ = Apinv * b
        elseif (ls_method == :least_squares)
            δxλ = least_squares(A,b)
        else
            U,S,V = svd(A)
            tol = rtol*maximum(S)
            ksig = 100.
            Sinv = 1. ./ (1. .+ exp.(-ksig*(S .- tol)/tol)) .* (1. ./ S)
            Sinv[isinf.(Sinv)] .= 0.
            Apinv = V * (Diagonal(Sinv) * U')
            δxλ = Apinv * b
        end

        δx = δxλ[1:num_x+num_ineqs]
        δλ = δxλ[num_x+num_ineqs+1:num_x+num_ineqs+num_eqs+num_ineqs]

        x += δx
        λ += δλ

        c *= c_sos
    end

    x[1:num_x], "Finished successfully"
end


export polyinterp!

function polyinterp!(ydes::VecO{Ty}, xdes::VecO{Tx}, ysrc::VecI, xsrc::VecI; order::Int=3) where {Tx<:Real,Ty<:Real}
    lx = 1                     # left index
    rx = sz = order + 1        # right index
    ia = sz >> 1 + 1           # anchor index
    mx = size(xsrc, 1)         # upper bound of window
    xb = Vector{Tx}(undef, sz) # xsrc buffer
    yb = Vector{Ty}(undef, sz) # ysrc buffer
    tb = Vector{Ty}(undef, sz) # interpolation buffer
    @inbounds for i in 1:sz
        xb[i] = xsrc[i]
        yb[i] = ysrc[i]
    end
    @inbounds xanchor = xb[ia]
    @inbounds for i in eachindex(xdes)
        xi = xdes[i]
        #### moving interpolation window
        @inbounds while xi > xanchor && rx < mx
            lx += 1
            rx += 1
            ix  = 1
            jx  = 2
            while ix < sz
                @inbounds xb[ix] = xb[jx]
                @inbounds yb[ix] = yb[jx]
                ix += 1
                jx += 1
            end
            xb[ix]  = xsrc[rx]
            yb[ix]  = ysrc[rx]
            xanchor = xb[ia]
        end
        ydes[i] = polyinterp!(xi, xb, yb, tb, sz)
    end
    return nothing
end

function polyinterp!(x::Real, xv::VecI, yv::VecI, bv::VecB, n::Int)
    one2n = eachindex(1:n)
    @inbounds for i in one2n
        x == xv[i] && return yv[i]
    end
    Δx = Inf
    yp = 0.0
    @inbounds for i in one2n
        δx = abs(x - xv[i])
        δx < Δx && (Δx = δx; yp = yv[i])
    end
    @simd for i in one2n
        @inbounds bv[i] = yv[i] - yp
    end
    @inbounds for k in 1:n-1, i in 1:n-k
        bv[i] += (bv[i] - bv[i+1]) * (x - xv[i]) / (xv[i] - xv[i+k])
    end
    @inbounds return yp + bv[1]
end


using PtFEM

data = Dict(
  # Frame(nels, nn, ndim, np_types, nip, finite_element(nod, nodof))
  :struc_el => Frame(10, 6, 2, 1, 1, Line(2, 2)),
  :properties => [2.0e5;],
  :x_coords => [0.0, 4.0, 4.0, 8.0, 8.0, 12.0],
  :y_coords => [3.0, 0.0, 3.0, 3.0, 0.0, 0.0],
  :g_num => [
    1 1 3 3 3 2 2 5 4 5;
    2 3 4 5 2 4 5 4 6 6
  ],
  :support => [
    (1, [0 0]),
    (2, [1 0])
    ],
  :loaded_nodes => [
    (6, [0.0 -10.0])],
  :penalty => 1e19
)

data |> display
println()

@time fem, dis_df, fm_df = p42(data)
println()

display(dis_df)
println()
display(fm_df)


# Adapted from the maple scripts in libxc
# Look at scripts/maple2c_new.pl in particular
#         maple/util.jl

using SymPy

r0, r1 = symbols("ρ1 ρ2")
s0, s1, s2 = symbols("σ0 σ1 σ2")

# unpol, 3D
DIMENSIONS = 3
RS_FACTOR    = (3/(4*PI))^(1//3)
X_FACTOR_C   = 3//8*(3/PI)^(1//3)*4^(2//3)
r_ws(n) = RS_FACTOR/n^(1//DIMENSIONS)
dens(r0, r1) = r0
zeta(r0, r1) = 0
xs0(r0, r1, sigma0, sigma2) = sqrt(sigma0/4)/((r0/2)^(1 + 1//DIMENSIONS))
xs1(r0, r1, sigma0, sigma2) = sqrt(sigma0/4)/((r0/2)^(1 + 1//DIMENSIONS))
xt(r0, r1, sigma0, sigma1, sigma2) = sqrt(sigma0)/r0^(1 + 1//DIMENSIONS)

f_zeta(z) = ((1 + z)^(4//3) + (1 - z)^(4//3) - 2)/(2^(4//3) - 2)

# Select case to print code for
case = "lda_x"
case = "lda_c_vwn"

include(case * ".jl")

println("#")
println("#-- " * case)
println("#")

# Energy per unit particle
f_expr = f(r_ws(dens(r0, r1)), zeta(r0, r1), xt(r0, r1, s0, s1, s2), xs0(r0, r1, s0, s2), xs1(r0, r1, s0, s2))

# Energy per unit volume
E_expr = f_expr * dens(r0, r1)

# Derivative with respect to r0
dE_expr = diff(E_expr, r0)

symbol = ["ε", "Vρ"]
f_expr = simplify(f_expr)
E_expr = simplify(E_expr)

sexpr, simplified = sympy.cse([f_expr, dE_expr])

for (var, expr) in sexpr
    println(sympy.julia_code(expr, assign_to=var))
end
println()
for (i, simp) in enumerate(simplified[1])
    println(symbol[i], " = ", sympy.julia_code(simp))
end


# # In-place transforms

using FFTW
using MPI
using PencilFFTs

using Random: randn!

const INPUT_DIMS = (64, 32, 64)

MPI.Init()

dims = INPUT_DIMS

# Combine r2r and c2c in-place transforms.
transforms = (
    Transforms.R2R!(FFTW.REDFT01),
    Transforms.FFT!(),
    Transforms.R2R!(FFTW.DHT),
)

# MPI topology information
comm = MPI.COMM_WORLD
Nproc = MPI.Comm_size(comm)
rank = MPI.Comm_rank(comm)

# Let's do a 1D decomposition.
proc_dims = (Nproc, )

# Create in-place plan
plan = PencilFFTPlan(dims, transforms, proc_dims, comm)
rank == 0 && println(plan)
@assert Transforms.is_inplace(plan)

# Allocate data for the plan.
# This returns a `ManyPencilArray` container that holds multiple
# `PencilArray` views.
A = allocate_input(plan) :: PencilArrays.ManyPencilArray

# The input output `PencilArray`s are recovered using `first` and
# `last`.
u_in = first(A) :: PencilArray
u_out = last(A) :: PencilArray

# Initialise input data.
randn!(u_in)

# Apply in-place forward transform on the `ManyPencilArray` container.
plan * A

# After the transform, operations should be performed on the output view
# `u_out`. For instance, let's compute the global sum of the transformed data.
sum_global = sum(u_out)  # note that `sum` reduces over all processes

# Apply in-place backward transform.
plan \ A

# Now we can again perform operations on the input view `u_in`...


#
#  This example is based on demo_m1.m from LAYPACK 1.0 and includes
#  many of the same comments.
#
#  MODEL REDUCTION BY THE ALGORITHMS LRSRM AND DSPMR. THE GOAL IS TO
#  GENERATE A REDUCED SYSTEM OF VERY SMALL ORDER.
#
#  This demo program shows the use of 'lp_lradi' and 'lp_para'

# -----------------------------------------------------------------------
# Generate test problem
# -----------------------------------------------------------------------
#
# This is an artificial test problem of a system, whose Bode plot shows
# "spires".
workspace()
include("../src/LargeMatrixEquations.jl")
using LargeMatrixEquations

A = speye(408,408); B = ones(408,1); C = ones(408,1)
A[1:2,1:2] = [-.01 -200; 200 .001]
A[3:4,3:4] = [-.2 -300; 300 -.1]
A[5:6,5:6] = [-.02 -500; 500 0]
A[7:8,7:8] = [-.01 -520; 520 -.01]
A[9:408,9:408] = diagm(-(1:400))

println("Problem dimensions:")

n = size(A,1);   # problem order (number of states)
m = size(B,2);  # number of inputs
q = size(C,1);   # number of outputs

# -----------------------------------------------------------------------
# Initialization/generation of data structures used in user-supplied
# functions and computation of ADI shift parameters
# -----------------------------------------------------------------------
#
# Note that A is a tridiagonal matrix. No preprocessing needs to be done.

println("Parameters for heuristic algorithm which computes ADI parameters:")
l0 = 10;   # desired number of distinct shift parameters
kp = 30;   # number of steps of Arnoldi process w.r.t. A
km = 15;   # number of steps of Arnoldi process w.r.t. inv(A)

b0 = ones(n,1);   # This is just one way to choose the Arnoldi start
                  # vector.
Bf = []#zeros(408,2)#
Kf = []
# computation of ADI shift parameters
p = lp_para(A,Bf,Kf,l0,kp,km,b0)

println("Actual number of ADI shift parameters:")
l = length(p)

println("ADI shift parameters:")
p

# -----------------------------------------------------------------------
# Solution of Lyapunov equations A*X+X*A' = -B*B' and
# A'*X+X*A = -C'*C
# -----------------------------------------------------------------------

println("Parameters for stopping criteria in LRCF-ADI iteration:")
max_it = 40;   # (will stop the iteration)
min_res = 1e-100;   # (avoided, but the residual history is shown)
with_rs = :N;   # (avoided)
min_in = 0;   # (avoided)

zk = :Z
rc = :R
Bf = ones(408,2)#[]
Kf = zeros(408,2)#[]
K = ones(408,408)
info = 2

println("... solving A*XB+XB*A'' = - B*B''...")
tp = :B

ZB,flag,res =lp_lradi(A,B,p=p,max_it=200)

X=ZB*ZB'
println(norm(A*X+X*A'+B*B'))


using Distributions, Random

function simulateGaussian(;μ=0, σ=1, Nd, kwargs...)
 (y = rand(Normal(μ,σ), Nd), N = Nd)
end


using Base.Test

using TALib
using TALib: basepath

using DataFrames

@testset "level 2 dataframes" begin

    epsilon = 1e-10
    eps_price = 1e-6

    filename = joinpath(basepath(), "ford_2012.csv")
    df = readtable(filename)
    df[:Date] = Date(df[:Date])

    indic = MA(df)
    @test indic[end, :Real] ≈ 11.546 atol=eps_price

    indic = MA(df, price=:Open)
    @test indic[end, :Real] ≈ 11.4613333 atol=eps_price

    indic = BBANDS(df)
    @test indic[end, :UpperBand] ≈ 13.13191 atol=eps_price  # UpperBand
    @test indic[end, :MiddleBand] ≈ 12.75400 atol=eps_price  # MiddleBand
    @test indic[end, :LowerBand] ≈ 12.37609 atol=eps_price  # LowerBand

end


#=
in this example, we just put a batch of the data into the
network at once over and over again to optimizes the
parameters of network.
=#

using Gadfly
using Random

# first class data and label
input1 = 3*randn(2,2000);
label1 = zeros(2,2000);
label1[1,:] .= 1.0;

# second class data and label
t = 0:pi/1000:2*pi;
input2 = zeros(2,length(t));
label2 = zeros(2,length(t));
label2[2,:] .= 1.0;
for i = 1:length(t)
    r = randn() + 16.0;
    input2[1,i] = r * cos(t[i]);
    input2[2,i] = r * sin(t[i]);
end

input = hcat(input1,input2)
label = hcat(label1,label2)

batch = 2048
epoch = 1000
lrate = 1e-4
topology = [2, 32,32,32, 2]
operator = [relu, relu, relu, relu]
indexs   = collect(1:batch)

mlpmodel = MLP(topology, operator)
paramter = paramsof(mlpmodel)
lossval  = zeros(epoch,1)
tic = time()
for i = 1:epoch
    ids  = shuffle!(indexs)
    outs = forward(mlpmodel, Variable(input[:,ids[1:batch]]))
    outs = softmax(outs; dims=1)
    COST = crossEntropyCost(outs, Variable(label[:,ids[1:batch]]))
    backward()
    update!(paramter,lrate)
    zerograds!(paramter)
    lossval[i] = COST
end
toc = time()
println("\n------- sgd optimizer loss: ", lossval[end])
l1 = layer(y=log.(lossval), Geom.line,Theme(default_color=colorant"cyan"))

mlpmodel  = MLP(topology, operator)
paramter  = paramsof(mlpmodel)
optimizer = Momentum(paramter,lr=lrate,inertia=0.9)
lossval   = zeros(epoch,1)
tic = time()
for i = 1:epoch
    ids  = shuffle!(indexs)
    outs = forward(mlpmodel, Variable(input[:,ids[1:batch]]))
    outs = softmax(outs; dims=1)
    COST = crossEntropyCost(outs, Variable(label[:,ids[1:batch]]))
    backward()
    update!(optimizer,paramter)
    zerograds!(paramter)
    lossval[i] = COST
end
toc = time()
println("-- Momentum optimizer loss: ", lossval[end])
l2 = layer(y=log.(lossval), Geom.line,Theme(default_color=colorant"yellow"))

mlpmodel  = MLP(topology, operator)
paramter  = paramsof(mlpmodel)
optimizer = Adam(paramter,lr=1e-3)
lossval   = zeros(epoch,1)
tic = time()
for i = 1:epoch
    ids  = shuffle!(indexs)
    outs = forward(mlpmodel, Variable(input[:,ids[1:batch]]))
    outs = softmax(outs; dims=1)
    COST = crossEntropyCost(outs, Variable(label[:,ids[1:batch]]))
    backward()
    update!(optimizer,paramter)
    zerograds!(paramter)
    lossval[i] = COST
end
toc = time()
println("------ Adma optimizer loss: ", lossval[end])
l3 = layer(y=log.(lossval), Geom.line,Theme(default_color=colorant"red"))

plot(l1,l2,l3)

mlpmodel  = Chain(
    Maxout(2,32,k=3),
    Residual(Dense(32,32),Dense(32,32),Linear(32,32)),
    Residual(Dense(32,32),Dense(32,32),Linear(32,32)),
    Residual(Dense(32,32),Dense(32,32),Linear(32,32)),
    Dense(32,2)
)

epoch = 70
parameter = paramsof(mlpmodel)
lossval  = zeros(epoch,1)
for i = 1:epoch
    ids  = shuffle!(indexs)
    outs = forward(mlpmodel, Variable(input[:,ids[1:batch]]))
    outs = softmax(outs; dims=1)
    COST = crossEntropyCost(outs, Variable(label[:,ids[1:batch]]))
    backward()
    update!(parameter,lrate)
    zerograds!(parameter)
    lossval[i] = COST
end
println(" sgd optimizer loss: ", lossval[end])


import DataStructures: BinaryMinHeap
using LinearAlgebra: norm

struct LSHBase
    alphas::Matrix{Float64} # Columns are linear equations.
    betas::Array{Float64, 2} # One per hash / one per alpha column.
    w::Array{Float64, 2} # Normalizing constant per hash.
end

"""
A gaussian LSH initialization method.

args:
    - n_dims: the number of dimensions in the input data.
    - w: the denominator and the max adjust value. May be an array of numbers.
    - n_hashes: number of hashes.
"""
function LSHBase(n_dims::Int, w=1.0, n_hashes=0)
    if n_hashes == 0
        n_hashes = Int(max(round(sqrt(n_dims)), 2))
    end
    alphas = randn(n_dims, n_hashes)

    # Normalize the alphas.
    alpha_norms = reshape([norm(col, 2) for col in eachcol(alphas)], 1, size(alphas, 2))
    unit_alphas = alphas ./ alpha_norms
    alphas = unit_alphas

    # Make sure that w is a column vector or an int.
    if isa(w, Float64) || isa(w, Int)
        w = reshape([w for _ in 1:n_hashes], 1, n_hashes)
    else
        w = reshape(w, 1, n_hashes)
    end
    betas = rand(1, n_hashes) .* w

    LSHBase(alphas, betas, w)
end

"""
The hash function.
"""
function hash(lsh::LSHBase, data)
    if length(size(data)) == 1
        data = reshape(data, 1, maximum(size(data)))
        hash_address = map(floor, ((data * lsh.alphas) .+ lsh.betas) ./ lsh.w)
        return vec(map(Int, hash_address))
    else
        hash_address = map(floor, ((data * lsh.alphas) .+ lsh.betas) ./ lsh.w)
        try
            return map(Int, hash_address)
        catch
            throw(error(join(["EXCEPTION", hash_address, data, lsh.alphas, lsh.betas, lsh.w], ", ")))
        end
    end
end

"""
Function to aggregate raw data into hash buckets.
"""
function reservoir_function(lshb::LSHBase, data::Array{Float64, 2})
    # Update vals
    vals = Dict{Array{Int, 1}, Any}()
    hashed = hash(lshb, data)
    for i in 1:size(hashed, 1)
        if !haskey(vals, hashed[i, :])
            vals[hashed[i, :]] = Array{Any, 1}()
        end
        push!(vals[hashed[i, :]], data[i, :])
    end
    return vals
end

"""
Function to aggregate raw data into hash buckets as average and count lists.
"""
function mean_count_function(lshb::LSHBase, data::Array{Float64, 2})
    vals = Dict{Array{Int, 1}, Any}()
    hashed = hash(lshb, data)
    for i in 1:size(hashed, 1)
        if !haskey(vals, hashed[i, :])
            vals[hashed[i, :]] = Array{Union{Int, Array{Float64, 1}}, 1}([[0. for _ in data[1, :]], 0])
        end
        vals[hashed[i, :]][1] += data[i, :]
        vals[hashed[i, :]][2] += 1
    end
    return vals
end

mutable struct LSH
    lshb::LSHBase
    aggregate_function::Function
    # Hash address to value.
    vals::Dict{Array{Int, 1}, Any}
end

"""
Create a new randomized LSH
args:
    - data: raw input data.
    - aggregate_function: a function to aggregate raw data points that fall
      into the same hash bucket.
    - n_hashes: the number of hashes.

"""
function LSH(data::Array{Float64, 2}, aggregate_function::Function=reservoir_function, n_hashes=0, w=.1)
    best = LSHBase(size(data, 2), w, n_hashes)

    # Update vals
    vals = aggregate_function(best, data)
    return LSH(best, aggregate_function, vals)
end

"""
Create a new LSH using the values in the provided LSHBase.
"""
function LSH(lshb::LSHBase, data::Array{Float64, 2}, aggregate_function::Function=reservoir_function)
    vals = aggregate_function(lshb, data)
    return LSH(lshb, aggregate_function, vals)
end

function hash(lsh::LSH, v)
    return hash(lsh.lshb, v)
end

"""
Calculate the score of a probe as the sum over the distance of the point from
the hash boundary + delta.

args:
    - query: a hash address.
    - delta: array of ints to apply to the hash address to get the probe
             location.
"""
function probe_score(lsh::LSHBase, query, delta)
    hashed_float = (query' * lsh.alphas .+ lsh.betas) ./ lsh.w
    if delta < 0
        return (map(floor, hashed_float) - hashed_float .+ [delta for _ in lsh.w]).^2
    else
        return (map(ceil, hashed_float) - hashed_float .+ [delta for _ in lsh.w]).^2
    end
end

"""
Generate query deltas in order of their probe_score and evaluate the probes for
neighbors. Return the neighbor hash addresses.

args
    lsh: mplsh object
    query: vector of data in original space
    target_neighbors: how many neighbors to return
    n_probes: the max number of probes to try
    probe_distance: positive int how many buckets away from query to include
                    in search.
"""
function get_neighbor_addresses(lsh::LSH, query, target_neighbors=1, probe_distance=1, n_probes=typemax(Int))
    neighbor_addresses = []
    query_address = hash(lsh, query)

    # Compute the probe_scores for all nonzero integer deltas in
    # [-probe_distance, probe_distance] (for each pair (i, delta_i))
    # perturbation_set will be (score, variable_index, probe_distance)
    perturbation_set = []
    for pd in 1:probe_distance
        perturbation_set = vcat(perturbation_set, vec([(score, hash_index, pd) for (hash_index, score) in enumerate(probe_score(lsh.lshb, query, pd))]))
        perturbation_set = vcat(perturbation_set, vec([(score, hash_index, -pd) for (hash_index, score) in enumerate(probe_score(lsh.lshb, query, -pd))]))
    end

    # Sort the pairs by score.
    scores = [pair[1] for pair in perturbation_set]
    perturbation_set = perturbation_set[sortperm(scores)]

    # Initialize the heap. The value [1] indicates that we'll start by trying
    # the first element in perturbation_set, which is sorted by cost.
    h = BinaryMinHeap{Tuple{Float64, Array{Int, 1}}}()
    push!(h, (perturbation_set[1][1], [1]))

    deltas = []
    count = 0
    n_probes_so_far = 0
    while length(deltas) < n_probes && length(h) > 0
        n_probes_so_far += 1
        count += 1
        # Get the next best probe.
        ai = pop!(h)

        # Create the hash address for this probe.
        probe = copy(query_address)
        for ind in ai[2]
            probe[perturbation_set[ind][2]] += perturbation_set[ind][3]
        end

        # If the bucket this delta indicates is not empty, add it to the neighbors set.
        if haskey(lsh.vals, probe)
            push!(neighbor_addresses, probe)
        end

        # If we've got the desired number of actual nonempty neighbors, return them.
        if length(neighbor_addresses) >= target_neighbors
            return neighbor_addresses
        end

        push!(deltas, ai[2])

        # Shift
        set = copy(ai[2])
        score = ai[1]
        # Remove the score for the current end perturbation.
        score -= perturbation_set[ai[2][end]][1]
        # Add the score for the new end perturbation.
        next = ai[2][end] + 1
        if next <= length(perturbation_set)
            score += perturbation_set[next][1]
            set[end] += 1
            push!(h, (score, set))
        end

        # Expand
        set = copy(ai[2])
        set = vcat(set, set[end] + 1)

        # Already updated the set at this point, offsets reflect.
        if length(set) <= length(lsh.lshb.w) - 1 && set[end] <= length(perturbation_set)
            score = ai[1] + perturbation_set[set[end]][1]
            push!(h, (score, set))
        end
    end

    return neighbor_addresses
end

"""
Find the approximate nearest neighbors with some constraints.

use:
Retrieve the first non-empty neighboring bucket's contents that can be found
adjacent to the query location.
    nearest_neighbors(lsh, [.23, -1, 3])
Retrieve the first non-empty neighboring bucket's contents within +-100 buckets
of the query location.
    nearest_neighbors(lsh, [.23, -1, 3], 1, probe_distance=100)
Retrieve the first 10 non-empty neighboring bucket's contents within +-100 buckets
of the query location.
    nearest_neighbors(lsh, [.23, -1, 3], 10, probe_distance=100)
And so on....

args
    lsh: mplsh object
    query: vector of data in original space
    n_neighbors: how many neighbors to return
    n_probes: the max number of probes to try
    probe_distance: positive int how many buckets away from query to include
                    in search.
"""
function nearest_neighbors(lsh::LSH, query, n_neighbors=1, probe_distance=3, n_probes=typemax(Int))
    nns = get_neighbor_addresses(lsh, query, n_neighbors, probe_distance, n_probes)
    neighbors = []
    for nn in nns
        neighbors = vcat(neighbors, lsh.vals[nn])
    end
    return Array{Float64, 2}(hcat(neighbors...)')
end


## Grid generation function.
function make_grid(zsize, ktot)
    dz = zsize / ktot
    z = range(0.5*dz, step=dz, length=ktot) |> collect
    return z
end

## Settings.
float_type = Float32

settings_grid = Dict{String, Any}(
    "itot" => 256,
    "jtot" => 256,
    "ktot" => 256,

    "xsize" => 3200.,
    "ysize" => 3200.,
    "zsize" => 3200.)

settings_fields = Dict(
    "visc" => 5.,
    "alpha" => 9.81 / 300)

settings_boundary = Dict(
    "mom_bot_type" => "Dirichlet",
    "mom_top_type" => "Neumann",
    "s_bot_type" => "Neumann",
    "s_top_type" => "Neumann")

settings_timeloop = Dict(
    "start_time" => 0.,
    "end_time" => 3600.,
    "save_time" => 900.,
    "check_time" => 2.,
    "dt" => 2.)

settings_multidomain = Dict(
    "enable_nudge" => false)

settings_d01 = Dict(
    "grid" => settings_grid,
    "fields" => settings_fields,
    "boundary" => settings_boundary,
    "timeloop" => settings_timeloop,
    "multidomain" => settings_multidomain)

settings_d01["grid"]["z"] = make_grid(settings_d01["grid"]["zsize"], settings_d01["grid"]["ktot"])

settings = [ settings_d01 ]


@testset "Primality check (zerodim subroutine)" begin

    R, (x, y) = Singular.PolynomialRing(Singular.QQ, ["x", "y"])

    @test check_primality_zerodim(Singular.Ideal(R, x^2 - 1, y^2 - 4)) == false

    @test check_primality_zerodim(Singular.Ideal(R, (x + 5) * (x^3 - 7), y - 3)) == false

    @test check_primality_zerodim(Singular.Ideal(R, x^3 - 5, y - 1)) == true

    @test check_primality_zerodim(Singular.Ideal(R, x^2 + 1, y^3 - 3 * x + x + 5)) == true

end


########################################################################
#
# Unsat.jl
#
# by Walt McNab (May 2018; updated in October 2019)
#
# A numerical model for steady-state, unsaturated flow
# in porous media, assuming a single fluid phase and passive gas phase
#
# Model domain is a 2-D vertical cross-section
#
########################################################################

using SparseArrays
using DelimitedFiles
using Statistics

### defined types ###

mutable struct Material
    name::AbstractString
    Kh::Float64                         # (saturated) hydraulic conductivity components
    Kz::Float64
    alpha::Float64                      # Van Genuchten unsaturated hydraulic parameters
    N::Float64
    m::Float64
    sr::Float64                         # residual saturation
end

mutable struct Cell
    x::Float64                          # grid cell center coordinates
    z::Float64
    P::Float64                          # pressure head (psi)
    Q::Float64                          # source/sink flux
    matNum::Int64                       # associated material index number
    fixed::Bool                         # fixed-head cell, or not (boolean)
    xConnect::Array{Int64, 1}           # list of connecting cells
    zConnect::Array{Int64, 1}
end

mutable struct Domain
    xLength::Float64                    # model discretization
    zLength::Float64
    nx::Int64
    nz::Int64
    dx::Float64
    dz::Float64
    Ax::Float64
    Az::Float64
    defaultMat::AbstractString          # default material assignment (heterogeneities handle in blocks.txt file)
end

mutable struct Solver
    psi0::Float64                       # initial pressure across entire model domain (to start iterations)
    f::Float64                          # pressure correction weighting factor, per iteration
    maxIter::Int64                      # number of iterations
end

### support functions (physics) ###

function kr(psi::Float64, material::Material)::Float64
    # relative permeability (Van Genuchten/Mualem model)
    if psi < 0.0
        y1 = 1.0 + abs(material.alpha * psi)^material.N
        y2 = 1.0 - abs(material.alpha * psi)^(material.N - 1.0) * y1^(-material.m)
        k = y2^2 / y1^(material.m/2.0)
    else
        k = 1.0
    end
    return k
end

function Se(psi::Float64, material::Material)::Float64
    # effective saturation (Van Genuchten/Mualem model)
    if psi < 0.0
        satEff = (1.0 + abs(material.alpha * psi)^material.N)^(-material.m)
    else
        satEff = 1.0
    end
    return satEff
end

function S(psi::Float64, material::Material)::Float64
    # saturation vs relative saturation
    return Se(psi, material) * (1.0 - material.sr) + material.sr
end

### support functions (utility) ###

function HMean2(a::Float64, b::Float64)::Float64
    # harmonic mean of two numbers
    return 2.0/(1.0/a + 1.0/b)
end

function GetIndex(x::Float64, z::Float64, domain::Domain)::Int64
    # find the index number of nearest cell object associated with location (x, y)
    row = maximum([round(Int64, z/domain.dz + 0.5), 1])
    col = maximum([round(Int64, x/domain.dx + 0.5), 1])
    return (row-1)*domain.nx + col
end

function GetMatNum(material::Array{Material, 1}, name::AbstractString)::Int64
    # return the index number of name in the materials list
    index = 0
    for (i, mat) in enumerate(material)
        if mat.name == name
            index = i
            break
        end
    end
    return index
end

function CreateCells(domain::Domain, matNum::Int64, psi0::Float64)::Array{Cell, 1}
    # create cell objects with default material assignment and initial pressure head
    cell = Cell[]
    for j = 1:domain.nz, i = 1:domain.nx
        x = (i-0.5) * domain.dx
        z = (j-0.5) * domain.dz
        xConnect = Int64[]
        zConnect = Int64[]
        push!(cell, Cell(x, z, psi0, 0.0, matNum, false, xConnect, zConnect))
    end
    # populate cell connection arrays
    for j = 0:domain.nz-1, i = 1:domain.nx-1
        push!(cell[i+j*domain.nx].xConnect, i+j*domain.nx+1)
        push!(cell[i+j*domain.nx+1].xConnect, i+j*domain.nx)
    end
    for j = 0:domain.nz-2, i = 1:domain.nx
        push!(cell[i+j*domain.nx].zConnect, i+(j+1)*domain.nx)
        push!(cell[i+(j+1)*domain.nx].zConnect, i+j*domain.nx)
    end
    return cell
end

function UpdatePressures(cell::Array{Cell, 1}, newP::Array{Float64, 1}, solverParams::Solver)::Array{Cell, 1}
    for (i, ce) in enumerate(cell)
        # weighting between old modeled head and new modeled head, per iteration
        ce.P = solverParams.f*newP[i] + (1.0-solverParams.f)*ce.P
    end
    return cell
end

function SumFluxes(cell::Array{Cell, 1}, material::Array{Material, 1}, domain::Domain)
    # compute fluxes (for model output and final flow balance estimates)
    flowSum = Float64[]
    Qx = Float64[]
    Qz = Float64[]
    errorF = Float64[]
    for ce in cell
        push!(flowSum, ce.Q)                # initialize summations
        push!(Qx, 0.0)
        push!(Qz, 0.0)
    end
    for (i, ce) in enumerate(cell)
        matI = material[ce.matNum]
        for xcon in ce.xConnect         # for each horizontal connection (up to 2 columns)
            matJ = material[cell[xcon].matNum]
            conduct = HMean2(kr(ce.P, matI)*matI.Kh, kr(cell[xcon].P, matJ)*matJ.Kh) * domain.Ax / domain.dx
            flux = conduct * (cell[xcon].P - ce.P)
            flowSum[i] += flux
            Qx[i] += sign(ce.x - cell[xcon].x) * flux
        end
        Qx[i] = Qx[i]/length(ce.xConnect)
        for zcon in ce.zConnect         # for each vertical connection (up to two columns)
            matJ = material[cell[zcon].matNum]
            conduct = HMean2(kr(ce.P, matI)*matI.Kz, kr(cell[zcon].P, matJ)*matJ.Kz) * domain.Az / domain.dz
            flux = conduct * ((cell[zcon].P + cell[zcon].z) - (ce.P + ce.z))
            flowSum[i] += flux
            Qz[i] += sign(ce.z - cell[zcon].z) * flux
        end
        Qz[i] = Qz[i]/length(ce.zConnect)
        flowThru = abs(Qx[i]) + abs(Qz[i])
        push!(errorF, flowSum[i]/flowThru)
    end
    return Qx, Qz, errorF
end

### matrix operations functions

function LHS_matrix(cell::Array{Cell, 1}, material::Array{Material, 1}, domain::Domain)
    # fill out the LHS of the equation matrix
    row_index = Int64[]                     # indexing system for sparse matrix
    col_index = Int64[]
    data = Float64[]
    for (i, ce) in enumerate(cell)          # for each row
        diag = 0.0
        if ce.fixed == false
            matI = material[ce.matNum]
            for xcon in ce.xConnect         # for each horizontal connection (up to 2 columns)
                matJ = material[cell[xcon].matNum]
                conduct = HMean2(kr(ce.P, matI)*matI.Kh, kr(cell[xcon].P, matJ)*matJ.Kh) * domain.Ax / domain.dx
                push!(row_index, i)
                push!(col_index, xcon)
                push!(data, conduct)
                diag -= conduct
            end
            for zcon in ce.zConnect         # for each vertical connection (up to two columns)
                matJ = material[cell[zcon].matNum]
                conduct = HMean2(kr(ce.P, matI)*matI.Kz, kr(cell[zcon].P, matJ)*matJ.Kz) * domain.Az / domain.dz
                push!(row_index, i)
                push!(col_index, zcon)
                push!(data, conduct)
                diag -= conduct
            end
        else
            diag = 1.0             # fixed pressure head cell
        end
        # set diagonal term
        push!(row_index, i)
        push!(col_index, i)
        push!(data, diag)
    end
    return data, row_index, col_index
end

function RHS_vector(cell::Array{Cell, 1}, material::Array{Material, 1}, domain::Domain)::Array{Float64, 1}
    # construct explicit matrix
    b = Float64[]
    for (i, ce) in enumerate(cell)          # for each row
        if ce.fixed == false
            bSum = -ce.Q
            matI = material[ce.matNum]
            for zcon in ce.zConnect         # for each vertical connection (up to two columns)
                matJ = material[cell[zcon].matNum]
                conduct = HMean2(kr(ce.P, matI)*matI.Kz, kr(cell[zcon].P, matJ)*matJ.Kz) * domain.Az / domain.dz
                bSum -= conduct * (cell[zcon].z - ce.z)
            end
            push!(b, bSum)
        else
            push!(b, ce.P)                  # fixed-head cell
        end
    end
    return b
end

### input and output functions

function ReadMaterials()::Array{Material, 1}
    # read various model parameters from file
    material = Material[]
    data = readdlm("materials.txt", '\t', header=true)
    for i = 1:size(data[1], 1)
        name = data[1][i, 1]
        Kh = Float64(data[1][i, 2])
        Kz = Float64(data[1][i, 3])
        alpha = Float64(data[1][i, 4])
        N = Float64(data[1][i, 5])
        m = 1-1/N
        sr = Float64(data[1][i, 6])
        push!(material, Material(name, Kh, Kz, alpha, N, m, sr))
    end
    println("Read material properties.")
    return material
end

function ReadSolverParams()::Solver
    # read numerical model "knobs" from file
    data = readdlm("solver.txt", '\t', header=false)
    psi0 = Float64(data[1, 2])
    f = Float64(data[2, 2])
    maxIter = Int64(data[3, 2])
    solverParams = Solver(psi0, f, maxIter)
    println("Read solver parameters.")
    return solverParams
end

function ReadDomain()::Domain
    # model geometry
    data = readdlm("domain.txt", '\t', header=false)
    xLength = Float64(data[1, 2])
    zLength = Float64(data[2, 2])
    nx = Int64(data[3, 2])
    nz = Int64(data[4, 2])
    dx = xLength/nx
    dz = zLength/nz
    Ax = dz             # placeholders for later expansion to 3-D
    Az = dx
    defaultMat = data[5, 2]
    domain = Domain(xLength, zLength, nx, nz, dx, dz, Ax, Az, defaultMat)
    println("Read model domain properties.")
    return domain
end

function ReadBlocks(cell::Array{Cell, 1}, material::Array{Material, 1})::Array{Cell, 1}
    # read and process rectangular material property heterogeneities
    data = readdlm("blocks.txt", '\t', header=true)
    for i = 1:size(data[1], 1)
        x0 = Float64(data[1][i, 1])
        z0 = Float64(data[1][i, 2])
        xf = Float64(data[1][i, 3])
        zf = Float64(data[1][i, 4])
        matBlock = data[1][i, 5]
        matNum = GetMatNum(material, matBlock)
        for ce in cell
            if ce.x >= x0 && ce.x <= xf && ce.z >= z0 && ce.z <= zf
                ce.matNum = matNum
            end
        end
    end
    println("Read material blocks.")
    return cell
end

function ReadSources(cell::Array{Cell, 1}, domain::Domain)::Array{Cell, 1}
    # read and process line sources/boundary conditions
    data = readdlm("sources.txt", '\t', header=true)
    for i = 1:size(data[1], 1)
        BC = data[1][i, 1]
        direction = data[1][i, 2]
        pos = Float64(data[1][i, 3])
        start = Float64(data[1][i, 4])
        fin = Float64(data[1][i, 5])
        value = Float64(data[1][i, 6])
        if direction == "x"
            cellStart = GetIndex(start, pos, domain)
            cellEnd = GetIndex(fin, pos, domain)
            stepSize = 1
            A = domain.Az
        else
            cellStart = GetIndex(pos, start, domain)
            cellEnd = GetIndex(pos, fin, domain)
            stepSize = domain.nx
            A = domain.Ax
        end
        if BC == "flux"
            for i = cellStart:stepSize:cellEnd
                cell[i].Q = value * A
            end
        else
            for i = cellStart:stepSize:cellEnd
                cell[i].fixed = true
                cell[i].P = value
            end
        end
    end
    println("Read sources/boundary condition constraints.")
    return cell
end

function ReadCells(cell::Array{Cell, 1})::Array{Cell, 1}
	# read in cell properties line by line
    data = readdlm("cells.txt", '\t', header=true)
    for i = 1:size(data[1], 1)
        indx = Int64(data[1][i, 1])
        cell[indx].P = Float64(data[1][i, 4])
        cell[indx].Q = Float64(data[1][i, 5])
        cell[indx].matNum = Int64(data[1][i, 6])
        cell[indx].fixed = Bool(data[1][i, 7])
    end
    println("Read individual cell properties.")
    return cell
end

function Spatial(domain::Domain, material::Array{Material, 1}, solverParams::Solver, readMode::Int64)::Array{Cell, 1}
    # set up the model spatial domain
    matNum = GetMatNum(material, domain.defaultMat)                 # default material index number
    cell = CreateCells(domain, matNum, solverParams.psi0)           # create grid cells
    if readMode == 1
		# apply simple heterogeneities (blocks)
		cell = ReadBlocks(cell, material)                           # read blocks file and update cell material assignments, as warranted
		cell = ReadSources(cell, domain)                            # read and assign sources/boundary conditions
	else
		# complex heterogeneities (cell-by-cell)
		cell = ReadCells(cell)
	end
    println("Set up model cells and distributed properties.")
    return cell
end

function WriteOutput(cell::Array{Cell, 1}, material::Array{Material, 1}, domain::Domain, fileName::AbstractString)
    # write property distribution to file
    csvfile = open(fileName,"w")
    flowImbalance = Float32[]
    Qx, Qz, errorF = SumFluxes(cell, material, domain)              # update flow balance estimates
    line_out = "cell" * "," * "x" * "," * "z" * "," * "psi" * "," * "h" * "," *
        "K_eff_h" * "," * "K_eff_z" * "," * "S" * "," * "Qx" * "," * "Qz" * "," * "error"
    println(csvfile, line_out)
    for (i, ce) in enumerate(cell)
        sat = S(ce.P, material[ce.matNum])
        k_rel = kr(ce.P, material[ce.matNum])
        K_eff_h = material[ce.matNum].Kh * k_rel
        K_eff_z = material[ce.matNum].Kz * k_rel
        line_out = string(i) * "," * string(ce.x) * "," * string(ce.z) * "," * string(ce.P) * "," * string(ce.z+ce.P) * "," *
            string(K_eff_h) * "," * string(K_eff_z) * "," * string(sat) * "," * string(Qx[i]) * "," * string(Qz[i]) * "," * string(errorF[i])
        println(csvfile,line_out)
    end
    close(csvfile)
    println("Wrote " * fileName * ".")
    if fileName == "final_estimates.csv"
        quants = [0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99]
        errors = quantile(errorF, quants)
        return quants, errors
    else
        return [0.], [0.]
    end
end

### main ###

function Unsat(readMode::Int64)

    # problem setup
    domain = ReadDomain()
    material = ReadMaterials()
    solverParams = ReadSolverParams()
    cell = Spatial(domain, material, solverParams, readMode)

    # write initial conditions to file
    fileName = "starting_estimates.csv"
    WriteOutput(cell, material, domain, fileName)

    println("Solving ...")

    # iterative solution of coupled non-linear flow balance PDEs
    numCells = length(cell)
    iter = 0
    while iter < solverParams.maxIter
        data, row_index, col_index = LHS_matrix(cell, material, domain)         # assemble and solve matrix of flux balance equations
        A = sparse(row_index, col_index, data, numCells, numCells)
        b = RHS_vector(cell, material, domain)
        newP = \(A, b)
        cell = UpdatePressures(cell, newP, solverParams)
        iter += 1
    end

    # write model results to file
    fileName = "final_estimates.csv"
    quants, errors = WriteOutput(cell, material, domain, fileName)
    println("Quantile", "\t", "Error (norm.)")
    for i = 1:length(quants)
        println(quants[i], "\t", errors[i])
    end
    println("Finished.")

end

### run model ###
readMode = 2
Unsat(readMode)


min_mass = 0.6
max_mass = 1.4
const g = -9.8
#z = [0;0;1]
sigma = 1.0

mass_resolution = 0.2
#mass_resolution = 0.05
masses = collect(min_mass:mass_resolution:max_mass)

np = 20
thetas = 2*pi * (1/np:1/np:1)
circle_ps = map(theta->[cos(theta);sin(theta);0], thetas)
circle_ws = to_wrench_matrix(circle_ps)

# move to separate attachment points
na = 10
thetas = 2*pi * (1/na:1/na:1)
r_attach = 0.8
attachment_ps = map(theta->r_attach*[cos(theta);sin(theta);0], thetas)
attachment_ws = map(point_to_wrench, attachment_ps)

# set up the occupancy prior
resolution = 0.1

# note resolution
interior_q(x) = norm(x) <= 1 - 0.1 * 3/2

normalized_error(err) = sqrt(sum((err.*[0.5;0.5;1/(max_mass - min_mass)]).^2))