text:
# export
# 	GoalConstraint,
# 	BoundConstraint,
# 	CircleConstraint,
# 	SphereConstraint,
# 	NormConstraint,
# 	LinearConstraint,
# 	VariableBoundConstraint,
# 	QuatNormConstraint,
# 	QuatSlackConstraint

import RobotDynamics: state_dim, control_dim

Base.copy(con::AbstractConstraint) = con

############################################################################################
#                              GOAL CONSTRAINTS 										   #
############################################################################################

"""
	GoalConstraint{P,T}

Constraint of the form ``x_g = a``, where ``x_g`` can be only part of the state
vector.

# Constructors:
```julia
GoalConstraint(xf::AbstractVector)
GoalConstraint(xf::AbstractVector, inds)
```
where `xf` is an n-dimensional goal state. If `inds` is provided,
only `xf[inds]` will be used.
"""
struct GoalConstraint{P,T} <: StateConstraint
	n::Int
	xf::MVector{P,T}
	inds::SVector{P,Int}
	function GoalConstraint(xf::AbstractVector{T}, inds::SVector{p,Int}) where {p,T}
		new{p,T}(length(xf), xf[inds], inds)
	end
end

function GoalConstraint(xf::AbstractVector, inds=1:length(xf))
	p = length(inds)
	inds = SVector{p}(inds)
	GoalConstraint(xf, inds)
end

Base.copy(con::GoalConstraint) = GoalConstraint(copy(con.xf), con.inds)

@inline sense(::GoalConstraint) = Equality()
@inline Base.length(con::GoalConstraint{P}) where P = P
@inline state_dim(con::GoalConstraint) = con.n
@inline is_bound(::GoalConstraint) = true
function primal_bounds!(zL,zU,con::GoalConstraint)
	for i in con.inds
		zL[i] = con.xf[i]
		zU[i] = con.xf[i]
	end
	return true
end

evaluate(con::GoalConstraint, x::StaticVector) = x[con.inds] - con.xf
function jacobian!(∇c, con::GoalConstraint, z::KnotPoint)
	T = eltype(∇c)
	for (i,j) in enumerate(con.inds)
		∇c[i,j] = one(T)
	end
	return true
end

∇jacobian!(G, con::GoalConstraint, z::AbstractKnotPoint, λ::AbstractVector) = true # zeros

function change_dimension(con::GoalConstraint, n::Int, m::Int, xi=1:n, ui=1:m)
	GoalConstraint(con.xf, xi[con.inds])
end

function set_goal_state!(con::GoalConstraint, xf::AbstractVector)
	if length(xf) != length(con.xf)
		for (i,j) in enumerate(con.inds)
			con.xf[i] = xf[j]
		end
	else
		con.xf .= xf
	end
end

############################################################################################
#                              LINEAR CONSTRAINTS 										   #
############################################################################################
"""
	LinearConstraint{S,P,W,T}

Linear constraint of the form ``Ay - b \\{\\leq,=\\} 0`` where ``y`` may be either the
state or controls (but not a combination of both).

# Constructor: ```julia
LinearConstraint{S,W}(n,m,A,b)
```
where `W <: Union{State,Control}`.
"""
struct LinearConstraint{S,P,W,T} <: StageConstraint
	n::Int
	m::Int
	A::SizedMatrix{P,W,T,2}
	b::SVector{P,T}
	sense::S
	inds::SVector{W,Int}
	function LinearConstraint(n::Int, m::Int, A::StaticMatrix{P,W,T}, b::StaticVector{P,T},
			sense::ConstraintSense, inds=1:n+m) where {P,W,T}
		@assert length(inds) == W
		inds = SVector{W}(inds)
		new{typeof(sense),P,W,T}(n,m,A,b,sense,inds)
	end
end

function LinearConstraint(n::Int, m::Int, A::AbstractMatrix, b::AbstractVector,
		sense::S, inds=1:n+m) where {S<:ConstraintSense}
	@assert size(A,1) == length(b)
	p,q = size(A)
	A = SizedMatrix{p,q}(A)
	b = SVector{p}(b)
	LinearConstraint(n,m, A, b, sense, inds)
end

Base.copy(con::LinearConstraint{S}) where S =
	LinearConstraint(con.n, con.m, copy(con.A), copy(con.b), S(), con.inds)

@inline sense(con::LinearConstraint) = con.sense
@inline Base.length(con::LinearConstraint{<:Any,P}) where P = P
@inline state_dim(con::LinearConstraint) = con.n
@inline control_dim(con::LinearConstraint) = con.m
evaluate(con::LinearConstraint, z::AbstractKnotPoint) = con.A*z.z[con.inds] .- con.b
function jacobian!(∇c, con::LinearConstraint, z::AbstractKnotPoint)
	∇c[:,con.inds] .= con.A
	return true
end

function change_dimension(con::LinearConstraint, n::Int, m::Int, ix=1:n, iu=1:m)
	inds0 = [ix; n .+ iu]  # indices of original z in new z
	inds = inds0[con.inds] # indices of elements in new z
	LinearConstraint(n, m, con.A, con.b, con.sense, inds)
end

############################################################################################
#                              CIRCLE/SPHERE CONSTRAINTS 								   #
############################################################################################
"""
	CircleConstraint{P,T}

Constraint of the form
`` (x - x_c)^2 + (y - y_c)^2 \\leq r^2 ``
where ``x``, ``y`` are given by `x[xi]`,`x[yi]`, ``(x_c,y_c)`` is the center
of the circle, and ``r`` is the radius.

# Constructor:
```julia
CircleConstraint(n, xc::SVector{P}, yc::SVector{P}, radius::SVector{P}, xi=1, yi=2)
```
"""
struct CircleConstraint{P,T} <: StateConstraint
	n::Int
	x::SVector{P,T}
	y::SVector{P,T}
	radius::SVector{P,T}
	xi::Int  # index of x-state
	yi::Int  # index of y-state
	function CircleConstraint{P,T}(n::Int, xc::AbstractVector, yc::AbstractVector, radius::AbstractVector,
			xi=1, yi=2) where {P,T}
    	@assert length(xc) == length(yc) == length(radius) == P "Lengths of xc, yc, and radius must be equal. Got lengths ($(length(xc)), $(length(yc)), $(length(radius)))"
        new{P,T}(n, xc, yc, radius, xi, yi)
    end
end
function CircleConstraint(n::Int, xc::AbstractVector, yc::AbstractVector, radius::AbstractVector,
		xi=1, yi=2)
    T = promote_type(eltype(xc), eltype(yc), eltype(radius))
    P = length(xc)
    CircleConstraint{P,T}(n, xc, yc, radius, xi, yi)
end
state_dim(con::CircleConstraint) = con.n

function evaluate(con::CircleConstraint, X::StaticVector)
	xc = con.x
	yc = con.y
	r = con.radius
	x = X[con.xi]
	y = X[con.yi]
	-(x .- xc).^2 - (y .- yc).^2 + r.^2
end

function jacobian!(∇c, con::CircleConstraint{P}, X::SVector) where P
	xc = con.x; xi = con.xi
	yc = con.y; yi = con.yi
	x = X[xi]
	y = X[yi]
	r = con.radius
	∇f(x,xc) = -2*(x - xc)
	for i = 1:P
		∇c[i,xi] = ∇f(x, xc[i])
		∇c[i,yi] = ∇f(y, yc[i])
	end
	return false
end

@inline Base.length(::CircleConstraint{P}) where P = P
@inline sense(::CircleConstraint) = Inequality()

function change_dimension(con::CircleConstraint, n::Int, m::Int, ix=1:n, iu=1:m)
	CircleConstraint(n, con.x, con.y, con.radius, ix[con.xi], ix[con.yi])
end

"""
	SphereConstraint{P,T}

Constraint of the form
`` (x - x_c)^2 + (y - y_c)^2 + (z - z_c)^2 \\leq r^2 ``
where ``x``, ``y``, ``z`` are given by `x[xi]`,`x[yi]`,`x[zi]`, ``(x_c,y_c,z_c)`` is the center
of the sphere, and ``r`` is the radius.

# Constructor:
```
SphereConstraint(n, xc::SVector{P}, yc::SVector{P}, zc::SVector{P},
	radius::SVector{P}, xi=1, yi=2, zi=3)
```
"""
struct SphereConstraint{P,T} <: StateConstraint
	n::Int
	x::SVector{P,T}
	y::SVector{P,T}
	z::SVector{P,T}
	xi::Int
	yi::Int
	zi::Int
	radius::SVector{P,T}
	function SphereConstraint{P,T}(n::Int, xc::AbstractVector, yc::AbstractVector,
            zc::AbstractVector, radius::AbstractVector,
			xi=1, yi=2, zi=3) where {P,T}
    	@assert length(xc) == length(yc) == length(radius) == length(zc) == P "Lengths of xc, yc, zc, and radius must be equal. Got lengths ($(length(xc)), $(length(yc)), $(length(zc)), $(length(radius)))"
        new{P,T}(n, xc, yc, zc, xi, yi, zi, radius)
    end
end
function SphereConstraint(n::Int, xc::AbstractVector, yc::AbstractVector,
        zc::AbstractVector, radius::AbstractVector,
		xi=1, yi=2, zi=3)
    T = promote_type(eltype(xc), eltype(yc), eltype(zc), eltype(radius))
    P = length(xc)
    SphereConstraint{P,T}(n, xc, yc, zc, radius, xi, yi, zi)
end

@inline state_dim(con::SphereConstraint) = con.n
@inline sense(::SphereConstraint) = Inequality()
@inline Base.length(::SphereConstraint{P}) where P = P

function evaluate(con::SphereConstraint, x::SVector)
	xc = con.x; xi = con.xi
	yc = con.y; yi = con.yi
	zc = con.z; zi = con.zi
	r = con.radius

	-((x[xi] .- xc).^2 + (x[yi] .- yc).^2 + (x[zi] .- zc).^2 - r.^2)
end

function jacobian!(con::SphereConstraint, X::SVector)
	xc = con.x; xi = con.xi
	yc = con.y; yi = con.yi
	zc = con.z; zi = con.zi
	x = X[xi]
	y = X[yi]
	z = X[zi]
	r = con.radius
	∇f(x,xc) = -2*(x - xc)
	for i = 1:P
		∇c[i,xi] = ∇f(x, xc[i])
		∇c[i,yi] = ∇f(y, yc[i])
		∇c[i,zi] = ∇f(z, zc[i])
	end
	return false
end

function change_dimension(con::SphereConstraint, n::Int, m::Int, ix=1:n, iu=1:m)
	SphereConstraint(n, con.x, con.y, con.z, con.radius, ix[con.xi], ix[con.yi], ix[con.zi])
end

############################################################################################
#  								SELF-COLLISION CONSTRAINT 								   #
############################################################################################

"""
    CollisionConstraint

Enforces a pairwise non self-collision constraint on the state, such that
    `norm(x[x1] - x[x2]).^2 > r^2`,
    where `x1` and `x2` are the indices of the positions of the respective bodies and `r`
    is the collision radius.

# Constructor
CollisionConstraint(n::Int, x1::AbstractVector{Int}, x2::AbstractVector{Int}, r::Real)
"""
struct CollisionConstraint{D} <: StateConstraint
	n::Int
    x1::SVector{D,Int}
    x2::SVector{D,Int}
    radius::Float64
    function CollisionConstraint(n::Int, x1::AbstractVector{Int}, x2::AbstractVector{Int}, r::Real)
        @assert length(x1) == length(x2) "Position dimensions must be of equal length, got $(length(x1)) and $(length(x2))"
        D = length(x1)
        new{D}(n, x1, x2, r)
    end
end

@inline state_dim(con::CollisionConstraint) = con.n
@inline sense(::CollisionConstraint) = Inequality()
@inline Base.length(::CollisionConstraint) = 1

