# example of neural net functions in Econometrics.jl
using Econometrics, Glob

# generate draws from linear regression model, and
# fitted coefficients from correct model, plus
# quadratic and cubic models (irrelevant regressors)
# and 5 pure noise statistics
function make_simdata(reps=100000)
    n = 30
    simdata = zeros(reps, 35+6)
    for rep = 1:reps
        # draw the regressors
        x = randn(n,4)
        z = [ones(n,1) x]
        # draw the parameters from prior
        b = randn(5,1)
        sig = exp.(randn(1))
        # generate dependent variable
        y = z*b + sig.*randn(n,1)
        # linear model
        bhat1 = z\y
        uhat = y-z*bhat1
        sighat1 = sqrt.(uhat'*uhat/(n-size(z,2)))
        # quadratic model
        z = [ones(n,1) x 0.1*x.^2.0]
        bhat2 = z\y
        uhat = y-z*bhat2
        sighat2 = sqrt.(uhat'*uhat/(n-size(z,2)))
        # cubic model
        z = [ones(n,1) x 0.1*x.^2.0 0.01*x.^3.0]
        bhat3 = z\y
        uhat = y-z*bhat3
        sighat3 = sqrt.(uhat'*uhat/(n-size(z,2)))
        # pure noise
        z = randn(1,5)
        # assemble:
        simdata[rep,:] = [b' log.(sig) bhat1' log.(sighat1) bhat2' log.(sighat2) bhat3' log.(sighat3) z]
    end
    return simdata
end

# fit a neural net to the linear model data, and show influence of statistics
function main()
    data = make_simdata()
    noutputs = 6
    trainsize = 80000
    savefile = "olsnet"
    layerconfig = [20, 10, 0, 0]
    epochs = 30
    TrainNet(data, trainsize, noutputs, layerconfig, 512, epochs, savefile)
    Y = data[trainsize+1:end,1:noutputs]
    olsfit = data[trainsize+1:end,(noutputs+1):(2*noutputs)]
    title = "linear regression example"
    params = ["constant", "x1","x2","x3","x4","x5"]
    fit = AnalyzeNet(savefile, epochs, data, trainsize, noutputs, title=title, params=params, doplot=true);
    rm.(glob("olsnet-*"))
end
main();


#   Unit "tools.jl" of the PosDefManifoldML Package for Julia language

#   MIT License
#   Copyright (c) 2019-2021
#   Marco Congedo, CNRS, Grenoble, France:
#   https://sites.google.com/site/marcocongedo/home

# ? CONTENTS :
#   This unit implements tools that are useful for building
#   Riemannian and Euclidean machine learning classifiers.

"""
```julia
function tsMap(	metric :: Metric,
		𝐏      :: ℍVector;
		w    	  :: Vector	= [],
		✓w   	  :: Bool = true,
		⏩   	 :: Bool = true,
		meanISR   :: Union{ℍ, Nothing}  = nothing,
		meanInit  :: Union{ℍ, Nothing}  = nothing,
		tol       :: Real               = 0.,
		transpose :: Bool   			 = true,
		vecRange  :: UnitRange          = 1:size(𝐏[1], 1))
```

The [tangent space mapping](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#PosDefManifold.logMap)
of positive definite matrices ``P_i``, *i=1...k* with mean *G*, once
those points have been parallel transported to the identity matrix,
is given by:

``S_i=\\textrm{log}(G^{-1/2} P_i G^{-1/2})``.

Given a vector of *k* matrices `𝐏` flagged by julia as `Hermitian`,
return a matrix *X* with such tangent vectors of the matrices in `𝐏`
vectorized as per the [vecP](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#PosDefManifold.vecP)
operation.

The mean *G* of the matrices in `𝐏` is found according to the
specified `metric`, of type
[Metric](https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#Metric::Enumerated-type-1).
A natural choice is the
[Fisher metric](https://marco-congedo.github.io/PosDefManifold.jl/dev/introToRiemannianGeometry/#Fisher-1).
If the metric is Fisher, logdet0 or Wasserstein the mean is found with an iterative
algorithm with tolerance given by optional keyword argument `tol`.
By default `tol` is set by the function
[mean](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#Statistics.mean).
For those iterative algorithms a particular initialization can be provided
as an Hermitian matrix by optional keyword argument `meanInit`.

A set of *k* optional non-negative weights `w` can be provided
for computing a weighted mean *G*, for any metrics.
If `w` is non-empty and optional keyword argument `✓w` is true (default),
the weights are normalized so as to sum up to 1,
otherwise they are used as they are passed and should be already normalized.
This option is provided to allow calling this function
repeatedly without normalizing the same weights vector each time.

If an Hermitian matrix is provided as optional keyword argument `meanISR`,
then the mean *G* is not computed, intead this matrix is used
directly in the formula as the inverse square root (ISR) ``G^{-1/2}``.
If `meanISR` is provided, arguments `tol` and `meanInit` have no effect
whatsoever.

If `meanISR` is not provided, return the 2-tuple ``(X, G^{-1/2})``,
otherwise return only matrix *X*.

If an `UnitRange` is provided with the optional keyword argument `vecRange`,
the vectorization concerns only the columns (or rows) of the matrices `𝐏`
specified by the range.

If optional keyword argument `transpose` is true (default),
*X* holds the *k* vectorized tangent vectors in its rows,
otherwise they are arranged in its columns.
The dimension of the rows in the former case and of the columns is the latter
case is *n(n+1)÷2* (integer division), where *n* is the size of the
matrices in `𝐏`, unless a `vecRange` spanning a subset of the columns or rows
of the matrices in `𝐏` has been provided, in which case the dimension will
be smaller. (see [vecP](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#PosDefManifold.vecP)
).

if optional keyword argument `⏩` if true (default),
the computation of the mean and the projection on the tangent space
are multi-threaded. Multi-threading is automatically disabled if the
number of threads Julia is instructed to use is *<2* or *<2k*.

**Examples**:
```julia
using PosDefManifoldML

# generate four random symmetric positive definite 3x3 matrices
Pset = randP(3, 4)

# project and vectorize in the tangent space
X, G⁻½ = tsMap(Fisher, Pset)

# X is a 4x6 matrix, where 6 is the size of the
# vectorized tangent vectors (n=3, n*(n+1)/2=6)

# If repeated calls have to be done, faster computations are obtained
# providing the inverse square root of the matrices in Pset, e.g.,
X1 = tsMap(Fisher, ℍVector(Pset[1:2]); meanISR = G⁻½)
X2 = tsMap(Fisher, ℍVector(Pset[3:4]); meanISR = G⁻½)
```

**See**: [the ℍVector type](https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#%E2%84%8DVector-type-1).

"""
function tsMap(metric :: Metric,
               𝐏      :: ℍVector;
         w    	   :: Vector 			 = [],
         ✓w   	   :: Bool   			 = true,
         ⏩   	  :: Bool   		    = true,
		 meanISR   :: Union{ℍ, Nothing}  = nothing,
		 meanInit  :: Union{ℍ, Nothing}  = nothing,
	  	 tol       :: Real               = 0.,
		 transpose :: Bool   			 = true,
		 vecRange  :: UnitRange          = 1:size(𝐏[1], 2))

	k, n, getMeanISR = dim(𝐏, 1), dim(𝐏, 2), meanISR===nothing
    getMeanISR ? G⁻½ = pow(mean(metric, 𝐏;
	                            w=w,
								✓w=✓w,
								init=meanInit,
								tol=tol,
								⏩=⏩), -0.5) : G⁻½ = meanISR

	# length of the tangent vectors for the given vecRange
	m=_triNum(𝐏[1], vecRange)

	if transpose
		V = Array{eltype(𝐏[1]), 2}(undef, k, m)
	    ⏩==true ? (@threads  for i = 1:k V[i, :] = vecP(log(cong(G⁻½, 𝐏[i], ℍ)); range=vecRange) end) :
	                (@inbounds for i = 1:k V[i, :] = vecP(log(cong(G⁻½, 𝐏[i], ℍ)); range=vecRange) end)
	else
		V = Array{eltype(𝐏[1]), 2}(undef, m, k)
		⏩==true ? (@threads  for i = 1:k V[:, i] = vecP(log(cong(G⁻½, 𝐏[i], ℍ)); range=vecRange) end) :
	                (@inbounds for i = 1:k V[:, i] = vecP(log(cong(G⁻½, 𝐏[i], ℍ)); range=vecRange) end)
	end
    return getMeanISR ? (V, G⁻½) : V
end

"""
```julia
function tsWeights(y::Vector{Int}; classWeights=[])
```

Given an [IntVector](@ref) of labels `y`, return a vector
of weights summing up to 1 such that the overall
weight is the same for all classes (balancing).
This is useful for machine learning models in the tangent
space with unbalanced classes for computing the mean,
that is, the base point to map PD matrices onto the tangent space.
For this mapping, giving equal weights to all observations
actually overweights the larger classes
and downweight the smaller classes.

Class labels for *n* classes must be the first *n* natural numbers,
that is, `1` for class 1, `2` for class 2, etc.
The labels in `y` can be provided in any order.

if a vector of *n* weights is specified as optional
keyword argument `classWeights`, the overall weights
for each class will be first balanced (see here above),
then weighted by the `classWeights`.
This allow user-defined control of weighting independently
from the number of observations in each class.
The weights in `classWeights` can be any integer or real
non-negative numbers. The returned weight vector will
nonetheless sum up to 1.

When you invoke the [`fit`](@ref) function for tangent space models
you don't actually need this function, as you can invoke it
implicitly passing symbol `:balanced` (or just `:b`) or a tuple
with the class weights as optional keyword argument `w`.

**Examples**
```julia
# generate some data; the classes are unbalanced
PTr, PTe, yTr, yTe=gen2ClassData(10, 30, 40, 60, 80, 0.1)

# Fit an ENLR lasso model and find the best model by cross-validation
# balancing the weights for tangent space mapping
m=fit(ENLR(), PTr, yTr; w=tsWeights(yTr))

# A simpler syntax is
m=fit(ENLR(), PTr, yTr; w=:balanced)

# to balance the weights and then give overall weight 0.5 to class 1
# and 1.5 to class 2:
m=fit(ENLR(), PTr, yTr; w=(0.5, 1.5))

# which is equivalent to
m=fit(ENLR(), PTr, yTr; w=tsWeights(yTr; classWeights=(0.5, 1.5)))

```

This is how it works:

julia> y=[1, 1, 1, 1, 2, 2]
6-element Array{Int64,1}:
 1
 1
 1
 1
 2
 2

We want the four observations of class 1 to count as much
as the two observations of class 2.

julia> tsWeights(y)
6-element Array{Float64,1}:
 0.125
 0.125
 0.125
 0.125
 0.25
 0.25

i.e., 0.125*4 = 1.25*2
and all weights sum up to 1

Now, suppose we want to give to class 2 a weight
four times bigger as compared to class 1:

julia> tsWeights(y, classWeights=[1, 4])
6-element Array{Float64,1}:
 0.05
 0.05
 0.05
 0.05
 0.4
 0.4

and, again, all weights sum up to 1

"""
function tsWeights(y::Vector{Int}; classWeights=[])

    Nobs = length(y)
    Nclass=length(unique(y))
    Nobsxclass=[count(i->(i==j), y) for j=1:Nclass]
    NobsxclassProp=[n/Nobs for n in Nobsxclass]
    minProportion=minimum(Nobsxclass)/Nobs
    minProportion<(1/(Nclass*10)) && @warn "the smallest class contains les then 10% of the observation as compared to a balances design" minProportion

    w=[1/(Nclass*Nobsxclass[l]) for l ∈ y]

    if !isempty(classWeights)
       if length(classWeights)≠Nclass
          @warn "the number of elements in argument ClassWeights is different from the number of unique classes in label vector y. Class weights have not been applied" length(classWeights)
       else
         for i=1:length(w) w[i]*= classWeights[y[i]] end
         w./=sum(w)
       end
    end

    return w
end

"""
```julia
function gen2ClassData(n        ::  Int,
                       k1train  ::  Int,
                       k2train  ::  Int,
                       k1test   ::  Int,
                       k2test   ::  Int,
                       separation :: Real = 0.1)
```

Generate a *training set* of `k1train`+`k2train`
and a *test set* of `k1test`+`k2test`
symmetric positive definite matrices.
All matrices have size *n*x*n*.

The training and test sets can be used to train and test any [MLmodel](@ref).

`separation` is a coefficient determining how well the two classs are
separable; the higher it is, the more separable the two classes are.
It must be in [0, 1] and typically a value of 0.5 already
determines complete separation.

Return a 4-tuple with

- an [ℍVector](https://marco-congedo.github.io/PosDefManifold.jl/dev/MainModule/#%E2%84%8DVector-type-1) holding the `k1train`+`k2train` matrices in the training set,
- an ℍVector holding the `k1test`+`k2test` matrices in the test set,
- a vector holding the `k1train`+`k2train` labels (integers) corresponding to the matrices of the training set,
- a vector holding the `k1test`+`k2test` labels corresponding to the matrices of the test set (*1* for class 1 and *2* for class 2).

**Examples**

```julia
using PosDefManifoldML

PTr, PTe, yTr, yTe=gen2ClassData(10, 30, 40, 60, 80, 0.25)

# PTr=training set: 30 matrices for class 1 and 40 matrices for class 2
# PTe=testing set: 60 matrices for class 1 and 80 matrices for class 2
# all matrices are 10x10
# yTr=a vector of 70 labels for the training set
# yTe=a vector of 140 labels for the testing set

```
"""
function gen2ClassData(n        ::  Int,
                       k1train  ::  Int,
                       k2train  ::  Int,
                       k1test   ::  Int,
                       k2test   ::  Int,
                       separation :: Real = 0.1)
	if separation<0 || separation >1
		@error 📌*", function "*gen2ClassData*": argument `separation` must be in range [0, 1]."
		return
	end

	G1=randP(n)
    G2=randP(n)

    # Create a set of k1+k2 random matrices and move the along
    # the Fisher Geodesic with arclength (1-a) the first k1 toward G1
    # and the last k2 toward G2. Geodesics are computed with the Schur method
    function getMatrices(k1::Int, k2::Int, n::Int, a::Real, G1::ℍ, G2::ℍ)

        function getChol(G::ℍ)
            L = cholesky(G, check=false)
            U⁻¹ = inv(L.U)
            return L, U⁻¹
        end

        k=k1+k2
        𝐗=ℍVector(undef, k)
            L, U⁻¹=getChol(G1)
            for i=1:k1
                F = schur(U⁻¹' * randP(n) * U⁻¹)
                𝐗[i]=ℍ(L.U' * (F.Z * F.T^a* F.Z') * L.U)
            end
            L, U⁻¹=getChol(G2)
            for i=k1+1:k
                F = schur(U⁻¹' * randP(n) * U⁻¹)
                𝐗[i]=ℍ(L.U' * (F.Z * F.T^a* F.Z') * L.U)
            end
        return 𝐗
    end

    𝐗train=getMatrices(k1train, k2train, n, 1-separation, G1, G2)
    𝐗test=getMatrices(k1test, k2test, n, 1-separation, G1, G2)
    y𝐗train=IntVector([repeat([1], k1train); repeat([2], k2train)])
    y𝐗test=IntVector([repeat([1], k1test); repeat([2], k2test)])

    return 𝐗train, 𝐗test, y𝐗train, y𝐗test
end

"""
```julia
function confusionMat(yTrue::IntVector, yPred::IntVector)
```

Return the *confusion matrix* given integer vectors of true label `yTrue`
and predicted labels `yPred`.

The length of `yTrue` and `yPred` must be equal. Furthermore,
the `yTrue` vector must comprise all natural numbers
in between 1 and *z*, where *z* is the number of classes.

The confusion matrix will have size *z*x*z*. It is computed
starting from a matrix filled everywhere with zeros and
adding, for each label, 1 at entry [i, j] of the matrix, where
i is the true label and j the predicted label, and finally
dividing the matrix by the sum of all its elements.
Therefore, the entries of the confusion matrix sum up to 1.0.

**See** [`predict`](@ref), [`predictAcc`](@ref), [`predictErr`](@ref).

**Examples**

```julia
using PosDefManifoldML
julia> confusionMat([1, 1, 1, 2, 2], [1, 1, 1, 1, 2])
# return: [0.6 0.0; 0.2 0.2]
```
"""
function confusionMat(yTrue::IntVector, yPred::IntVector)

	n1=length(yTrue)
	n2=length(yPred)
	if n1≠n2
		@error 📌*", function ConfusionMat: the length of the two argument vectors must be equal." n1 n2
		return
	end

	cTrue=sort(unique(yTrue))
	z = length(cTrue)
	if cTrue≠[i for i∈1:z]
		@error 📌*", function ConfusionMat: the `yTrue` vector must contains all natural numbers from 1 to the number of classes. It contains instead: " cTrue
		return
	end

	CM = zeros(Float64, z, z)
	for i=1:n1 CM[yTrue[i], yPred[i]]+=1. end
	return CM/=sum(CM)
end

"""
```julia
(1)
function predictAcc(yTrue::IntVector, yPred::IntVector;
		scoring:: Symbol = :b,
		digits::Int=3)

(2)
function predictAcc(CM:: Matrix{R};
		scoring:: Symbol = :b,
		digits::Int=3) where R<:Real
```

Return the prediction accuracy as a proportion, that is, ∈[0, 1],
given

- (1) the integer vectors of true labels `yTrue` and of predicted labels `yPred`,
- (2) a confusion matrix.

If `scoring`=:b (default) the **balanced accuracy** is computed.
Any other value will make the function returning the regular **accuracy**.
Balanced accuracy is to be preferred for unbalanced classes.
For balanced classes the balanced accuracy reduces to the
regular accuracy, therefore there is no point in using regular accuracy
if not to avoid a few unnecessary computations when the class are balanced.

The error is rounded to the number of optional keyword argument
`digits`, 3 by default.

**Maths**

The regular *accuracy* is given by sum of the diagonal elements
of the confusion matrix.

For the *balanced accuracy*, the diagonal elements
of the confusion matrix are divided by the respective row sums
and their mean is taken.

**See** [`predict`](@ref), [`predictErr`](@ref), [`confusionMat`](@ref)

**Examples**

```julia
using PosDefManifoldML
julia> predictAcc([1, 1, 1, 2, 2], [1, 1, 1, 1, 2]; scoring=:a)
# regular accuracy, return: 0.8
julia> predictAcc([1, 1, 1, 2, 2], [1, 1, 1, 1, 2])
# balanced accuracy, return: 0.75
```
"""
function predictAcc(yTrue::IntVector, yPred::IntVector;
					scoring:: Symbol = :b,
	          		digits::Int=3)
	n1=length(yTrue)
	n2=length(yPred)
	if n1≠n2
		@error 📌*", function `predictAcc` or `predictErr`: the length of the two argument vectors must be equal." n1 n2
		return
	end

	if scoring≠:b # regular accuracy
		return round(sum(y1==y2 for (y1, y2) ∈ zip(yTrue, yPred))/n1; digits=digits)
	else # balanced accuracy
		CM=confusionMat(yTrue, yPred)
		z=size(CM, 1)
		return round(sum(CM[i, i]/sum(CM[i, :]) for i=1:z) / z; digits=digits)
	end
end

function predictAcc(CM:: Matrix{R};
					scoring:: Symbol = :b,
					digits::Int=3) where R<:Real
					num_of_rows, num_of_cols = size(CM)

	num_of_rows, num_of_cols = size(CM)
	if num_of_rows≠num_of_cols
		@error 📌*", function predictAcc or predictErr: the `CM` argument must be square as this must be a confusion matrix." num_of_rows num_of_cols
		return
	end

	sum_of_elements=sum(CM)
	if sum_of_elements≉  1.0
		@error 📌*", function predictAcc or predictErr: the elements of `CM` matrix argument must sum up to 1.0 as this must be a confusion matrix." sum_of_elements
		return
	end

	return scoring==:b ? round(sum(CM[i, i]/sum(CM[i, :]) for i=1:size(CM, 1)) / size(CM, 1);
								digits=digits) :
						 round(tr(CM);
						 		digits=digits)
end

"""
```julia
(1)
function predictErr(yTrue::IntVector, yPred::IntVector;
		scoring:: Symbol = :b,
		digits::Int=3)
(2)
function predictErr(CM:: Matrix{R};
		scoring:: Symbol = :b,
		digits::Int=3) where R<:Real
```

Return the complement of the predicted accuracy, that is, 1.0 minus
the result of [`predictAcc`](@ref), given

- (1) the integer vectors of true labels `yTrue` and of predicted labels `yPred`,
- (2) a confusion matrix.

**See** [`predictAcc`](@ref).
"""
predictErr(yTrue::IntVector, yPred::IntVector;
			scoring::Symbol = :b,
	        digits::Int=3) =
	return (acc=predictAcc(yTrue, yPred;
				scoring=scoring, digits=8))≠nothing ? round(1.0-acc;
													  digits=digits) : nothing
predictErr(CM:: Matrix{R};
			scoring:: Symbol = :b,
			digits::Int=3) where R<:Real =
	return (acc=predictAcc(CM;
				scoring=scoring, digits=8))≠nothing ? round(1.0-acc;
													  digits=digits) : nothing
"""
```julia
function rescale!(X::Matrix{T},	bounds::Tuple=(-1, 1);
		dims::Int=1) where T<:Real
```
Rescale the columns or the rows of real matrix `X` to be in range [a, b],
where a and b are the first and seconf elements of tuple `bounds`.

By default rescaling apply to the columns. Use `dims=2` for rescaling
the rows.

This function is used, for instance, by the SVM fit and predict functions.
"""
function rescale!(X::Matrix{T},
	              bounds::Tuple=(-1, 1);
				  dims::Int=1) where T<:Real
 dims∉(1, 2) && throw(ArgumentError, "rescale! function: the `dims` keyword argument must be either 1 or 2; dims=$dims")
 length(bounds) ≠ 2 && throw(ArgumentError, "rescale! function: tuple `bounds` must contain two elements; bounds=$bounds")
 a=first(bounds)
 b=last(bounds)
 a isa Number && b isa Number || throw(ArgumentError, "rescale! function: the two elements of tuple `bounds` must be numbers; bounds=$bounds")
 c=b-a
 c<=0 && throw(ArgumentError, "rescale! function: the two elements (a, b) of tuple `bounds` must verify b-a>0; bounds=$bounds")
 @inbounds dims==1 ? (for j=1:size(X, 2)
			 	        minX, maxX=extrema(X[:, j])
				        range=maxX-minX
			 	        for i=1:size(X, 1) X[i, j]=a+(c*(X[i, j]-minX)/range) end
					end) :
			        (for i=1:size(X, 1)
			 			minX, maxX=extrema(X[i, :])
						range=maxX-minX
			 			for j=1:size(X, 2) X[i, j]=a+(c*(X[i, j]-minX)/range) end
					end)
end

# -------------------------------------------------------- #
# INTERNAL FUNCTIONS #

# return a vector of ranges partitioning lineraly and
# as much as possible evenly `n` elements in `threads` ranges.
# `threads` is the number of threads to which the ranges are to be
# dispatched. If `threads` is not provided, it is set to the number
# of threads Julia is currently instructed to use.
# For example, for `k`=99
# and `threads`=4, return Array{UnitRange{Int64},1}:[1:25, 26:50, 51:75, 76:99].
function _partitionLinRange4threads(n::Int, threads::Int=0)
    threads==0 ? thr=nthreads() : thr=threads
    n<thr ? thr = n : nothing
    d = max(round(Int64, n / thr), 1)
    return [(r<thr ? (d*r-d+1:d*r) : (d*thr-d+1:n)) for r=1:thr]
end

function _GetThreads(n::Int, callingFunction::String)
	threads=Threads.nthreads()
	threads==1 && @warn 📌*", function "*callingFunction*": Julia is instructed to use only one thread."
	if n<threads && n<3
		@warn 📌*", function "*callingFunction*": the number of operations (n) is too low for taking advantage of multi-threading" threads n
		threads=1
	end
	return threads
end

function _GetThreadsAndLinRanges(n::Int, callingFunction::String)
	threads = _GetThreads(n, callingFunction)
	ranges=_partitionLinRange4threads(n, threads)
	return threads, ranges
end

# checks for `fit function`
function _check_fit(model       :: MLmodel,
              		 dim𝐏Tr     :: Int,
              		 dimyTr     :: Int,
           			 dimw  	    :: Int,
					 dimWeights :: Int,
					 modelName  :: String)
    errMsg1="the number of data do not match the number of labels."
	errMsg2="the number of data do not match the number of elements in `w`."
	errMsg3="the number of data do not match the number of elements in `weights`."
    if dim𝐏Tr ≠ dimyTr
		@error 📌*", fit function, model "*modelName*": "*errMsg1
		return false
	end
    if dimw ≠ 0 && dimw ≠ dimyTr
		@error 📌*", fit function, model "*modelName*": "*errMsg2
		return false
	end
	if dimWeights ≠ 0 && dimWeights ≠ dimyTr
		@error 📌*", fit function, model "*modelName*": "*errMsg3
		return false
	end

	return true
end

# check for argument `what` in `predict` function
_whatIsValid(what::Symbol, funcName::String) =
	if what ∉ (:l, :labels, :p, :probabilities, :f, :functions)
		@error 📌*", "*funcName*" function: the `what` symbol is not supported."
		return false
	else
		return true
	end

# translate the `what` argument to a string to give feedback
_what2Str(what::Symbol) =
	if      what ∈(:f, :fucntions)      return "functions"
	elseif  what ∈(:l, :labels)         return "labels"
	elseif  what ∈(:p, :probabilities)  return "prob. of belonging to each class"
	end

# return the dimension of the manifold of PD matrices: (n*n+1)/2
_triNum(P::ℍ, vecRange::UnitRange) =
	if length(vecRange)==size(P, 1)
		result=_triNum(P)
	else
		m=0; for j=vecRange, i=j:size(P, 1) m+=1 end
		result=m
	end

_triNum(P::ℍ) = ( size(P, 1) * (size(P, 1)+1) ) ÷ 2

# dimension of the manifold if 𝐏Tr is an ℍVector,
# dimension of the tagent(feature) vectors if 𝐏Tr is a Matrix
_getDim(𝐏Tr :: Union{ℍVector, Matrix{Float64}}, vecRange::UnitRange) =
	𝐏Tr isa ℍVector ? _triNum(𝐏Tr[1], vecRange) : length(vecRange)

_getDim(𝐏Tr :: Union{ℍVector, Matrix{Float64}}) =
	𝐏Tr isa ℍVector ? _triNum(𝐏Tr[1]) : size(𝐏Tr, 2)

# convert a ML model in a atring to release information
_modelStr(model::MLmodel) =
  if 		model isa MDMmodel
	  		return "MDM"

  elseif    model isa ENLRmodel
    		if     model.alpha≈1. return "Lasso logistic regression"
    		elseif model.alpha≈0. return "Ridge logistic regression"
    		else                  return "El. Net (α=$(round(model.alpha; digits=2))) log. reg."
			end

  elseif    model isa SVMmodel
			if 		model.svmType==SVC 			return "SVC"
			elseif  model.svmType==NuSVC 		return "NuSVC"
			elseif  model.svmType==EpsilonSVR 	return "EpsilonSVR"
			elseif  model.svmType==OneClassSVM 	return "OneClassSVM"
			elseif  model.svmType==NuSVR 		return "NuSVR"
			else    return  "Warning: the SVM type is unknown"
			end

  else      return "unknown"
  end

# check on `what` argument of `fit` function
_fitTypeIsValid(what::Symbol, funcName::String) =
	if what ∉ (:best, :path)
		@error 📌*", "*funcName*" function: the `fitType` symbol must be `:best` or `:path`."
		return false
	else
		return true
	end

# check on `onWhich` argument of `predict` function
function _ENLRonWhichIsValid(model::ENLRmodel, fitType::Symbol,
                    onWhich::Int, funcName::String)
    if fitType==:best
		return true
	else #fitType==:path
		i=length(model.path.lambda)
		if !(0<=onWhich<=i)
			@error 📌*", "*funcName*" function: the `onWhich` integer argument must be comprised between 0 (all models) and $i."
			return false
		else
			return true
		end
	end
end

# return a string to give information on what model is used to predict
# based on arguments `fitType` and `onWhich`
_ENLRonWhichStr(model::ENLRmodel, fitType::Symbol, onWhich::Int) =
	if 		fitType==:best
		return "from the best "*_modelStr(model)*" model (λ=$(round(model.best.lambda[1]; digits=5)))"
	else # :path
		if onWhich == 0
			return "from all "*_modelStr(model)*" models"
		else
			return "from "*_modelStr(model)*" model $(onWhich) (λ=$(round(model.path.lambda[onWhich]; digits=5)))"
		end
	end

# create a copy of optional keyword arguments `args`
# removing all arguments with names listed in tuple `remove`.
# Examples:
# fitArgs✔=_rmArgs((:meanISR, :fitType); fitArgs...)
# fitArgs✔=_rmArgs((:meanISR,); fitArgs...) # notice the comma after `meanISR`
# note: named tuples are immutable, that's why a copy must be created
function _rmArgs(remove::Tuple; args...)
	D = Dict(args)
	for key ∈ remove delete!(D, key) end
    return Tuple(D)
end

# given optional keyword arguments `args`,
# return the value of the argument with key `key`.
# If the argument does not exist, return `nothing`
function _getArgValue(key::Symbol; args...)
   D = Dict(args)
   return haskey(D, key) ? D[key] : nothing
end

# get a valid weigts `w` object and perform check given
# the user-defined `w` argument. Used in `fit` and `cvAcc` functions.
function _getWeights(w :: Union{Symbol, Tuple, Vector}, y::IntVector, funcName::String)
	if 		(w isa Vector && isempty(w)) || w==:uniform || w==:u return []
	elseif	w isa Vector && !isempty(w)
		    nObs, length_w = length(y), length(w)
		    if length_w==nObs return w
			else @error 📌*", "*funcName*"invalid vector `w`. `w` must contain as many elements as there are observations" length_w nObs
			end
	elseif  w==:balanced || w==:b return tsWeights(y)
	elseif  w isa Tuple
		    nClasses, length_w = length(unique(y)), length(w)
			if length_w==nClasses return tsWeights(y; classWeights=collect(w))
			else @error 📌*", "*funcName*"invalid tuple `w`. `w` must contain as many elements as there are classes" length_w nClasses
			end
	else
			@error 📌*", "*funcName*"invalid argument `w`. `w` must be a vector, an empty vector, a tuple of as many real numbers as classes, or symbol `:balanced`, or symbol `:uniform`"
			return nothing
	end
end

# Get the feature matrix for fit functions of ML model in the tangent space:
# if `𝐏Tr` is a matrix just return the columns in `vecRange` (by default all).
# if `𝐏Tr` is vector of Hermitian matrices, they are projected onto the
# tangent space. If the inversesquare root of a base point `meanISR`
# is provided, the projection is obtained at this base point, otherwise the
# mean of all points is computed and used as base point.
# If the mean is to be computed by an iterative algorithm (e.g., if the metric
# of the model is the Fisher metric), an initialization `meanInit`, weights
# `w` and a tolerance `tol` are used.
# Once projected onto the tangent space, the matrices in `𝐏Tr` are vectorized
# using only the rows (or columns) specified by `vecRange`.
# if `verbose` is true, print "Projecting data onto the tangent space..."
# if `transpose` the feature vectors are in the rows of `X`, otherwise in the
# columns of `X`.
# if ⏩ is true, the projection onto the tangent space
# and the algorithm to compute the mean are multi-threaded
function _getFeat_fit!(ℳ	    :: TSmodel,
					   𝐏Tr	  	  :: Union{ℍVector, Matrix{Float64}},
			 		   meanISR	 :: Union{ℍ, Nothing},
					   meanInit	 :: Union{ℍ, Nothing},
					   tol		 :: Real,
			 	 	   w		 :: Union{Symbol, Tuple, Vector},
			 	 	   vecRange	 :: UnitRange,
					   transpose :: Bool,
					   verbose	 :: Bool,
					   ⏩	   :: Bool)
	if 𝐏Tr isa ℍVector
		verbose && println(greyFont, "Projecting data onto the tangent space...")
		if meanISR==nothing
			(X, G⁻½)=tsMap(ℳ.metric, 𝐏Tr;
			               w=w,
						   ⏩=⏩,
						   vecRange=vecRange,
						   meanInit=meanInit,
						   tol=tol,
						   transpose=transpose)
			ℳ.meanISR = G⁻½
		else
			X=tsMap(ℳ.metric, 𝐏Tr;
			        w=w,
					⏩=⏩,
					vecRange=vecRange,
					meanISR=meanISR,
					transpose=transpose)
			ℳ.meanISR = meanISR
		end
	else X=𝐏Tr[:, vecRange]
	end
	return X
end

# Get the feature matrix for predict functions of ML model in the tangent space:
# if `𝐏Te` is a matrix just return the columns in `vecRange` (by default all).
# if `𝐏Te` is vector of Hermitian matrices, they are projected onto the
# tangent space. If an inverse square root 5ISR) `tranfer` is provided
# (typically the ISR of the mean of the matrices in `𝐏Te`), this is used
# as ISR of the base point, otherwise the ISR of the base point stored in the
# model ℳ is used. The latter is the classical approach, the former realizes the
# adaptation (transfer learning) explained in Barachant et al. (2013).
# Once projected onto the tangent spave, the matrces in `𝐏Te` are vectorized
# using only the rows (or columns) specified by `vecRange`.
# if `transpose` the feature vectors are in the rows of `X`, otherwise in the
# columns of `X`.
# if `verbose` is true, print "Projecting data onto the tangent space..."
# if ⏩ is true, the projection onto the tangent space is multi-threaded
_getFeat_Predict!(ℳ	   		:: TSmodel,
			      𝐏Te		  :: Union{ℍVector, Matrix{Float64}},
				  transfer   :: Union{ℍ, Nothing},
				  vecRange	 :: UnitRange,
				  transpose  :: Bool,
				  verbose	 :: Bool,
				  ⏩	       :: Bool) =
	if 𝐏Te isa ℍVector
		verbose && println(greyFont, "Projecting data onto the tangent space...")
		return tsMap(ℳ.metric, 𝐏Te;
				     meanISR = transfer==nothing ? ℳ.meanISR : transfer,
				     ⏩=⏩,
				     vecRange=vecRange,
					 transpose=transpose)
	else
		return 𝐏Te[:, vecRange]
	end


# # Custom structs

#=
In this example, we demonstrate implicit differentiation through functions that manipulate `NamedTuple`s, as a first step towards compatibility with general structs.
=#

using ComponentArrays
using ImplicitDifferentiation
using Zygote

using ChainRulesCore  #src
using ChainRulesTestUtils  #src
using Test  #src

# ## Implicit function wrapper

#=
We replicate a componentwise square function with `NamedTuple`s, taking `a=(x,y)` as input and returning `b=(u,v)`.
=#

forward(a::ComponentVector) = ComponentVector(u=a.x .^ 2, v=a.y .^ 2)

function conditions(a::ComponentVector, b::ComponentVector)
    return vcat(b.u .- a.x .^ 2, b.v .- a.y .^ 2)
end

implicit = ImplicitFunction(forward, conditions);

#=
In order to be able to call `Zygote.gradient`, we use `implicit` to define a convoluted version of the squared Euclidean norm, which takes a `ComponentVector` as input and returns a real number.
=#

function fullnorm(a::ComponentVector)
    b = implicit(a)
    return sum(b)
end

function halfnorm(a::ComponentVector)  #src
    b = implicit(a)  #src
    return sum(b.u)  #src
end;  #src

# ## Testing

a = ComponentVector(x=rand(3), y=rand(3))

# Let us first check that our weird squared norm returns the correct result.

fullnorm(a) ≈ sum(abs2, a)
halfnorm(a) ≈ sum(abs2, a.x)  #src

# Now we go one step further and compute its gradient, which involves the reverse rule for `implicit`.

Zygote.gradient(fullnorm, a)[1] ≈ 2a

#-

Zygote.gradient(halfnorm, a)[1] ≈ vcat(2a.x, zero(a.y))  #src

# The following tests are not included in the docs.  #src

@test fullnorm(a) ≈ sum(abs2, a)  #src
@test halfnorm(a) ≈ sum(abs2, a.x)  #src

@test Zygote.gradient(fullnorm, a)[1] ≈ 2a  #src
@test Zygote.gradient(halfnorm, a)[1] ≈ vcat(2a.x, zero(a.y))  #src


using FileIO
using SparseArrays

@enum ExecutionPolicy_t begin
    serial          # uses Gauss-Seidel smoother and coarse solver, and no ParallelStencil code at all
    parallel        # uses Jacobi smoother and coarse solver, and ParallelStencil code
    parallel_shmem  # uses Jacobi smoother and coarse solver, and ParallelStencil code + shared memory for residual_2DPoisson
end

""" Loads double-precision matrices in binary format, leading with nx and ny. """
function load(f)
    nx = Array{Int32}(undef, 1, 1)
    ny = Array{Int32}(undef, 1, 1)
    read!(f, nx)
    read!(f, ny)
    sim_results = Array{Float64}(undef, nx[1, 1], ny[1, 1])
    read!(f, sim_results)
    return sim_results
end

""" Applies Dirichlet and Neumann boundary conditions on all 4 edges of the rectangular domain. """
@views function apply_boundary_conditions!(T)
    apply_boundary_conditions_dirichlet!(T)
    apply_boundary_conditions_neumann!(T)
end

""" Applies Dirichlet boundary conditions on bottom and top edges of the rectangular domain. """
@views function apply_boundary_conditions_dirichlet!(T)
    nx, ny = size(T)
    T[:, 1]  .= 1.0  # High temperature on the bottom edge
    T[:, ny] .= 0.0  # Low temperature on the top edge
end

""" Applies Neumann boundary conditions on left and right edges of the rectangular domain. """
@views function apply_boundary_conditions_neumann!(T)
    nx, ny = size(T)
    T[1, :]  .= T[2, :]
    T[nx, :] .= T[nx-1, :]
end

""" Creates sparse coefficient matrix as defined by 5-point stencil. """
@views function stencil_5pt(nx, ny)
    Δx = spdiagm(-1 => 1 * ones(nx - 1), 0 => -2 * ones(nx), 1 => 1 * ones(nx-1))
    Δy = spdiagm(-1 => 1 * ones(ny - 1), 0 => -2 * ones(ny), 1 => 1 * ones(ny-1))
    Ix = spdiagm(0 => ones(nx))
    Iy = spdiagm(0 => ones(ny))
    Δ = kron(Δy, Ix) + kron(Iy, Δx)
    return Δ
end


"""
    SCHash

Separate chaining hash table, using dynamic arrays.
"""
mutable struct SCHash{K, V}
    table::Vector{Vector{Pair{K, V}}}
    nkeys::Int # number of items
    tablesize::Int # table size
    function SCHash{K, V}(m::Int = 10) where {K, V}
        table = Vector{Vector{Pair{K, V}}}(undef, m)
        for i in 1:m
            table[i] = []
        end
        new(table, 0, m)
    end
end

"""
    SCLHash

Separate chaining hash table, using linked list.
"""
mutable struct SCLHash{K, V}
    table::Vector{DLinkedList{Pair{K, V}}}
    nkeys::Int
    tablesize::Int
    function SCLHash{K, V}(m::Int = 10) where {K, V}
        table = Vector{DLinkedList{Pair{K, V}}}(undef, m)
        for i in 1:m
            table[i] = DLinkedList{Pair{K, V}}()
        end
        new(table, 0, m)
    end
end

"""
    hash(map, key)

Hash function for key, implement for allowed key types, for now Numbers and Strings.
"""
function hash(map::Union{SCHash, SCLHash}, key::K) where K <: Number
    return compress(Int(ceil(abs(key))), map.tablesize)
end

function hash(map::Union{SCHash, SCLHash}, key::K) where K <: AbstractString
    h = 0
    for c in key
        h  = 37h + UInt8(c)
    end
    return compress(h, map.tablesize)
end

"""
    compress(h, tablesize)

Compress the hash, quite simple.
"""
function compress(h::Int, tablesize::Int)
    return h % tablesize + 1
end

"""
    get(map, key)

Get value from key.
"""
function get(map::SCHash, key::K) where K
    if map.nkeys == 0 return nothing end

    bucket = map.table[hash(map, key)]
    for (i, v) in enumerate(bucket)
        if key == v[1]
            return v[2]
        end
    end

    return nothing
end

function get(map::SCLHash, key::K) where K
    if map.nkeys == 0 return nothing end

    bucket = map.table[hash(map, key)]
    n = bucket.head
    while n != n.next
        if key == n.value[1]
            return n.value[2]
        end
    end

    return nothing
end

"""
    put(map, key, value)

Add value to hash table with key.
"""
function put!(map::SCHash, key::K, value::V) where {K, V}
    bucket = map.table[hash(map, key)]
    for (i, v) in enumerate(bucket)
        if key == v[1]
            bucket[i] = Pair(key, value)
            return nothing
        end
    end
    push!(bucket, Pair(key, value))
    map.nkeys += 1

    #if map.nkeys > 8 * map.tablesize resize(map) end

    return nothing
end

function put!(map::SCLHash, key::K, value::V) where {K, V}
    bucket = map.table[hash(map, key)]
    n = bucket.head
    while n != n.next
        if key == n.value[1]
            n.value = Pair(key, value)
            return nothing
        end
        n = n.next
    end
    addlast!(bucket, Pair(key, value))
    map.nkeys += 1
    map.tablesize = length(map.table)

    return nothing
end

"""
    delete!(map, key)

Delete kay value pair from map.
"""
function delete!(map::SCHash, key::K)  where K
    bucket = map.table[hash(map, key)]
    for (i, v) in enumerate(bucket)
        if key == v[1]
            bucket[i][2] = value
            return nothing
        end
    end

    return nothing
end

function delete!(map::SCLHash, key::K)  where K
    h = hash(map, key)
    bucket = map.table[h]
    n = bucket.head
    while n != n.next
        if key == n.value[1]
            n.prev.next = n.next
            n.next.prev = n.prev
            return nothing
        end
    end

    return nothing
end

"""
    resize(map, factor)

Resize the map. Resizing should be to a close prime (not implemented).
"""
function resize(map::SCHash{K, V}, factor) where {K,V}
    oldtable = map.table
    map.table = Vector{Vector{Pair{K, V}}}(undef, factor * map.tablesize)
    for bucket in oldtable
        for i in bucket
            put!(map, i)
        end
    end
end

function resize(map::SCLHash{K, V}, factor) where {K,V}
    oldtable = map.table
    map.table = Vector{DLinkedList{Pair{K, V}}}(undef, factor * map.tablesize)
    for bucket in oldtable
        n = bucket.head
        while n != n.next
            put!(map, n.value)
            n = n.next
        end
    end
end

"""
    OAHash

Open addressing hash table.
"""
struct OAHash
    table
end


@testset "Bounding boxes" begin
  # sphere
  @test boundingbox(Sphere(P2(0,0),T(1))) == Box(P2(-1,-1),P2(1,1))
  @test boundingbox(Sphere(P2(1,1),T(1))) == Box(P2(0,0),P2(2,2))

  # boxes
  b1 = Box(P2(0,0), P2(1,1))
  b2 = Box(P2(-1,-1), P2(0.5,0.5))
  @test boundingbox([b1, b2]) == Box(P2(-1,-1), P2(1,1))

  # geometries
  b = Box(P2(-3,-1), P2(0.5,0.5))
  s = Sphere(P2(0,0), T(2))
  @test boundingbox([b, s]) == Box(P2(-3,-2), P2(2,2))

  # point set
  @test boundingbox(PointSet(T[0 1 2; 0 2 1])) == Box(P2(0,0), P2(2,2))
  @test boundingbox(PointSet(T[1 2; 2 1])) == Box(P2(1,1), P2(2,2))

  # cartesian grid
  @test boundingbox(CartesianGrid{T}(10,10)) == Box(P2(0,0), P2(10,10))
  @test boundingbox(CartesianGrid{T}(100,200)) == Box(P2(0,0), P2(100,200))
  @test boundingbox(CartesianGrid((10,10), T.((1,1)), T.((1,1)))) == Box(P2(1,1), P2(11,11))

  # point set views
  d = PointSet(T[0 1 2; 0 2 1])
  v = view(d, 1:2)
  @test boundingbox(v) == Box(P2(0,0), P2(1,2))

  # cartesian grid views
  d = CartesianGrid{T}(10,10)
  v = view(d, 1:2)
  @test boundingbox(v) == Box(P2(0,0), P2(2,1))
end


abstract type AbstractTransition end

"""
    DiscreteTransition{RT, GT, IT⁻, IT⁺, WT} <: AbstractTransition

Type that encodes a discrete transition with an affine assignment of the form:

```math
    post_d(X) = (R(X ∩ G ∩ I⁻) ⊕ W) ∩ I⁺
```
where ``I⁻`` and ``I⁺``  are invariants at the source and the target locations
respectively, ``G ⊆ \\mathbb{R}^n`` is the guard set of the transition, the assignment
is of the form ``x^+ := Rx + w``, ``w ∈ W``, ``x^+ ∈ \\mathbb{R}^m`` are the values after
the transition, ``R ∈ \\mathbb{R}^{m\\times n}`` is the assignment map (or reset map)
and ``W ⊆ \\mathbb{R}^m`` is a closed and bounded convex set of non-deterministic inputs.

### Fields

- `R`  -- assignment map
- `W`  -- non-deterministic inputs
- `G`  -- guard set of the transition from the source location to the target location
- `I⁻` -- invariant at the source location
- `I⁺` -- invariant at the target location
"""
struct DiscreteTransition{RT, WT, GT, IT⁻, IT⁺} <: AbstractTransition
    R::RT
    W::WT
    G::GT
    I⁻::IT⁻
    I⁺::IT⁺
end

reset_map(tr::DiscreteTransition{<:AbstractMatrix})  = ConstrainedAffineMap(tr.R, tr.W)
reset_map(tr::DiscreteTransition{<:IdentityMap}) = ConstrainedIdentityMap(dim(tr.W), tr.W)
HybridSystems.guard(tr::DiscreteTransition) = tr.G
source_invariant(tr::DiscreteTransition) = tr.I⁻
target_invariant(tr::DiscreteTransition) = tr.I⁺

# constructor when the reset map is the identity
function DiscreteTransition(; guard::GT, source_invariant::IT⁻, target_invariant::IT⁺) where {GT, IT⁻, IT⁺}
    n = dim(source_invariant)
    return DiscreteTransition(IdentityMap(n), ZeroSet(n), guard, source_invariant, target_invariant)
end

# constructor from hybrid system given source and target modes id's and the transition t
function DiscreteTransition(H::HybridSystem, transition)
    assignment = resetmap(H, transition)
    R = _state_matrix(assignment)
    W = _affine_term(assignment)
    G = stateset(assignment) # HybridSystems.guard(H, transition) doesn't work
    I⁻ = stateset(H, source(H, transition))
    I⁺ = stateset(H, target(H, transition))
    return DiscreteTransition(R, W, G, I⁻, I⁺)
end

# -----------------------------------------------
# Getter functions for MathematicalSystems maps
# (see also MathematicalSystems#209)
# -----------------------------------------------

_state_matrix(m::IdentityMap) = m       # no-op, to dispatch on identity assignments
_state_matrix(m::ConstrainedIdentityMap) = IdentityMap(statedim(m))
_state_matrix(m::MathematicalSystems.LinearMap) = m.A
_state_matrix(m::ConstrainedLinearMap) = m.A
_state_matrix(m::MathematicalSystems.AffineMap) = m.A
_state_matrix(m::ConstrainedAffineMap) = m.A

_affine_term(m::IdentityMap) = ZeroSet(statedim(m)) # there is no numeric type in these maps; use ZeroSet default
_affine_term(m::ConstrainedIdentityMap) = ZeroSet(statedim(m))
_affine_term(m::MathematicalSystems.LinearMap{N}) where {N} = ZeroSet{N}(statedim(m))
_affine_term(m::ConstrainedLinearMap{N}) where {N} = ZeroSet{N}(statedim(m))
_affine_term(m::MathematicalSystems.AffineMap{N, <:AbstractVector}) where {N} = Singleton(m.c)
_affine_term(m::MathematicalSystems.AffineMap{N, <:LazySet}) where {N} = m.c
_affine_term(m::ConstrainedAffineMap{N, MT, VT, <:AbstractVector}) where {N, MT, VT} = Singleton(m.c)
_affine_term(m::ConstrainedAffineMap{N, MT, VT, <:LazySet}) where {N, MT, VT} = Singleton(m.c)

LazySets.dim(c::ConstrainedResetMap) = statedim(c)

function _state_matrix(m::Union{<:MathematicalSystems.ResetMap{N},
                                <:MathematicalSystems.ConstrainedResetMap{N}}) where {N}
    n = dim(m)
    v = ones(N, n)
    for i in keys(m.dict)
        v[i] = zero(N)
    end
    return Diagonal(v)
end

function _affine_term(m::Union{<:MathematicalSystems.ResetMap{N},
                               <:MathematicalSystems.ConstrainedResetMap{N}}) where {N}
    n = dim(m)
    b = sparsevec(Int[], N[], n)
    for (i, val) in m.dict
        b[i] = val
    end
    return Singleton(b)
end

# -------------------------------------------
# Applying reset maps
# -------------------------------------------

@inline function _apply_reset(R::AbstractMatrix, P::HPolytope, W::ZeroSet, ::HRepIntersection)
    return linear_map(R, P) # TODO pass options here
end

@inline function _apply_reset(R::AbstractMatrix, P::HPolytope, W::AbstractVector, ::HRepIntersection)
    return affine_map(R, P, W) # TODO pass options here
end

@inline function _apply_reset(R::AbstractMatrix, P::HPolytope, W::LazySet, ::HRepIntersection)
    Ylm = linear_map(R, P) # TODO pass options here
    return minkowski_sum(Ylm, W)
end

# we assume that the reach-set is bounded: it's either a polytopic set
# or a union of polytopic sets
# => the return type is either the empty set or an HPolytope unless otherwise stated
# for examples in cases involving unions this method may return a union of polytopes
apply(tr::DiscreteTransition,
      R::AbstractLazyReachSet,
      method::AbstractIntersectionMethod) = apply(tr, set(R), method)

# ==========================================
# Intersection method: HRepIntersection
# ==========================================

# -------------------------------------------
# Cases where the reset map is the identity
# -------------------------------------------

# all sets are polyhedral
function  apply(tr::DiscreteTransition{<:IdentityMap, <:ZeroSet, GT, IT⁻, IT⁺},
                X::AbstractPolytope{N},
                method::HRepIntersection) where {N,
                                GT<:AbstractPolyhedron{N},
                                IT⁻<:AbstractPolyhedron{N},
                                IT⁺<:AbstractPolyhedron{N}}

    success, out = _intersection(X, tr.G, tr.I⁻, tr.I⁺, method)
    return success ? HPolytope(out) : EmptySet(dim(X))
end

# the source invariant is the set union of half-spaces, the other sets are polyhedral
function apply(tr::DiscreteTransition{<:IdentityMap, <:ZeroSet, GT, IT⁻, IT⁺},
               X::AbstractPolytope{N},
               method::HRepIntersection) where {N,
                                GT<:AbstractPolyhedron{N},
                                VN,
                                IT⁻<:UnionSetArray{N, LinearConstraint{N, VN}},
                                IT⁺<:AbstractPolyhedron{N}}

    success, aux = _intersection(X, tr.G, tr.I⁺, method)
    !success && return EmptySet(dim(X))

    clist_I⁻ = tr.I⁻.array
    out = Vector{HPolytope{N, Vector{N}}}()

    # intersect Y := X ∩ G ∩ I⁺ with each half-space ci in I⁻
    # only store those polytpes for which Y ∩ ci is non-empty
    for ci in clist_I⁻
        Y_ci = vcat(aux, ci)
        success = remove_redundant_constraints!(Y_ci)
        !success && continue
        push!(out, HPolytope(copy(Y_ci)))
    end
    return isempty(out) ? EmptySet(dim(X)) : UnionSetArray(out)
end

# the source invariant is the set union of half-spaces, the other sets are polyhedral
# and the set X is the convex hull of polytopes
function apply(tr::DiscreteTransition{<:IdentityMap, <:ZeroSet, GT, IT⁻, IT⁺},
               X::ConvexHullArray{N, PT},
               method::HRepIntersection) where {N,
                                PT<:AbstractPolytope{N},
                                GT<:AbstractPolyhedron{N},
                                IT⁻<:AbstractPolyhedron{N},
                                IT⁺<:AbstractPolyhedron{N}}

    vlist = vertices_list(X) # TODO pass the backend as an option
    Xhrep = tohrep(vlist)
    return apply(tr, Xhrep, method)
end

# the guard is the set union of half-spaces, the source invariant is polyhedral,
# and the target invariant is universal
function apply(tr::DiscreteTransition{<:IdentityMap, <:ZeroSet, GT, IT⁻, IT⁺},
                  X::AbstractPolytope{N},
                  method::HRepIntersection) where {N, VN,
                                                   GT<:UnionSetArray{N, HalfSpace{N, VN}},
                                                   IT⁻<:AbstractPolyhedron{N},
                                                   IT⁺<:Universe{N}}

    success, aux = _intersection(X, tr.I⁻, method)
    !success && return EmptySet(dim(X))

    clist_G = tr.G.array
    out = Vector{HPolytope{N, Vector{N}}}()

    # intersect Y := X ∩ I⁻ with each half-space ci in G
    # only store those polytpes for which Y ∩ ci is non-empty
    for ci in clist_G
        Y_ci = vcat(aux, ci)
        success = remove_redundant_constraints!(Y_ci)
        !success && continue
        push!(out, HPolytope(copy(Y_ci)))
    end

    return !isempty(out) ? UnionSetArray(out) : EmptySet(dim(X))
end

# -------------------------------------------
# Cases where the reset map is RX ⊕ W
# -------------------------------------------

function _apply_hrep(R, W, G, I⁻, I⁺, X, method::HRepIntersection)
    success, aux = _intersection(X, G, I⁻, method)
    !success && return EmptySet(dim(X))

    P = HPolytope(aux)
    Y = _apply_reset(R, P, W, method)

    success, out = _intersection(Y, I⁺, method)
    return success ? HPolytope(out) : EmptySet(dim(X))
end

# cases with W = 0
function apply(tr::DiscreteTransition{<:AbstractMatrix, <:ZeroSet, GT, IT⁻, IT⁺},
               X::AbstractPolytope{N},
               method::HRepIntersection) where {N, GT<:AbstractPolyhedron{N},
                                                   IT⁻<:AbstractPolyhedron{N},
                                                   IT⁺<:AbstractPolyhedron{N}}
    return _apply_hrep(tr.R, tr.W, tr.G, tr.I⁻, tr.I⁺, X, method)
end

# cases where the inputs are deterministic (W is a vector)
function apply(tr::DiscreteTransition{<:AbstractMatrix, <:AbstractVector, GT, IT⁻, IT⁺},
               X::AbstractPolytope{N},
               method::HRepIntersection) where {N, GT<:AbstractPolyhedron{N},
                                                   IT⁻<:AbstractPolyhedron{N},
                                                   IT⁺<:AbstractPolyhedron{N}}
    return _apply_hrep(tr.R, tr.W, tr.G, tr.I⁻, tr.I⁺, X, method)
end

# cases where the inputs are non-deterministic (W is a set)
function apply(tr::DiscreteTransition{<:AbstractMatrix, <:LazySet, GT, IT⁻, IT⁺},
               X::AbstractPolytope{N},
               method::HRepIntersection) where {N, GT<:AbstractPolyhedron{N},
                                                   IT⁻<:AbstractPolyhedron{N},
                                                   IT⁺<:AbstractPolyhedron{N}}
    return _apply_hrep(tr.R, tr.W, tr.G, tr.I⁻, tr.I⁺, X, method)
end

# ==========================================
# Intersection method: BoxIntersection
# ==========================================

function apply(tr::DiscreteTransition{<:IdentityMap, <:ZeroSet, GT, IT⁻, IT⁺},
                X::AbstractPolytope{N},
                ::BoxIntersection) where {N,
                                       GT<:AbstractPolyhedron{N},
                                       VN,
                                       IT⁻<:UnionSetArray{N, LinearConstraint{N, VN}},
                                       IT⁺<:AbstractPolyhedron{N}}

    out = apply(tr, X, HRepIntersection())
    return overapproximate(out, Hyperrectangle)
end

# the guard is the set union of half-spaces, the source invariant is polyhedral,
# and the target invariant is the set union of half-spaces
function apply(tr::DiscreteTransition{<:IdentityMap, <:ZeroSet, GT, IT⁻, IT⁺},
                  X::AbstractPolytope{N},
                  method::BoxIntersection) where {N, VN,
                                                  GT<:UnionSetArray{N, HalfSpace{N, VN}},
                                                  IT⁻<:AbstractPolyhedron{N},
                                                  IT⁺<:Universe{N}}

    out = apply(tr, X, HRepIntersection())
    return overapproximate(out, Hyperrectangle)
end

#=
function _apply(tr::DiscreteTransition{<:IdentityMap, <:ZeroSet, GT, IT⁻, IT⁺},
                X::ConvexHullArray{N, PT},
                ::BoxEnclosure) where {N,
                                       PT<:AbstractPolyhedron{N},
                                       GT<:AbstractPolyhedron{N},
                                       VN,
                                       IT⁻<:UnionSetArray{N, LinearConstraint{N, VN}},
                                       IT⁺<:AbstractPolyhedron{N}}
    Xbox = overapproximate(X, Hyperrectangle)
    out = _apply(tr, Xbox, HRepIntersection())
    return overapproximate(out, Hyperrectangle)
end
=#

#=

function _apply(tr::DiscreteTransition{<:Universe, <:ZeroSet, GT, IT⁻, IT⁺},
                X::AbstractPolyhedron{N},
                ::BoxEnclosure) where {N,
                                       GT<:AbstractPolyhedron{N},
                                       VN,
                                       IT⁻<:UnionSetArray{N, LinearConstraint{N, VN}},
                                       IT⁺<:AbstractPolyhedron{N}}
    out = _apply(tr, X, HRepIntersection())
    return overapproximate(out, Hyperrectangle)
end

function _apply(tr::DiscreteTransition{<:Universe, <:ZeroSet, GT, IT⁻, IT⁺},
                X::ConvexHullArray{N, PT},
                ::BoxEnclosure) where {N,
                                       PT<:AbstractPolyhedron{N},
                                       GT<:AbstractPolyhedron{N},
                                       IT⁻<:AbstractPolyhedron{N},
                                       IT⁺<:AbstractPolyhedron{N}}

    Xbox = overapproximate(X, Hyperrectangle)
    out = _apply(tr, Xbox, HRepIntersection())
    return overapproximate(out, Hyperrectangle)
end
=#

# ===============================================
# Intersection method: TemplateHullIntersection
# ===============================================

# compute (R(X ∩ G ∩ I⁻) ⊕ W) ∩ I⁺ first computing the exact intersection
# of Y := X ∩ G ∩ I⁻ all of which are polyhedral, using their list of constraints
# then we compute (RY ⊕ W) ∩ I⁺ using again the given template, Z := [(RY ⊕ W) ∩ I⁺]_dirs
function apply(tr::DiscreteTransition{<:AbstractMatrix, <:LazySet, GT, IT⁻, IT⁺},
               X::PT,
               method::TemplateHullIntersection) where {N,
                                PT<:AbstractPolyhedron{N},
                                GT<:AbstractPolyhedron{N},
                                IT⁻<:AbstractPolyhedron{N},
                                IT⁺<:AbstractPolyhedron{N}}

    # intersect X with the guard and the source invariant (NOTE can be cached)
    success, Y = _intersection(X, tr.G, tr.I⁻, HRepIntersection())
    !success && return EmptySet(dim(X))

    # compute the intersection [(RY ⊕ W) ∩ I⁺]_dirs using the template
    K = (tr.R * HPolyhedron(Y)) ⊕ tr.W # lazy affine map

    Z = _intersection(K, tr.I⁺, method)
    return Z
end

# compute (R(X ∩ G ∩ I⁻) ⊕ W) ∩ I⁺ first computing the exact intersection
# of Y := G ∩ I⁻ all of which are polyhedral, then compute K := [X ∩ Y]_dirs,
# finally we compute (RK ⊕ W) ∩ I⁺ using again the given template, Z := [(RK ⊕ W) ∩ I⁺]_dirs
function apply(tr::DiscreteTransition{<:AbstractMatrix, <:LazySet, GT, IT⁻, IT⁺},
               X::PT,
               method::TemplateHullIntersection) where {N,
                                PT<:LazySet{N},
                                GT<:AbstractPolyhedron{N},
                                IT⁻<:AbstractPolyhedron{N},
                                IT⁺<:AbstractPolyhedron{N}}

    # intersect the guard and the source invariant
    success, Y = _intersection(tr.G, tr.I⁻, HRepIntersection())
    !success && return EmptySet(dim(X))

    K = _intersection(X, HPolyhedron(Y), method)
    isempty(K) && return EmptySet(dim(X))

    # compute Z := [(RK ⊕ W) ∩ I⁺]_dirs using the template
    Km = (tr.R * K) ⊕ tr.W # lazy affine map
    Z = _intersection(Km, tr.I⁺, method)
    return Z
end

# compute X ∩ G ∩ I⁻ ∩ I⁺ first computing the exact intersection
# of Y := G ∩ I⁻ ∩ I⁺ all of which are polyhedral, then compute K := [X ∩ Y]_dirs
function apply(tr::DiscreteTransition{<:IdentityMap, <:ZeroSet, GT, IT⁻, IT⁺},
               X::PT,
               method::TemplateHullIntersection) where {N,
                                PT<:LazySet{N},
                                GT<:AbstractPolyhedron{N},
                                IT⁻<:AbstractPolyhedron{N},
                                IT⁺<:AbstractPolyhedron{N}}

    # intersect the guard and the source invariant
    success, Y = _intersection(tr.G, tr.I⁻, tr.I⁺, HRepIntersection())
    !success && return EmptySet(dim(X))

    # compute the intersection K := [X ∩ Y]_dirs using the template
    K = _intersection(X, HPolyhedron(Y), method)
    return K
end

# ======================================
# Template hulls with set unions
#
# TODO: now X is a vector of reach-sets,
# we may dispatch on UnionSetArray later
# ==============================

# general case, requires constraints list of each Xi
function apply(tr::DiscreteTransition{<:AbstractMatrix, <:LazySet, GT, IT⁻, IT⁺},
               X::AbstractVector{RT},
               method::TemplateHullIntersection) where {N,
                                RT<:AbstractReachSet{N},
                                GT<:AbstractPolyhedron{N},
                                IT⁻<:AbstractPolyhedron{N},
                                IT⁺<:AbstractPolyhedron{N}}

    @unpack R, W, G, I⁻, I⁺ = tr
    got_inv = LazySets.Arrays.isinvertible(R)

    m = length(X)
    Y = Vector{HPolytope{N, Vector{N}}}(undef, m)
    for i in 1:m
        Xi = set(X[i])

        # concrete Xi ∩ G ∩ I⁻
        success, Qi = _intersection(Xi, G, I⁻, HRepIntersection())
        !success && return EmptySet(dim(Xi))

        if got_inv
            # concretely compute R*Qi then lazily add W and lazily intersect with I⁺
            E = linear_map(R, HPolytope(Qi)) ⊕ W
            Y[i] = _intersection(E, I⁺, method)
        else
            Y[i] = _intersection(R * HPolytope(Qi) ⊕ W, I⁺, method)
        end
    end
    return overapproximate(ConvexHullArray(Y), method.dirs)
end

# case with R generic and W = 0
# requires constraints list of X
function apply(tr::DiscreteTransition{<:AbstractMatrix, <:ZeroSet, GT, IT⁻, IT⁺},
               X::AbstractVector{RT},
               method::TemplateHullIntersection) where {N,
                                RT<:AbstractReachSet{N},
                                GT<:AbstractPolyhedron{N},
                                IT⁻<:AbstractPolyhedron{N},
                                IT⁺<:AbstractPolyhedron{N}}

    @unpack R, W, G, I⁻, I⁺ = tr
    got_inv = LazySets.Arrays.isinvertible(R)

    m = length(X)
    Y = Vector{HPolytope{N, Vector{N}}}(undef, m)
    for i in 1:m
        Xi = set(X[i])

        # concrete Xi ∩ G ∩ I⁻
        success, Qi = _intersection(Xi, G, I⁻, HRepIntersection())
        !success && return EmptySet(dim(Xi))

        # lazily or concretely compute R*Qi ∩ I⁺, depending on the invertibility of R
        if got_inv
            E = linear_map(R, HPolytope(Qi))
            success, Yi = _intersection(E, I⁺, HRepIntersection())
            !success && return EmptySet(dim(Xi))
            Y[i] = HPolytope(Yi)
        else
            Y[i] = _intersection(R * HPolytope(Qi), I⁺, method)
        end
    end
    return overapproximate(ConvexHullArray(Y), method.dirs)
end

# the R = Id and W = 0
function apply(tr::DiscreteTransition{<:IdentityMap, <:ZeroSet, GT, IT⁻, IT⁺},
               X::AbstractVector{RT},
               method::TemplateHullIntersection) where {N,
                                RT<:AbstractReachSet{N},
                                GT<:AbstractPolyhedron{N},
                                IT⁻<:AbstractPolyhedron{N},
                                IT⁺<:AbstractPolyhedron{N}}

    @unpack R, W, G, I⁻, I⁺ = tr

    # precompute L = G ∩ I⁻ ∩ I⁺
    success, L = _intersection(G, I⁻, I⁺, HRepIntersection())
    !success && return EmptySet(dim(X))

    m = length(X)
    Y = Vector{HPolytope{N, Vector{N}}}(undef, m)

    for i in 1:m
        Xi = set(X[i])

        # concrete intersection of Xi ∩ L
        success, Qi = _intersection(Xi, L, HRepIntersection())
        !success && return EmptySet(dim(Xi))
        Y[i] = HPolytope(Qi)
    end
    return overapproximate(ConvexHullArray(Y), method.dirs)
end

# the R = Id and W = 0
# and AbstractZonotope => use template directions for the intersection
function apply(tr::DiscreteTransition{<:IdentityMap, <:ZeroSet, GT, IT⁻, IT⁺},
               X::AbstractVector{RT},
               method::TemplateHullIntersection) where {N,
                                ZT<:AbstractZonotope{N},
                                RT<:ReachSet{N, ZT},
                                GT<:AbstractPolyhedron{N},
                                IT⁻<:AbstractPolyhedron{N},
                                IT⁺<:AbstractPolyhedron{N}}

    @unpack R, W, G, I⁻, I⁺ = tr

    # precompute L = G ∩ I⁻ ∩ I⁺
    success, L = _intersection(G, I⁻, I⁺, HRepIntersection())
    !success && return EmptySet(dim(X))
    L = HPolyhedron(L)

    m = length(X)
    Y = Vector{HPolytope{N, Vector{N}}}(undef, m)

    for i in 1:m
        Xi = set(X[i])

        # concrete intersection of Xi ∩ L
        Y[i] = _intersection(Xi, L, method)
    end

    return overapproximate(ConvexHullArray(Y), method.dirs)
end


using Plots
plotlyjs()

# Rossler space equation
function rossler(t, state, a, b, c)
    x, y, z = state
    dx = -y - z
    dy = x + a*y
    dz = b + z*(x-c)
    return [dx, dy, dz]
end

"""
    rk4(f, x0, time_span[, h, args])

4-th order Runge-Kutta. ``f`` is a handler to the
space state funtion, ``x0`` is a vector of initial
conditions, ``time_span`` is a tuple containing
the start and end simulation times and ``h`` is the
integration interval.

The function ``f`` is called as:

    dx = f(t, x, args...)
"""
function rk4_old(f, x0, time_span; h=(time_span[2]-time_span[1])/1e4, args=())
    t = time_span[1]:h:time_span[2]
    N = length(t)
    n = length(x0)
    # Allocate buffers
    x = zeros(N, n)
    xd = zeros(4, n)
    # Simulate
    x[1, :] = x0
    for i = 1:N-1
        # 1-st evaluation
        xd[1, :] = f(t[i], x[i, :], args...)
        # 2-nd evaluation
        xd[2, :] = f(t[i]+h/2, x[i, :]+h/2*xd[1, :], args...)
        # 3-rd evaluation
        xd[3, :] = f(t[i]+h/2, x[i, :]+h/2*xd[2, :], args...)
        # 4-th evaluation
        xd[4, :] = f(t[i]+h, x[i, :]+h*xd[3, :], args...)
        # Next point
        x[i+1, :] = x[i, :]
        x[i+1, :] += (xd[1, :] + 2*xd[2, :] + 2*xd[3, :] + xd[4, :]) * h/6
    end
    return x
end

# Simulate Rossler
time_span = (0, 1000)
x0 = [1.0, 1.0, 2.0]
args = (0.2, 0.2, 5.7)

@time x = rk4_old(rossler, x0, time_span, h=0.001, args=args)

plot(x)

plot(x[:, 1], x[:,2], x[:, 3])


#=
    run_tests
    Copyright © 2019 Mark Wells <mwellsa@gmail.com>

    Distributed under terms of the MIT license.
=#

using InterpolatedRejectionSampling
using Test

import Random.seed!
seed!(1234)

@testset "irsample" begin

    X = range(0,π,length=5)
    Y = range(0,π/4,length=4)
    knots = (X,Y)
    prob = [sin(x)+tan(y) for x in X, y in Y]

    n = 100
    xy = irsample(knots, prob, n)
    @test isa(xy, Matrix{Float64})
    @test size(xy) == (2,n)

    xy = Matrix{Union{Float64,Missing}}(missing,2,n)
    irsample!(xy, knots, prob)
    @test isa(xy, Matrix{Union{Missing,Float64}})
    @test size(xy) == (2,n)
    @test iszero(count(ismissing.(xy)))

    x = irsample(X, sin.(X), n)
    @test isa(x, Vector{Float64})
    @test length(x) == n
end


using ManifoldsBase

using LinearAlgebra
using DoubleFloats
using ForwardDiff
using ReverseDiff
using StaticArrays
using Test

struct CustomDefinedRetraction <: ManifoldsBase.AbstractRetractionMethod end
struct CustomUndefinedRetraction <: ManifoldsBase.AbstractRetractionMethod end

function ManifoldsBase.injectivity_radius(
    ::ManifoldsBase.DefaultManifold,
    ::CustomDefinedRetraction,
)
    return 10.0
end

struct MatrixVectorTransport{T} <: AbstractVector{T}
    m::Matrix{T}
end

Base.getindex(x::MatrixVectorTransport, i) = x.m[:, i]

Base.size(x::MatrixVectorTransport) = (size(x.m, 2),)

@testset "Testing Default (Euclidean)" begin
    M = ManifoldsBase.DefaultManifold(3)
    types = [
        Vector{Float64},
        SizedVector{3,Float64},
        MVector{3,Float64},
        Vector{Float32},
        SizedVector{3,Float32},
        MVector{3,Float32},
        Vector{Double64},
        MVector{3,Double64},
        SizedVector{3,Double64},
    ]

    @test repr(M) == "DefaultManifold(3; field = ℝ)"
    @test isa(manifold_dimension(M), Integer)
    @test manifold_dimension(M) ≥ 0
    @test base_manifold(M) == M
    @test number_system(M) == ℝ
    @test ManifoldsBase.representation_size(M) == (3,)

    @test injectivity_radius(M) == Inf

    rm = ManifoldsBase.ExponentialRetraction()
    irm = ManifoldsBase.LogarithmicInverseRetraction()

    rm2 = CustomDefinedRetraction()
    rm3 = CustomUndefinedRetraction()

    for T in types
        @testset "Type $T" begin
            pts = convert.(Ref(T), [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])

            @test injectivity_radius(M, pts[1]) == Inf
            @test injectivity_radius(M, pts[1], rm) == Inf
            @test injectivity_radius(M, rm) == Inf
            @test injectivity_radius(M, rm2) == 10
            @test injectivity_radius(M, pts[1], rm2) == 10
            @test_throws ErrorException injectivity_radius(M, rm3)
            @test_throws ErrorException injectivity_radius(M, pts[1], rm3)

            tv1 = log(M, pts[1], pts[2])

            for pt in pts
                @test is_point(M, pt)
            end
            @test is_vector(M, pts[1], tv1; atol = eps(eltype(pts[1])))

            tv2 = log(M, pts[2], pts[1])
            @test isapprox(M, pts[2], exp(M, pts[1], tv1))
            @test isapprox(M, pts[1], exp(M, pts[1], tv1, 0))
            @test isapprox(M, pts[2], exp(M, pts[1], tv1, 1))
            @test isapprox(M, pts[1], exp(M, pts[2], tv2))
            @test is_point(M, retract(M, pts[1], tv1))
            @test isapprox(M, pts[1], retract(M, pts[1], tv1, 0))

            @test is_point(M, retract(M, pts[1], tv1, rm))
            @test isapprox(M, pts[1], retract(M, pts[1], tv1, 0, rm))

            new_pt = exp(M, pts[1], tv1)
            retract!(M, new_pt, pts[1], tv1)
            @test is_point(M, new_pt)
            for x in pts
                @test isapprox(M, zero_vector(M, x), log(M, x, x); atol = eps(eltype(x)))
                @test isapprox(
                    M,
                    zero_vector(M, x),
                    inverse_retract(M, x, x);
                    atol = eps(eltype(x)),
                )
                @test isapprox(
                    M,
                    zero_vector(M, x),
                    inverse_retract(M, x, x, irm);
                    atol = eps(eltype(x)),
                )
            end
            zero_vector!(M, tv1, pts[1])
            @test isapprox(M, pts[1], tv1, zero_vector(M, pts[1]))
            log!(M, tv1, pts[1], pts[2])
            @test norm(M, pts[1], tv1) ≈ sqrt(inner(M, pts[1], tv1, tv1))

            @test isapprox(M, exp(M, pts[1], tv1, 1), pts[2])
            @test isapprox(M, exp(M, pts[1], tv1, 0), pts[1])

            @test distance(M, pts[1], pts[2]) ≈ norm(M, pts[1], tv1)

            @test mid_point(M, pts[1], pts[2]) == [0.5, 0.5, 0.0]
            midp = allocate(pts[1])
            @test mid_point!(M, midp, pts[1], pts[2]) === midp
            @test midp == [0.5, 0.5, 0.0]

            @testset "Geodesic interface test" begin
                @test isapprox(M, geodesic(M, pts[1], tv1)(0.0), pts[1])
                @test isapprox(M, geodesic(M, pts[1], tv1)(1.0), pts[2])
                @test isapprox(M, geodesic(M, pts[1], tv1, 1.0), pts[2])
                @test isapprox(M, geodesic(M, pts[1], tv1, 1.0 / 2), (pts[1] + pts[2]) / 2)
                @test isapprox(M, shortest_geodesic(M, pts[1], pts[2])(0.0), pts[1])
                @test isapprox(M, shortest_geodesic(M, pts[1], pts[2])(1.0), pts[2])
                @test isapprox(M, shortest_geodesic(M, pts[1], pts[2], 0.0), pts[1])
                @test isapprox(M, shortest_geodesic(M, pts[1], pts[2], 1.0), pts[2])
                @test all(
                    isapprox.(
                        Ref(M),
                        geodesic(M, pts[1], tv1, [0.0, 1.0 / 2, 1.0]),
                        [pts[1], (pts[1] + pts[2]) / 2, pts[2]],
                    ),
                )
                @test all(
                    isapprox.(
                        Ref(M),
                        shortest_geodesic(M, pts[1], pts[2], [0.0, 1.0 / 2, 1.0]),
                        [pts[1], (pts[1] + pts[2]) / 2, pts[2]],
                    ),
                )
            end

            @testset "basic linear algebra in tangent space" begin
                @test isapprox(
                    M,
                    pts[1],
                    0 * tv1,
                    zero_vector(M, pts[1]);
                    atol = eps(eltype(pts[1])),
                )
                @test isapprox(M, pts[1], 2 * tv1, tv1 + tv1)
                @test isapprox(M, pts[1], 0 * tv1, tv1 - tv1)
                @test isapprox(M, pts[1], (-1) * tv1, -tv1)
            end

            @testset "Hat and vee in the tangent space" begin
                X = log(M, pts[1], pts[2])
                a = vee(M, pts[1], X)
                b = similar(a)
                vee!(M, b, pts[1], X)
                Y = hat(M, pts[1], a)
                Z = similar(Y)
                hat!(M, Z, pts[1], a)
                @test a == b
                @test X == Y
                @test Z == X
                @test a == vec(X)
            end

            @testset "broadcasted linear algebra in tangent space" begin
                @test isapprox(M, pts[1], 3 * tv1, 2 .* tv1 .+ tv1)
                @test isapprox(M, pts[1], -tv1, tv1 .- 2 .* tv1)
                @test isapprox(M, pts[1], -tv1, .-tv1)
                v = similar(tv1)
                v .= 2 .* tv1 .+ tv1
                @test v ≈ 3 * tv1
            end

            @testset "project test" begin
                # point
                @test isapprox(M, pts[1], project(M, pts[1]))
                pt = similar(pts[1])
                project!(M, pt, pts[1])
                @test isapprox(M, pt, pts[1])

                @test isapprox(M, pts[1], embed(M, pts[1]))
                pt = similar(pts[1])
                embed!(M, pt, pts[1])
                @test isapprox(M, pt, pts[1])

                # tangents
                @test isapprox(M, pts[1], tv1, project(M, pts[1], tv1))
                tv = similar(tv1)
                project!(M, tv, pts[1], tv1)
                @test isapprox(M, pts[1], tv, tv1)

                @test isapprox(M, pts[1], tv1, embed(M, pts[1], tv1))
                tv = similar(tv1)
                embed!(M, tv, pts[1], tv1)
                @test isapprox(M, pts[1], tv, tv1)
            end

            @testset "vector transport" begin
                # test constructor and alias
                @test DifferentiatedRetractionVectorTransport(ExponentialRetraction()) ==
                      ParallelTransport()
                v1 = log(M, pts[1], pts[2])
                v2 = log(M, pts[1], pts[3])
                v1t1 = vector_transport_to(M, pts[1], v1, pts[3])
                v1t2 = zero(v1t1)
                vector_transport_to!(M, v1t2, pts[1], v1, v2, ProjectionTransport())
                v1t3 = vector_transport_direction(M, pts[1], v1, v2)
                @test is_vector(M, pts[3], v1t1)
                @test is_vector(M, pts[3], v1t3)
                @test isapprox(M, pts[3], v1t1, v1t3)
                # along a `Vector` of points
                c = [pts[1]]
                v1t4 = vector_transport_along(M, pts[1], v1, c)
                @test isapprox(M, pts[1], v1, v1t4)
                v1t5 = allocate(v1)
                vector_transport_along!(M, v1t5, pts[1], v1, c)
                @test isapprox(M, pts[1], v1, v1t5)
                # along a custom type of points
                T = eltype(pts[1])
                c2 = MatrixVectorTransport{T}(reshape(pts[1], length(pts[1]), 1))
                v1t4c2 = vector_transport_along(M, pts[1], v1, c2)
                @test isapprox(M, pts[1], v1, v1t4c2)
                v1t5c2 = allocate(v1)
                vector_transport_along!(M, v1t5c2, pts[1], v1, c2)
                @test isapprox(M, pts[1], v1, v1t5c2)
                # On Euclidean Space Schild & Pole are identity
                @test vector_transport_to(
                    M,
                    pts[1],
                    v2,
                    pts[2],
                    SchildsLadderTransport(),
                ) == v2
                @test vector_transport_to(M, pts[1], v2, pts[2], PoleLadderTransport()) ==
                      v2
                @test vector_transport_to(
                    M,
                    pts[1],
                    v2,
                    pts[2],
                    ScaledVectorTransport(ParallelTransport()),
                ) == v2

                # along is also the identity
                c = [0.5 * (pts[1] + pts[2]), pts[2], 0.5 * (pts[2] + pts[3]), pts[3]]
                @test vector_transport_along(M, pts[1], v2, c, SchildsLadderTransport()) ==
                      v2
                @test vector_transport_along(M, pts[1], v2, c, PoleLadderTransport()) == v2
                @test vector_transport_along(M, pts[1], v2, c, ParallelTransport()) == v2
                # check mutating ones with defaults
                p = allocate(pts[1])
                ManifoldsBase.pole_ladder!(M, p, pts[1], pts[2], pts[3])
                # -log_p3 p == log_p1 p2
                @test isapprox(M, -log(M, pts[3], p), log(M, pts[1], pts[2]))
                ManifoldsBase.schilds_ladder!(M, p, pts[1], pts[2], pts[3])
                @test isapprox(M, log(M, pts[3], p), log(M, pts[1], pts[2]))

                @test repr(ParallelTransport()) == "ParallelTransport()"
                @test repr(ScaledVectorTransport(ParallelTransport())) ==
                      "ScaledVectorTransport(ParallelTransport())"
            end

            @testset "ForwardDiff support" begin
                exp_f(t) = distance(M, pts[1], exp(M, pts[1], t * tv1))
                d12 = distance(M, pts[1], pts[2])
                for t in 0.1:0.1:0.9
                    @test d12 ≈ ForwardDiff.derivative(exp_f, t)
                end

                retract_f(t) = distance(M, pts[1], retract(M, pts[1], t * tv1))
                for t in 0.1:0.1:0.9
                    @test ForwardDiff.derivative(retract_f, t) ≥ 0
                end
            end

            isa(pts[1], Union{Vector,SizedVector}) && @testset "ReverseDiff support" begin
                exp_f(t) = distance(M, pts[1], exp(M, pts[1], t[1] * tv1))
                d12 = distance(M, pts[1], pts[2])
                for t in 0.1:0.1:0.9
                    @test d12 ≈ ReverseDiff.gradient(exp_f, [t])[1]
                end

                retract_f(t) = distance(M, pts[1], retract(M, pts[1], t[1] * tv1))
                for t in 0.1:0.1:0.9
                    @test ReverseDiff.gradient(retract_f, [t])[1] ≥ 0
                end
            end
        end
    end

    @testset "mid_point on 0-index arrays" begin
        M = ManifoldsBase.DefaultManifold(1)
        p1 = fill(0.0)
        p2 = fill(1.0)
        @test isapprox(M, fill(0.5), mid_point(M, p1, p2))
    end

    @testset "Retracion" begin
        a = NLsolveInverseRetraction(ExponentialRetraction())
        @test a.retraction isa ExponentialRetraction
    end

    @testset "copy of points and vectors" begin
        M = DefaultManifold(2)
        p = [2.0, 3.0]
        q = similar(p)
        copyto!(M, q, p)
        @test p == q
        r = copy(M, p)
        @test r == p
        X = [4.0, 5.0]
        Y = similar(X)
        copyto!(M, Y, p, X)
        @test Y == X
        Z = copy(M, p, X)
        @test Z == X
    end
end


abstract type AbstractDomain end

abstract type RealDomain <: AbstractDomain end

struct RealNumbers <: RealDomain end

const ℝ = RealNumbers()

Base.minimum(::RealNumbers) = static(-Inf)
Base.maximum(::RealNumbers) = static(Inf)

Base.in(x, ::RealNumbers) = isreal(x)

Base.show(io::IO, ::typeof(ℝ)) = print(io, "ℝ")

struct BoundedReals{L,U} <: RealDomain
    lower :: L
    upper :: U
end

Base.in(x, b::BoundedReals) = b.lower ≤ x ≤ b.upper

export ℝ, ℝ₊, 𝕀, ℤ

const ℝ₊ = BoundedReals(static(0.0), static(Inf))
const 𝕀 = BoundedReals(static(0.0), static(1.0))

Base.minimum(b::BoundedReals) = b.lower
Base.maximum(b::BoundedReals) = b.upper

Base.show(io::IO, ::typeof(ℝ₊)) = print(io, "ℝ₊")
Base.show(io::IO, ::typeof(𝕀)) = print(io, "𝕀")

testvalue(::typeof(ℝ)) = 0.0
testvalue(::typeof(ℝ₊)) = 1.0
testvalue(::typeof(𝕀)) = 0.5

abstract type IntegerDomain <: AbstractDomain end

struct IntegerNumbers <: IntegerDomain end

Base.in(x, ::IntegerNumbers) = isinteger(x)

const ℤ = IntegerNumbers()

Base.show(io::IO, ::typeof(ℤ)) = print(io, "ℤ")

Base.minimum(::IntegerNumbers) = static(-Inf)
Base.maximum(::IntegerNumbers) = static(Inf)
struct BoundedInts{L,U} <: IntegerDomain
    lower :: L
    upper :: U
end

Base.in(x, b::BoundedInts) = x ∈ ℤ && b.lower ≤ x ≤ b.upper

Base.minimum(b::BoundedInts) = b.lower
Base.maximum(b::BoundedInts) = b.upper

function Base.show(io::IO, b::BoundedInts)
    io = IOContext(io, :compact => true)
    print(io, "ℤ[", b.lower, ":", b.upper, "]")
end

testvalue(b::BoundedInts) = min(b.lower, 0)

function Base.getindex(::typeof(ℤ), r::AbstractUnitRange)
    BoundedInts(extrema(r)...)
end


abstract Initializer

type NormalInit <: Initializer
    mu :: AbstractFloat
    sd :: AbstractFloat
    function NormalInit(mu=0.0, sd=1.0)
        new(mu, sd)
    end
end

type UniformInit <: Initializer
    min :: AbstractFloat
    max :: AbstractFloat
    function UniformInit(min=0.0, max=1.0)
        new(min, max)
    end
end

function initialize(T::DataType, mapdim::Tuple, init::UniformInit)
    data = rand(T, mapdim)
    data -= T(init.min)
    data /= T(init.max-init.min)
end

function initialize(T::DataType, mapdim::Tuple, init::NormalInit)
    data = convert(Array{T},randn(mapdim))
    data *= T(init.sd)
    data += T(init.mu)
end


struct SymmetriesGraph
    graph::SimpleGraph{Int64}
    node_lookup
    reverse_node_lookup
end

"""
    random_symmetric_2players_game(dist::Distributions.Sampleable,p1_size::Int,p2_size::Int)

Returns random symmetric payoff matrices for 2 players with given number of actions

**Input parameters**
* `dist::Distributions.Sampleable` - distribution from which payoffs are sampled
* `p_size::Int` - number of actions of a player
"""
function random_symmetric_2players_game(dist::Distributions.Sampleable,
                               p_size::Int)
    random_data = rand(dist, p_size, p_size)
    data = [random_data, transpose(random_data)]
    return Dict("player"*string(i)=>data[i] for i in 1:2)
end

"""
    find_symmetric_nash_equilibrium_2players_game(game::Dict{String,<:AbstractArray}, s::Array{SymPy.Sym,1})

Finds necessary values to find symmetric Nash equilibrium

**Input parameters**
* `game::Dict{String,<:AbstractArray}` - dictionary of players and their payoff matrices
* `s::Array{SymPy.Sym,1}` - collection of actions probabilities
"""
function find_symmetric_nash_equilibrium_2players_game(game::Dict{String,<:AbstractArray},
               s::Array{SymPy.Sym,1})
    n = size(game["player1"])[1]
    id_m = Matrix(1I, n, n)
    res = 0
    for i in 1:n
        res = res + max(get_payoff(game, [map(x->sympify(x), id_m[i, :]), s])["player1"] - get_payoff(game, [s, s])["player1"], 0)^2
    end
    return find_zeros(res, 0, 1)
end

"""
    create_symmetries_graph(actions_no_vector::Array{Int64})

Return graph of needed symmetries in the game

**Input parameters**
* `actions_no_vector::Array{Int64}` - array of number of actions of consecutive players
"""
function create_symmetries_graph(actions_no_vector::Array{Int64})
    payout_tensor = [()]
    for k in actions_no_vector
        payout_tensor = collect(Iterators.flatten([[(j...,i) for i in 1:k] for j in payout_tensor]))
    end
    payout_tensor = collect(Iterators.flatten([[(j,i) for i in 1:length(actions_no_vector)] for j in payout_tensor]))
    @info "Created payout tensor: $(payout_tensor)"

    p = permutations(collect(1:length(actions_no_vector)))
    equalities = unique(collect(Iterators.flatten([[(b, (b[1][i], i[b[2]])) for i in p] for b in payout_tensor])))
    equalities = filter(x -> x[1]!=x[2], equalities)
    G₁ = Graph(length(payout_tensor))
    node_lookup = Dict(payout_tensor[i] => i for i in 1:length(payout_tensor))
    for i in equalities
        add_edge!(G₁, node_lookup[i[1]], node_lookup[i[2]])
    end
    return SymmetriesGraph(G₁, node_lookup, Dict(value => key for (key, value) in node_lookup))
end

"""
    check_equality_condition(g::SymmetriesGraph, payout_1, payout_2)

Checks if there needs to be an equality between two payouts

**Input parameters**
* `g::SymmetriesGraph` - symmetries graph
* `payout_1` - first payout
* `payout_2` - second payout
"""
function check_equality_condition(g::SymmetriesGraph, payout_1, payout_2)
    return has_path(g.graph, g.node_lookup[payout_1], g.node_lookup[payout_2])
end

"""
    find_all_equalities(g::SymmetriesGraph)

Finds all payout tuples, that need to be equal

**Input parameters**
* `g::SymmetriesGraph` - symmetries graph
"""
function find_all_equalities(g::SymmetriesGraph)
    tmp = Dict{Int64,Vector{Tuple{Int64, Int64}}}()
    Threads.@threads for i in vertices(g.graph)
        tmp[i] = collect(unique(filter(x->x[2]>x[1],map(x->(dst(x),i), edges(dfs_tree(g.graph, i))))))
    end
    return map(x -> (g.reverse_node_lookup[x[1]], g.reverse_node_lookup[x[2]]), collect(Iterators.flatten(values(tmp))))
end


@testset "ilog2 $U(2^$e)" for e in (0, 7, 63, 64, 128, 200), U in (UInt8, UInt16, UInt32, UInt64, UInt128)

    a = U(2)^e
    @test iszero(a) || ilog2(a) == e
    @test iszero(a) || ilog2(a - 1) == e - 1
end

@testset "ilog2 $U(2^$e)" for e in (0, 7, 63, 64, 128, 200), U in (Int8, Int16, Int32, Int64, Int128, BigInt)

    a = U(2)^e
    @test a <= 0 || ilog2(a) == e
    @test a <= 0 || ilog2(a - 1) == e - 1
    @test a <= 0 || ilog2(-a) == e
    @test a <= 0 || ilog2(-a + 1) == e - 1
end

@testset "ilog2 $U(0)" for U in (UInt8, UInt16, UInt32, UInt64, UInt128, BigInt)
    a = zero(U)
    @test ilog2(a) == -1
    @test ilog2(signed(a)) == - 1
end


using PkgSRA, Test, Random
Random.seed!(13)

# include("../utils/sindy_turing_utils.jl")
# include("../utils/sparse_regression_functions.jl")
# include("../utils/sindy_statistics_utils.jl")

# Generate test data: Lorenz
include("../examples/example_lorenz.jl")
dat = Array(solve_lorenz_system())
numerical_grad = numerical_derivative(dat, ts)

# Test new interface
alg = slstQuantile(4, 0.2) # Will get extra terms
sindy_library = Dict("cross_terms"=>2,"constant"=>nothing);
test_model = sindy(dat, numerical_grad, ts,
                        library=sindy_library,
                        optimizer=alg)

# Built correct terms: names
named_terms = build_term_names(test_model)
@testset "Term Metadata" begin
    @test named_terms ==
        ["x", "y", "z", "", "xx", "xy", "xz", "yy", "yz", "zz"]
end

# Built correct terms: values
library_functions = convert_string2function(test_model.library)
terms = calc_augmented_data(dat, library_functions)

@testset "Data Augmentation" begin
    @test size(terms,1) == length(named_terms)
    @test terms[1:3,:] == dat
    @test terms[4,:] == ones(size(dat,2))
    @test terms[5:10,:] == calc_cross_terms(dat, 2)
end;

# Nonzero term and type tests
t1 = get_nonzero_terms(test_model)
t2 = get_nonzero_terms(core_dyn_true)
# typeof(test_model) <: sindycModel
# @testset "Type Tests" begin
@test typeof(test_model) <: sindyModel
@test issubset(t2, t1)
# end

###
### Cross validation functions
###

# First, retraining functions
test_model2 = sindy_retrain(test_model, dat, numerical_grad)
t1 = get_nonzero_terms(test_model)
t2 = get_nonzero_terms(test_model2)

@testset "Retrained model" begin
    @test typeof(test_model2) <: sindyModel
    @test t1==t2
    @test calc_coefficient_error(test_model, test_model2) ≈ 0.0
end

# Also test cross validation function
null_model = sindy(dat, numerical_grad, ts,
                        library=sindy_library,
                        optimizer=slstHard(100)) # i.e. not sparse
err_test = sindy_cross_validate(test_model, dat, numerical_grad)
err_null = sindy_cross_validate(null_model, dat, numerical_grad)

@testset "Cross Validation" begin
    @test err_test < err_null
end

## One more cross validation test
include("../examples/example_lotkaVolterra.jl")
dat = Array(solve_lv_system())
numerical_grad = numerical_derivative(dat, ts)

# First, a decent model on one variable
sindy_library = Dict("cross_terms"=>2,"constant"=>nothing);
lib = convert_string2function(sindy_library)
#      x   y   c   xx xy yy
A1 = [[0.7 0   0   0 -1.3 0 ];
      [0  0   0   0   0   1e-6 ]]
n = size(A, 1)
m1 = sindyModel(ts, A1, lib ,["x", "y"])

# Second, a decent model on both variables
#      x   y   c   xx xy yy
A2 = [[0.7 0   0   0 -1.3 1e-6 ];
      [0  -0.7 0  1e-6 0.7  0]]
m2 = sindyModel(ts, A2, lib ,["x", "y"])

# Third, another decent model on both variables
#      x   y   c   xx xy yy
A3 = [[0.7 0   0   0 -1.3 0 ];
      [1e-6 -0.9 0 0 0.8  0]]
m3 = sindyModel(ts, A3, lib ,["x", "y"])

## Calculate errors
# Coefficients
cerr1 = calc_coefficient_error(m1, core_dyn_true)
cerr2 = calc_coefficient_error(m2, core_dyn_true)
# Residual
ind = 1:500
rss0 = rss_sindy_derivs(core_dyn_true, dat[:,ind], numerical_grad[:,ind])
rss1 = rss_sindy_derivs(m1, dat[:,ind], numerical_grad[:,ind])
rss2 = rss_sindy_derivs(m2, dat[:,ind], numerical_grad[:,ind])
rss3 = rss_sindy_derivs(m3, dat[:,ind], numerical_grad[:,ind])

#
@testset "Error calculations" begin
    @test cerr1 > cerr2
    @test rss1 > rss0
    @test rss2 > rss0
    @test rss1 > rss2
end

using Plots
pyplot()
u0 = dat[:,1]
plot_sindy_model(m1, u0, which_dim=2)
    plot!(m1.ts, dat[2,:])
    title!("Model 1")
plot_sindy_model(m2, u0, which_dim=2)
    plot!(m2.ts, dat[2,:])
    title!("Model 2")

# SINDY tests: up to cubed terms

# Generate test data: Van der Pol
include("../examples/example_vanDerPol.jl")
dat = Array(solve_vdp_system())
numerical_grad = numerical_derivative(dat, ts)

sindy_library = Dict("cross_terms"=>[2,3],
                    "constant"=>nothing);
test_model = sindy(dat, numerical_grad, ts,
                        library=sindy_library,
                        optimizer=alg,
                        var_names=["x", "y"])

named_terms = build_term_names(test_model)
@testset "Cubic term names" begin
    @test named_terms ==
        ["x", "y", "", "xx", "xy", "yy", "xxx", "xxy", "xyy", "yyy"]
end


using Test
using LogicCircuits
using ProbabilisticCircuits

# This tests are supposed to test queries on the circuits
@testset "Logistic Circuit Class Conditional" begin
    # Uses a Logistic Circuit with 4 variables, and tests 3 of the configurations to
    # match with python version.

    EPS = 1e-7;
    my_opts = (max_factors= 2,
            compact⋀=false,
            compact⋁=false)

    logistic_circuit = zoo_lc("little_4var.circuit", 2);
    @test logistic_circuit isa Vector{<:LogisticΔNode};

    flow_circuit = FlowΔ(logistic_circuit, 16, Float64, my_opts)
    @test flow_circuit isa Vector{<:FlowΔNode};

    # Step 1. Check Probabilities for 3 samples
    data = XData(Bool.([0 0 0 0; 0 1 1 0; 0 0 1 1]));

    true_prob = [3.43147972 4.66740416;
                4.27595352 2.83503504;
                3.67415087 4.93793472]

    CLASSES = 2
    calc_prob = class_conditional_likelihood_per_instance(flow_circuit, CLASSES, data)

    for i = 1:3
        for j = 1:2
            @test true_prob[i,j] ≈ calc_prob[i,j] atol= EPS;
        end
    end

    # 2. Testing different API
    fc2, calc_prob2 = class_conditional_likelihood_per_instance(logistic_circuit, CLASSES, data)
    for i = 1:3
        for j = 1:2
            @test true_prob[i,j] ≈ calc_prob2[i,j] atol= EPS;
        end
    end

end


using DelimitedFiles

function step_fish(input_fishnum)
    if input_fishnum == 0
        return 6
    else
        return input_fishnum - 1
    end
end

function step_day(input_list)
    num_fish_to_add = count(x->(x==0), input_list)
    output_list = map(x->step_fish(x), input_list)
    append!(output_list, ones(Int, num_fish_to_add)*8)
    return output_list
end

function test_step_day()
    test_list = [3,2,1]
    @assert step_day(test_list) == [2,1,0]
    test_list = [0,2,1]
    @assert step_day(test_list) == [6,1,0,8]
    test
end

function step_multiple_days(fish_list, num_days)
    for i in 1:num_days
        println("Processing day ", i, " num_fish is ", length(fish_list))
        fish_list = step_day(fish_list)
    end
    return fish_list
end

# read input file, note single quotes around the delimiter
input = readdlm("input.txt", ',',Int)
# convert to a column vector
input = vec(input)

# step 80 days
output = step_multiple_days(input, 80)
print("number of fish is ", length(output))