function evaluate(con::CollisionConstraint, x::SVector)
    x1 = x[con.x1]
    x2 = x[con.x2]
    d = x1 - x2
    @SVector [con.radius^2 - d'd]
end

function jacobian!(∇c, con::CollisionConstraint, x::SVector)
    x1 = x[con.x1]
    x2 = x[con.x2]
    d = x1 - x2
	∇x1 = -2d
	∇x2 =  2d
	∇c[1,con.x1] .= ∇x1
	∇c[1,con.x2] .= ∇x2
	return false
end

function change_dimension(con::CollisionConstraint, n::Int, m::Int, ix=1:n, iu=1:m)
	CollisionConstraint(n, ix[con.x1], ix[con.x2], con.radius)
end

############################################################################################
#								NORM CONSTRAINT											   #
############################################################################################

"""
	NormConstraint{S,D,T}

Constraint of the form
``\\|y\\|^2 \\{\\leq,=\\} a``
where ``y`` is made up of elements from the state and/or control vectors.

# Constructor:
```
NormConstraint(n, m, a, sense, [inds])
```
where `n` is the number of states,
    `m` is the number of controls,
    `a` is the constant on the right-hand side of the equation,
    `sense` is either `Inequality()` or `Equality()`, and
    `inds` can be a `UnitRange`, `AbstractVector{Int}`, or either `:state` or `:control`

# Examples:
```julia
NormConstraint(3, 2, 4, Equality(), :control)
```
creates a constraint equivalent to
``\\|u\\|^2 = 4.0`` for a problem with 2 controls.

```julia
NormConstraint(3, 2. 3, Inequality(), :state
```
creates a constraint equivalent to
``\\|x\\|^2 \\leq 2.3`` for a problem with 3 states.

"""
struct NormConstraint{S,D,T} <: StageConstraint
	n::Int
	m::Int
	val::T
	sense::S
	inds::SVector{D,Int}
	function NormConstraint(n::Int, m::Int, val::T, sense::ConstraintSense,
			inds=SVector{n+m}(1:n+m)) where T
		if inds == :state
			inds = SVector{n}(1:n)
		elseif inds == :control
			inds = SVector{m}(n .+ (1:m))
		end
		@assert val ≥ 0 "Value must be greater than or equal to zero"
		new{typeof(sense),length(inds),T}(n,m,val,sense,inds)
	end
end

@inline state_dim(con::NormConstraint) = con.n
@inline control_dim(con::NormConstraint) = con.m
@inline sense(con::NormConstraint) = con.sense
@inline Base.length(::NormConstraint) = 1

function evaluate(con::NormConstraint, z::AbstractKnotPoint)
	x = z.z[con.inds]
	return @SVector [x'x - con.val]
end

function jacobian!(∇c, con::NormConstraint, z::AbstractKnotPoint)
	x = z.z[con.inds]
	∇c[1,con.inds] .= 2*x
	return false
end

function change_dimension(con::NormConstraint, n::Int, m::Int, ix=1:n, iu=1:m)
	NormConstraint(n, m, con.val, con.sense, ix[con.inds])
end

############################################################################################
# 								COPY CONSTRAINT 										   #
############################################################################################

# struct CopyConstraint{K,W,S,P,N,M} <: AbstractConstraint{W,S,P}
# 	con::AbstractConstraint{W,S,P}
#     xinds::Vector{SVector{N,Int}}
#     uinds::Vector{SVector{M,Int}}
# end
#
# function evaluate(con::CopyConstraint{K}, z::KnotPoint)
# 	c = evaluate(con,)
# 	for 2 = 1:K
# 	end
# end

############################################################################################
# 								BOUND CONSTRAINTS 										   #
############################################################################################
"""
	BoundConstraint{P,NM,T}

Linear bound constraint on states and controls
# Constructors
```julia
BoundConstraint(n, m; x_min, x_max, u_min, u_max)
```
Any of the bounds can be ±∞. The bound can also be specifed as a single scalar, which applies the bound to all state/controls.
"""
struct BoundConstraint{P,NM,T} <: StageConstraint
	n::Int
	m::Int
	z_max::SVector{NM,T}
	z_min::SVector{NM,T}
	i_max::Vector{Int}
	i_min::Vector{Int}
	inds::SVector{P,Int}
end

Base.copy(bnd::BoundConstraint{P,nm,T}) where {P,nm,T} =
	BoundConstraint(bnd.n, bnd.m, bnd.z_max, bnd.z_min,
		copy(bnd.i_max), copy(bnd.i_min), bnd.inds)

function BoundConstraint(n, m; x_max=Inf*(@SVector ones(n)), x_min=-Inf*(@SVector ones(n)),
		u_max=Inf*(@SVector ones(m)), u_min=-Inf*(@SVector ones(m)))
	nm = n+m

	# Check and convert bounds
	x_max, x_min = checkBounds(n, x_max, x_min)
	u_max, u_min = checkBounds(m, u_max, u_min)

	# Concatenate bounds
	z_max = [x_max; u_max]
	z_min = [x_min; u_min]
	b = [-z_max; z_min]
	inds = findall(isfinite, b)
	inds = SVector{length(inds)}(inds)

	# Get linear indices of 1s of Jacobian
	a_max = findall(isfinite, z_max)
	a_min = findall(isfinite, z_min)
	u = length(a_max)
	l = length(a_min)
	carts_u = [CartesianIndex(i,   j) for (i,j) in enumerate(a_max)]
	carts_l = [CartesianIndex(i+u, j) for (i,j) in enumerate(a_min)]
	∇c = zeros(u+l, n+m)
	linds_u = LinearIndices(zeros(u+l,n+m))[carts_u]
	linds_l = LinearIndices(zeros(u+l,n+m))[carts_l]

	BoundConstraint(n, m, z_max, z_min, linds_u, linds_l, inds)
end

function con_label(con::BoundConstraint, ind::Int)
	i = con.inds[ind]
	n,m = state_dim(con), control_dim(con)
	if 1 <= i <= n
		return "x max $i"
	elseif n < i <= n + m
		j = i - n
		return "u max $j"
	elseif n + m < i <= 2n+m
		j = i - (n+m)
		return "x min $j"
	elseif 2n+m < i <= 2n+2m
		j = i - (2n+m)
		return "u min $j"
	else
		throw(BoundsError())
	end
end

function checkBounds(n::Int, u::AbstractVector, l::AbstractVector)
	if all(u .>= l)
		return SVector{n}(u), SVector{n}(l)
	else
		throw(ArgumentError("Upper bounds must be greater than or equal to lower bounds"))
	end
end

checkBounds(n::Int, u::Real, l::Real) =
	checkBounds(n, (@SVector fill(u,n)), (@SVector fill(l,n)))
checkBounds(n::Int, u::AbstractVector, l::Real) = checkBounds(n, u, (@SVector fill(l,N)))
checkBounds(n::Int, u::Real, l::AbstractVector) = checkBounds(n, (@SVector fill(u,N)), l)

@inline state_dim(con::BoundConstraint) = con.n
@inline control_dim(con::BoundConstraint) = con.m
@inline is_bound(::BoundConstraint) = true
@inline lower_bound(bnd::BoundConstraint) = bnd.z_min
@inline upper_bound(bnd::BoundConstraint) = bnd.z_max
@inline sense(::BoundConstraint) = Inequality()
@inline Base.length(con::BoundConstraint) = length(con.i_max) + length(con.i_min)

function primal_bounds!(zL, zU, bnd::BoundConstraint)
	for i = 1:length(zL)
		zL[i] = max(bnd.z_min[i], zL[i])
		zU[i] = min(bnd.z_max[i], zU[i])
	end
	return true
end

function evaluate(bnd::BoundConstraint, z::AbstractKnotPoint)
	[(z.z - bnd.z_max); (bnd.z_min - z.z)][bnd.inds]
end

function jacobian!(∇c, bnd::BoundConstraint{U,L}, z::AbstractKnotPoint) where {U,L}
	for i in bnd.i_max
		∇c[i]  = 1
	end
	for i in bnd.i_min
		∇c[i] = -1
	end
	return true
end

∇jacobian!(G, con::BoundConstraint, z::AbstractKnotPoint, λ::AbstractVector) = true # zeros

function change_dimension(con::BoundConstraint, n::Int, m::Int, ix=1:n, iu=1:m)
	n0,m0 = con.n, con.m
	x_max = fill(Inf,n)
	x_min = fill(Inf,n)
	u_max = fill(Inf,m)
	u_min = fill(Inf,m)
	x_max[ix] = con.z_max[1:n0]
	x_min[ix] = con.z_min[1:n0]
	u_max[iu] = con.z_max[n0 .+ (1:m0)]
	u_min[iu] = con.z_min[n0 .+ (1:m0)]
	BoundConstraint(n, m, x_max=x_max, x_min=x_min, u_max=u_max, u_min=u_min)
end

############################################################################################
#  							VARIABLE BOUND CONSTRAINT 									   #
############################################################################################

# struct VariableBoundConstraint{T,P,NM,PNM} <: AbstractConstraint{Inequality,Stage,P}
# 	n::Int
# 	m::Int
# 	z_max::Vector{SVector{NM,T}}
# 	z_min::Vector{SVector{NM,T}}
# 	b::Vector{SVector{P,T}}
# 	B::SMatrix{P,NM,T,PNM}
# 	function VariableBoundConstraint(n::Int,m::Int,
# 			z_max::Vector{<:SVector{NM,T}}, z_min::Vector{<:SVector{NM,T}},
# 			b::Vector{<:SVector{P}}, B::SMatrix{P,NM,T,PNM}) where {T,P,PN,NM,PNM}
# 		new{T,P,NM,PNM}(n,m,z_max,z_min,b,B)
# 	end
# end
#
# state_dim(con::VariableBoundConstraint) = con.n
# control_dim(con::VariableBoundConstraint) = con.m
# is_bound(::VariableBoundConstraint) = true
#
# function evaluate!(vals::Vector{<:AbstractVector},
# 		con::VariableBoundConstraint, Z::Traj, inds=1:length(Z)-1)
# 	for (i,k) in enumerate(inds)
# 		vals[i] = con.B*Z[k].z + con.b[k]
# 	end
# end
#
# function jacobian(con::VariableBoundConstraint, z::KnotPoint)
# 	return con.B
# end
#
# function VariableBoundConstraint(n, m, N;
# 		x_max=[Inf*(@SVector ones(n)) for k = 1:N], x_min=[-Inf*(@SVector ones(n)) for k = 1:N],
# 		u_max=[Inf*(@SVector ones(m)) for k = 1:N], u_min=[-Inf*(@SVector ones(m)) for k = 1:N])
# 	@assert length(x_max) == N
# 	@assert length(u_max) == N
# 	@assert length(x_min) == N
# 	@assert length(u_min) == N
#
# 	# Check and convert bounds
# 	for k = 1:N
# 		x_max[k], x_min[k] = checkBounds(Val(n), x_max[k], x_min[k])
# 		u_max[k], u_min[k] = checkBounds(Val(m), u_max[k], u_min[k])
# 	end
#
# 	# Concatenate bounds
# 	z_max = [SVector{n+m}([x_max[k]; u_max[k]]) for k = 1:N]
# 	z_min = [SVector{n+m}([x_min[k]; u_min[k]]) for k = 1:N]
# 	b = [[-z_max[k]; z_min[k]] for k = 1:N]
#
# 	active = map(x->isfinite.(x), b)
# 	equal_active = all(1:N-2) do k
# 		active[k] == active[k+1]
# 	end
# 	if !equal_active
# 		throw(ArgumentError("All bounds must have the same active constraints"))
# 	end
# 	active = active[1]
# 	p = sum(active)
#
# 	inds = SVector{p}(findall(active))
#
# 	b = [bi[inds] for bi in b]
# 	B = SMatrix{2(n+m), n+m}([1.0I(n+m); -1.0I(n+m)])
#
# 	VariableBoundConstraint(n, m, z_max, z_min, b, B[inds,:])
# end

############################################################################################
#  								INDEXED CONSTRAINT 	 									   #
############################################################################################
"""
	IndexedConstraint{C,N,M}

Compute a constraint on an arbitrary portion of either the state or control,
or both. Useful for dynamics augmentation. e.g. you are controlling two models, and have
individual constraints on each. You can define constraints as if they applied to the individual
model, and then wrap it in an `IndexedConstraint` to apply it to the appropriate portion of
the concatenated state. Assumes the indexed state or control portion is contiguous.

# Type params:
* S - Inequality or Equality
* W - ConstraintType
* P - Constraint length
* N,M - original state and control dimensions
* NM - N+M
* Bx - location of the first element in the state index
* Bu - location of the first element in the control index
* C - type of original constraint

# Constructors:
```julia
IndexedConstraint(n, m, con)
IndexedConstraint(n, m, con, ix::UnitRange, iu::UnitRange)
```
where the arguments `n` and `m` are the state and control dimensions of the new dynamics.
`ix` and `iu` are the indices into the state and control vectors. If left out, they are
assumed to start at the beginning of the vector.

NOTE: Only part of this functionality has been tested. Use with caution!
"""
struct IndexedConstraint{C,N,M} <: StageConstraint
	n::Int  # new dimension
	m::Int  # new dimension
	n0::Int # old dimension
	m0::Int # old dimension
	con::C
	ix::SVector{N,Int}  # index of old x in new z
	iu::SVector{M,Int}  # index of old u in new z
	∇c::Matrix{Float64}
	A::SubArray{Float64,2,Matrix{Float64},Tuple{UnitRange{Int},UnitRange{Int}},false}
	B::SubArray{Float64,2,Matrix{Float64},Tuple{UnitRange{Int},UnitRange{Int}},false}
end

@inline state_dim(con::IndexedConstraint) = con.n
@inline control_dim(con::IndexedConstraint) = con.m
@inline Base.length(con::IndexedConstraint) = length(con.con)
@inline sense(con::IndexedConstraint) = sense(con.con)

function Base.copy(c::IndexedConstraint{C,n0,m0}) where {C,n0,m0}
	IndexedConstraint{C,n0,m0,}(c.n, c.m, c.n0, c.m0, copy(c.con), c.ix, c.iu,
		copy(∇c),copy(A),copy(B))
end

function IndexedConstraint(n,m,con::AbstractConstraint,
		ix::UnitRange{Int}, iu::UnitRange{Int})
	p = length(con)
	n0,m0 = length(ix), length(iu)
	iu = iu .+ n
	iz = SVector{n0+m0}([ix; iu])
	w = widths(con)[1]
	∇c = zeros(p,w)
	if con isa StageConstraint
		if con isa ControlConstraint
			A = view(∇c, 1:p, 1:0)
			B = view(∇c, 1:p, 1:m0)
		else
			A = view(∇c, 1:p, 1:n0)
			if con isa StateConstraint
				B = view(∇c, 1:p, n0:n0-1)
			else
				B = view(∇c, 1:p, n0 .+ (1:m0))
			end
		end
	else
		throw(ArgumentError("IndexedConstraint not support for CoupledConstraint yet"))
	end
	IndexedConstraint{typeof(con),n0,m0}(n,m,n0,m0,con,ix,iu,∇c,A,B)
end

function IndexedConstraint(n,m,con::AbstractConstraint)
	if con isa Union{StateConstraint, CoupledStateConstraint}
		m0 = m
	else
		m0 = control_dim(con)
	end
	if con isa Union{ControlConstraint, CoupledControlConstraint}
		n0 = n
	else
		n0 = state_dim(con)
	end
	ix = 1:n0
	iu = 1:m0
	IndexedConstraint(n, m, con, ix, iu)
end

function evaluate(con::IndexedConstraint, z::AbstractKnotPoint)
	x0 = z.z[con.ix]
	u0 = z.z[con.iu]
	z_ = StaticKnotPoint(x0, u0, z.dt, z.t)
	evaluate(con.con, z_)
end

@generated function jacobian!(∇c, con::IndexedConstraint{C}, z::AbstractKnotPoint) where C
	if C <: StateConstraint
		assignment = quote
			∇c_ = uview(∇c, :, con.ix)
			isconst = jacobian!(∇c_, con.con, z_)
		end
	elseif C <: ControlConstraint
		assignment = quote
			∇c_ = uview(∇c, :, con.iu)
			isconst = jacobian!(∇c_, con.con, z_)
		end
	else
		assignment = quote
			∇c_ = con.∇c
			isconst = jacobian!(∇c_, con.con, z_)
			uview(∇c, :, con.ix) .= con.A
			uview(∇c, :, con.iu) .= con.B
		end
	end
	quote
		x0 = z.z[con.ix]
		u0 = z.z[con.iu]
		z_ = StaticKnotPoint(x0, u0, z.dt, z.t)
		$assignment
		return isconst
	end
end

@inline is_bound(idx::IndexedConstraint) = is_bound(idx.con)
@inline upper_bound(idx::IndexedConstraint) = upper_bound(idx.con)
@inline lower_bound(idx::IndexedConstraint) = lower_bound(idx.con)

function change_dimension(con::AbstractConstraint, n::Int, m::Int, ix=1:n, iu=1:m)
	IndexedConstraint(n, m, con, ix, iu)
end
#
# # TODO: define higher-level evaluate! function instead
# @generated function evaluate(con::IndexedConstraint{<:Any,<:Stage,<:Any,N,M}, z::KnotPoint) where {N,M}
# 	ix = SVector{N}(1:N)
# 	iu = N .+ SVector{M}(1:M)
# 	return quote
# 		x0 = state(z)[con.ix]
# 		u0 = control(z)[con.iu]
# 		z_ = StaticKnotPoint([x0; u0], $ix, $iu, z.dt, z.t)
# 		evaluate(con.con, z_)
# 	end
# end
#
# # TODO: define higher-leel jacobian! function instead
# @generated function jacobian!(∇c, con::IndexedConstraint{<:Any,Stage,P,N0,M0},
# 		z::KnotPoint{<:Any,N}) where {P,N0,M0,N}
# 	iP = 1:P
# 	ix = SVector{N0}(1:N0)
# 	iu = SVector{M0}(N0 .+ (1:M0))
# 	if eltype(∇c) <: SizedMatrix
# 		assignment = quote
# 			uview(∇c.data,$iP,iA) .= con.A
# 			uview(∇c.data,$iP,iB) .= con.B
# 		end
# 	else
# 		assignment = quote
# 			uview(∇c,$iP,iA) .= con.A
# 			uview(∇c,$iP,iB) .= con.B
# 		end
# 	end
# 	quote
# 		x0 = state(z)[con.ix]
# 		u0 = control(z)[con.iu]
# 		z_ = StaticKnotPoint([x0;u0], $ix, $iu, z.dt, z.t)
# 		jacobian!(con.∇c, con.con, z_)
# 		iA = con.ix
# 		iB = N .+ con.iu
# 		$assignment
# 	end
# end
#
# @generated function jacobian!(∇c, con::IndexedConstraint{<:Any,State,P,N0,M0},
# 		z::KnotPoint{<:Any,N}) where {P,N0,M0,N}
# 	iP = 1:P
# 	ix = SVector{N0}(1:N0)
# 	iu = SVector{M0}(N0 .+ (1:M0))
# 	if eltype(∇c) <: SizedArray
# 		assignment = :(uview(∇c.data,$iP,iA) .= con.∇c)
# 	else
# 		assignment = :(uview(∇c,$iP,iA) .= con.∇c)
# 	end
# 	quote
# 		x0 = state(z)[con.ix]
# 		u0 = control(z)[con.iu]
# 		z_ = StaticKnotPoint([x0;u0], $ix, $iu, z.dt, z.t)
# 		jacobian!(con.∇c, con.con, z_)
# 		iA = con.ix
# 		$assignment
# 	end
# end
#
# @generated function jacobian!(∇c, con::IndexedConstraint{<:Any,Control,P,N0,M0},
# 		z::KnotPoint{<:Any,N}) where {P,N0,M0,N}
# 	iP = 1:P
# 	ix = SVector{N0}(1:N0)
# 	iu = SVector{M0}(N0 .+ (1:M0))
# 	if eltype(∇c) <: SizedArray
# 		assignment = :(uview(∇c.data,$iP,iB) .= con.∇c)
# 	else
# 		assignment = :(uview(∇c,$iP,iB) .= con.∇c)
# 	end
# 	quote
# 		x0 = state(z)[con.ix]
# 		u0 = control(z)[con.iu]
# 		z_ = StaticKnotPoint([x0;u0], $ix, $iu, z.dt, z.t)
# 		jacobian!(con.∇c, con.con, z_)
# 		iB = con.iu
# 		$assignment
# 	end
# end


text:
# ============================================================================
# Thermostat model
# See https://fenix.tecnico.ulisboa.pt/downloadFile/3779579688470/lygeros.pdf,
# Section 1.3.4.
# ============================================================================

using Reachability, HybridSystems, MathematicalSystems, LazySets, LinearAlgebra
using LazySets: HalfSpace  # resolve name-space conflicts with Polyhedra

function thermostat()
    c_a = 0.1

    # automaton structure
    automaton = LightAutomaton(2)

    # mode on
    A = hcat(-c_a)
    B = hcat(30.)
    U = Singleton([c_a])
    inv = HalfSpace([1.0], 22.0)  # x ≤ 22
    # @vars x v
    # @set x ≤ 22
    m_on = ConstrainedLinearControlContinuousSystem(A, B, inv, ConstantInput(U))

    # mode off
    A = hcat(-c_a)
    B = hcat(0.0)
    U = Singleton([0.0])
    inv = HalfSpace([-1.0], -18.0)  # x ≥ 18
    m_off = ConstrainedLinearControlContinuousSystem(A, B, inv, ConstantInput(U))

    # modes
    modes = [m_on, m_off]

    # transition from on to off
    add_transition!(automaton, 1, 2, 1)
    guard = HalfSpace([-1.0], -21.0)  # x ≥ 21
    t_on2off = ConstrainedIdentityMap(2, guard)

    # transition from off to on
    add_transition!(automaton, 2, 1, 2)
    guard = HalfSpace([1.0], 19.0)  # x ≤ 19
    t_off2on = ConstrainedIdentityMap(2, guard)

    # transition annotations
    resetmaps = [t_on2off, t_off2on]

    # switching
    switchings = [AutonomousSwitching()]

    ℋ = HybridSystem(automaton, modes, resetmaps, switchings)

    # initial condition in mode off
    X0 = Singleton([18.0])
    initial_condition = [(2, X0)]

    system = InitialValueProblem(ℋ, initial_condition)

    options = Options(:mode=>"reach", :T=>5.0,
                      :max_jumps=>1, :verbosity=>1)

    return (system, options)
end

ext: jl


text:
abstract type ConstrainedOptimizer{T}  <: AbstractOptimizer end
abstract type IPOptimizer{T} <: ConstrainedOptimizer{T} end # interior point methods


text:
# polynomial drift
include("common.jl")
import SDELab2
#
T=1 #  final time
x0=1. # initial data
alp=2. # parameter
sig=1 #
quad0=SDELab2.delta_measure(x0)
###########################################
function drift_ff(u,alp)
  x=u[1]; mf=u[2]; mf_sq=u[3]
  return alp*x+mf-x*mf_sq
end
##########################################
function diff_gg(u,sig)
  x=u[1]
  return sig*x
end
#########################################
function mf_r(x)
  return [x; x^2]
end
#########################################
fcn=SDELab2.set_fcn_mf(u->drift_ff(u,alp),
          u->diff_gg(u,sig),
          x->mf_r(x))
######################################
function odef(t, r)
    # Extract the coordinates from the r vector
    (x, y) = r
    dx_dt = (alp+1)*x- x*y
    dy_dt = (2*alp+1)*y + 2*x*x -2*y*y
    # Return the derivatives as a vector
    return [dx_dt; dy_dt]
end
####################################
using ODE
const odedt = 0.1
tvec = 0:odedt:T
tic()
tvec, pos = ode45(odef, [x0;x0], tvec,abstol=1e-9,reltol=1e-9)
exact=pos[end];
println("ode soln ",exact,"; cpu time ",toq()," secs.")
#############################################

dt=1e-2
solver="MFEM0"
params=SDELab2.set_opt(dt,solver)
params["cutoff"]=40 # range R for Alg 3.1 cutoff
params["const"]=1 # consts for no_Gauss_points
params["length_scale"]=10 # for no_Gauss_points
params["skiphalf"]=true
params["periodic"]=false
params["no_partition"]=1
dtpts=[0.05,0.025,0.01,0.005,0.0025]#,0.001]
errors=zeros(length(dtpts),2)
errors_no=deepcopy(errors);errors_nd=deepcopy(errors)
cpu_times=zeros(length(dtpts));
cpu_times_no=deepcopy(cpu_times);cpu_times_nd=deepcopy(cpu_times)
phi=x->[x;x.^2]
params["Solver"]="MFEM0"
dtpts_no=dtpts/4
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts_no[i]
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,false)
    errors_no[i,:],cpu_times_no[i]=helper(qoi, exact, cpu_time, "normal")
end
########################
params["Solver"]="MFEM0"
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts[i]
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,true)
    errors[i,:],cpu_times[i]=helper(qoi, exact, cpu_time, "normal")
end
################
params["Solver"]="MFEM2"
for i=1:(length(dtpts))
    params["MaxStepSize"]=dtpts[i]
    qoi,cpu_time,stats,q=get_soln(fcn,quad0,params,phi,false)
    errors_nd[i,:],cpu_times_nd[i]=helper(qoi, exact, cpu_time, "normal")
end
#########################################
str1="--" # extrap
str2="k-." # 2nd-order
figure(1);
# first functional
i=1
time_step_err(dtpts,errors[:,i],
  dtpts_no,errors_no[:,i],
  dtpts,errors_nd[:,i],[1e0,1e0],utf8("mom1.pdf"))
err_cpu(errors[:,i],cpu_times[:],
  errors_no[:,i],cpu_times_no[:],
  errors_nd[:,i],cpu_times_nd[:],[1e0,1e0], utf8("mom1_.pdf"))
# second functional
i=2
time_step_err(dtpts,errors[:,i],
      dtpts_no,errors_no[:,i],
      dtpts,errors_nd[:,i],[1e0,1],utf8("mom2.pdf"))
err_cpu(errors[:,i], cpu_times[:],
      errors_no[:,i],cpu_times_no[:],
      errors_nd[:,i],cpu_times_nd[:],[1e0,1], utf8("mom2_.pdf"))
#########################################
println("finished preparing plots")


text:
using Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MathOptInterface.Test
const MOIU = MathOptInterface.Utilities
const MOIB = MathOptInterface.Bridges

include("../utilities.jl")

mock = MOIU.MockOptimizer(MOIU.UniversalFallback(MOIU.Model{Float64}()))
config = MOIT.TestConfig()

bridged_mock = MOIB.Variable.Free{Float64}(mock)

@testset "solve_multirow_vectoraffine_nonpos" begin
    MOIU.set_mock_optimize!(
        mock,
        (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(
            mock,
            MOI.OPTIMAL,
            (MOI.FEASIBLE_POINT, [0.5, 0.0]),
        ),
        (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(
            mock,
            MOI.OPTIMAL,
            (MOI.FEASIBLE_POINT, [0.25, 0.0]),
        ),
    )
    MOIT.solve_multirow_vectoraffine_nonpos(bridged_mock, config)

    MOI.set(
        bridged_mock,
        MOI.ConstraintName(),
        MOI.get(
            mock,
            MOI.ListOfConstraintIndices{
                MOI.VectorAffineFunction{Float64},
                MOI.Nonpositives,
            }(),
        ),
        ["c"],
    )

    @testset "Mock model" begin
        var_names = ["xpos", "xneg"]
        MOI.set(
            mock,
            MOI.VariableName(),
            MOI.get(mock, MOI.ListOfVariableIndices()),
            var_names,
        )
        MOI.set(
            bridged_mock,
            MOI.ConstraintName(),
            MOI.get(
                mock,
                MOI.ListOfConstraintIndices{
                    MOI.VectorOfVariables,
                    MOI.Nonnegatives,
                }(),
            ),
            ["nonneg"],
        )
        s = """
        variables: xpos, xneg
        nonneg: [xpos, xneg] in MathOptInterface.Nonnegatives(2)
        c: [4.0xpos + -4.0xneg + -1.0, 3.0xpos + -3.0xneg + -1.0] in MathOptInterface.Nonpositives(2)
        maxobjective: xpos + -1.0xneg
        """
        model = MOIU.Model{Float64}()
        MOIU.loadfromstring!(model, s)
        MOIU.test_models_equal(mock, model, var_names, ["nonneg", "c"])
    end
    @testset "Bridged model" begin
        s = """
        variables: x
        c: [4.0x + -1.0, 3.0x + -1.0] in MathOptInterface.Nonpositives(2)
        maxobjective: 1.0x
        """
        model = MOIU.Model{Float64}()
        MOIU.loadfromstring!(model, s)
        MOIU.test_models_equal(bridged_mock, model, ["x"], ["c"])
    end
end

@testset "Linear6" begin
    MOIU.set_mock_optimize!(
        mock,
        (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [0, 0, 0, 0]),
        (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [100, 0, 0, 0]),
        (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(
            mock,
            [100, 0, 0, 100],
            (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) =>
                [1.0],
            (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}) =>
                [-1.0],
        ),
    )
    MOIT.linear6test(bridged_mock, config)

    loc = MOI.get(bridged_mock, MOI.ListOfConstraints())
    @test length(loc) == 2
    @test !((MOI.VectorOfVariables, MOI.Reals) in loc)
    @test (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) in loc
    @test (MOI.ScalarAffineFunction{Float64}, MOI.LessThan{Float64}) in loc
    @test MOI.get(mock, MOI.NumberOfVariables()) == 4
    @test MOI.get(bridged_mock, MOI.NumberOfVariables()) == 2
    vis = MOI.get(bridged_mock, MOI.ListOfVariableIndices())
    @test vis == MOI.VariableIndex.([-1, -2])

    cx = MOI.ConstraintIndex{MOI.VectorOfVariables,MOI.Reals}(vis[1].value)
    @test MOI.get(bridged_mock, MOI.ConstraintPrimal(), cx) == [100.0]
    @test MOI.get(bridged_mock, MOI.ConstraintDual(), cx) == [0.0]
    cy = MOI.ConstraintIndex{MOI.VectorOfVariables,MOI.Reals}(vis[2].value)
    @test MOI.get(bridged_mock, MOI.ConstraintPrimal(), cy) == [-100.0]
    @test MOI.get(bridged_mock, MOI.ConstraintDual(), cy) == [0.0]

    test_delete_bridged_variable(
        bridged_mock,
        vis[1],
        MOI.Reals,
        2,
        (
            (MOI.VectorOfVariables, MOI.Nonnegatives, 0),
            (MOI.VectorOfVariables, MOI.Nonpositives, 0),
        ),
    )
    test_delete_bridged_variable(
        bridged_mock,
        vis[2],
        MOI.Reals,
        1,
        (
            (MOI.VectorOfVariables, MOI.Nonnegatives, 0),
            (MOI.VectorOfVariables, MOI.Nonpositives, 0),
        ),
    )
end

@testset "Linear7" begin
    function set_mock_optimize_linear7Test!(mock)
        return MOIU.set_mock_optimize!(
            mock,
            (mock::MOIU.MockOptimizer) ->
                MOIU.mock_optimize!(mock, [0, 0, 0, 0]),
            (mock::MOIU.MockOptimizer) ->
                MOIU.mock_optimize!(mock, [100, 0, 0, 0]),
            (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(
                mock,
                [100, 0, 0, 100],
                (MOI.VectorAffineFunction{Float64}, MOI.Nonnegatives) =>
                    [[1.0]],
                (MOI.VectorAffineFunction{Float64}, MOI.Nonpositives) =>
                    [[-1.0]],
            ),
        )
    end
    set_mock_optimize_linear7Test!(mock)
    MOIT.linear7test(bridged_mock, config)

    x, y = MOI.get(bridged_mock, MOI.ListOfVariableIndices())

    cx = MOI.ConstraintIndex{MOI.VectorOfVariables,MOI.Reals}(x.value)
    @test MOI.get(bridged_mock, MOI.ConstraintPrimal(), cx) == [100.0]
    @test MOI.get(bridged_mock, MOI.ConstraintDual(), cx) == [0.0]
    cy = MOI.ConstraintIndex{MOI.VectorOfVariables,MOI.Reals}(y.value)
    @test MOI.get(bridged_mock, MOI.ConstraintPrimal(), cy) == [-100.0]
    @test MOI.get(bridged_mock, MOI.ConstraintDual(), cy) == [0.0]

    @test MOI.supports(
        bridged_mock,
        MOI.VariablePrimalStart(),
        MOI.VariableIndex,
    )
    @test MOI.supports(
        bridged_mock,
        MOI.VariablePrimalStart(),
        typeof(MOIB.bridge(bridged_mock, x)),
    )
    MOI.set(bridged_mock, MOI.VariablePrimalStart(), [x, y], [1.0, -1.0])
    xa, xb, ya, yb = MOI.get(mock, MOI.ListOfVariableIndices())
    @test MOI.get(mock, MOI.VariablePrimalStart(), [xa, xb, ya, yb]) ==
          [1.0, 0.0, 0.0, 1.0]
    @test MOI.get(bridged_mock, MOI.VariablePrimalStart(), x) == 1
    @test MOI.get(bridged_mock, MOI.VariablePrimalStart(), y) == -1
end

@testset "Linear11" begin
    MOIU.set_mock_optimize!(
        mock,
        (mock::MOIU.MockOptimizer) ->
            MOIU.mock_optimize!(mock, [1.0, 1.0, 0.0, 0.0]),
        (mock::MOIU.MockOptimizer) ->
            MOIU.mock_optimize!(mock, [0.5, 0.5, 0.0, 0.0]),
    )
    MOIT.linear11test(bridged_mock, config)

    vis = MOI.get(bridged_mock, MOI.ListOfVariableIndices())
    @test vis == MOI.VariableIndex.([-1, -2])

    MOI.set(
        bridged_mock,
        MOI.ConstraintName(),
        MOI.get(
            mock,
            MOI.ListOfConstraintIndices{
                MOI.ScalarAffineFunction{Float64},
                MOI.GreaterThan{Float64},
            }(),
        ),
        ["c1"],
    )
    MOI.set(
        bridged_mock,
        MOI.ConstraintName(),
        MOI.get(
            mock,
            MOI.ListOfConstraintIndices{
                MOI.ScalarAffineFunction{Float64},
                MOI.LessThan{Float64},
            }(),
        ),
        ["c2"],
    )

    @testset "Mock model" begin
        var_names = ["v1pos", "v2pos", "v1neg", "v2neg"]
        MOI.set(
            mock,
            MOI.VariableName(),
            MOI.get(mock, MOI.ListOfVariableIndices()),
            var_names,
        )
        MOI.set(
            bridged_mock,
            MOI.ConstraintName(),
            MOI.get(
                mock,
                MOI.ListOfConstraintIndices{
                    MOI.VectorOfVariables,
                    MOI.Nonnegatives,
                }(),
            ),
            ["nonneg"],
        )
        s = """
        variables: v1pos, v2pos, v1neg, v2neg
        nonneg: [v1pos, v2pos, v1neg, v2neg] in MathOptInterface.Nonnegatives(4)
        c1: v1pos + -1.0v1neg + v2pos + -1.0v2neg >= 1.0
        c2: v1pos + -1.0v1neg + v2pos + -1.0v2neg <= 2.0
        minobjective: v1pos + -1.0v1neg + v2pos + -1.0v2neg
        """
        model = MOIU.Model{Float64}()
        MOIU.loadfromstring!(model, s)
        MOIU.test_models_equal(mock, model, var_names, ["nonneg", "c1", "c2"])
    end
    @testset "Bridged model" begin
        var_names = ["v1", "v2"]
        MOI.set(bridged_mock, MOI.VariableName(), vis, var_names)
        s = """
        variables: v1, v2
        c1: v1 + v2 >= 1.0
        c2: v1 + v2 <= 2.0
        minobjective: v1 + v2
        """
        model = MOIU.Model{Float64}()
        MOIU.loadfromstring!(model, s)
        MOIU.test_models_equal(bridged_mock, model, var_names, ["c1", "c2"])
    end

    test_delete_bridged_variable(
        bridged_mock,
        vis[1],
        MOI.Reals,
        2,
        (
            (MOI.VectorOfVariables, MOI.Nonnegatives, 0),
            (MOI.VectorOfVariables, MOI.Nonpositives, 0),
        ),
        used_bridges = 0,
        used_constraints = 0,
    )
    test_delete_bridged_variable(
        bridged_mock,
        vis[2],
        MOI.Reals,
        1,
        (
            (MOI.VectorOfVariables, MOI.Nonnegatives, 0),
            (MOI.VectorOfVariables, MOI.Nonpositives, 0),
        ),
    )
end


text:
function optimize(model, loss, cb, opt, AD, train_dl; normalize=false)
  pp = initial_params(model)
  lb, ub = get_bounds(model)

  @show length(train_dl)
  @show size(first(train_dl)[1])
  @show size(first(train_dl)[1][1])
  @show length(pp)
  @show length(lb)
  @show length(ub)

  optfun = GalacticOptim.OptimizationFunction((θ,p,x,y) -> loss(θ,model,x,y), AD)
  optfunc = GalacticOptim.instantiate_function(optfun, pp, AD, nothing)
  optprob = GalacticOptim.OptimizationProblem(optfunc, pp, lb=lb, ub=ub,
                                #grad = true, hess = true, sparse = true,
                                #parallel=ModelingToolkit.MultithreadedForm()
                                )

  if normalize
    solve_normalized(optprob, opt, train_dl, scale=true, cb = cb)
  else
    GalacticOptim.solve(optprob, opt, train_dl, cb = cb)
  end
end

function optimize(model::Flux.Chain, loss, cb, opt, AD, train_dl)
  pp, re = LTC.destructure(model)
  lb, ub = get_bounds(model)

  @show length(train_dl)
  @show size(first(train_dl)[1])
  @show size(first(train_dl)[1][1])
  @show length(pp)
  @show length(lb)
  @show length(ub)
  @show sum(length.(pp))

  optfun = GalacticOptim.OptimizationFunction((θ,p,x,y) -> loss(θ,re,x,y), AD)
  optfunc = GalacticOptim.instantiate_function(optfun, pp, AD, nothing)
  optprob = GalacticOptim.OptimizationProblem(optfunc, pp, lb=lb, ub=ub,
                                #grad = true, hess = true, sparse = true,
                                #parallel=ModelingToolkit.MultithreadedForm()
                                )
  GalacticOptim.solve(optprob, opt, train_dl, cb = cb)
end

# https://github.com/SciML/GalacticOptim.jl/issues/146
# Skip the DiffEqBase handling
struct InverseScale{T}
    scale::T
end

(inv_scale::InverseScale)(x, compute_inverse::Bool = false) =
    compute_inverse ? x .* inv_scale.scale : x ./ inv_scale.scale

function solve_normalized(prob::OptimizationProblem, opt, args...;
               scale::Bool = false, scaling_function = nothing,
               cb = (args...) -> (false), kwargs...)
    !scale && return GalacticOptim.solve(prob, opt, args...; kwargs...)

    θ_start = copy(prob.u0)
    for (i,p) in enumerate(θ_start)
      if iszero(p)
        @show i
        @show p
        θ_start[i] += 0.00001f0
      end
    end

    if isnothing(scaling_function) && any(iszero.(θ_start))
        error("Default Inverse Scaling is not compatible with `0` as initial guess")
    end

    scaling_function = if isnothing(scaling_function)
            InverseScale(θ_start)
        else
            # Check if arguments are compatible
            # First arg is the parameter
            # 2nd one denotes inverse computation or not
            scaling_function(θ_start, false)
            scaling_function
        end

    normalized_f(α, args...) = prob.f.f(scaling_function(α, true), args...)
    normalized_cb(α, args...) = cb(scaling_function(α, true), args...)

    lb = isnothing(prob.lb) ? nothing : scaling_function(prob.lb, false)
    ub = isnothing(prob.ub) ? nothing : scaling_function(prob.ub, false)

    u0s = scaling_function(prob.u0, false)
    optfun = OptimizationFunction(normalized_f, prob.f.adtype,
                             grad = prob.f.grad, hess = prob.f.hess,
                             hv = prob.f.hv, cons = prob.f.cons,
                             cons_j = prob.f.cons_j, cons_h = prob.f.cons_h)
    optfunc = GalacticOptim.instantiate_function(optfun, u0s, prob.f.adtype, nothing)

    _prob = remake(prob, u0 = u0s, lb = lb,
                   ub = ub,
                   f = optfunc)

    optsol = GalacticOptim.solve(_prob, opt, args...; cb = normalized_cb, kwargs...)
    optsol.u .= scaling_function(optsol.u, true)
    optsol
end


text:
using CUDA.NVML

macro maybe_unsupported(ex)
    quote
        try
            $(esc(ex))
        catch err
            (isa(err, NVML.NVMLError) && err.code == NVML.ERROR_NOT_SUPPORTED) || rethrow()
        end
    end
end

@testset "system" begin
    @test NVML.version() isa VersionNumber
    @test NVML.driver_version() isa VersionNumber
    @test NVML.cuda_driver_version() == CUDA.version()
end

@testset "devices" begin
    let dev = NVML.Device(0)
        @test dev == first(NVML.devices())
        @test NVML.index(dev) == 0

        str = sprint(io->show(io, "text/plain", dev))
        @test occursin("NVML.Device(0)", str)
    end

    cuda_dev = CuDevice(0)
    mig = uuid(cuda_dev) != parent_uuid(cuda_dev)

    # tests for the parent device
    let dev = NVML.Device(parent_uuid(cuda_dev))
        @test NVML.uuid(dev) == parent_uuid(cuda_dev)
        NVML.brand(dev)
        @test occursin(NVML.name(dev), name(cuda_dev))
        @maybe_unsupported NVML.serial(dev)

        @maybe_unsupported NVML.power_usage(dev)
        @maybe_unsupported NVML.energy_consumption(dev)

        @maybe_unsupported NVML.utilization_rates(dev)

        NVML.compute_mode(dev)
        @test NVML.compute_capability(dev) == capability(cuda_dev)
    end

    # tests for the compute instance
    let dev = NVML.Device(uuid(cuda_dev); mig)
        @test NVML.uuid(dev) == uuid(cuda_dev)
        @test NVML.name(dev) == name(cuda_dev)

        NVML.memory_info(dev)

        context()
        # FIXME: https://github.com/NVIDIA/gpu-monitoring-tools/issues/63
        #@test getpid() in keys(NVML.compute_processes(dev))
        NVML.compute_processes(dev)
    end
end


text:
include("../mjWrap.jl")
using ..mjWrap

T = 1#0
K = 1 #50

modelname = "models/humanoid.xml"

myskip = 1
mjsys = mjw.load_model(modelname, myskip, "normal")
mjsys2 = mjw.load_model(modelname, myskip, "normal")
samples = mjw.allocateTrajSamples(mjsys, T, K)
samples2 = mjw.allocateTrajSamples(mjsys, T, K)

samples.ctrl[:] .= 0.5
samples2.ctrl[:] .= 0.5

#@time mjw.rollout(mjsys, samples)
@time mjw.rollout2(mjsys2, samples2)
#Profile.clear_malloc_data()
println()
#J@time mjw.rollout(mjsys, samples)
@time mjw.rollout2(mjsys2, samples2)
println()
#@time mjw.rollout(mjsys, samples)
@time mjw.rollout2(mjsys2, samples2)

# actual test??


text:
export
    set_zero!,
    set_one!,
    used_variables,
    size

function set_zero!(hobo::HOBO, v::Vector{Int})
    for k = keys(values(hobo))
        if v ∈ k
            delete!(values(hobo), k)
        end
    end
    hobo
end

function set_one!(hobo::HOBO, v::Vector{Int})
    for k = keys(values(hobo))
        if v ∈ k
            val = hobo[k...]
            new_k = setdiff(Set{Vector{Int}}(k), Set([v]))
            hobo[new_k...] += val
            delete!(values(hobo), k)
        end
    end
    hobo
end

function used_variables(hobo::BinaryOptimizationModel)
    sort(unique(vcat(keys(values(hobo))...)))
end


text:
"""
sinkhorn_dvg(α, β, K, a₀, b₀, d₁₀, d₂₀, SP, caches)
"""
function sinkhorn_dvg(  α::Vector{T},
                        β::Vector{V},
                        a₀, b₀,
                        d₁₀, d₂₀,
                        K,
                        SP, caches   ) where {T,V}

    VC = caches.VC

    VC[:a,V] .= a₀
    VC[:b,V] .= b₀
    if SP.debias
        VC[:d₁,V] .= d₁₀
        VC[:d₂,V] .= d₂₀
    end

    for l in 1:SP.L
        mul!(VC[:t1,V], K, VC[:b,V])
        VC[:a,V] .= α ./ VC[:t1,V]
        mul!(VC[:t1,V], K', VC[:a,V])
        VC[:b,V] .= β ./ VC[:t1,V]

        if SP.debias # && (l%2==0 || l==SP.L)
            mul!(VC[:t1,V], K, VC[:d₁,V])
            VC[:d₁,V] .= sqrt.( α .* VC[:d₁,V] ./ VC[:t1,V] )
            mul!(VC[:t1,V], K, VC[:d₂,V])
            VC[:d₂,V] .= sqrt.( β .* VC[:d₂,V] ./ VC[:t1,V] )
        end
    end

    S_ε = 0

    if SP.debias
        for i in eachindex(α)
            S_ε += α[i] * (_safe_log( VC[:a,V][i] ) - _safe_log( VC[:d₁,V][i] )) +
                    β[i] * (_safe_log( VC[:b,V][i] ) - _safe_log( VC[:d₂,V][i] ))
        end
    else
        for i in eachindex(α)
            S_ε += α[i] * _safe_log( VC[:a,V][i] ) +
                    β[i] * _safe_log( VC[:b,V][i] )
        end
    end

    return SP.ε * S_ε
end

function sinkhorn_barycenter(  λ::AbstractVector{T}, α::AbstractVector{AT},
                                b_₀, d₀,
                                K, SP, caches,
                                ) where {T, V, AT <: AbstractArray{V}}

    VC = caches.VC
    VVC = caches.VVC

    for s in 1:VVC.S
        VVC[:b,T][s] .= b_₀[s]
    end
    if SP.debias
        VC[:d,T] .= d₀
    end

    for l in 1:SP.L
        for s in 1:VVC.S
            mul!(VC[:t1,V], K, VVC[:b,V][s])
            VC[:t1,V] .= α[s] ./ VC[:t1,V]
            mul!(VVC[:φ,T][s], K', VC[:t1,V])
        end

        SP.debias ? VC[:μ,T] .= VC[:d,T] : VC[:μ,T] .= 1

        for s in 1:VC.S
            VC[:μ,T] .*= VVC[:φ,T][s] .^ λ[s]
        end

        for s in 1:VC.S
            VVC[:b,T][s] .= VC[:μ,T] ./ VVC[:φ,T][s]
        end

        if SP.debias
            mul!(VC[:t1,V], K, VC[:d,T])
            VC[:d,T] .= sqrt.( VC[:d,T] .* VC[:μ,T] ./ VC[:t1,T] )
        end
    end

    if SP.update_potentials
        for s in 1:VVC.S
            b_₀[s] .= ForwardDiff.value.(VVC[:b,T][s])
        end
        if SP.debias
            d₀ .= ForwardDiff.value.(VC[:d,T])
        end
    end

    return copy(VC[:μ,T])
end

max_stars_repo_name: JuliaRCM/WassersteinDictionaries

max_issues_repo_name: JuliaRCM/WassersteinDictionaries

max_forks_repo_name: JuliaRCM/WassersteinDictionaries


text:
using Test
using MPI
using CLIMA.Mesh.Topologies

function main()
  MPI.Init()
  comm = MPI.COMM_WORLD
  crank = MPI.Comm_rank(comm)
  csize = MPI.Comm_size(comm)

  @assert csize == 5

  topology = BrickTopology(comm, (0:10,); boundary=((1,2),),
                           periodicity=(true,))

  elems       = topology.elems
  realelems   = topology.realelems
  ghostelems  = topology.ghostelems
  sendelems   = topology.sendelems
  elemtocoord = topology.elemtocoord
  elemtoelem  = topology.elemtoelem
  elemtoface  = topology.elemtoface
  elemtoordr  = topology.elemtoordr
  elemtobndy  = topology.elemtobndy
  nabrtorank  = topology.nabrtorank
  nabrtorecv  = topology.nabrtorecv
  nabrtosend  = topology.nabrtosend

  globalelemtoelem = [10 1 2 3 4 5 6 7  8 9
                       2 3 4 5 6 7 8 9 10 1]

  globalelemtoface = [2 2 2 2 2 2 2 2 2 2
                      1 1 1 1 1 1 1 1 1 1]

  globalelemtoordr = ones(Int, size(globalelemtoface))
  globalelemtobndy = zeros(Int, size(globalelemtoface))

  globalelemtocoord = Array{Int}(undef, 1, 2, 10)
  globalelemtocoord[:,:, 1] = [0  1]
  globalelemtocoord[:,:, 2] = [1  2]
  globalelemtocoord[:,:, 3] = [2  3]
  globalelemtocoord[:,:, 4] = [3  4]
  globalelemtocoord[:,:, 5] = [4  5]
  globalelemtocoord[:,:, 6] = [5  6]
  globalelemtocoord[:,:, 7] = [6  7]
  globalelemtocoord[:,:, 8] = [7  8]
  globalelemtocoord[:,:, 9] = [8  9]
  globalelemtocoord[:,:,10] = [9 10]

  @assert csize == 5
  nrealelem = 2
  if crank == 0
    globalelems = [1,2,3,10]
    elemtoelem_expect = [4 1 3 4; 2 3 3 4]
    nabrtorank_expect = [1, 4]
  elseif crank == 1
    globalelems = [3,4,2,5]
    elemtoelem_expect = [3 1 3 4; 2 4 3 4]
    nabrtorank_expect = [0, 2]
  elseif crank == 2
    globalelems = [5,6,4,7]
    elemtoelem_expect = [3 1 3 4; 2 4 3 4]
    nabrtorank_expect = [1, 3]
  elseif crank == 3
    globalelems = [7,8,6,9]
    elemtoelem_expect = [3 1 3 4; 2 4 3 4]
    nabrtorank_expect = [2, 4]
  elseif crank == 4
    globalelems = [9,10,1,8]
    elemtoelem_expect = [4 1 3 4; 2 3 3 4]
    nabrtorank_expect = [0, 3]
  end
  nabrtorecv_expect = UnitRange{Int}[1:1, 2:2]
  nabrtosend_expect = UnitRange{Int}[1:1, 2:2]

  @test elems == 1:length(globalelems)
  @test realelems == 1:nrealelem
  @test ghostelems == nrealelem+1:length(globalelems)
  @test elemtocoord == globalelemtocoord[:,:,globalelems]
  @test elemtoface[:,realelems] == globalelemtoface[:,globalelems[realelems]]
  @test elemtoelem == elemtoelem_expect
  @test elemtobndy == globalelemtobndy[:,globalelems]
  @test elemtoordr == globalelemtoordr[:,globalelems]
  @test nabrtorank == nabrtorank_expect
  @test nabrtorecv == nabrtorecv_expect
  @test nabrtosend == nabrtosend_expect
end

main()

max_stars_repo_name: akshaysridhar/CLIMA

max_issues_repo_name: akshaysridhar/CLIMA

max_forks_repo_name: akshaysridhar/CLIMA


text:
struct SVGCmd{T}
    args::Vector{<:Float64}
    SVGCmd{T}(args::Vararg{<:Float64}) where T = new{T}(Float64[args...])
end

mutable struct SVGPath
    path::String
    commands::Vector{SVGCmd}
end

function SVGPath(path::String)
    svg_path = SVGPath(path, SVGCmd[])
    δ = ""
    Δ = Float64[]
    m = missing
    chars = Char[]

    foreach(path) do char
        if isletter(char)
            push!(chars, ' ')
            push!(chars, char)
            push!(chars, ' ')
        elseif char == '-'
            push!(chars, ' ')
            push!(chars, char)
        else
            push!(chars, char)
        end
    end

    push!(chars, ' ')
    push!(chars, 'E')

    for char in chars
        if isletter(char)
            # println("char is a letter '$char'")
            if ismissing(m)
                # println("Previous letter is missing.")
                m = Symbol(uppercase(char))
                # println("Setting char to ':$m'")
            else
                # println("Previous letter is not missing")
                # @show m Δ
                push!(svg_path.commands, SVGCmd{m}(Δ...))
                empty!(Δ)
                m = Symbol(uppercase(char))
            end
        elseif isdigit(char) || char in ".-"
            # println("Char is part of a number")
            δ *= char
            # @show δ
        elseif (isspace(char) || char == ',') && δ != ""
            # println("char is a space")
            push!(Δ, parse(Float64, δ))
            # @show δ Δ
            δ = ""
        else
            # println("Skipping $char")
            continue
        end
        #println("-" ^ 32)
    end

    return svg_path
end

# Quadratic bezier curve
qB(t, P...) = @. P[2] + (1 - t)^2 * (P[1] - P[2]) + t^2 * (P[3] - P[2])
# Cubic bezier curve
cB(t, P...) = @. (1 - t) * qB(t, P[1], P[2], P[3]) + t * qB(t, P[2], P[3], P[4])

bezier_points(curve::Function, P, Δt = 1/20) = [curve(t, P...) for t in 0:Δt:1]

Base.parse(cmd::SVGCmd{:M}, points) = push!(points, cmd.args)

Base.parse(cmd::SVGCmd{:L}, points) = push!(points, cmd.args)

Base.parse(cmd::SVGCmd{:V}, points) = begin
    x = cmd.args[1]
    y = points[end][2]
    push!(points, Float64[x, y])
end

Base.parse(cmd::SVGCmd{:H}, points) = begin
    x = points[end][1]
    y = cmd.args[1]
    push!(points, Float64[x, y])
end

Base.parse(::SVGCmd{:Z}, points) = push!(points, points[1])

Base.parse(cmd::SVGCmd{:C}, points) = begin
    P = [last(points), [cmd.args[i:i+1] for i in 1:2:6]...]
    for point in bezier_points(cB, P)
        push!(points, point)
    end
end

Base.parse(cmd::SVGCmd{:S}, points) = begin
    args = cmd.args
    P = [points[end-1:end]..., args[1:2], args[3:4]]
    for point in bezier_points(cB, P)
        push!(points, point)
    end
end

Base.parse(cmd::SVGCmd{:Q}, points) = begin
    args = cmd.args
    P = [points[end], args[1:2], args[3:4]]
    for point in bezier_points(cB, P)
        push!(points, point)
    end
end

Base.parse(cmd::SVGCmd{:T}, points) = begin
    args = cmd.args
    P = [points[end-1:end]..., args[1:2], args[3:4]]
    for point in bezier_points(qB, P)
        push!(points, point)
    end
end

Base.parse(::SVGCmd{T}, points) where T = throw(error("SVG Path command of type '$T' not yet supported."))

function getpoints(path::SVGPath)::Vector{Point}
    points = Vector{<:Float64}[]
    for cmd in path.commands
        parse(cmd, points)
    end
    return point.(points)
end

getpoints(path::String) = getpoints(SVGPath(path))

# TODO: implement svg 'arc' command, but i'm kinda lazy
# TODO: implement a way of getting the path from a svg file, without the need of get it manually


text:
using Base: Int64, Float64, NamedTuple
using Printf
using Glob

# LaMEM I/O
#
# These are routines that help to create a LaMEM marker files from a ParaviewData structure, which can be used to perform geodynamic simulations
# We also include routines with which we can read LaMEM *.pvtr files into julia

export LaMEM_grid, ReadLaMEM_InputFile
export Save_LaMEMMarkersParallel, Save_LaMEMTopography
export GetProcessorPartitioning, ReadData_VTR, ReadData_PVTR, CreatePartitioningFile

"""
Structure that holds information about the LaMEM grid (usually read from an input file).
"""
struct LaMEM_grid
    nmark_x :: Int64
    nmark_y :: Int64
    nmark_z :: Int64
    nump_x  :: Int64
    nump_y  :: Int64
    nump_z  :: Int64

    nel_x   :: Int64
    nel_y   :: Int64
    nel_z   :: Int64

    W       ::  Float64
    L       ::  Float64
    H       ::  Float64

    coord_x
    coord_y
    coord_z

    x1D_c
    y1D_c
    z1D_c

    X
    Y
    Z

end

"""
    ParaviewData(Grid::LaMEM_grid, fields::NamedTuple)

Creates a `ParaviewData` struct from a LaMEM grid and from fields stored on that grid. Note that one needs to have a field `Phases` and optionally a field `Temp` to create LaMEM marker files.
"""
ParaviewData(Grid::LaMEM_grid, fields::NamedTuple) = ParaviewData(Grid.X, Grid.Y, Grid.Z, fields)

"""
    CartData(Grid::LaMEM_grid, fields::NamedTuple)

Creates a `CartData` struct from a LaMEM grid and from fields stored on that grid. Note that one needs to have a field `Phases` and optionally a field `Temp` to create LaMEM marker files.
"""
CartData(Grid::LaMEM_grid, fields::NamedTuple) = CartData(Grid.X, Grid.Y, Grid.Z, fields)

"""
    Below = BelowSurface(Data_LaMEM::LaMEM_grid, DataSurface_Cart::CartData)

Determines if points within the 3D `LaMEM_grid` structure are below the Cartesian surface DataSurface_Cart
"""
function BelowSurface(Grid::LaMEM_grid, DataSurface_Cart::CartData)
    return AboveSurface(CartData(Grid,(Z=Grid.Z,)), DataSurface_Cart; above=false)
end

"""
    Above = AboveSurface(Data_LaMEM::LaMEM_grid, DataSurface_Cart::CartData)

Determines if points within the 3D `LaMEM_grid` structure are above the Cartesian surface DataSurface_Cart
"""
function AboveSurface(Grid::LaMEM_grid, DataSurface_Cart::CartData)
    return AboveSurface(CartData(Grid,(Z=Grid.Z,)), DataSurface_Cart; above=true)
end

"""
    value = ParseValue_LaMEM_InputFile(file,keyword,type)

Extracts a certain `keyword` from a LaMEM input `file` and convert it to a certain type

# Example
```julia
julia> nmark_z = ParseValue_LaMEM_InputFile("SaltModels.dat","nmark_z",Int64)
```
"""
function ParseValue_LaMEM_InputFile(file,keyword,type)
    value = nothing
    for line in eachline(file)
        line_strip = lstrip(line)       # strip leading tabs/spaces

        # Strip comments
        ind        = findfirst("#", line)
        if isnothing(ind)
            # no comments
        else
            line_strip = line_strip[1:ind[1]-2];
        end
        line_strip = rstrip(line_strip)       # strip last tabs/spaces

        if startswith(line_strip, keyword)
            ind = findfirst("=", line_strip)
            if type==String
                value = split(line_strip)[3:end]
            else
                value = parse.(type,split(line_strip)[3:end])

                if length(value)==1
                    value=value[1];
                end
            end
        end

    end

    return value
end

"""
    Grid::LaMEM_grid = ReadLaMEM_InputFile(file)

Parses a LaMEM input file and stores grid information in the `Grid` structure.

# Example
```julia
julia> Grid = ReadLaMEM_InputFile("SaltModels.dat")
LaMEM Grid:
nel         : (32, 32, 32)
marker/cell : (3, 3, 3)
markers     : (96, 96, 96)
x           ϵ [-3.0 : 3.0]
y           ϵ [-2.0 : 2.0]
z           ϵ [-2.0 : 0.0]
```
"""
function ReadLaMEM_InputFile(file)

    nmark_x = ParseValue_LaMEM_InputFile(file,"nmark_x",Int64)
    nmark_y = ParseValue_LaMEM_InputFile(file,"nmark_y",Int64)
    nmark_z = ParseValue_LaMEM_InputFile(file,"nmark_z",Int64)

    nel_x   = ParseValue_LaMEM_InputFile(file,"nel_x",Int64)
    nel_y   = ParseValue_LaMEM_InputFile(file,"nel_y",Int64)
    nel_z   = ParseValue_LaMEM_InputFile(file,"nel_z",Int64)

    coord_x = ParseValue_LaMEM_InputFile(file,"coord_x",Float64)
    coord_y = ParseValue_LaMEM_InputFile(file,"coord_y",Float64)
    coord_z = ParseValue_LaMEM_InputFile(file,"coord_z",Float64)

    if (length(coord_x)>2) || (length(coord_y)>2) || (length(coord_z)>2)
        error("Routine currently not working for variable grid spacing")
    end

    W       = coord_x[end]-coord_x[1];
    L       = coord_y[end]-coord_y[1];
    H       = coord_z[end]-coord_z[1];

    nump_x  = nel_x*nmark_x;
    nump_y  = nel_y*nmark_y;
    nump_z  = nel_z*nmark_z;

    dx      =   W/nump_x;
    dy      =   L/nump_y;
    dz      =   H/nump_z;

    # these lines should be replaced with a separate routine for variable spacing
    x       =   coord_x[1]+dx/2: dx : coord_x[end]-dx/2;
    y       =   coord_y[1]+dy/2: dy : coord_y[end]-dy/2;
    z       =   coord_z[1]+dz/2: dz : coord_z[end]-dz/2;

    X,Y,Z   =   XYZGrid(x,y,z); # create 3D grid using regular spacing

    Grid    =  LaMEM_grid(  nmark_x,    nmark_y,    nmark_z,
                            nump_x,     nump_y,     nump_z,
                            nel_x,      nel_y,      nel_z,
                            W,          L,          H,
                            coord_x,    coord_y,    coord_z,
                            x,          y,          z,
                            X,          Y,          Z);

    return Grid
end

# Print an overview of the LaMEM Grid struct:
function Base.show(io::IO, d::LaMEM_grid)
    println(io,"LaMEM Grid: ")
    println(io,"  nel         : ($(d.nel_x), $(d.nel_y), $(d.nel_z))")
    println(io,"  marker/cell : ($(d.nmark_x), $(d.nmark_y), $(d.nmark_z))")
    println(io,"  markers     : ($(d.nump_x), $(d.nump_x), $(d.nump_x))")
    println(io,"  x           ϵ [$(d.coord_x[1]) : $(d.coord_x[2])]")
    println(io,"  y           ϵ [$(d.coord_y[1]) : $(d.coord_y[2])]")
    println(io,"  z           ϵ [$(d.coord_z[1]) : $(d.coord_z[2])]")
end

"""
    Save_LaMEMMarkersParallel(Grid::CartData; PartitioningFile=empty, directory="./markers", verbose=true)

Saves a LaMEM marker file from the `CartData` structure `Grid`. It must have a field called `Phases`, holding phase information (as integers) and optionally a field `Temp` with temperature info.
It is possible to provide a LaMEM partitioning file `PartitioningFile`. If not, output is assumed to be for one processor.

The size of `Grid` should be consistent with what is provided in the LaMEM input file. In practice, the size of the mesh can be retrieved from a LaMEM input file using `ReadLaMEM_InputFile`.

# Example

```
julia> Grid    = ReadLaMEM_InputFile("LaMEM_input_file.dat")
julia> Phases  = zeros(Int32,size(Grid.X));
julia> Temp    = ones(Float64,size(Grid.X));
julia> Model3D = CartData(Grid, (Phases=Phases,Temp=Temp))
julia> Save_LaMEMMarkersParallel(Model3D)
Writing LaMEM marker file -> ./markers/mdb.00000000.dat
```
If you want to create a LaMEM input file for multiple processors:
```
julia> Save_LaMEMMarkersParallel(Model3D, PartitioningFile="ProcessorPartitioning_4cpu_1.2.2.bin")
Writing LaMEM marker file -> ./markers/mdb.00000000.dat
Writing LaMEM marker file -> ./markers/mdb.00000001.dat
Writing LaMEM marker file -> ./markers/mdb.00000002.dat
Writing LaMEM marker file -> ./markers/mdb.00000003.dat
```

"""
function Save_LaMEMMarkersParallel(Grid::CartData; PartitioningFile=empty, directory="./markers", verbose=true)

    x = ustrip.(Grid.x.val[:,1,1]);
    y = ustrip.(Grid.y.val[1,:,1]);
    z = ustrip.(Grid.z.val[1,1,:]);

    if haskey(Grid.fields,:Phases)
        Phases = Grid.fields[:Phases];
    else
        error("You must provide the field :Phases in the structure")
    end

    if haskey(Grid.fields,:Temp)
        Temp = Grid.fields[:Temp];
    else
        if verbose
            println("Field :Temp is not provided; setting it to zero")
        end
        Temp = zeros(size(Phases));
    end

    if PartitioningFile==empty
        # in case we run this on 1 processor only
        Nprocx  =   1;
        Nprocy  =   1;
        Nprocz  =   1;
        xc,yc,zc = x,y,z;
    else
        Nprocx,Nprocy,Nprocz,
        xc,yc,zc,
        nNodeX,nNodeY,nNodeZ = GetProcessorPartitioning(PartitioningFile)
    end

    Nproc                       =   Nprocx*Nprocy*Nprocz;
    num, num_i, num_j, num_k    =   get_numscheme(Nprocx, Nprocy, Nprocz);

    xi,ix_start,ix_end          =   get_ind(x,xc,Nprocx);
    yi,iy_start,iy_end          =   get_ind(y,yc,Nprocy);
    zi,iz_start,iz_end          =   get_ind(z,zc,Nprocz);

    x_start                     =   ix_start[num_i[:]];
    y_start                     =   iy_start[num_j[:]];
    z_start                     =   iz_start[num_k[:]];
    x_end                       =   ix_end[num_i[:]];
    y_end                       =   iy_end[num_j[:]];
    z_end                       =   iz_end[num_k[:]];

    # Loop over all processors partition
    for n=1:Nproc
        # Extract coordinates for current processor

        part_x   = ustrip.(Grid.x.val[x_start[n]:x_end[n],y_start[n]:y_end[n],z_start[n]:z_end[n]]);
        part_y   = ustrip.(Grid.y.val[x_start[n]:x_end[n],y_start[n]:y_end[n],z_start[n]:z_end[n]]);
        part_z   = ustrip.(Grid.z.val[x_start[n]:x_end[n],y_start[n]:y_end[n],z_start[n]:z_end[n]]);
        part_phs = Phases[x_start[n]:x_end[n],y_start[n]:y_end[n],z_start[n]:z_end[n]];
        part_T   =   Temp[x_start[n]:x_end[n],y_start[n]:y_end[n],z_start[n]:z_end[n]];
        num_particles = size(part_x,1)* size(part_x,2) * size(part_x,3);

        # Information vector per processor
        num_prop        =   5;      # number of properties we save [x/y/z/phase/T]
        lvec_info       =   num_particles;

        lvec_prtcls     =   zeros(Float64,num_prop*num_particles);

        lvec_prtcls[1:num_prop:end] = part_x[:];
        lvec_prtcls[2:num_prop:end] = part_y[:];
        lvec_prtcls[3:num_prop:end] = part_z[:];
        lvec_prtcls[4:num_prop:end] = part_phs[:];
        lvec_prtcls[5:num_prop:end] = part_T[:];

        # Write output files
        if ~isdir(directory); mkdir(directory); end         # Create dir if not existent
        fname = @sprintf "%s/mdb.%1.8d.dat"  directory (n-1);   # Name
        if verbose
            println("Writing LaMEM marker file -> $fname")                   # print info
        end
        lvec_output    = [lvec_info; lvec_prtcls];          # one vec with info about length

        PetscBinaryWrite_Vec(fname, lvec_output)            # Write PETSc vector as binary file

    end
end

# Internal routine to retrieve indices of local portion of the grid
function get_ind(x,xc,Nprocx)
    if Nprocx == 1
        xi       = length(x);
        ix_start = [1];
        ix_end   = [length(x)];
    else

        xi = zeros(Int64,Nprocx)
        for k= 1:Nprocx
            if k==1
                xi[k] = length(x[ (x .>=xc[k]) .& (x .<=xc[k+1]) ]);
            else
                xi[k] = length(x[ (x.>xc[k]) .& (x.<=xc[k+1])]);
            end
        end
        ix_start = cumsum( [0; xi[1:end-1]] ) .+ 1;
        ix_end   = cumsum(xi[1:end]);
    end

    return xi,ix_start,ix_end
end

# Internal routine
function get_numscheme(Nprocx,Nprocy,Nprocz)
    n   = zeros(Int64, Nprocx*Nprocy*Nprocz)
    nix = zeros(Int64, Nprocx*Nprocy*Nprocz)
    njy = zeros(Int64, Nprocx*Nprocy*Nprocz)
    nkz = zeros(Int64, Nprocx*Nprocy*Nprocz)

    num=0;
    for k=1:Nprocz
        for j=1:Nprocy
            for i=1:Nprocx
                num=num+1;
                n[num]   = num;
                nix[num]= i;
                njy[num]= j;
                nkz[num]= k;
            end
        end
    end

    return n,nix,njy,nkz
end

# Internal routine, to write a PETSc vector (as Float64)
"""
    PetscBinaryWrite_Vec(filename, A)

Writes a vector `A` to disk, such that it can be read with `PetscBinaryRead` (which assumes a Big Endian type)

"""
function PetscBinaryWrite_Vec(filename, A)

    # Note: use "hton" to transfer to Big Endian type, which is what PETScBinaryRead expects
    open(filename,"w+") do f
        n               =   length(A);
        nummark         =   A[1];           # number of markers

        write(f,hton(Float64(1211214)));    # header (not actually used)
        write(f,hton(Float64(nummark)));    # info about # of markers written

        for i=2:n
             write(f,hton(Float64(A[i])));  # Write data itself
        end

    end

end

"""
    nProcX,nProcY,nProcZ, xc,yc,zc, nNodeX,nNodeY,nNodeZ = GetProcessorPartitioning(filename)

Reads a LaMEM processor partitioning file, used to create marker files, and returns the parallel layout

"""
function GetProcessorPartitioning(filename)

    io = open(filename, "r")

    nProcX = ntoh(read(io,Int32))
    nProcY = ntoh(read(io,Int32))
    nProcZ = ntoh(read(io,Int32))

    nNodeX = ntoh(read(io,Int32))
    nNodeY = ntoh(read(io,Int32))
    nNodeZ = ntoh(read(io,Int32))

    iX = [ntoh(read(io,Int32)) for i=1:nProcX+1];
    iY = [ntoh(read(io,Int32)) for i=1:nProcY+1];
    iZ = [ntoh(read(io,Int32)) for i=1:nProcZ+1];

    CharLength = ntoh(read(io,Float64))
    xcoor = [ntoh(read(io,Float64)) for i=1:nNodeX].*CharLength;
    ycoor = [ntoh(read(io,Float64)) for i=1:nNodeY].*CharLength;
    zcoor = [ntoh(read(io,Float64)) for i=1:nNodeZ].*CharLength;

    xc = xcoor[iX .+ 1]
    yc = ycoor[iY .+ 1]
    zc = zcoor[iZ .+ 1]

    close(io)

    return  nProcX,nProcY,nProcZ,
            xc,yc,zc,
            nNodeX,nNodeY,nNodeZ

end

"""
    coord, Data_3D_Arrays, Name_Vec = ReadData_VTR(fname)

Reads a VTR (structured grid) VTK file `fname` and extracts the coordinates, data arrays and names of the data.
In general, this only contains a piece of the data, and one should open a `*.pvtr` file to retrieve the full data
"""
function ReadData_VTR(fname, FullSize)
    file = open(fname, "r")

    header = true
    num = 1;
    CoordOffset = zeros(Int64,3);
    Offset_Vec  =   [];
    Name_Vec    =   [];     Type_Vec = [];
    NumComp_Vec =   [];     PieceExtent=[]; WholeExtent=[];
    while header==true

        line        = readline(file)
        line_strip  = lstrip(line)
        if startswith(line_strip, "<RectilinearGrid WholeExtent")
            id_start    = findfirst("\"", line_strip)[1]+1
            id_end      = findlast("\"", line_strip)[1]-1
            WholeExtent = parse.(Int64,split(line_strip[id_start:id_end]))
        end
        if startswith(line_strip, "<Piece Extent=")
            id_start    = findfirst("\"", line_strip)[1]+1
            id_end      =  findlast("\"", line_strip)[1]-1
            PieceExtent =  parse.(Int64,split(line_strip[id_start:id_end]))

        end
        if startswith(line_strip, "<Coordinates>")
            # Read info where the coordinates are stored
            Type, Name, NumberOfComponents, CoordOffset[1]  = Parse_VTR_Line(readline(file)); num += 1
            Type, Name, NumberOfComponents, CoordOffset[2]  = Parse_VTR_Line(readline(file)); num += 1
            Type, Name, NumberOfComponents, CoordOffset[3]  = Parse_VTR_Line(readline(file)); num += 1
        end

        if startswith(line_strip, "<PointData>")
            line_strip  = lstrip(readline(file))
            while ~startswith(line_strip, "</PointData>")
                Type, Name, NumberOfComponents, Offset  = Parse_VTR_Line(line_strip);  num += 1

                Offset_Vec  = [Offset_Vec;  Offset];
                Name_Vec    = [Name_Vec;    Name];
                Type_Vec    = [Type_Vec;   Type];
                NumComp_Vec = [NumComp_Vec; NumberOfComponents];
                line_strip  = lstrip(readline(file))
            end
        end

        if startswith(line_strip, "<CellData>")
            line_strip  = lstrip(readline(file))
            while ~startswith(line_strip, "</CellData>")
                Type, Name, NumberOfComponents, Offset  = Parse_VTR_Line(line_strip);  num += 1

                Offset_Vec  = [Offset_Vec;  Offset];
                Name_Vec    = [Name_Vec;    Name];
                Type_Vec    = [Type_Vec;   Type];
                NumComp_Vec = [NumComp_Vec; NumberOfComponents];
                line_strip  = lstrip(readline(file))
                 # if we have cell Data, for some reason we need to increment this by one.
                PieceExtent[1:2:end] .+= 1
            end

        end

        if startswith(line_strip, "<AppendedData ")
            header=false
        end

        num += 1
    end

    # Skip to beginning of raw data (linebreak)
    skip(file, 5)
    start_bin = position(file);     # start of binary data

    # Determine the end of the raw data
    seekend(file);
    skip(file, -29)

    end_bin = position(file);

    # Start with reading the coordinate arrays:
    coord_x     =   ReadBinaryData(file, start_bin, CoordOffset[1],   (PieceExtent[2]-PieceExtent[1]+1)*sizeof(Float32))
    coord_y     =   ReadBinaryData(file, start_bin, CoordOffset[2],   (PieceExtent[4]-PieceExtent[3]+1)*sizeof(Float32))
    coord_z     =   ReadBinaryData(file, start_bin, CoordOffset[3],   (PieceExtent[6]-PieceExtent[5]+1)*sizeof(Float32))

    # Read data arrays:
    Data_3D_Arrays  = [];
    ix = PieceExtent[1]:PieceExtent[2];
    iy = PieceExtent[3]:PieceExtent[4];
    iz = PieceExtent[5]:PieceExtent[6];
    numPoints       = length(ix)*length(iy)*length(iz);

    coord_x_full = zeros(Float64, FullSize[1]);
    coord_y_full = zeros(Float64, FullSize[2]);
    coord_z_full = zeros(Float64, FullSize[3]);

    coord_x_full[ix] = coord_x[1:length(ix)];
    coord_y_full[iy] = coord_y[1:length(iy)];
    coord_z_full[iz] = coord_z[1:length(iz)];

    for i=1:length(Name_Vec)-1

        data3D      =   ReadBinaryData(file, start_bin, Offset_Vec[i],    numPoints*NumComp_Vec[i]*sizeof(Float32) )
        data3D      =   getArray(data3D, PieceExtent, NumComp_Vec[i]);

        data3D_full =   zeros(Float64,NumComp_Vec[i],FullSize[1],FullSize[2],FullSize[3])  # Generate full data

        # uggly hack to make it work with parallel files
        ix_left = ix; ix_right = 1:length(ix_left);
        iy_left = iy; iy_right = 1:length(iy_left);
        iz_left = iz; iz_right = 1:length(iz_left);
        if ix_left[1]>1; ix_left = ix_left[2:end]; ix_right=ix_right[2:end];    end
        if iy_left[1]>1; iy_left = iy_left[2:end]; iy_right=iy_right[2:end];    end
        if iz_left[1]>1; iz_left = iz_left[2:end]; iz_right=iz_right[2:end];    end

        data3D_full[1:NumComp_Vec[i], ix_left, iy_left, iz_left]        = data3D[1:NumComp_Vec[i],ix_right, iy_right, iz_right];
        #data3D_full[1:NumComp_Vec[i], ix, iy, iz]        = data3D;

        Data_3D_Arrays = [Data_3D_Arrays; data3D_full]
    end
    i=length(Name_Vec);

    if Type_Vec[i]=="UInt8"
        data3D   =   ReadBinaryData(file, start_bin, Offset_Vec[i],    numPoints*NumComp_Vec[i]*sizeof(UInt8), DataType=UInt8)
    else
        data3D   =   ReadBinaryData(file, start_bin, Offset_Vec[i],    numPoints*NumComp_Vec[i]*sizeof(Float32) )
    end

    data3D   =   getArray(data3D, PieceExtent, NumComp_Vec[i]);
    data3D_full =   zeros(Float64,NumComp_Vec[i],FullSize[1],FullSize[2],FullSize[3])  # Generate full d
    data3D_full[1:NumComp_Vec[i], ix, iy, iz]        = data3D[1:NumComp_Vec[i],1:length(ix),1:length(iy),1:length(iz)];

    Data_3D_Arrays = [Data_3D_Arrays; data3D_full]

    return coord_x_full, coord_y_full, coord_z_full, Data_3D_Arrays, Name_Vec, NumComp_Vec, ix, iy, iz
end

 # Parses a line of a *.vtr file & retrieve Type/Name/NumberOfComponents/Offset
 function Parse_VTR_Line(line)
    line_strip  = lstrip(line)

    # Retrieve Type
    if findfirst("type", line_strip) != nothing
        id_start    = findfirst("type", line_strip)[1]+6

        line_strip  = line_strip[id_start:end]
        id_end      = findfirst("\"", line_strip)[1]-1
        Type        = line_strip[1:id_end]
        line_strip  = line_strip[id_end:end]
    else
        Type=nothing;
    end

    # Retrieve Name
    if findfirst("Name", line_strip) != nothing
        id_start    = findfirst("Name", line_strip)[1]+6
        line_strip  = line_strip[id_start:end]
        id_end      = findfirst("\"", line_strip)[1]-1
        Name        = line_strip[1:id_end]
        line_strip  = line_strip[id_end:end]
    else
        Name=nothing
    end

    # Retrieve number of components
    if findfirst("NumberOfComponents", line_strip) != nothing
        id_start    = findfirst("NumberOfComponents", line_strip)[1]+20
        line_strip  = line_strip[id_start:end]
        id_end      = findfirst("\"", line_strip)[1]-1
        NumberOfComponents     = parse(Int64,line_strip[1:id_end])
        line_strip  = line_strip[id_end:end]
    else
        NumberOfComponents=nothing
    end

    # Offset
    if findfirst("offset", line_strip) != nothing
        id_start    = findfirst("offset", line_strip)[1]+8
        line_strip  = line_strip[id_start:end]
        id_end      = findfirst("\"", line_strip)[1]-1
        Offset     = parse(Int64,line_strip[1:id_end])
    else
        Offset=nothing;
    end
    return Type, Name, NumberOfComponents, Offset
end

function getArray(data, PieceExtent, NumComp)
    data        =   reshape(data, (NumComp, PieceExtent[2]-PieceExtent[1]+1,  PieceExtent[4]-PieceExtent[3]+1,  PieceExtent[6]-PieceExtent[5]+1))
    return data
end

function ReadBinaryData(file::IOStream, start_bin::Int64, Offset::Int64, BytesToRead; DataType=Float32)

    seekstart(file);                                # go to start
    skip(file, start_bin+Offset)                    # move to beginning of raw binary data
    buffer      =   read(file,BytesToRead)          # Read necesaary bytes
    data        =   reinterpret(DataType,buffer)    # Transfer to buffer

    data        =   Float64.(data[1:end]);        # Transfer to Float64
    return data
end

"""
    Data::ParaviewData = ReadData_PVTR(fname, dir)

Reads a parallel, rectilinear, `*.vts` file with the name `fname` and located in `dir` and create a 3D `Data` struct from it.

# Example
```julia
julia> Data = ReadData_PVTR("Haaksbergen.pvtr", "./Timestep_00000005_3.35780500e-01/")
ParaviewData
  size  : (33, 33, 33)
  x     ϵ [ -3.0 : 3.0]
  y     ϵ [ -2.0 : 2.0]
  z     ϵ [ -2.0 : 0.0]
  fields: (:phase, :density, :visc_total, :visc_creep, :velocity, :pressure, :temperature, :dev_stress, :strain_rate, :j2_dev_stress, :j2_strain_rate, :plast_strain, :plast_dissip, :tot_displ, :yield, :moment_res, :cont_res)
```
"""
function  ReadData_PVTR(fname, dir)
    file = open(joinpath(dir,fname), "r")

    header = true
    num = 1;
    FullSize= (1,1,1);
    num_data_sets = 1;
    Data_3D=[]; coord_x=[]; coord_y=[]; coord_z=[]; NumComp=[]; Names=[]
    while header==true

        line        = readline(file)
        line_strip  = lstrip(line)
        if startswith(line_strip, "<PRectilinearGrid")
            id_start    = findfirst("WholeExtent=", line_strip)[1]+13
            line_strip  = line_strip[id_start:end]
            id_end      = findfirst("\"", line_strip)[1]-1
            line_piece  = line_strip[1:id_end]

            WholeExtent = parse.(Int64,split(line_piece))
            FullSize    = (WholeExtent[2],WholeExtent[4],WholeExtent[6])
        end

        if startswith(line_strip, "<Piece")
            id_start    = findfirst("Source=", line_strip)[1]+8
            line_strip  = line_strip[id_start:end]
            id_end      = findfirst("\"", line_strip)[1]-1
            fname_piece = line_strip[1:id_end]

            if num_data_sets==1
                coord_x, coord_y, coord_z, Data_3D, Names, NumComp, ix,iy,iz = ReadData_VTR(joinpath(dir,fname_piece), FullSize);
            else
                coord_x1, coord_y1, coord_z1, Data_3D1, Names, NumComp, ix,iy,iz  = ReadData_VTR(joinpath(dir,fname_piece), FullSize);
                coord_x[ix]   = coord_x1[ix];
                coord_y[iy]   = coord_y1[iy];
                coord_z[iz]   = coord_z1[iz];

                Data_3D = Data_3D+Data_3D1;

            end
            num_data_sets += 1
        end

        if startswith(line_strip, "</PRectilinearGrid")
            header=false;
        end
    end

    # Create a named-Tuple out of the fields
    NamesSymbol = [];
    for i=1:length(Names)
        id = findfirst(" ", Names[i])
        if id == nothing
            Names_Strip = Names[i]
        else
            Names_Strip = Names[i][1:findfirst(" ", Names[i])[1]-1];
        end
        NamesSymbol = [NamesSymbol; Names_Strip]
    end

  #  NamesSymbol =   [Names[i][1:findfirst(" ", Names[i])[1]-1] for i=1:length(Names)]
    Names1      =   Symbol.(NamesSymbol)

    Data_Array = [];
    num     =   1;
    for i=1:length(NumComp)
        data        =   Data_3D[num:num+NumComp[i]-1,:,:,:];
        data_arrays =   [data[i,:,:,:] for i=1:size(data,1)]
        data_tuple  =   tuple(data_arrays...)

        if size(data,1)>1
            Data_NamedTuple = NamedTuple{(Names1[i],)}((data_tuple,))
        else
            Data_NamedTuple = NamedTuple{(Names1[i],)}((data_tuple[1],))
        end
        Data_Array = [Data_Array; Data_NamedTuple]

        num = num+NumComp[i];
    end

    # Merge vector with tuples into a NamedTuple
    fields = Data_Array[1];
    for i=2:length(Data_Array)
        fields = merge(fields, Data_Array[i])
    end

    # Create a ParaviewData struct from it.
    X,Y,Z       =   XYZGrid(coord_x, coord_y, coord_z)
    DataC       =   ParaviewData(X,Y,Z, fields);

    return DataC
end

"""
    Save_LaMEMTopography(Topo::CartData, filename::String)

This writes a topography file `Topo` for use in LaMEM, which should have size `(nx,ny,1)` and contain the field `:Topography`
"""
function Save_LaMEMTopography(Topo::CartData, filename::String)

    if (size(Topo.z.val,3) != 1)
        error("Not a valid `CartData' Topography file (size in 3rd dimension should be 1)")
    end
    if !haskey(Topo.fields,:Topography)
        error("The topography `CartData` structure requires a field :Topography")
    end

    # Code the topograhic data into a vector
    nx = Float64(size(Topo.fields.Topography,1));
    ny = Float64(size(Topo.fields.Topography,2));
    x0 = ustrip(Topo.x.val[1,1,1])
    y0 = ustrip(Topo.y.val[1,1,1])
    dx = ustrip(Topo.x.val[2,2,1]) - x0
    dy = ustrip(Topo.y.val[2,2,1]) - y0
    Topo_vec = [ nx;ny;x0;y0;dx;dy; ustrip.(Topo.fields.Topography[:])]

    # Write as PetscBinary file
    PetscBinaryWrite_Vec(filename, Topo_vec)

    println("Written LaMEM topography file: $(filename)")

    return nothing
end

"""
    CreatePartitioningFile(LaMEM_input::String, NumProc::Int64; LaMEM_dir::String=pwd(), LaMEM_options::String="", MPI_dir="")

This executes LaMEM for the input file `LaMEM_input` & creates a parallel partitioning file for `NumProc` processors.
The directory where the LaMEM binary is can be specified; if not it is assumed to be in the current directory.
Likewise for the `mpiexec` directory (if not specified it is assumed to be available on the command line).

"""
function CreatePartitioningFile(LaMEM_input::String,NumProc::Int64; LaMEM_dir::String=pwd(), LaMEM_options="", MPI_dir="")

    # Create string to execute LaMEM
    mpi_str     =  MPI_dir*"mpiexec -n $(NumProc) "
    LaMEM_str   =  LaMEM_dir*"/"*"LaMEM -ParamFile "*LaMEM_input*" -mode save_grid "
    str         =  mpi_str*LaMEM_str

    println("Executing command: $str")

    # Run
    exit=run(`sh -c $str`, wait=false);

    # Retrieve newest file
    if success(exit)
        files=readdir(glob"ProcessorPartitioning_*.bin")
        time_modified = zeros(length(files))
        for (i,file) in enumerate(files)
            time_modified[i] = stat(file).mtime
        end
        id          = findall(time_modified.==maximum(time_modified))   # last modified
        PartFile    = files[id]

        println("Successfuly generated PartitioningFile: $(PartFile[1])")
    else
        error("Something went wrong with executing command ")
    end

    return PartFile[1]

end


text:
# Doing epsilon greedy
# https://youtu.be/PnHCvfgC_ZA?t=3013

using Game2048
using Flux

const EPSILON = 0.5
const LAMBDA = 0.9 # the discount fcator

function mirrors_rotations(board)
    tboard = transpose(board) |> collect

    return board,  rotl90(board),  rot180(board),  rotr90(board),
           tboard, rotl90(tboard), rot180(tboard), rotr90(tboard)
end

actor_model = Chain(
    board -> Float32.(board ./ maximum(board)),
    normalised_board -> reshape(normalised_board, 4, 4, 1, 1),
    Conv((2,2), 1=>64, relu),
    Conv((2,2), 64=>128, relu),
    x->reshape(x, 2*2*128),
    Dense(128*2*2, 4),
    softmax
)

critic_model1 = Chain(
    # boardmove -> (boardmove[1] ./ maximum(boardmove[1]), boardmove[2]),
    boardmove -> (reshape(boardmove[1], 4, 4, 1, 1), boardmove[2]),
    boardmove -> (Conv((2,2), 1=>64, relu)(boardmove[1]), boardmove[2]),
    boardmove -> (Conv((2,2), 64=>128, relu)(boardmove[1]), boardmove[2]),
    boardmove -> vcat(reshape(boardmove[1], 2*2*128), boardmove[2]),
    Dense(2*2*128+4, 256, relu),
    # Dense(256, 1, exp)
)

critic_model2 = Chain(
    Dense(2*2*128+4, 256, relu),
    Dense(256, 1, exp)
)

function action2vec(dir)
    a = zeros(Float32, 4)
    a[Int(dir)+1] = 1
    a
end

board = initboard()

p = Flux.params(critic_model1)
normalised_board(board) = Float32.(board ./ maximum(board))
critic_model1([board |> normalised_board, action2vec(left)])

loss_qa(boardmove, expected_value) = begin
    #sum(sum.(boardmove))
    oldSA = critic_model1(boardmove)
    sum(oldSA)
    #sum(oldSA)
    #sum((oldSA .- expected_value).^2)
    # sum((oldSA .- (reward .+ LAMBDA.*newSA)).^2)

end

opt = ADAM()
normalised_board(board) = Float32.(board ./ maximum(board))
X = [board |> normalised_board, action2vec(left)]
Y = Float32[10.0]
loss_qa(X, Y)

Flux.train!(loss_qa, p, [(X, Y)], opt)

# S,A
# R = Reward
# S'
# A'
# = SARSA

res = move.(Ref(board), DIRS)
possible = map(r->r[2][2], res)
prob_dir = actor_model(board)[possible]
_, greedy_pos = findmax(prob_dir)
reward = map(r->r[2][1], res)[possible][greedy_pos]
dir = rand() < EPSILON ? rand(DIRS[possible]) : DIRS[possible][prob_dir]

features = critic_model1([board, action2vec(dir)])
SA = critic_model2(vcat(features, action2vec(dir)))

#  up to here I have the SAR
new_state = res[possible][greedy_pos][1]
gen_new_tile!(new_state)
res = move.(Ref(new_state), DIRS)
possible = map(r->r[2][2], res)
prob_dir = actor_model(new_state)[possible]
_, greedy_pos = findmax(prob_dir)
reward = map(r->r[2][1], res)[possible][greedy_pos]
dir = rand() < EPSILON ? rand(DIRS[possible]) : DIRS[possible][greedy_pos]
features = critic_model1(new_state)
newSA = critic_model2(vcat(features, action2vec(dir)))

action = dir
loss_qa(boardmove, expected_value) = begin
    oldSA = critic_model1(boardmove)
    sum((oldSA .- estimate_reward).^2)
    # sum((oldSA .- (reward .+ LAMBDA.*newSA)).^2)
end

estimate_reward = reward + LAMBDA*newSA[1]

p = Flux.params(critic_model1)
opt = ADAM()
loss_qa(([board, action2vec(dir)], estimate_reward)...)
Flux.train!(loss_qa, p, [([board, action2vec(dir)], estimate_reward)], opt)

@time actor_model(board)
@time critic_model.(map(res->res[1], res))

using Flux:crossentropy
crossentropy([0.97, 0.01, 0.01, 0.01],[ 1, 0, 0, 0])


text:
using SparseDiffTools
using ForwardDiff: Dual, jacobian
using SparseArrays, Test
using LinearAlgebra
using BlockBandedMatrices, ArrayInterfaceBlockBandedMatrices
using BandedMatrices, ArrayInterfaceBandedMatrices
using StaticArrays
using ArrayInterfaceStaticArrays

fcalls = 0
function f(dx, x)
    global fcalls += 1
    for i in 2:length(x)-1
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    dx[end] = x[end-1] - 2x[end]
    nothing
end

function oopf(x)
    global fcalls += 1
    dx = zero(x)
    for i in 2:length(x)-1
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    dx[end] = x[end-1] - 2x[end]
    dx
end

function nsqf(x)#length(dx)<length(x)
    global fcalls += 1
    dx = zero(x)[1:div(length(x), 2)]
    for i in 2:length(dx)
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    dx
end

function nsqf2(x)#length(dx)>length(x)
    global fcalls += 1
    dx = zeros(eltype(x), length(x) * 2)
    for i in 2:length(x)-1
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    dx
end

function nsqf!(dx, x)
    global fcalls += 1
    for i in 2:length(dx)
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    nothing
end

function nsqf2!(dx, x)
    global fcalls += 1
    for i in 2:length(x)-1
        dx[i] = x[i-1] - 2x[i] + x[i+1]
    end
    dx[1] = -2x[1] + x[2]
    nothing
end

function staticf(x, N=length(x))
    global fcalls += 1
    SVector{N}([i == 1 ? -2x[1] + x[2] : (i == N ? x[N-1] - 2x[N] : x[i-1] - 2x[i] + x[i+1]) for i in 1:N])
end

function staticnsqf(x, N=div(length(x), 2))
    global fcalls += 1
    SVector{N}(vcat([-2x[1] + x[2]], [x[i-1] - 2x[i] + x[i+1] for i in 2:N]))
end

function second_derivative_stencil(N)
    A = zeros(N, N)
    for i in 1:N, j in 1:N
        (j - i == -1 || j - i == 1) && (A[i, j] = 1)
        j - i == 0 && (A[i, j] = -2)
    end
    A
end

@info "ended definitions"

x = rand(30)
dx = rand(30)

J = jacobian(f, dx, x)
@test J ≈ second_derivative_stencil(30)
_J = sparse(J)
@test fcalls == 3

fcalls = 0
_J1 = similar(_J)
forwarddiff_color_jacobian!(_J1, f, x, colorvec=repeat(1:3, 10))
@test _J1 ≈ J
@test fcalls == 1

@info "second passed"

fcalls = 0
_J1 = forwarddiff_color_jacobian(oopf, x, colorvec=repeat(1:3, 10), sparsity=_J, jac_prototype=_J)
@test _J1 ≈ J
@test typeof(_J1) == typeof(_J)
@test fcalls == 1

@info "third passed"

fcalls = 0
_J1 = forwarddiff_color_jacobian(oopf, x, colorvec=repeat(1:3, 10), sparsity=_J)
@test _J1 ≈ J
@test fcalls == 1

#oop with in-place Jacobian
fcalls = 0
_oop_jacout = sparse(1.01 .* J) # want to be nonzero to check that the pre-allocated matrix is overwritten properly
forwarddiff_color_jacobian(_oop_jacout, oopf, x; colorvec=repeat(1:3, 10), sparsity=_J, jac_prototype=_J)
@test _oop_jacout ≈ J
@test typeof(_oop_jacout) == typeof(_J)
@test fcalls == 1

# BandedMatrix
_oop_jacout = BandedMatrix(-1 => diag(J, -1) .* 1.01, 0 => diag(J, 0) .* 1.01,
    1 => diag(J, 1) .* 1.01) # check w/BandedMatrix instead of sparse
fcalls = 0
forwarddiff_color_jacobian(_oop_jacout, oopf, x; colorvec=repeat(1:3, 10), sparsity=_J)
@test _oop_jacout ≈ J
@test isa(_oop_jacout, BandedMatrix)
@test fcalls == 1

@info "4th passed"

fcalls = 0
_J1 = forwarddiff_color_jacobian(staticf, SVector{30}(x), colorvec=repeat(1:3, 10), sparsity=_J, jac_prototype=SMatrix{30,30}(_J))
@test _J1 ≈ J
@test fcalls == 1

@info "5"

_J1 = forwarddiff_color_jacobian(staticf, SVector{30}(x), jac_prototype=SMatrix{30,30}(_J))
@test _J1 ≈ J
_J1 = forwarddiff_color_jacobian(oopf, x, jac_prototype=similar(_J))
@test _J1 ≈ J
_J1 = forwarddiff_color_jacobian(oopf, x)
@test _J1 ≈ J

#Non-square Jacobian
#length(dx)<length(x)
nsqJ = jacobian(nsqf, x)
spnsqJ = sparse(nsqJ)
_nsqJ = forwarddiff_color_jacobian(nsqf, x, dx=nothing)
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf, x, colorvec=repeat(1:3, 10), sparsity=spnsqJ)
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf, x, jac_prototype=similar(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf, x, colorvec=repeat(1:3, 10), sparsity=spnsqJ, jac_prototype=similar(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf, x, jac_prototype=SMatrix{15,30}(nsqJ))
@test _nsqJ ≈ nsqJ
@test typeof(_nsqJ) == typeof(SMatrix{15,30}(nsqJ))
_nsqJ = forwarddiff_color_jacobian(staticnsqf, SVector{30}(x), jac_prototype=SMatrix{15,30}(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(staticnsqf, SVector{30}(x), jac_prototype=SMatrix{15,30}(nsqJ), colorvec=repeat(1:3, 10), sparsity=spnsqJ)
@test _nsqJ ≈ nsqJ
_nsqJ = similar(nsqJ)
forwarddiff_color_jacobian!(_nsqJ, nsqf!, x)
@test _nsqJ ≈ nsqJ
_nsqJ = similar(nsqJ)
forwarddiff_color_jacobian!(_nsqJ, nsqf!, x, colorvec=repeat(1:3, 10), sparsity=spnsqJ)
@test _nsqJ ≈ nsqJ

#length(dx)>length(x)
nsqJ = jacobian(nsqf2, x)
spnsqJ = sparse(nsqJ)
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, dx=nothing)
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, colorvec=repeat(1:3, 10), sparsity=spnsqJ)
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, jac_prototype=similar(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, colorvec=repeat(1:3, 10), sparsity=spnsqJ, jac_prototype=similar(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = forwarddiff_color_jacobian(nsqf2, x, jac_prototype=SMatrix{60,30}(nsqJ))
@test _nsqJ ≈ nsqJ
_nsqJ = similar(nsqJ)
forwarddiff_color_jacobian!(_nsqJ, nsqf2!, x)
@test _nsqJ ≈ nsqJ
_nsqJ = similar(nsqJ)
forwarddiff_color_jacobian!(_nsqJ, nsqf2!, x, colorvec=repeat(1:3, 10), sparsity=spnsqJ)
@test _nsqJ ≈ nsqJ

fcalls = 0
_J1 = similar(_J)
jac_cache = ForwardColorJacCache(f, x, colorvec=repeat(1:3, 10), sparsity=_J1)
forwarddiff_color_jacobian!(_J1, f, x, jac_cache)
@test _J1 ≈ J
@test fcalls == 1

fcalls = 0
_J1 = similar(_J)
_denseJ1 = collect(_J1)
forwarddiff_color_jacobian!(_denseJ1, f, x, colorvec=repeat(1:3, 10), sparsity=_J1)
@test _denseJ1 ≈ J
@test fcalls == 1

fcalls = 0
_J1 = similar(_J)
_denseJ1 = collect(_J1)
jac_cache = ForwardColorJacCache(f, x, colorvec=repeat(1:3, 10), sparsity=_J1)
forwarddiff_color_jacobian!(_denseJ1, f, x, jac_cache)
@test _denseJ1 ≈ J
@test fcalls == 1

_Jt = similar(Tridiagonal(J))
forwarddiff_color_jacobian!(_Jt, f, x, colorvec=repeat(1:3, 10), sparsity=_Jt)
@test _Jt ≈ J

#https://github.com/JuliaDiff/FiniteDiff.jl/issues/67#issuecomment-516871956
function f(out, x)
    x = reshape(x, 100, 100)
    out = reshape(out, 100, 100)
    for i in 1:100
        for j in 1:100
            out[i, j] = x[i, j] + x[max(i - 1, 1), j] + x[min(i + 1, size(x, 1)), j] + x[i, max(j - 1, 1)] + x[i, min(j + 1, size(x, 2))]
        end
    end
    return vec(out)
end
x = rand(10000)
J = BandedBlockBandedMatrix(Ones(10000, 10000), fill(100, 100), fill(100, 100), (1, 1), (1, 1))
Jsparse = sparse(J)
colors = matrix_colors(J)
forwarddiff_color_jacobian!(J, f, x, colorvec=colors)
forwarddiff_color_jacobian!(Jsparse, f, x, colorvec=colors)
@test J ≈ Jsparse

# Non vector input
x = rand(2, 2)
oopf(x) = x
iipf(fx, x) = (fx .= x)
J = forwarddiff_color_jacobian(oopf, x)
@test J ≈ Matrix(I, 4, 4)
J = zero(J)
forwarddiff_color_jacobian!(J, iipf, x, dx=similar(x))
@test J ≈ Matrix(I, 4, 4)

#1x1 SVector test
x = SVector{1}([1.0])
f(x) = x
J = forwarddiff_color_jacobian(f, x)
@test J isa SArray
@test J ≈ SMatrix{1,1}([1.0])


text:
@testset "Distances" begin
  comps = rand(Composition{5}, 100)
  d = CoDaDistance()

  # Test CoDaDistance
  @test d(comps[1], comps[2]) ≈ distance(comps[1], comps[2])

  # Test pairwise with CoDaDistance
  @test pairwise(d, comps) ≈ [distance(comps[i], comps[j]) for i in 1:100, j in 1:100]
end