using ClusterManagers
using Distributed

OnCluster = true #set to false to run locally
addWorkers = true #set to false to run serially
println("OnCluster = $(OnCluster)")

# Current number of workers
#--------------------------
currentWorkers = nworkers()
println("Initial number of workers = $(currentWorkers)")

# Increase the number of workers available
#-----------------------------------------
maxNumberWorkers = 10
if addWorkers == true
	if OnCluster == true
	  addprocs(SlurmManager(maxNumberWorkers))
	else
	  addprocs(maxNumberWorkers)
	end
end

# Sanity checks
#-------------
hosts = []
pids = []
for i in workers()
	host, pid = fetch(@spawnat i (gethostname(), getpid()))
	println("Hello I am worker $(i), my host is $(host)")
	push!(hosts, host)
	push!(pids, pid)
end

currentWorkers = nworkers()
println("Number of workers = $(currentWorkers)")

using Plots
using LaTeXStrings
using ParallelDataTransfer

@everywhere using MSM
@everywhere using DataStructures
@everywhere using OrderedCollections
@everywhere using Distributions
@everywhere using Random
@everywhere using DataStructures
@everywhere using Statistics
@everywhere using LinearAlgebra

# Generate simulated data
Random.seed!(1234)  #for replicability reasons
T = 100000          #number of periods
P = 2               #number of dependent variables
beta0 = rand(P)     #choose true coefficients by drawing from a uniform distribution on [0,1]
alpha0 = rand(1)[]  #intercept
theta0 = 0.0        #coefficient to create serial correlation in the error terms
println("True intercept = $(alpha0)")
println("True coefficient beta0 = $(beta0)")
println("Serial correlation coefficient theta0 = $(theta0)")

# Generation of error terms
# row = individual dimension
# column = time dimension
U = zeros(T)
d = Normal()
U[1] = rand(d, 1)[] #first error term
# loop over time periods
for t = 2:T
    U[t] = rand(d, 1)[] + theta0*U[t-1]
end

# Let's simulate the dependent variables x_t
x = zeros(T, P)

d = Uniform(0, 5)
for p = 1:P
    x[:,p] = rand(d, T)
end

# Let's calculate the resulting y_t
y = zeros(T)

for t=1:T
    y[t] = alpha0 + x[t,1]*beta0[1] + x[t,2]*beta0[2] + U[t]
end

# Send simulated data to workers
sendto(workers(), y=y)
sendto(workers(), x=x)

# Visualize data
p1 = scatter(x[1:100,1], y[1:100], xlabel = "x1", ylabel = "y", legend=:none, smooth=true)
p2 = scatter(x[1:100,2], y[1:100], xlabel = "x2", ylabel = "y", legend=:none, smooth=true)
plot(p1, p2)

# Define locally
optionsSMM = MSMOptions(maxFuncEvals=2000, globalOptimizer = :dxnes, localOptimizer = :NelderMead)
myProblem = MSMProblem(options = optionsSMM);

# Send to workers
sendto(workers(), optionsSMM=optionsSMM)
sendto(workers(), myProblem=myProblem)

# Priors
dictPriors = OrderedDict{String,Array{Float64,1}}()
dictPriors["alpha"] = [0.5, 0.001, 1.0]
dictPriors["beta1"] = [0.5, 0.001, 1.0]
dictPriors["beta2"] = [0.5, 0.001, 1.0]
sendto(workers(), dictPriors=dictPriors)

# Empirical moments
dictEmpiricalMoments = OrderedDict{String,Array{Float64,1}}()
dictEmpiricalMoments["mean"] = [mean(y)] #informative on the intercept
dictEmpiricalMoments["mean^2"] = [mean(y.^2)] #informative on the intercept
dictEmpiricalMoments["mean_x1y"] = [mean(x[:,1] .* y)] #informative on betas
dictEmpiricalMoments["mean_x2y"] = [mean(x[:,2] .* y)] #informative on betas
dictEmpiricalMoments["mean_x1y^2"] = [mean((x[:,1] .* y).^2)] #informative on betas
dictEmpiricalMoments["mean_x2y^2"] = [mean((x[:,2] .* y).^2)] #informative on betas

W = Matrix(1.0 .* I(length(dictEmpiricalMoments)))#initialization
#Special case: diagonal matrix
#(you may choose something else)
for (indexMoment, k) in enumerate(keys(dictEmpiricalMoments))
    W[indexMoment,indexMoment] = 1.0/(dictEmpiricalMoments[k][1])^2
end

# Send to workers
sendto(workers(), dictPriors=dictPriors)
sendto(workers(), dictEmpiricalMoments=dictEmpiricalMoments)
sendto(workers(), W=W)

@everywhere set_priors!(myProblem, dictPriors)
@everywhere set_empirical_moments!(myProblem, dictEmpiricalMoments)
@everywhere set_weight_matrix!(myProblem, W)

# x[1] corresponds to the intercept, x[2] corresponds to beta1, x[3] corresponds to beta2
@everywhere function functionLinearModel(x; uniform_draws::Array{Float64,1}, simX::Array{Float64,2}, nbDraws::Int64 = length(uniform_draws), burnInPerc::Int64 = 0)
    T = nbDraws
    P = 2       #number of dependent variables

    alpha = x[1]
    beta = x[2:end]
    theta = 0.0     #coefficient to create serial correlation in the error terms

    # Creation of error terms
    # row = individual dimension
    # column = time dimension
    U = zeros(T)
    d = Normal()
    # Inverse cdf (i.e. quantile)
    gaussian_draws = quantile.(d, uniform_draws)
    U[1] = gaussian_draws[1] #first error term

    # loop over time periods
    for t = 2:T
        U[t] = gaussian_draws[t] + theta*U[t-1]
    end

    # Let's calculate the resulting y_t
    y = zeros(T)

    for t=1:T
        y[t] = alpha + simX[t,1]*beta[1] + simX[t,2]*beta[2] + U[t]
    end

	# Get rid of the burn-in phase:
    #------------------------------
    startT = max(1, Int(nbDraws * (burnInPerc / 100)))

    # Moments:
    #---------
    output = OrderedDict{String,Float64}()
    output["mean"] = mean(y[startT:nbDraws])
    output["mean^2"] = mean(y[startT:nbDraws].^2)
    output["mean_x1y"] = mean(simX[startT:nbDraws,1] .* y[startT:nbDraws])
    output["mean_x2y"] = mean(simX[startT:nbDraws,2] .* y[startT:nbDraws])
    output["mean_x1y^2"] = mean((simX[startT:nbDraws,1] .* y[startT:nbDraws]).^2)
    output["mean_x2y^2"] = mean((simX[startT:nbDraws,2] .* y[startT:nbDraws]).^2)

    return output
end

# Let's freeze the randomness during the minimization
d_Uni = Uniform(0,1)
nbDraws = 2*T #number of draws in the simulated data
uniform_draws = rand(d_Uni, nbDraws)
simX = zeros(length(uniform_draws), 2)
d = Uniform(0, 5)
for p = 1:2
  simX[:,p] = rand(d, length(uniform_draws))
end

# Send to workers
sendto(workers(), simX=simX)
sendto(workers(), uniform_draws=uniform_draws)

# Construct the objective function everywhere
@everywhere set_simulate_empirical_moments!(myProblem, x -> functionLinearModel(x, uniform_draws = uniform_draws, simX = simX))
@everywhere construct_objective_function!(myProblem)

# Safety check: value on the master node == values on slave nodes?
using Test
val_local = myProblem.objective_function(ones(3)); #local execution
val_workers = [];
for w in workers() #Execution on workers
    push!(val_workers, @fetchfrom w myProblem.objective_function(ones(3)))
end
for (wIndex, w) in enumerate(workers())
    @test abs(val_local - val_workers[wIndex]) < 10e-10
end

println("Global Algorithm")
# Choose a global optimizer that supports parallel evaluations (e.g. xnes or dxnes)
# (see the documentation: https://github.com/robertfeldt/BlackBoxOptim.jl)
msm_optimize!(myProblem, verbose = false)

# Access the results using best_candidate
minimizer = msm_minimizer(myProblem)
minimum_val = msm_minimum(myProblem)

println("Minimum objective function = $(minimum_val)")
println("Estimated value for alpha = $(minimizer[1]). True value for beta1 = $(alpha0[1]) \n")
println("Estimated value for beta1 = $(minimizer[2]). True value for beta1 = $(beta0[1]) \n")
println("Estimated value for beta2 = $(minimizer[3]). True value for beta2 = $(beta0[2]) \n")

println("Multistart Algorithm")
# Start several local optimization algorithms in parallel
# Choose algorithms from the package Optim.jl (https://github.com/JuliaNLSolvers/Optim.jl)
# The "global" mininimum is the minimum of the local minima.
msm_multistart!(myProblem, nums = nworkers(), verbose = false)

minimizer_multistart = msm_multistart_minimizer(myProblem)
minimum_multistart = msm_multistart_minimum(myProblem)

println("Minimum objective function = $(minimum_multistart)")
println("Estimated value for alpha = $(minimizer_multistart[1]). True value for beta1 = $(alpha0[1]) \n")
println("Estimated value for beta1 = $(minimizer_multistart[2]). True value for beta1 = $(beta0[1]) \n")
println("Estimated value for beta2 = $(minimizer_multistart[3]). True value for beta2 = $(beta0[2]) \n")

# Empirical Series
#-----------------
X = zeros(T, length(dictEmpiricalMoments))
X[:,1] = y
X[:,2] = y.^2
X[:,3] = (x[:,1] .* y)
X[:,4] = (x[:,2] .* y)
X[:,5] = (x[:,1] .* y).^2
X[:,6] = (x[:,2] .* y).^2

# "Distance Matrix" (see Duffie and Singleton, 1993)
Sigma0 = cov(X)

# Heatmap to visualize correlation
xs = [string("x", i) for i = 1:length(dictEmpiricalMoments)]
ys = [string("x", i) for i = 1:length(dictEmpiricalMoments)]
z = cor(X)
hh = heatmap(xs, ys, z, aspect_ratio = 1)

set_Sigma0!(myProblem, Sigma0)
# nbDraws = number of draws in the simulated data
# To decrease standard errors, increase nbDraws
calculate_Avar!(myProblem, minimizer_multistart, tau = T/nbDraws)

df = summary_table(myProblem, minimizer_multistart, T, 0.05)
println(df)

# Check the rank condition
# Local identification requires D to be full column rank (in a neighborhood of the solution)
D = calculate_D(myProblem, minimizer_multistart)
println("number of parameters: $(size(D,2))")
println("rank of D is: $(rank(D))")

# Slices
vXGrid, vYGrid = msm_slices(myProblem, minimizer_multistart, nbPoints = 7);

p1 = plot(vXGrid[:, 1],vYGrid[:, 1],title = L"\alpha", label = "",linewidth = 3, xrotation = 45)
plot!(p1, [minimizer_multistart[1]], seriestype = :vline, label = "",linewidth = 1)
p2 = plot(vXGrid[:, 2],vYGrid[:, 2],title = L"\beta_1", label = "",linewidth = 3, xrotation = 45)
plot!(p2, [minimizer_multistart[2]], seriestype = :vline, label = "",linewidth = 1)
p3 = plot(vXGrid[:, 3],vYGrid[:, 3],title = L"\beta_2", label = "",linewidth = 3, xrotation = 45)
plot!(p3, [minimizer_multistart[3]], seriestype = :vline, label = "",linewidth = 1)
plot_combined = plot(p1, p2, p3)
display(plot_combined)

# Compare results with GLM
using DataFrames, GLM
data = DataFrame(x1=x[:,1], x2=x[:,2], y= y[:]);
ols = lm(@formula(y ~ x1 + x2), data)
println(ols)


# # Vehicle Platoon Model
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/models/Platoon.ipynb)
#
#md # !!! note "Overview"
#md #     System type: affine with hybrid dynamics\
#md #     State dimension: 9 + 1\
#md #     Application domain: Autonomous Driving
#
# ## Model description
#
# This benchmark considers a platoon of three vehicles following each
# other. This benchmark considers loss of communication between vehicles. The model
# has two modes, connected and disconnected.
# Three scenarios are considered for the loss of communication:
#
# **PLAA01 (arbitrary loss)** The loss of communication can occur at any time. This
#       includes the possibility of no communication at all.
#
# **PLADxy (loss at deterministic times)** The loss of communication occurs at fixed
#       points in time, which are determined by clock constraints ``c_1`` and ``c_2``.
#       The clock t is reset when communication is lost and when it is re-established.
#       Note that the transitions have must-semantics, i.e., they take place as soon
#       as possible. We will consider `PLAD01: c1 = c_2 = 5`.
#
# **PLANxy (loss at nondeterministic times)** The loss of communication occurs at
#       any time ``t ∈ [t_b, t_c]``. The clock t is reset when communication is lost
#       and when it is reestablished. Communication is reestablished at any time
#       ``t ∈ [0, t_r]``. This scenario covers loss of communication after an
#       arbitrarily long time ``t ≥ t_c`` by reestablishing communication in zero time.
#       We will consider `PLAN01: tb = 10, tc = 20, tr = 20`.

# The Julia model is developed next. It's convenient to create two independent
# functions, `platoon_connected` and `platoon_disconnected`, which describe the
# dynamics of the connected (resp. disconnected) modes. These functions are
# connected in the `platoon` function using a hybrid automaton with two modes
# and two discrete transitions, see the following diagram.

# ![](https://github.com/JuliaReach/ReachabilityAnalysis.jl/blob/master/examples/Platoon/platoon_hybrid.png?raw=true)

# On the left, the system can switch arbitrarily between the modes.
# In the middle, mode switches are only possible at given points in time, i.e. deterministic switching.
# On the right, mode switches are only possible during given time intervals.
# In this notebook we will only consider the case of deterministic switching.

# ### Dynamics of the "connected" platoon

using ReachabilityAnalysis, SparseArrays, Symbolics

const var = @variables x[1:9] t

function platoon_connected(; deterministic_switching::Bool=true,
                             c1=5.0)  # clock constraints
    n = 9 + 1

    ## x' = Ax + Bu + c
    A = Matrix{Float64}(undef, n, n)
    A[1, :] = [0, 1.0, 0, 0, 0, 0, 0, 0, 0, 0]
    A[2, :] = [0, 0, -1.0, 0, 0, 0, 0, 0, 0, 0]
    A[3, :] = [1.6050, 4.8680, -3.5754, -0.8198, 0.4270, -0.0450, -0.1942,  0.3626, -0.0946, 0.]
    A[4, :] = [0, 0, 0, 0, 1.0, 0, 0, 0, 0, 0,]
    A[5, :] = [0, 0, 1.0, 0, 0, -1.0, 0, 0, 0, 0]
    A[6, :] = [0.8718, 3.8140, -0.0754,  1.1936, 3.6258, -3.2396, -0.5950,  0.1294, -0.0796, 0.]
    A[7, :] = [0, 0, 0, 0, 0, 0, 0, 1.0, 0, 0]
    A[8, :] = [0, 0, 0, 0, 0, 1.0, 0, 0, -1.0, 0]
    A[9, :] = [0.7132, 3.5730, -0.0964,  0.8472, 3.2568, -0.0876,  1.2726,  3.0720, -3.1356, 0.]
    A[10, :] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0]; # t' = 1

    if deterministic_switching
        invariant = HalfSpace(t <= c1, var)
    else
        invariant = Universe(n)
    end

    ## acceleration of the lead vehicle + time
    B = sparse([2], [1], [1.0], n, 1)
    U = Hyperrectangle(low=[-9.], high=[1.])
    c = [0, 0, 0, 0, 0, 0, 0, 0, 0, 1.0]
    @system(x' = Ax + Bu + c, x ∈ invariant, u ∈ U)
end

# ### Dynamics of the "disconnected" platoon

function platoon_disconnected(; deterministic_switching::Bool=true,
                                c2=5.0)  # clock constraints
    n = 10 # 9 dimensions + time

    ## x' = Ax + Bu + c
    A = Matrix{Float64}(undef, n, n)
    A[1, :] = [0, 1.0, 0, 0, 0, 0, 0, 0, 0, 0]
    A[2, :] = [0, 0, -1.0, 0, 0, 0, 0, 0, 0, 0]
    A[3, :] = [1.6050, 4.8680, -3.5754, 0, 0, 0, 0, 0, 0, 0]
    A[4, :] = [0, 0, 0, 0, 1.0, 0, 0, 0, 0, 0,]
    A[5, :] = [0, 0, 1.0, 0, 0, -1.0, 0, 0, 0, 0]
    A[6, :] = [0, 0, 0,  1.1936, 3.6258, -3.2396, 0, 0, 0, 0.]
    A[7, :] = [0, 0, 0, 0, 0, 0, 0, 1.0, 0, 0]
    A[8, :] = [0, 0, 0, 0, 0, 1.0, 0, 0, -1.0, 0]
    A[9, :] = [0.7132, 3.5730, -0.0964,  0.8472, 3.2568, -0.0876,  1.2726,  3.0720, -3.1356, 0.]
    A[10, :] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0]; # t' = 1

    if deterministic_switching
        invariant = HalfSpace(t <= c2, var)
    else
        invariant = Universe(n)
    end

    ## acceleration of the lead vehicle + time
    B = sparse([2], [1], [1.0], n, 1)
    U = Hyperrectangle(low=[-9.], high=[1.])
    c = [0, 0, 0, 0, 0, 0, 0, 0, 0, 1.0]
    @system(x' = Ax + Bu + c, x ∈ invariant, u ∈ U)
end

# ### Hybrid system

function platoon(; deterministic_switching::Bool=true,
                   c1=5.0,  # clock constraints
                   c2=5.0,  # clock constraints
                   tb=10.0,  # lower bound for loss of communication
                   tc=20.0, tr=20.0) # upper bound for loss of communication (tc) and reset time (tr)

    ## three variables for each vehicle, (ei, d(et)/dt, ai) for
    ## (spacing error, relative velocity, speed), and the last dimension is time
    n = 9 + 1

    ## transition graph
    automaton = LightAutomaton(2)
    add_transition!(automaton, 1, 2, 1)
    add_transition!(automaton, 2, 1, 2)

    ## modes
    mode1 = platoon_connected(deterministic_switching=deterministic_switching, c1=c1)
    mode2 = platoon_disconnected(deterministic_switching=deterministic_switching, c2=c2)
    modes = [mode1, mode2]

    ## common reset
    reset = Dict(n => 0.)

    ## transition l1 -> l2
    if deterministic_switching
        guard = Hyperplane(t == c1, var)
    else
        guard = HPolyhedron([tb <= t, t <= tc], var)
    end
    t1 = ConstrainedResetMap(n, guard, reset)

    ## transition l2 -> l1
    if deterministic_switching
        guard = Hyperplane(t == c2, var)
    else
        guard = HalfSpace(t <= tr, var)
    end
    t2 = ConstrainedResetMap(n, guard, reset)
    resetmaps = [t1, t2]

    H = HybridSystem(automaton, modes, resetmaps, [AutonomousSwitching()])

    ## initial condition is at the orgin in mode 1
    X0 = BallInf(zeros(n), 0.0)
    initial_condition = [(1, X0)]

    return IVP(H, initial_condition)
end

# ## Safety specifications
#
# The verification goal is to check whether the minimum distance between
# vehicles is preserved. The choice of the coordinate system is such that the
# minimum distance is a negative value.
#
# **BNDxy**: Bounded time (no explicit bound on the number of transitions): For all ``t ∈ [0, 20] [s]``,
#
#    - ``x_1(t) ≥ −d_{min} [m]``
#    - ``x_4(t) ≥ −d_{min} [m]``
#    - ``x_7(t) ≥ −d_{min} [m]``
#
# **BND42**: ``d_{min} = 42``.
#
# **BND30**: ``d_{min} = 30``.

function dmin_specification(sol, dmin)
    return (-ρ(sparsevec([1], [-1.0], 10), sol) > -dmin) &&
           (-ρ(sparsevec([4], [-1.0], 10), sol) > -dmin) &&
           (-ρ(sparsevec([7], [-1.0], 10), sol) > -dmin)
end

# ## Results

# We will only consider the case of deterministic switching in this notebook. We consider two
# scenarios of increasing difficulty, BND42 and BND30.

prob_PLAD01 = platoon();

const boxdirs = BoxDirections(10)
const octdirs = OctDirections(10);

#  ### PLAD01 - BND42

# This scenario can be solved using a hyperrectangular set representation with step
# size ``δ = 0.01``.  We use a template that contains all box (i.e., canonical)
# directions in the ambient space of the state-space, ``\mathbb{R}^{10}``.
# There are ``20`` such directions, two for each coordinate:

length(boxdirs)

#-

alg = BOX(δ=0.01)
sol_PLAD01_BND42 = solve(prob_PLAD01,
                         alg=alg,
                         clustering_method=BoxClustering(1),
                         intersection_method=TemplateHullIntersection(boxdirs),
                         intersect_source_invariant=false,
                         tspan = (0.0 .. 20.0));

# Let's verify that the specification holds:

dmin_specification(sol_PLAD01_BND42, 42)

# In more detail we can check how is the flowpipe from violating the property.
# The specification requires that each of the follwing quantities is greater
# than `-dmin = -42`.

# Minimum of ``x_1(t)``:
-ρ(sparsevec([1], [-1.0], 10), sol_PLAD01_BND42)
#-
# Minimum of ``x_4(t)``:
-ρ(sparsevec([4], [-1.0], 10), sol_PLAD01_BND42)
#-
# Minimum of ``x_7(t)``:
-ρ(sparsevec([7], [-1.0], 10), sol_PLAD01_BND42)

# We plot variable ``x_1`` vs time.

using Plots, LaTeXStrings

plot(sol_PLAD01_BND42, vars=(0, 1), xlab=L"t", ylab=L"x_1", title="PLAD01 - BND42", lw=0.1)
plot!(x->x, x->-42., 0., 20., linewidth=2, color="red", linestyle=:dash, leg=nothing)

#  ### PLAD01 - BND30

# Note that the previous solution obtained for `PLAD01 - BND42` does not verify the
# `BND30` specifications since, for example,
# the minimum of variable ``x_4(t)`` is ``≈ -35.52``, which
# is below the given bound ``-d_{min} = -30``. As a consequence, to prove the safety
# properties for this scenario we have to use a solver with more precision.
# Instead of the `BOX` algorithm, we propose to use `LGG09` with step-size
# ``δ=0.03`` and octagonal template directions. There are 200 such directions:

length(octdirs)

#md # !!! tip "Performance tip"
#md # The increase in the number of directions implies an increase in runtime.
#md # Since evaluating the 200 directions of the template is quite expensive,
#md # we try using a concrete set after the discretization, instead of using a
#md # lazy discretization. This is achieved by passing the option
#md # `approx_model=Forward(setops=octdirs)` to the `LGG09` algorithm, specifying that
#md # we want to oveapproximate the initial set of the set-based recurrence, that we call `Ω0`,
#md # with an octagonal template. It turns out in this example that the option gives a
#md # gain in runtime of ``~30\%``, without a noticeable loss in precision.

alg = LGG09(δ=0.03, template=octdirs, approx_model=Forward(setops=octdirs));
sol_PLAD01_BND30 = solve(prob_PLAD01,
                         alg=alg,
                         clustering_method=LazyClustering(1),
                         intersection_method=TemplateHullIntersection(octdirs),
                         intersect_source_invariant=false,
                         tspan = (0.0 .. 20.0));

# Verifying that the specification holds:

dmin_specification(sol_PLAD01_BND30, 30)

# Let's check in more detail how close is the flowpipe to the safety conditions:

# Minimum of ``x_1(t)``:
-ρ(sparsevec([1], [-1.0], 10), sol_PLAD01_BND30)
#-
# Minimum of ``x_4(t)``:
-ρ(sparsevec([4], [-1.0], 10), sol_PLAD01_BND30)
#-
# Minimum of ``x_7(t)``:
-ρ(sparsevec([7], [-1.0], 10), sol_PLAD01_BND30)

#-

plot(sol_PLAD01_BND30, vars=(0, 1), xlab=L"t", ylab=L"x_1", title="PLAD01 - BND30", lw=0.1)
plot!(x->x, x->-30., 0., 20., linewidth=2, color="red", linestyle=:dash, leg=nothing)

# ## References

# [^]:


using Base.Iterators
using Test

using Arbitrary

const alltypes = [Nothing, Bool, Char,
                  Int8, Int16, Int32, Int64, Int128,
                  UInt8, UInt16, UInt32, UInt64, UInt128,
                  Float16, Float32, Float64,
                  BigInt, BigFloat,
                  Rational{BigInt},
                  Tuple{},
                  Tuple{Int8, Int16},
                  Tuple{Int8, Int16, Int32},
                  Tuple{Int8, Int16, Int32, Int64}]
hasonevalue(T) = T in [Nothing, Tuple{}]

const Array0{T} = Array{T, 0}
const Array1{T} = Array{T, 1}
const Array2{T} = Array{T, 2}
const allcontainers = [nothing, Ref, Tuple, Array0, Array1, Array2]
hasfixedshape(C) = C in [nothing, Ref, Tuple, Array0]

myequal(x::T, y::T) where {T} = isequal(x, y)
myequal(x::Ref{T}, y::Ref{T}) where {T} = isequal(x[], y[])

for C in allcontainers
    for E in alltypes
        T = C === nothing ? E : C{E}
        @testset "Basic functionality for type $T" begin
            # Generate arbitrary values
            arb = arbitrary(T)
            values = collect(take(arb, 100))
            # Generate some other arbitrary values
            arb2 = arbitrary(T)
            values2 = collect(take(arb2, 100))
            # Ensure they are different
            @test all(myequal.(values, values2)) ==
                (hasonevalue(E) && hasfixedshape(C))
            # Generate values from a known RNG
            arb3 = arbitrary(T, UInt(42))
            arb4 = arbitrary(T, UInt(42))
            @test all(myequal.(collect(take(arb3, 100)),
                               collect(take(arb4, 100))))
        end
    end
end

# Floating-point numbers do NOT satisfy the usual arithmetic laws
const arithmetic_types = [Int8, Int16, Int32, Int64, Int128,
                          UInt8, UInt16, UInt32, UInt64, UInt128,
                          # Float16, Float32, Float64,
                          BigInt, BigFloat,
                          Rational{BigInt}]
const division_types = [# Float16, Float32, Float64,
                        # BigFloat,
                        Rational{BigInt}]
for T in arithmetic_types
    @testset "Arithmetic identities for type $T" begin
        # Generate arbitrary values
        xs = collect(take(arbitrary(T), 100))
        ys = collect(take(arbitrary(T), 100))
        zs = collect(take(arbitrary(T), 100))
        ds = collect(take(Iterators.filter(x -> !isequal(x, T(0)),
                                           arbitrary(T)), 100))
        # Addition:
        # Commutativity
        @test all(isequal.(xs .+ ys, ys .+ xs))
        # Associativity
        @test all(isequal.((xs .+ ys) .+ zs, xs .+ (ys .+ zs)))
        # Neutral element
        @test all(isequal.(xs .+ T(0), xs))
        @test all(isequal.(T(0) .+ xs, xs))
        # Inverse
        @test all(isequal.(xs .+ (-xs), T(0)))
        @test all(isequal.(xs .- ys, xs .+ (-ys)))
        # Multiplication:
        # Commutativity
        @test all(isequal.(xs .* ys, ys .* xs))
        # Associativity
        @test all(isequal.((xs .* ys) .* zs, xs .* (ys .* zs)))
        # Neutral element
        @test all(isequal.(xs .* T(1), xs))
        @test all(isequal.(T(1) .* xs, xs))
        # Inverse
        if T in division_types
            @test all(isequal.(ds .* inv.(ds), T(1)))
            @test all(isequal.(xs ./ ds, xs .* inv.(ds)))
        end
    end
end

const function_types = [Int, Float64, Char, UInt]

@testset "Function identities" begin
    T, U, V, R = function_types
    fs = collect(take(arbitrary(Fun{T, U}), 100))
    gs = collect(take(arbitrary(Fun{U, V}), 100))
    hs = collect(take(arbitrary(Fun{V, R}), 100))

    # Neutral element
    for f in fs
        @test fun_isequal(T, f ∘ identity, f)
        @test fun_isequal(T, identity ∘ f, f)
    end

    # Associativity
    for (f, g, h) in zip(fs, gs, hs)
        @test fun_isequal(T, (h ∘ g) ∘ f, h ∘ (g ∘ f))
    end
end


export iauTporv
"""
In the tangent plane projection, given the rectangular coordinates
of a star and its direction cosines, determine the direction
cosines of the tangent point.

This function is part of the International Astronomical Union's
SOFA (Standards of Fundamental Astronomy) software collection.

Status:  support function.

Given:
   xi,eta   double    rectangular coordinates of star image (Note 2)
   v        double[3] star's direction cosines (Note 3)

Returned:
   v01      double[3] tangent point's direction cosines, Solution 1
   v02      double[3] tangent point's direction cosines, Solution 2

Returned (function value):
               int     number of solutions:
                     0 = no solutions returned (Note 4)
                     1 = only the first solution is useful (Note 5)
                     2 = both solutions are useful (Note 5)

Notes:

   1. The tangent plane projection is also called the "gnomonic
   projection" and the "central projection".

   2. The eta axis points due north in the adopted coordinate system.
   If the direction cosines represent observed (RA,Dec), the tangent
   plane coordinates (xi,eta) are conventionally called the
   "standard coordinates".  If the direction cosines are with
   respect to a right-handed triad, (xi,eta) are also right-handed.
   The units of (xi,eta) are, effectively, radians at the tangent
   point.

   3. The vector v must be of unit length or the result will be wrong.

   4. Cases where there is no solution can arise only near the poles.
   For example, it is clearly impossible for a star at the pole
   itself to have a non-zero xi value, and hence it is meaningless
   to ask where the tangent point would have to be.

   5. Also near the poles, cases can arise where there are two useful
   solutions.  The return value indicates whether the second of the
   two solutions returned is useful;  1 indicates only one useful
   solution, the usual case.

   6. The basis of the algorithm is to solve the spherical triangle
   PSC, where P is the north celestial pole, S is the star and C is
   the tangent point.  Calling the celestial spherical coordinates
   of the star and tangent point (a,b) and (a0,b0) respectively, and
   writing rho^2 = (xi^2+eta^2) and r^2 = (1+rho^2), and
   transforming the vector v into (a,b) in the normal way, side c is
   then (pi/2-b), side p is sqrt(xi^2+eta^2) and side s (to be
   found) is (pi/2-b0), while angle C is given by sin(C) = xi/rho
   and cos(C) = eta/rho;  angle P (to be found) is (a-a0).  After
   solving the spherical triangle, the result (a0,b0) can be
   expressed in vector form as v0.

   7. This function is a member of the following set:

      spherical      vector         solve for

      iauTpxes      iauTpxev         xi,eta
      iauTpsts      iauTpstv          star
      iauTpors    > iauTporv <       origin

References:

   Calabretta M.R. & Greisen, E.W., 2002, "Representations of
   celestial coordinates in FITS", Astron.Astrophys. 395, 1077

   Green, R.M., "Spherical Astronomy", Cambridge University Press,
   1987, Chapter 13.

This revision:   2018 January 2

SOFA release 2018-01-30

Copyright (C) 2018 IAU SOFA Board.  See notes at end.
"""

# int iauTporv(double xi, double eta, double v[3],
#              double v01[3], double v02[3])

function iauTporv(xi::Real, eta::Real, v::AbstractVector{<:Real})
   v01 = zeros(Float64, 3)
   v02 = zeros(Float64, 3)

   status = ccall((:iauTporv, libsofa_c), Cint,
       (Cdouble, Cdouble, Ref{Cdouble},
       Ptr{Cdouble}, Ptr{Cdouble}),
       convert(Float64, xi),
       convert(Float64, eta),
       convert(Array{Float64, 1}, v),
       v01, v02)

   return status, SVector{3}(v01), SVector{3}(v02)
end


using MAT
using ReachabilityAnalysis
using SparseArrays
using LinearAlgebra
using StructuralDynamicsODESolvers
using LazySets.Arrays

using Plots
using LaTeXStrings
using Plots.PlotMeasures

# utils for plots
__toL(x, digits) = L"%$(round(x, digits=digits))"
xticklatex(vec, digits) = (vec, __toL.(vec, Ref(digits)))

LazySets.set_rtol(Float64, eps(Float64))
LazySets.set_ztol(Float64, eps(Float64))
LazySets.set_atol(Float64, eps(Float64))

include("../krylov_concept.jl")

include("lambs_modelo.jl")

# cases:
#    - "constant" (Heaviside)
#    - "rickler" (step-load approximation of the rickler wavelet)
#
function sol_numerica_lamb(p::FEMProblem; case="heaviside", δt=nothing, δtfac=1.0,
                           alg=Trapezoidal, X0=nothing, kwargs...)

    if isnothing(δt)
        δt = step_size_lamb(p, δtfac)
    end

    Tmax = 1.0 # seconds
    NSTEPS = round(Int, Tmax/δt)
    tsp = range(0.0, Tmax, length=NSTEPS+1)

    m = size(p.s.K, 1)
    B = Matrix(1.0I, m, m)
    X = Universe(m)

    H = Heaviside
    if case == "constant" || case == "heaviside"
        Fval = -1e6
        Fval = [p.b * Fval for ti in tsp]

    elseif case == "ricker"
        ancho_escalon = 0.05 # 50 ms de ancho
        # por cuestion de simetria se usa 2e6 / 2 respecto del articulo de Accurate Wave Prop
        fr(t) = -1e6 * (H(ancho_escalon*3 - t) - 3*H(ancho_escalon*2 - t) + 3*H(ancho_escalon - t))
        Fval = [p.b * fr(ti) for ti in tsp] # "step-loading" en el articulo

    else
        throw(ArgumentError("case $case not defined"))
    end
    sys = SecondOrderConstrainedLinearControlContinuousSystem(sparse(p.s.M), p.s.C, p.s.K, B, Universe(m), Fval)

    if isnothing(X0)
        X0 = (zeros(m), zeros(m))
    end
    prob = @ivp(sys, u(0) ∈ X0)
    alg = alg(Δt=δt)
    sol = solve(prob, alg, NSTEPS=NSTEPS)

    # al final no invertimos
    # invertimos el signo de todos los desplazamientos para comparar con el articulo Accurate Wave prop
    #sol.U .= -1 .* sol.U
    # lo mismo con las velocidades
    #sol.U′ .= -1 .* sol.U′

    return sol
end

function sol_numerica_lamb_statistics(p::FEMProblem; case="ricker", δt=nothing, δtfac=1.0,
                                      alg=Trapezoidal, X0=nothing,
                                      sampler, idx, vecsamples,
                                      kwargs...)
    if isnothing(δt)
        δt = step_size_lamb(p, δtfac)
    end

    Tmax = 1.0 # seconds
    NSTEPS = round(Int, Tmax/δt)
    tsp = range(0.0, Tmax, length=NSTEPS+1)

    m = size(p.s.K, 1)
    B = Matrix(1.0I, m, m)
    X = Universe(m)

    H = Heaviside
    if case == "constant" || case == "heaviside"
        Fval = -1e6
        Fval = [p.b * Fval for ti in tsp]

    elseif case == "ricker"
        ancho_escalon = 0.05 # 50 ms de ancho
        # por cuestion de simetria se usa 2e6 / 2 respecto del articulo de Accurate Wave Prop
        fr(t) = -1e6 * (H(ancho_escalon*3 - t) - 3*H(ancho_escalon*2 - t) + 3*H(ancho_escalon - t))
        Fval = [p.F * fr(ti) for ti in tsp] # "step-loading" en el articulo

    else
        throw(ArgumentError("case $case not defined"))
    end
    sys = SecondOrderConstrainedLinearControlContinuousSystem(sparse(p.s.M), p.s.C, p.s.K, B, Universe(m), Fval)

    if isnothing(X0)
        X0 = (zeros(m), zeros(m))
    end
    prob = @ivp(sys, u(0) ∈ X0)
    alg = alg(Δt=δt)

    sol_statistics = StructuralDynamicsODESolvers._solve_statistics(alg, prob, NSTEPS, sampler, idx, vecsamples)

    # invertimos el signo de todos los desplazamientos para comparar con el articulo Accurate Wave prop
    #sol.U .= -1 .* sol.U

    return sol_statistics
end

function plot_nodo(sol; nodeidx=3, direction="Horizontal", method="Newmark")

    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{%$direction displacement} (\times 10^{-5} \textrm{m})",
               legend=:topright,
               legendfontsize=25,
               tickfont=font(20, "Times"),
               guidefontsize=10,
               xguidefont=font(20, "Times"),
               yguidefont=font(20, "Times"),
               xtick = xticklatex([0, 0.2, 0.4, 0.6, 0.8, 1.0], 2),
               ytick = ([-4e-5, -2e-5, 0, 2e-5, 4e-5, 6e-5], [L"-4", L"-2", L"0", L"2", L"4", L"6"]),
               xlims=(0.0, 1.0),
               ylims=(-5e-5, 8e-5),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    plot!(fig, sol, vars=(0, nodeidx), seriestype=:path, markershape=:none, lw=2.0, lab=L"\textrm{%$method}")
    fig
end

function plot_lambs(p::FEMProblem; nodeidx=3, direction="Horizontal", case="heaviside", kwargs...)

    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{%$direction displacement} (\times 10^{-5} \textrm{m})",
               legend=:topright,
               legendfontsize=25,
               tickfont=font(20, "Times"),
               guidefontsize=10,
               xguidefont=font(20, "Times"),
               yguidefont=font(20, "Times"),
               xtick = xticklatex([0, 0.2, 0.4, 0.6, 0.8, 1.0], 2),
               ytick = ([-4e-5, -2e-5, 0, 2e-5, 4e-5, 6e-5], [L"-4", L"-2", L"0", L"2", L"4", L"6"]),
               xlims=(0.0, 1.0),
               ylims=(-5e-5, 8e-5),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    # resolver numericas
    sol_newmark = sol_numerica_lamb(p; nodeidx=nodeidx, case=case, alg=Trapezoidal, kwargs...)
    sol_bathe = sol_numerica_lamb(p; nodeidx=nodeidx, case=case, alg=Bathe, kwargs...)

    # resolver flowpipe
    out = _sol_reachability_lamb(p; case=case, nodeidx=nodeidx, kwargs...)

    plot!(fig, out, vars=(0, nodeidx), lc=:blue, c=:blue, lab=L"\textrm{Flowpipe}")
    plot!(fig, sol_newmark, vars=(0, nodeidx), seriestype=:path, markershape=:none, lw=2.0, lab=L"\textrm{Newmark}")
    plot!(fig, sol_bathe, vars=(0, nodeidx), seriestype=:path, markershape=:none, lw=2.0, lab=L"\textrm{Bathe}")
    fig
end

function generar_figura_malla2_singleton_position()

    # =================
    # setear parametros
    # =================
    p = load_lambs(mesh=2);
    @show m = size(p.s.M, 1)
    nodeidx = 3
    δtfac = 1.0
    δt = step_size_lamb(p, δtfac) # 0.001
    case = "heaviside"
    condicion_inicial = "singleton"
    mk = 140
    tol = 1e-7

    # ====================
    # resolver numericas
    # ====================
    println("Resolviendo con Newmark... ")
    @time sol_newmark = sol_numerica_lamb(p; nodeidx=nodeidx, case=case, alg=Trapezoidal, δt=δt, condicion_inicial=condicion_inicial)

    println("Resolviendo con Bathe... ")
    @time sol_bathe = sol_numerica_lamb(p; nodeidx=nodeidx, case=case, alg=Bathe, δt=δt, condicion_inicial=condicion_inicial)

    # =================
    # resolver flowpipe
    # =================
    println("Resolviendo Flowpipe... ")
    @time out = _sol_reachability_lamb(p; case=case, nodeidx=nodeidx, δt=δt, condicion_inicial=condicion_inicial, mk=mk, tol=tol)

    # =================
    # graficar
    # =================
    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{Displacement} (\times 10^{-5} \textrm{m})",
               legend=:bottomright,
               legendfontsize=15,
               tickfont=font(25, "Times"),
               guidefontsize=15,
               xguidefont=font(25, "Times"),
               yguidefont=font(25, "Times"),
               xtick = xticklatex([0, 0.2, 0.4, 0.6, 0.8, 1.0], 2),
               ytick = ([-1e-5, 0, 1e-5, 2e-5, 3e-5, 4e-5], [L"-1", L"0", L"1", L"2", L"3", L"4"]),
               xlims=(0.0, 1.0),
               ylims=(-1.2e-5, 4.2e-5),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    plot!(fig, out, vars=(0, nodeidx), c=:lightblue, lc=:black, lw=0.5, alpha=1., lab=L"\textrm{Flowpipe}")

    # grafico continuo
    plot!(fig, sol_newmark, vars=(0, nodeidx), seriestype=:path, marker=:none, lw=2.0, color=:red, lab="")
    plot!(fig, sol_bathe, vars=(0, nodeidx), seriestype=:path, marker=:none, lw=2.0, color=:green, lab="")

    # con markers salteados
    paso = 15
    npoints = length(sol_newmark.U)
    plot!(fig, sol_newmark.t[1:paso:npoints], [sol_newmark.U[i][nodeidx] for i in 1:paso:npoints], vars=(0, nodeidx), linetype=:scatter, marker=:circle, markersize=8, lw=2.0, color=:red, lab=L"\textrm{Newmark}")
    plot!(fig, sol_bathe.t[1:paso:npoints], [sol_bathe.U[i][nodeidx] for i in 1:paso:npoints], vars=(0, nodeidx), linetype=:scatter, marker=:utriangle, markersize=8, lw=2.0, color=:green, lab=L"\textrm{Bathe}")

    # grafico bounding box donde hago zoom
    tzoom_left = 0.40
    tzoom_right = 0.50
    yzoom_bottom = 2e-5
    yzoom_top = 4e-5

    PBotLeft  = [tzoom_left, yzoom_bottom]
    PBotRight = [tzoom_right, yzoom_bottom]
    PTopLeft  = [tzoom_left, yzoom_top]
    PTopRight = [tzoom_right, yzoom_top]

    L1 = LineSegment(PBotLeft, PBotRight)
    L2 = LineSegment(PBotRight, PTopRight)
    L3 = LineSegment(PTopRight, PTopLeft)
    L4 = LineSegment(PTopLeft, PBotLeft)

    plot!(fig, L1, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)
    plot!(fig, L2, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)
    plot!(fig, L3, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)
    plot!(fig, L4, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)

    savefig(fig, "malla2_singleton_pos_nodeidx_3.pdf")

    # =============================
    # graficar parte zoomeada
    # =============================

    # como todos los solvers usan el mismo step size, puedo usar Newmark para encontrar
    # los indices donde hago zoom
    a = findfirst(>(0), sol_newmark.t .- tzoom_left)
    b = findfirst(>(0), sol_newmark.t .- tzoom_right)

    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{Displacement} (\times 10^{-5} \textrm{m})",
               legend=:bottomright,
               legendfontsize=15,
               tickfont=font(25, "Times"),
               guidefontsize=15,
               xguidefont=font(25, "Times"),
               yguidefont=font(25, "Times"),
               xtick = xticklatex([0.4, 0.425, 0.45, 0.475, 0.5], 3),
               ytick = ([2.5e-5, 2.75e-5, 3.0e-5, 3.25e-5, 3.5e-5], [L"2.5", L"2.75", L"3.0", L"3.25", L"3.5"]),
               xlims=(0.395, 0.505),
               ylims=(2.4e-5, 3.7e-5),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    plot!(fig, out[a:b], vars=(0, nodeidx), c=:lightblue, lc=:black, lw=0.5, alpha=1., lab=L"\textrm{Flowpipe}")

    plot!(fig, sol_newmark.t[a:b], [sol_newmark.U[i][nodeidx] for i in a:b], vars=(0, nodeidx), seriestype=:path, markershape=:none, lw=2.0, color=:red, lab="", alpha=1.)
    plot!(fig, sol_bathe.t[a:b], [sol_bathe.U[i][nodeidx] for i in a:b], seriestype=:path, markershape=:none, lw=2.0, color=:green, lab="", alpha=1.)

    # con markers salteados
    paso = 2
    plot!(fig, sol_newmark.t[a:paso:b], [sol_newmark.U[i][nodeidx] for i in a:paso:b], vars=(0, nodeidx), linetype=:scatter, marker=:circle, markersize=10, lw=2.0, color=:red, lab=L"\textrm{Newmark}")
    plot!(fig, sol_bathe.t[a:paso:b], [sol_bathe.U[i][nodeidx] for i in a:paso:b], vars=(0, nodeidx), linetype=:scatter, marker=:utriangle, markersize=10, lw=2.0, color=:green, lab=L"\textrm{Bathe}")

    savefig(fig, "malla2_singleton_pos_nodeidx_3_zoom.pdf")

    return nothing
end

function generar_figura_malla2_singleton_velocity()

    # =================
    # setear parametros
    # =================
    p = load_lambs(mesh=2);
    @show m = size(p.s.M, 1)
    nodeidx = 3
    δtfac = 1.0
    δt = step_size_lamb(p, δtfac) # 0.001
    case = "heaviside"
    condicion_inicial = "singleton"
    mk = 140
    tol = 1e-7

    # ====================
    # resolver numericas
    # ====================
    println("Resolviendo con Newmark... ")
    @time sol_newmark = sol_numerica_lamb(p; nodeidx=nodeidx, case=case, alg=Trapezoidal, δt=δt, condicion_inicial=condicion_inicial)

    println("Resolviendo con Bathe... ")
    @time sol_bathe = sol_numerica_lamb(p; nodeidx=nodeidx, case=case, alg=Bathe, δt=δt, condicion_inicial=condicion_inicial)

    # =================
    # resolver flowpipe
    # =================
    println("Resolviendo Flowpipe... ")
    @time out = _sol_reachability_lamb(p; case=case, nodeidx=nodeidx+m, δt=δt, condicion_inicial=condicion_inicial, mk=mk, tol=tol)

    # =================
    # graficar
    # =================
    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{Velocity} (\times 10^{-4} \textrm{m/s})",
               legend=:topright,
               legendfontsize=15,
               tickfont=font(25, "Times"),
               guidefontsize=15,
               xguidefont=font(25, "Times"),
               yguidefont=font(25, "Times"),
               xtick = xticklatex([0, 0.2, 0.4, 0.6, 0.8, 1.0], 2),
               ytick = ([-1e-4, 0, 1e-4, 2e-4, 3e-4, 4e-4, 5e-4], [L"-1", L"0", L"1", L"2", L"3", L"4", L"5"]),
               xlims=(-0.02, 1.0),
               ylims=(-2.1e-4, 5.5e-4),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    plot!(fig, out, vars=(0, nodeidx+m), c=:lightblue, lc=:black, lab=L"\textrm{Flowpipe}", alpha=1.)

    # grafico continuo
    plot!(fig, sol_newmark.t, [sol_newmark.U′[i][nodeidx] for i in 1:length(sol_newmark.U′)], vars=(0, nodeidx), seriestype=:path, markershape=:none, lw=2.0, lab="")
    plot!(fig, sol_bathe.t, [sol_bathe.U′[i][nodeidx] for i in 1:length(sol_bathe.U′)], seriestype=:path, markershape=:none, lw=2.0, lab="")

    # con markers salteados
    paso = 15
    npoints = length(sol_newmark.U)
    plot!(fig, sol_newmark.t[1:paso:npoints], [sol_newmark.U′[i][nodeidx] for i in 1:paso:npoints], vars=(0, nodeidx), linetype=:scatter, marker=:circle, markersize=10, lw=2.0, c=:red, lab=L"\textrm{Newmark}")
    plot!(fig, sol_bathe.t[1:paso:npoints], [sol_bathe.U′[i][nodeidx] for i in 1:paso:npoints], linetype=:scatter, marker=:utriangle, markersize=10, lw=2.0, c=:green, lab=L"\textrm{Bathe}")

    # grafico bounding box donde hago zoom
    tzoom_left = 0.8
    tzoom_right = 0.99
    yzoom_bottom = -1e-4
    yzoom_top = 1e-4

    PBotLeft  = [tzoom_left, yzoom_bottom]
    PBotRight = [tzoom_right, yzoom_bottom]
    PTopLeft  = [tzoom_left, yzoom_top]
    PTopRight = [tzoom_right, yzoom_top]

    L1 = LineSegment(PBotLeft, PBotRight)
    L2 = LineSegment(PBotRight, PTopRight)
    L3 = LineSegment(PTopRight, PTopLeft)
    L4 = LineSegment(PTopLeft, PBotLeft)

    plot!(fig, L1, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)
    plot!(fig, L2, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)
    plot!(fig, L3, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)
    plot!(fig, L4, color=:black, style=:dash, lw=3.0, seriestype=:solid, marker=:none, alpha=1.)

    savefig(fig, "malla2_singleton_vel_nodeidx_3.pdf")

    # =============================
    # graficar parte zoomeada
    # =============================

    # como todos los solvers usan el mismo step size, puedo usar Newmark para encontrar
    # los indices donde hago zoom
    a = findfirst(>(0), sol_newmark.t .- tzoom_left)
    b = findfirst(>(0), sol_newmark.t .- tzoom_right)

    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{Velocity} (\times 10^{-4} \textrm{m/s})",
               legend=:topright,
               legendfontsize=15,
               tickfont=font(25, "Times"),
               guidefontsize=15,
               xguidefont=font(25, "Times"),
               yguidefont=font(25, "Times"),
               xtick = xticklatex([0.8, 0.9, 1.0], 3),
               ytick = ([-1.0e-4, -0.5e-4, 0, 0.5e-4, 1e-4], [L"-1.0", L"-0.5", L"0.0", L"0.5", L"1.0"]),
               xlims=(0.79, 1.01),
               ylims=(-1.05e-4, 1.05e-4),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    plot!(fig, out[a:b-1], vars=(0, nodeidx+m), lc=:black, c=:lightblue, lw=0.5, alpha=1., lab=L"\textrm{Flowpipe}")
    plot!(fig, sol_bathe.t[a:b], [sol_bathe.U′[i][nodeidx] for i in a:b], seriestype=:path, markershape=:none, c=:red, lw=2.0, lab="", alpha=1.)
    plot!(fig, sol_newmark.t[a:b], [sol_newmark.U′[i][nodeidx] for i in a:b], vars=(0, nodeidx), seriestype=:path, markershape=:none, c=:green, lw=2.0, lab="", alpha=1.)

    # con markers salteados
    paso = 2
    plot!(fig, sol_newmark.t[a:paso:b], [sol_newmark.U′[i][nodeidx] for i in a:paso:b], vars=(0, nodeidx), linetype=:scatter, marker=:circle, markersize=10, lw=2.0, color=:red, alpha=1., lab=L"\textrm{Newmark}")
    plot!(fig, sol_bathe.t[a:paso:b], [sol_bathe.U′[i][nodeidx] for i in a:paso:b], vars=(0, nodeidx), linetype=:scatter, marker=:utriangle, markersize=10, lw=2.0, color=:green, alpha=1., lab=L"\textrm{Bathe}")

    savefig(fig, "malla2_singleton_vel_nodeidx_3_zoom.pdf")

    return nothing
end

# ========================================
# Metodo de reachability
# ========================================

function _sol_reachability_lamb(p::FEMProblem; δt=nothing, nodeidx=[3, 4],
                                case="ricker", δtfac=1.0, alg=Trapezoidal,
                                mk=150, tol=1e-12,
                                Eplus_optim=true,
                                kwargs...)

    if case == "heaviside"
        return _sol_reachability_lamb_heaviside(p; δt=δt, nodeidx=nodeidx, δtfac=δtfac, alg=alg, mk=mk, tol=tol, Eplus_optim=Eplus_optim, kwargs...)
    else
        error("ver lambs_krylov_old.jl")
    end
end

function _sol_reachability_lamb_heaviside(p::FEMProblem;
                                          δtfac=1.0,
                                          δt=step_size_lamb(p, δtfac),
                                          nodeidx=[3, 4],
                                          alg=Trapezoidal,
                                          mk=150, tol=1e-12,
                                          Eplus_optim=true,
                                          condicion_inicial="singleton",

                                          # para condicion inicial "distribuida"
                                          radio_desplazamientos = nothing, # 3.5e-7
                                          radio_velocidades = nothing, # 5e-6

                                          kwargs...)

    # ===================
    # Definir la carga
    # ===================

    # @assert δt <= δtnom

    Tmax = 1.0 # seconds
    NSTEPS = round(Int, Tmax/δt)

    H = Heaviside

    ancho_escalon = 0.05 # 50 ms de ancho
    # por cuestion de simetria se usa 2e6 / 2 respecto del articulo de Accurate Wave Prop
    Δ = ancho_escalon
    α = -1e6
    fload(t) = α * H(t)

    # ==============================
    # Homogeneizar
    # ==============================
    m = size(p.s.K, 1)
    Shom = _homogeneize(p)
    Aext = Shom.A
    @assert size(Aext, 1) == 2m+1

    # =====================================
    # Preparacion de la condicion inicial
    # =====================================

    if condicion_inicial == "singleton"
        X0_init = Singleton(zeros(2m+1))
        X0_init.element[end] = α

    elseif condicion_inicial == "distribuida"
        X0_init = Hyperrectangle(zeros(2m+1), zeros(2m+1))
        if isnothing(radio_desplazamientos) || isnothing(radio_velocidades)
            error("debe especificar `radio_desplazamientos` y `radio_velocidades`")
        end
        X0_init.center[end] = α # fuerza
        X0_init.radius[1:m] .= radio_desplazamientos # desplazamientos iniciales
        X0_init.radius[m+1:2m] .= radio_velocidades # velocidades iniciales

    else
        error("condicion_inicial = $condicion_inicial no existe")
    end

    # ==============================
    # Propagacion a tiempo continuo
    # ==============================
    Aᵀδ = copy(transpose(Aext .* δt))
    num_nodes = isa(nodeidx, Int) ? 1 : length(nodeidx)

    # alg = LGG09(δ=δt, n=2m+1, vars=nodeidx)
    dirs = Vector{LazySets.Arrays.SingleEntryVector{Float64}}()
    for j in 1:num_nodes
        node = nodeidx[j]
        push!(dirs, SingleEntryVector(node, 2m+1, 1.0))  # 0, 0, 1, ... , 0, 0, 0, 0, 0, 0
        push!(dirs, SingleEntryVector(node, 2m+1, -1.0))
    end
    cdirs = CustomDirections(dirs)

    println("Calculando el factor de bloating...")

    Eplus = _build_Eplus_optim(Aext, X0_init, δt; m=mk, tol=tol, cutoff=eps(Float64))

    println("Propagando a tiempo continuo...")

    if condicion_inicial == "singleton"
        ΦX0_init = _expv(Aext, X0_init.element, 1, δt, m=mk, tol=tol) |> Singleton
        Ω0 = CH(X0_init, ΦX0_init ⊕ Eplus)

        Mout = Matrix{Float64}(undef, 2*num_nodes, NSTEPS)
        for j in 1:num_nodes
            node = nodeidx[j]
            d = zeros(2m+1)
            d[node] = 1.0
            println("Propagando indice $node ...")
            idx = (2(j-1) + 1):2j
            _reach_homog_krylov_LGG09_modif!(view(Mout, idx, :), Ω0, Aᵀδ, d, NSTEPS; m=mk, tol=tol)
        end
        F = [TemplateReachSet(cdirs, view(Mout, :, k), (0.0 .. δt) + (k-1) * δt) for k in 1:NSTEPS]

    elseif condicion_inicial == "distribuida"
        # Ω0 = CH(X0_init, ΦX0_init * X0_init ⊕ Eplus)

        Mout = Matrix{Float64}(undef, 2*num_nodes, NSTEPS)
        for j in 1:num_nodes
            node = nodeidx[j]

            d = zeros(2m+1)
            d[node] = 1.0

            println("Propagando indice $node ...")
            idx = (2(j-1) + 1):2j
            _reach_homog_krylov_LGG09_modif!(view(Mout, idx, :), X0_init, Aext, δt, Eplus, d, NSTEPS; m=mk, tol=tol)
        end
        F = [TemplateReachSet(cdirs, view(Mout, :, k), (0.0 .. δt) + (k-1) * δt) for k in 1:NSTEPS]

    else
        error("condicion_inicial = $condicion_inicial no existe")
    end

    # niego para ajustar el sistema de coordenadas con el articulo
    proyectar = get(kwargs, :proyectar, false)
    if proyectar
        # OJO si uso esto, el flowpipe tiene una sola variable 1
        @assert length(nodeidx) == 1
        out = [ReachSet(Interval(-Mout[2, k], Mout[1, k]), (0.0 .. δt) + (k-1) * δt) for k in 1:size(Mout, 2)]
    else
        #α = -1.0
        α = 1.0
        out = [ReachSet(α * set(R), tspan(R)) for R in F]
    end

    return out
end

# computes the support functions of Ω0 = CH(X0, Φ*X0 ⊕ Eplus)
# along +d and -d for NSTEPS
#
# for k = 0, 1, ..., NSTEPS-1
#
# ρ(d, Φ^k Ω0) = ρ((Φ^T)^k d, CH(X0, Φ * X0 ⊕ Eplus))
#              = max(ρ((Φ^T)^k d, X0), ρ((Φ^T)^k d, Φ * X0) + ρ((Φ^T)^k d, Eplus))
#              = max(ρ((Φ^T)^k d, X0), ρ((Φ^T)^(k+1) d, X0) + ρ((Φ^T)^k d, Eplus))
#              = max(ρ(vk, X0), ρ((Φ^T)^k * vk, X0) + ρ(vk, Eplus)), vk := (Φ^T)^k d
#
function _reach_homog_krylov_LGG09_modif!(out, X0, A, δ, Eplus,
                                          d::AbstractVector, NSTEPS;
                                          hermitian=false, m=min(30, size(Aᵀδ, 1)), tol=1e-7)

    @assert size(out, 1) == 2
    @assert size(out, 2) == NSTEPS

    Aᵀδ = transpose(A) .* δ

    # initialization of the krylov subspace
    # esta parte genera una base de
    # {Aᵀδ * d, (Aᵀδ)^2 * d, (Aᵀδ)^3 * d, ..., (Aᵀδ)^m * d}
    TA, Tb = eltype(Aᵀδ), eltype(d)
    T = promote_type(TA, Tb)
    Ks = KrylovSubspace{T, real(T)}(length(d), m)
    arnoldi!(Ks, Aᵀδ, d; m=m, ishermitian=hermitian, tol=tol)

    # rᵢ stores is the cache for each vector: (Φᵀ)^i d and (Φᵀ)^{i+1} d
    rᵢ = deepcopy(d)
    rᵢ₊₁ = similar(d)

    #
    # TODO se puede optimizar el loop ya que el ρ(rᵢ₊₁, X0) se repite en ρ(rᵢ, X0)
    #
    @inbounds for i in 1:NSTEPS
        expv!(rᵢ₊₁, i*1.0, Ks) # rᵢ₊₁ <- exp(A^T δ * i) * d
        out[1, i] = max(ρ(rᵢ, X0), ρ(rᵢ₊₁, X0) + ρ(rᵢ, Eplus))
        out[2, i] = max(ρ(-rᵢ, X0), ρ(-rᵢ₊₁, X0) + ρ(-rᵢ, Eplus))
        copyto!(rᵢ, rᵢ₊₁)
    end
    return out
end

# ===========================
# estadisticas con numericas
# ===========================

using Sobol
using StructuralDynamicsODESolvers: SolutionExtrema

sol_numericas_random_statistics(p; m, vecsamples::Int, idx, radio_desplazamientos, radio_velocidades) = sol_numericas_random_statistics(p; m=m,
    vecsamples=[vecsamples], idx=idx, radio_desplazamientos=radio_desplazamientos, radio_velocidades=radio_velocidades)

# ejemplo: vecsamples = [10, 20, 100]
function sol_numericas_random_statistics(p; m, vecsamples::Vector{Int}, idx, δt, radio_desplazamientos, radio_velocidades)
    #X0rand_desplazamientos = [[2(rand() - 0.5) * 5e-7 for _ in 1:825] for i in 1:NTRAJ]
    #X0rand_velocidades = [[2(rand() - 0.5) * 5e-7 for _ in 1:825] for i in 1:NTRAJ]
    #remake = [(X0d, X0v) for (X0d, X0v) in zip(X0rand_desplazamientos, X0rand_velocidades)]

    H0 = Hyperrectangle(zeros(2m), vcat(fill(radio_desplazamientos, m), fill(radio_velocidades, m)))

    sampleado = :vertices

    if sampleado == :sobol
        # para samplear dentro de H0
        sampler = SobolSeq(H0)
    elseif sampleado == :vertices
        # para samplear en los vertices
        sampler = VertexSampler(H0)
    end

    @time out = sol_numerica_lamb_statistics(p; case="heaviside", alg=Trapezoidal, δt=δt, X0=nothing, sampler=sampler, idx=idx, vecsamples=vecsamples)
end

Sobol.SobolSeq(H::Hyperrectangle) = SobolSeq(low(H), high(H))

struct VertexSampler{T}
    H::T
end

function StructuralDynamicsODESolvers._next!(S::VertexSampler, unew)
    c = S.H.center
    r = S.H.radius
    for i in 1:length(unew)
        a = ifelse(rand() > 0.5, 1, -1)
        unew[i] = c[i] + a*r[i]
    end
end

StructuralDynamicsODESolvers._next!(s::ScaledSobolSeq, unew) = next!(s, unew)

function generar_figura_malla2_envolventes_position(; vecsamples = nothing)

    # =================
    # setear parametros
    # =================
    p = load_lambs(mesh=2);
    @show m = size(p.s.M, 1)
    nodeidx = 3
    δt = step_size_lamb(p, 1.0) # 0.001
    case = "heaviside"
    condicion_inicial = "distribuida"
    mk = 140
    tol = 1e-7

    radio_desplazamientos = 3.5e-7
    radio_velocidades = 5e-6

    # ====================
    # correr simulaciones
    # ====================
    out_extrema_malla2 = nothing
    if !isnothing(vecsamples)
        out_extrema_malla2 = sol_numericas_random_statistics(p, m=m, idx=nodeidx, vecsamples=vecsamples, δt=δt, radio_desplazamientos=radio_desplazamientos, radio_velocidades=radio_velocidades)
    end

    # ====================
    # calcular flowpipe
    # ====================
    @time out_flowpipe = _sol_reachability_lamb(p, case=case, condicion_inicial=condicion_inicial, nodeidx=nodeidx, mk=mk, tol=tol, δt=δt, proyectar=true, radio_desplazamientos=radio_desplazamientos, radio_velocidades=radio_velocidades)

    # ====================
    # graficar flowpipe
    # ====================
    # q = length(vecsamples)
    # c = distinguishable_colors(q);
    fig = plot()

    plot!(fig, out_flowpipe, vars=(0, 1), lw=0.0, lc=:lightblue, c=:lightblue, lab=L"\textrm{Flowpipe}")

    if !isnothing(vecsamples)
        [plot!(fig, out_extrema_malla2.t, out_extrema_malla2.hist[k][:U_min], lw=2.0, lab="$k") for (i, k) in enumerate(vecsamples)]
        [plot!(fig, out_extrema_malla2.t, out_extrema_malla2.hist[k][:U_max], lw=2.0, lab="") for (i, k) in enumerate(vecsamples)]
    end
    savefig(fig, "malla2_distribuida_pos_nodeidx_3.pdf")

    out_flowpipe, out_extrema_malla2, fig
end

function generar_figura_malla2_envolventes_velocity(; vecsamples = nothing)

    # =================
    # setear parametros
    # =================
    p = load_lambs(mesh=2);
    @show m = size(p.s.M, 1)
    nodeidx = 3
    δt = step_size_lamb(p, 1.0)
    case = "heaviside"
    condicion_inicial = "distribuida"
    mk = 140
    tol = 1e-7

    radio_desplazamientos = 3.5e-7
    radio_velocidades = 5e-6

    # ====================
    # correr simulaciones
    # ====================
    println("Corriendo simulaciones...")
    out_extrema_malla2 = nothing
    if !isnothing(vecsamples)
        out_extrema_malla2 = sol_numericas_random_statistics(p, m=m, idx=nodeidx, vecsamples=vecsamples, δt=δt, radio_desplazamientos=radio_desplazamientos, radio_velocidades=radio_velocidades)
    end

    # ====================
    # calcular flowpipe
    # ====================
    println("Calculando el flowpipe...")
    @time out_flowpipe = _sol_reachability_lamb(p, case=case, condicion_inicial=condicion_inicial, nodeidx=nodeidx+m, mk=mk, tol=tol, δt=δt, proyectar=true, radio_desplazamientos=radio_desplazamientos, radio_velocidades=radio_velocidades)

    # ====================
    # graficar flowpipe
    # ====================
    # q = length(vecsamples)
    # c = distinguishable_colors(q);
    fig = plot()

    plot!(fig, out_flowpipe, vars=(0, 1), lw=0.0, lc=:lightblue, c=:lightblue, lab=L"\textrm{Flowpipe}")
    [plot!(fig, out_extrema_malla2.t, out_extrema_malla2.hist[k][:Udot_min], lw=2.0, lab="$k") for (i, k) in enumerate(vecsamples)]
    [plot!(fig, out_extrema_malla2.t, out_extrema_malla2.hist[k][:Udot_max], lw=2.0, lab="") for (i, k) in enumerate(vecsamples)]

    savefig(fig, "malla2_distribuida_vel_nodeidx_3.pdf")

    out_flowpipe, out_extrema_malla2, fig
end

function _norma_inf_envolvente_numericas(sol, vecsamples)
    sol_norma_inf = Vector{Float64}()
    for k in vecsamples
        Udotmin_k = sol.hist[k][:Udot_min] .|> abs
        Udotmax_k = sol.hist[k][:Udot_max] .|> abs

        # norma infinito de las envolventes numericas
        sol_norma_inf_k = maximum(max.(Udotmin_k, Udotmax_k))
        push!(sol_norma_inf, sol_norma_inf_k)
    end
    return sol_norma_inf
end

function _norma_1_envolvente_numericas(sol, vecsamples)
    sol_norma_1 = Vector{Float64}()
    for k in vecsamples
        Udotmin_k = sol.hist[k][:Udot_min] .|> abs
        Udotmax_k = sol.hist[k][:Udot_max] .|> abs

        # norma infinito de las envolventes numericas
        sol_norma_1_k = sum(max(Udotmin_k[i], Udotmax_k[i]) for i in 1:length(Udotmin_k))
        push!(sol_norma_1, sol_norma_1_k)
    end
    δt = out_extrema_malla2_vel.alg.Δt
    return sol_norma_1 * δt
end

function generar_figura_malla2_envolvente_velocidad()

    results = matread("results.mat")

    vecsamples = "N" .* string.(results["vecsamples"])

    times = results["out_extrema_malla2_vel_t"]
    vmin = [results["out_extrema_malla2_vel_hist"][k]["NUdot_min"] for k in vecsamples]
    vmax = [results["out_extrema_malla2_vel_hist"][k]["NUdot_max"] for k in vecsamples];

    xx = [Interval(a, b) for (a, b) in zip(results["out_flowpipe_vel_lo"], results["out_flowpipe_vel_hi"])];
    tt = [IntervalArithmetic.Interval(a, b) for (a, b) in zip(results["out_flowpipe_vel_tstart"], results["out_flowpipe_vel_tend"])];
    F = [ReachSet(a, b) for (a, b) in zip(xx, tt)];

    cols = distinguishable_colors(length(vecsamples))

    fig = plot(xlab=L"\textrm{Time (s)}",
               ylab=L"\textrm{Velocity} (\times 10^{-4} \textrm{m/s})",
               legend=:outerright,
               legendfontsize=20,
               tickfont=font(20, "Times"),
               guidefontsize=10,
               xguidefont=font(20, "Times"),
               yguidefont=font(20, "Times"),
               xtick = xticklatex([0, 0.2, 0.4, 0.6, 0.8, 1.0], 2),
               ytick = ([-10e-4, -5e-4, 0, 5e-4, 10e-4], [L"-10", L"-5", L"0", L"5", L"10"]),
               xlims=(0.0, 1.0),
               ylims=(-1.3e-3, 1.2e-3),
               bottom_margin=10mm, left_margin=6mm, right_margin=12mm, top_margin=3mm, size=(900, 600))

    labs = [L"10^0", L"10^1", L"10^2", L"10^3", L"10^4", L"10^5"]

    plot!(fig, F, vars=(0, 1), lw=1., lc=:lightblue, c=:lightblue, lab=L"\textrm{Flowpipe}")
    [plot!(fig, times, vmin[i], lw=2.0, c=cols[i], lab=labs[i], ls=:solid) for i in 1:length(vmin)]
    [plot!(fig, times, vmax[i], lw=2.0, c=cols[i], lab="", ls=:solid) for i in 1:length(vmin)]

    savefig(fig, "fig_vel.pdf")

    fig
end


# To test a simplified quadrotor.
# Modeling the quadrotor frame as a cube and the props
# as another rigid body with a revolute joint between the two.

include("../src/plotSol.jl")
include("../src/simulateDiff.jl")
include("../src/OrientationConversion.jl")
##
clearconsole()
j1 = Nothing
j2 = Nothing
##
mb = 0.155; mp = 0.02;

Ib =  [ 0.0009    0.0000   -0.0000;
        0.0000    0.0053    0.0000;
       -0.0000    0.0000    0.0054]#*(10^-3)

Ip =    [ 0.0029    0.0000    0.0000;
          0.0000    0.5824    0.0000;
          0.0000    0.0000    0.5805]*(10^-4)

InFrame = InertialFrameAsRB()
QuadCube = RigidBody(mb,Ib,2)
Props = RigidBody(mp,Ip,3)

x0Cube = [zeros(3);[1;zeros(3)];zeros(3);zeros(4)]
initialiseRigidBody!(QuadCube,x0Cube)
x0Props = [[0.0;0.0;0.5];[1;zeros(3)];zeros(3);zeros(4)]
initialiseRigidBody!(Props,x0Props)

axis = [0.0 0.0 1.0][:]
rjCube = [0.0 0.0 0.5][:]
rjProp = [0.0 0.0 0.0][:]

j1 = Joint(InFrame,QuadCube,zeros(3),zeros(3))
j2 = Joint(QuadCube,Props,rjCube,rjProp,type = "Revolute",axis = axis);#,jointTorque = [0.0;0.0;0.0001])

j = [j1,j2];

## Linearization
# x0, u0 = getXU_0(j);
# A,B = linearizeDiff(x0, u0, j);
# println("A = ",A)
# println("B = ",B)
# sleep(1200)
## Simulation
tEnd = 0.01
tSpan = 0.01
g = [0.0,0.0,0.0]; # Gravity Vector.
@btime simulateDiff(tEnd, tInt, j, g=g);
sleep(1000);
tSim, solFinal = simulate(tEnd,tSpan,j,g=g)

solQuad = solFinal[1]
solProp = solFinal[2]
##
ωCube, ωProp, ωRel, jointLoc = checkRevJoint(solQuad, solProp, rjCube, rjProp);

##
plotErrNorm(tSim,solQuad.β)
plotErrNorm(tSim,solProp.β)
# Check if joint location has moved
# jointLoc = Matrix{Float64}(undef,length(tSim),3)
# ωCube = Matrix{Float64}(undef,length(tSim),3)
# ωProp = Matrix{Float64}(undef,length(tSim),3)
# ωPropInCube = Matrix{Float64}(undef,length(tSim),3)
# ωRel = Matrix{Float64}(undef,length(tSim),3)
# for i=1:length(tSim)
#     QuadCube.dcm = quat2dcm(solQuad.β[i,:])
#     ωCube[i,:] = angVel(solQuad.β[i,:],solQuad.βdot[i,:])
#     ωProp[i,:] = angVel(solProp.β[i,:],solProp.βdot[i,:])
#     ωPropInCube[i,:] = quat2dcm(solQuad.β[i,:])*transpose(quat2dcm(solProp.β[i,:]))*ωProp[i,:]
#     ωRel[i,:] = ωPropInCube[i,:] - ωCube[i,:]
#     jointLoc[i,:] = solQuad.r[i,:] + transpose(QuadCube.dcm)*rjCube - solProp.r[i,:]
# end
##
close("all");
plotPos(tSim,jointLoc)
plotPos(tSim, solQuad.r)
# plotVel(tSim,solQuad.v - solProp.v)
# plotAngVel(tSim,ωCube)
plotAngVel(tSim,ωProp)
# plotAngVel(tSim, ωRel)
# plotPos(tSim, solQuad.r)
# plotPos(tSim, solProp.r)
# plotQuat(tSim, solProp.β)
# plotPos(tSim, solQuad.r + solProp.r)
##

# plotAngVel(tSim,ωRel)
# plotAngVel(tSim,ωCube - ωProp)

# # Check revJoint axis
# axisSol = Matrix{Float64}(undef,length(tSim),3)
# errAxis = Matrix{Float64}(undef,length(tSim),3)
# errAxisNorm = Vector{Float64}(undef,length(tSim))
# for i=1:length(tSim)
#     axisSol[i,:] = solProp.β[i,2:4]/norm(solProp.β[i,2:4],2)
#
#     # β1inv = [solQuad.β[i,1];solQuad.β[i,2:4]]
#     # βRev = quaternionProduct(solProp.β[i,:],β1inv)
#     # axisSol[i,:] = βRev[2:4]/norm(βRev[2:4])
#
#     # axisSol[i,:] = (solQuad.β[i,:])*solProp.β[i,2:4]/norm(solProp.β[i,2:4],2)
#
#     # axisSol[i,:] = axixRev(solQuad.β[i,:],solProp.β[i,:])
#     errAxis[i,:] = axisSol[i,:] - axis
#     errAxisNorm[i] = norm(axisSol[i,:],2) - 1.0
# end
# plot(tSim[2:end],errAxisNorm[2:end])
# plotPos(tSim,errAxis)


# Test get/set objective coefficients in LP

using Gurobi, Test

@testset "LP 03" begin
	env = Gurobi.Env()
	setparam!(env, "OutputFlag", 0)

	# original model
	#
	#   maximize  2x + 2y
	#
	#	s.t. 0.2 <= x, y <= 1
	#

	model = gurobi_model(env;
		name="lp_03",
		sense=:maximize,
		f=[2.0, 2.0],
		lb=[0.2, 0.2],
		ub=[1.0, 1.0])

	lb_ = lowerbounds(model)
	ub_ = upperbounds(model)
	c_ = objcoeffs(model)

	@test lb_ == [0.2, 0.2]
	@test ub_ == [1.0, 1.0]
	@test c_ == [2.0, 2.0]

	optimize(model)

	@test get_solution(model) == [1.0, 1.0]
	@test get_objval(model) == 4.0

	# change objective (warm start)
	#
	#	maximize x - y
	#
	#	s.t. 0.2 <= x, y <= 1
	#

	set_objcoeffs!(model, [1, -1])
	update_model!(model)

	c_ = objcoeffs(model)
	@test c_ == [1.0, -1.0]

	optimize(model)

	@test get_solution(model) == [1.0, 0.2]
	@test get_objval(model) == 0.8

	GC.gc()
end


""" Test JuMP model with a variety of variable bounds and sizes. """
function test_model()
    model = Model()
    @variables(model, begin
        -5 <= x[1:5] <= 5
        -4 <= y[1:3] <= 1
        -30 <= z
        -2 <= a[1:2, 1:3]
    end)
    JuMP.set_upper_bound(a[2,2], 3)
    JuMP.set_upper_bound(a[1,1], 3)
    return model, x, y, z, a
end

""" Basic qp for testing. """
function test_qp(solver = OCTHaGOn.SOLVER_SILENT)
    m = JuMP.Model(with_optimizer(solver))
    @variable(m, 2 >= x[1:2] >= 0)
    @objective(m, Min, 3*x[1]^2 + x[2]^2 + 2*x[1]*x[2] + x[1] + 6*x[2] + 2)
    @constraint(m, 2*x[1] + 3*x[2] >= 4)
    return m
end

""" Basic gqp, same as above. """
function test_gqp(solver = OCTHaGOn.SOLVER_SILENT)
    m = JuMP.Model(with_optimizer(solver))
    @variable(m, 2 >= x[1:2] >= 0)
    @constraint(m, 2*x[1] + 3*x[2] >= 4)
    @variable(m, obj)
    @objective(m, Min, obj)
    gm = GlobalModel(model = m)
    add_nonlinear_constraint(gm, :(x -> 3*x[1]^2 + x[2]^2 + 2*x[1]*x[2] + x[1] + 6*x[2] + 2), dependent_var = obj)
    return gm
end

"""
    random_qp(dims::Int64, nconstrs::Int64, sparsity=dims, solver = OCTHaGOn.SOLVER_SILENT)

Generates a random quadratic program with the specified parameters.
"""
function random_qp(dims::Int64, nconstrs::Int64, sparsity=dims, solver = OCTHaGOn.SOLVER_SILENT)
    m = JuMP.Model(with_optimizer(solver))
    @variable(m, -1 <= x[1:dims] <= 1)
    m[:sparse_vars] = randperm(dims)[1:sparsity]
    m[:A] = rand(dims, sparsity).*2 .- 1
    m[:b] = rand(dims)
    m[:C] = rand(nconstrs, dims).*2 .- 1
    m[:d] = rand(nconstrs)
    @constraint(m, m[:C]*x .>= m[:d])
    @objective(m, Min, sum((m[:A]*x[m[:sparse_vars]] - m[:b]).^2))
    return m
end

""" Turns random_qp into a GlobalModel. """
function gmify_random_qp(m::JuMP.Model)
    @variable(m, obj)
    @objective(m, Min, obj)
    gm = GlobalModel(model = m)
    add_nonlinear_constraint(gm, :(x -> sum(($(m[:A])*x[$(m[:sparse_vars])] - $(m[:b])).^2)), vars = m[:x][m[:sparse_vars]],
                                    dependent_var = obj, name = "qp")
    return gm
end


#-------------------------------------------------------------------------------------------------
#--
#-- Ususal libraries
#--
using CSV, Dates
using DataFrames, DataFramesMeta
using Plots, PyPlot
using DifferentialEquations

#-------------------------------------------------------------------------------------------------
#--
#-- Date functions
#--
# The time unit is days (as floating point)
# Day 0 is taken at 1 March 2020
const BASE_DATE = Date(2020, 3, 1)
const BASE_DAYS = 0

function date2days(d)
    return convert(Float64, datetime2rata(d) - datetime2rata(BASE_DATE))
end

function days2date(d)
    return BASE_DATE + Day(d)
end

#%% R₀ calculations
# Default values for R_0
const BaseR₀ = 2.7

# Peak date
const peakDate = Dict(
    :north => date2days(Date(2020, 1, 1)),
    :tropical => date2days(Date(2020, 1, 1)),    # although no impact
    :south => date2days(Date(2020, 7, 1))
)

# Seasonal forcing parameter \epsilon
const ϵ = Dict(:north => 0.2, :tropical => 0.0, :south => 0.2)

# Gives R_0 at a given date
function R₀(d; r₀ = baseR₀, latitude = :north)
    eps = ϵ[latitude]
    peak = peakDate[latitude]

    return r₀ * (1 + eps * cos(2.0 * π * (d - peak) / 365.25))
end

#%% Epidemy mitigation
const DEFAULT_MITIGATION = [(0, 1.00), (30, 0.80), (60, 0.5), (90, 0.20)]

function getCurrentRatio(d; start = BASE_DAYS, schedule = DEFAULT_MITIGATION)
    l = length(schedule)

    # If l = 1, ratio will be the only one
    if l == 1
        return schedule[1][2]
    else
        for i in 2:l
            d1 = schedule[i-1][1]
            d2 = schedule[i  ][1]

            if d < d2
                deltaR = schedule[i][2] - schedule[i-1][2]
                return schedule[i-1][2] + deltaR * (d - d1) / (d2 - d1)
            end
        end

        # Last possible choice
        return schedule[l][2]
    end
end

#%% Commpartments
# The population will be modeled as a single vector.
# The vector will be a stack of several vectors, each of them represents a compartment.
# Each compartment vector has a size $nAgeGroup$ representing each age group.
# The compartments are: S, E, I, H, C, R, D, K, L
# We also track the hospital bed usage BED and ICU

# Population to compartments
function Pop2Comp(P)

    # To make copy/paste less prone to error
    g = 0

    S = P[ g*nAgeGroup + 1: (g+1)*nAgeGroup]; g += 1
    E = P[ g*nAgeGroup + 1: (g+1)*nAgeGroup]; g += 1
    I = P[ g*nAgeGroup + 1: (g+1)*nAgeGroup]; g += 1
    J = P[ g*nAgeGroup + 1: (g+1)*nAgeGroup]; g += 1
    H = P[ g*nAgeGroup + 1: (g+1)*nAgeGroup]; g += 1
    C = P[ g*nAgeGroup + 1: (g+1)*nAgeGroup]; g += 1
    R = P[ g*nAgeGroup + 1: (g+1)*nAgeGroup]; g += 1
    D = P[ g*nAgeGroup + 1: (g+1)*nAgeGroup]; g += 1
    K = P[ g*nAgeGroup + 1: (g+1)*nAgeGroup]; g += 1
    L = P[ g*nAgeGroup + 1: (g+1)*nAgeGroup]; g += 1

    h = 1
    BED = P[ g*nAgeGroup + h: g*nAgeGroup + h]; h += 1
    ICU = P[ g*nAgeGroup + h: g*nAgeGroup + h]; h += 1

    return S, E, I, J, H, C, R, D, K, L, BED, ICU
end

#-------------------------------------------------------------------------------------------------
#--
#-- EPIDEMIOLOGY
#-

#-- Susceptibility to contagion and transition from a compartment to another
const AgeGroup = ["0-9", "10-19", "20-29", "30-39", "40-49", "50-59", "60-69", "70-79", "80+"]
const zₐ =       [0.05,   0.05,   0.10,    0.15,    0.20,    0.25,    0.30,    0.40,    0.50]
const mₐ =       [0.01,   0.03,   0.03,    0.03,    0.06,    0.10,    0.25,    0.35,    0.50]
const cₐ =       [0.05,   0.10,   0.10,    0.15,    0.20,    0.25,    0.35,    0.45,    0.55]
const fₐ =       [0.30,   0.30,   0.30,    0.30,    0.30,    0.40,    0.40,    0.50,    0.50]

const nAgeGroup = length(AgeGroup)

#-- Contagion multiplier
# asymptomatic individuals (compartment E)
const γₑ = 0.50

# Infected / symptomatic individuals
const γᵢ=1.0

# Severe symptoms
const γⱼ=1.0

# Critical symptoms
const γₖ = 2.0

#-- Fatality Multiplier
# In ICU
const δᵤ = 1.0

# In hospital
const δₗ = 2.0

# Out of hospital
const δₖ = 3.0

#-- Transition times (all in days)
# Time to infectiousness (written t\_l)
const tₗ = 5.0

# Time to infectiousness (written t\_i)
const tᵢ = 3.0

# Time in hospital bed (not ICU)
const tₕ = 4.0

# Time in ICU
const tᵤ = 14.0

#-------------------------------------------------------------------------------------------------
#--
#-- MODELLING
#--

const compartments =  ["S", "E", "I", "J", "H", "C", "R", "D", "K", "L"];

# How many compartments before D in the list of solution
const nb4D = findfirst(compartments .== "D") - 1

# index of D0
const D0Index = nb4D * nAgeGroup + 1

# Helper function to never change the number of individuals in a compartment in a way that would
# make it below 0.1 (to avoid rounding errors around 0)
function ensurePositive(d,s)
    return max.(d .+ s, 0.1) .- s
end

# The dynamics of the epidemy is a function that mutates its argument with a precise signature
# Don't pay too much attetion to the print debugs/

function epiDynamics!(dP, P, params, t)

    S, E, I, J, H, C, R, D, K, L, BED, ICU = Pop2Comp(P)

    BED = BED[1]
    ICU = ICU[1]

    r₀, tₗ, tᵢ, tₕ, tᵤ, γₑ, γᵢ, γⱼ, γₖ, δₖ, δₗ, δᵤ, mitigation, BED_max, ICU_max = params
    # println(mitigation)

    ####################################
    # Arrows reflecting epidemiology - Check signs (just in case)
    EI = ones(nAgeGroup) .* E / tₗ;  EI = max.(EI, 0.0); IE = -EI
    IJ = mₐ              .* I / tᵢ;  IJ = max.(IJ, 0.0); JI = -IJ
    JK = cₐ              .* J / tₕ;  JK = max.(JK, 0.0); KJ = -JK
    HL = cₐ              .* H / tₕ;  HL = max.(HL, 0.0); LH = -HL

    # Recovery arrows
    IR = (1 .- mₐ)       .* I / tᵢ;  IR = max.(IR, 0.0); RI = -IR
    JR = (1 .- cₐ)       .* J / tₕ;  JR = max.(JR, 0.0); RJ = -JR
    HR = (1 .- cₐ)       .* H / tₕ;  HR = max.(HR, 0.0); RH = -HR
    KR = (1 .- δₖ .* fₐ) .* K / tᵤ;  KR = max.(KR, 0.0); RK = -KR
    LR = (1 .- δₗ .* fₐ) .* L / tᵤ;  LR = max.(LR, 0.0); RL = -LR
    CR = (1 .- δᵤ .* fₐ) .* C / tᵤ;  CR = max.(CR, 0.0); RC = -CR

    # Deaths
    KD = δₖ .* fₐ        .* K / tᵤ;  KD = max.(KD, 0.0); DK = -KD
    LD = δₗ .* fₐ        .* L / tᵤ;  LD = max.(LD, 0.0); DL = -LD
    CD = δᵤ .* fₐ        .* C / tᵤ;  CD = max.(CD, 0.0); DC = -CD

    ####################################
    # Bed transfers

    ####### Step 1:
    # Decrease in ICU usage after 14 days (recall that CD and CR are vectors over the age groups)
    dICU = - (sum(CD) + sum(CR))
    dICU = ensurePositive(dICU, ICU)

    # ICU beds available
    ICU_free = ICU_max - (ICU + dICU)

    # Move as many patients as possible from $L$ to $C$ in proportion of each group
    ICU_transfer = min(sum(L), ICU_free)
    LC = ICU_transfer / sum(L) .* L;                     CL = -LC

    # Overall change in ICU bed becomes
    dICU = dICU + ICU_transfer
    dICU = ensurePositive(dICU, ICU)

    # And some normal beds are freed
    dBED = -ICU_transfer
    dBED = ensurePositive(dBED, BED)
    # println(dBED); println(BED); println(floor(BED_max)); println(ICU_transfer)

    ####### Step 2:
    # Beds available
    BED_free = BED_max - (BED + dBED)

    # Move as many patients as possible from $K$ to $L$ in proportion of each group
    BED_transfer = min(sum(K), BED_free)
    KL = BED_transfer / sum(K) .* K;                     LK = -KL

    # Overall change in normal bed becomes
    dBED = dBED + BED_transfer
    dBED = ensurePositive(dBED, BED)

    ####### Step 3:
    # Beds available
    BED_free = BED_max - (BED + dBED)

    # Move as many patients as possible from $J$ to $H$ in proportion of each group
    BED_transfer = min(sum(J), BED_free)
    JH = BED_transfer / sum(J) .* J;                     HJ = -JH

    # Overall change in ICU bed becomes
    dBED = dBED + BED_transfer
    dBED = ensurePositive(dBED, BED)

    ####################################
    # Sum of all flows + Check never negative compartment

    # Susceptible
    # Calculation of β
    β = getCurrentRatio(t; start = BASE_DAYS, schedule = mitigation) .* zₐ .* r₀

    dS = -β .* (γₑ.*E + γᵢ.*I + γⱼ.*J + γₖ.*K)
    dS = min.(-0.01, dS)
    dS = ensurePositive(dS, S)
    # println(floor(sum(dS)))

    # Exposed
    dE = -dS + IE
    dE = ensurePositive(dE, E)

    # Infected.
    dI = EI + JI + RI
    dI = ensurePositive(dI, I)

    # Infected no hospital
    dJ = IJ + HJ + KJ + RJ
    dJ = ensurePositive(dJ, J)

    # Infected in hospital
    dH = JH + LH + RH
    dH = ensurePositive(dH, H)

    # Critical no hospital
    dK = JK + LK + DK + RK
    dK = ensurePositive(dK, K)

    # Critical in hospital
    dL = KL + HL + CL + DL + RL
    dL = ensurePositive(dL, L)

    # Critical in ICU
    dC = LC + DC + RC
    dC = ensurePositive(dC, C)

    # Recovery (can only increase)
    dR = IR + JR + HR + KR + LR + CR
    dR = max.(dR, 0.01)

    # Dead (can only increase)
    dD = KD + LD + CD
    dD = max.(dD, 0.01)

    # Vector change of population and update in place
    result = vcat(dS, dE, dI, dJ, dH, dC, dR, dD, dK, dL, [dBED], [dICU])
    for i = 1:length(result)
        dP[i] = result[i]
    end
end

function errorDeaths(actualD, predictedD)
    return sqrt(sum( (actualD .- predictedD) .^ 2 ))
end

function calculateTotalDeaths(sol)
    # The solutions are returned as a vector of vectors:
    #  - it is a vector of size the number of timesteps
    #  - each element of the vector is a vector of all the variables
    nSteps = length(sol.t);
    nVars  = length(sol.u[1]);

    solMat = zeros((nSteps, nVars));
    for i = 1:nSteps
        solMat[i, :] = sol.u[i]
    end;

    # Creates matrix (varialbles x time steps)
    solMat = transpose(solMat)

    # Select the rows of Dx and sums to have total deaths at each time period
    forecastDeaths = sum(solMat[D0Index:D0Index + nAgeGroup - 1, :]; dims=1)

    # Make a linear approximation of the forecast to match the actual days
    start = date2days(first(countryData[country][:cases].time))
    finish = date2days(last(countryData[country][:cases].time))

    return [ linearInterpolation(t, sol.t, forecastDeaths) for t in start:finish]
end

function my_loss_function(sol)
   tot_loss = 0.0
   if any((s.retcode != :Success for s in sol))
     tot_loss = Inf
   else

       # Calculate the total number of deaths interpolated to actual deaths
       forecastDeaths = calculateTotalDeaths(sol)

       # Calculate RMSE
       tot_loss = sqrt(sum( (ActualDeaths .- forecastDeaths) .^ 2 ))

   end
   tot_loss
end


using Test, Random, LSHFunctions, LinearAlgebra, SparseArrays

include(joinpath("..", "utils.jl"))

#==================
Tests
==================#
@testset "SignALSH tests" begin
    Random.seed!(RANDOM_SEED)

    @testset "Construct SignALSH" begin
        hashfn = SignALSH(; maxnorm=1)
        @test n_hashes(hashfn) == 1
        @test isa(hashfn, SignALSH{Float32})
        @test isa(hashfn, LSHFunctions.AsymmetricLSHFunction)
        @test hashtype(hashfn) == Bool
        @test similarity(hashfn) == inner_prod

        hashfn = SignALSH(32; maxnorm=1)
        @test n_hashes(hashfn) == 32

        hashfn = SignALSH(; dtype=Float64, maxnorm=1)
        @test isa(hashfn, SignALSH{Float64})

        # n_hashes must be positive
        @test_throws ErrorException SignALSH(-1; maxnorm=1)
        @test_throws ErrorException SignALSH( 0; maxnorm=1)

        # maxnorm must be specified and non-negative
        @test_throws ErrorException SignALSH()
        @test_throws ErrorException SignALSH(; maxnorm=-1)

        # m must be positive
        @test_throws ErrorException SignALSH(; m=-1)
        @test_throws ErrorException SignALSH(; m=0)
    end

    @test_skip @testset "Can hash inputs correctly with SignALSH" begin
        input_length = 5
        n_hashes = 8

        hashfn = SignALSH(n_hashes; m=3)
        simhash = SimHash([hashfn.coeff_A hashfn.coeff_B])

        # By default, hashfn should be using Float32 for its coefficients
        @test isa(hashfn, SignALSH{Float32})

        X = rand(Float32, input_length, 128)
        ihashes = index_hash(hashfn, X)
        qhashes = query_hash(hashfn, X)

        @test eltype(ihashes) == eltype(qhashes)
        @test hashtype(hashfn) == Bool

        # 1. Compute the indexing hashes manually
        norms = map(norm, eachcol(X))
        maxnorm = maximum(norms)
        norms .*= 1 / maxnorm
        Y = X ./ maxnorm

        @test all(map(norm,eachcol(Y)) .≤ 1)
        @test all(0 .≤ norms .≤ 1)

        norm_powers = -[norms.^2 norms.^4 norms.^8] .+ 1/2
        Px = [Y; norm_powers']

        @test simhash(Px) == ihashes

        # 2. Compute the querying hashes manually
        Qx = [X ./ norms'; zeros(1,128); zeros(1,128); zeros(1,128)]

        @test simhash(Qx) == qhashes
    end

    @testset "SignALSH can't hash inputs of norm > maxnorm" begin
        hashfn = SignALSH(; maxnorm=0)
        @test_throws ErrorException index_hash(hashfn, rand(4))
        @test_throws ErrorException query_hash(hashfn, rand(4))

        # Should have no issue if norm(x) == maxnorm
        @test index_hash(hashfn, zeros(4)) |> length == 1
        @test query_hash(hashfn, zeros(4)) |> length == 1
    end

    @testset "SignALSH generates collisions for large inner products" begin
        input_length = 5; n_hashes = 128;

        # Compare a random vector x against four other vectors:
        # a) 10 * x
        # b) x
        # c) A vector of all zeros
        # d) -x
        x = randn(input_length)
        x2, x3, x4 = 10*x, zero(x), -x

        maxnorm = (x, x2, x3, x4) .|> norm |> maximum
        hashfn = SignALSH(n_hashes; maxnorm=maxnorm)

        x_query_hashes = query_hash(hashfn, x)

        dataset = [x2 x x3 x4]
        p_hashes = index_hash(hashfn, dataset)

        # Each collection of hashes should be different from one another
        @test let result = true
            for (ii,jj) in Iterators.product(1:4, 1:4)
	            if ii != jj && p_hashes[:,ii] == p_hashes[:,jj]
		            result = false
		            break
	            end
            end
            result
        end

        # The number of collisions should be highest for x and 2*x, second-highest
        # for x and x, second-lowest for x and zeros, and lowest for x and -x
        n_collisions = [sum(x_query_hashes .== p) for p in eachcol(p_hashes)]
        @test n_collisions[1] > n_collisions[2] >
              n_collisions[3] > n_collisions[4]
    end

    @testset "Can hash sparse arrays" begin
        input_size = 100
        n_inputs = 150
        n_hashes = 2

        hashfn = SignALSH(n_hashes; maxnorm=4*input_size)
        x = sprandn(input_size, n_inputs, 0.2)

        # Mostly just need to test that the following lines don't crash
        ih = index_hash(hashfn, x)
        qh = query_hash(hashfn, x)

        @test size(ih) == (n_hashes, n_inputs)
        @test size(qh) == (n_hashes, n_inputs)
    end

    @testset "Can hash matrix adjoints" begin
        input_size = 100
        n_inputs = 150
        n_hashes = 2
        hashfn = SignALSH(n_hashes; maxnorm=4*input_size)

        ## Test 1: regular matrix adjoint
        x = randn(n_inputs, input_size)'
        @test index_hash(hashfn, x) == index_hash(hashfn, copy(x))
        @test query_hash(hashfn, x) == query_hash(hashfn, copy(x))

        ## Test 2: sparse matrix adjoint
        x = sprandn(n_inputs, input_size, 0.2)'
        @test index_hash(hashfn, x) == index_hash(hashfn, copy(x))
        @test query_hash(hashfn, x) == query_hash(hashfn, copy(x))
    end

    @testset "Hash inputs of different sizes" begin
        n_hashes = 42
        hashfn = SignALSH(n_hashes; maxnorm=100)

        @test size(hashfn.coeff_A) == (n_hashes, 0)

        index_hash(hashfn, rand(10))
        @test size(hashfn.coeff_A) == (n_hashes, 10)

        query_hash(hashfn, rand(20))
        @test size(hashfn.coeff_A) == (n_hashes, 20)

        index_hash(hashfn, rand(5))
        @test size(hashfn.coeff_A) == (n_hashes, 20)

        query_hash(hashfn, rand(15))
        @test size(hashfn.coeff_A) == (n_hashes, 20)

        index_hash(hashfn, rand(20))
        @test size(hashfn.coeff_A) == (n_hashes, 20)

        query_hash(hashfn, rand(20))
        @test size(hashfn.coeff_A) == (n_hashes, 20)
    end

    @testset "Hash inputs of different sizes with resize_pow2 = true" begin
        n_hashes = 25
        hashfn = SignALSH(n_hashes; maxnorm=100, resize_pow2=true)

        @test size(hashfn.coeff_A) == (n_hashes, 0)

        index_hash(hashfn, rand(4))
        @test size(hashfn.coeff_A) == (n_hashes, 4)

        query_hash(hashfn, rand(8))
        @test size(hashfn.coeff_A) == (n_hashes, 8)

        index_hash(hashfn, rand(9))
        @test size(hashfn.coeff_A) == (n_hashes, 16)

        query_hash(hashfn, rand(17))
        @test size(hashfn.coeff_A) == (n_hashes, 32)

        index_hash(hashfn, rand(15))
        @test size(hashfn.coeff_A) == (n_hashes, 32)

        query_hash(hashfn, rand(31))
        @test size(hashfn.coeff_A) == (n_hashes, 32)

        index_hash(hashfn, rand(32))
        @test size(hashfn.coeff_A) == (n_hashes, 32)

        query_hash(hashfn, rand(32))
        @test size(hashfn.coeff_A) == (n_hashes, 32)
    end
end


"""
    pmfN(data; tol = 1e-5) -> Vector

Compute the probability mass function for the number of guesses `N`, given a
strategy. The `n`th entry of the output vector gives the probability for
guessing the correct answer on the `n`th try, for `n = 1 : K`. If the number of
allowed guesses, `K`, is smaller than `length(p)`, then there is an additional
last entry which gives the probability of never guessing the correct answer.

* `data` should be a `NamedTuple` with entries for `p`, `ρBs`, `Es`, `K`, and
  `povm_outcomes`
* `tol` provides a tolerance above which to warn about imaginary
  or negative probabilities.

"""
function pmfN(data; tol = 1e-5)
    @unpack p, K = data
    T = eltype(p)

    probs = [_prob(n, data, tol, T) for n = 1:K]

    if K < length(p)
        prob_no_guess = 1 - sum(probs)

        if prob_no_guess < 0
            prob_no_guess < -tol && @warn "Probability negative " prob_no_guess
            prob_no_guess = zero(T)
        end

        push!(probs, prob_no_guess)
    end

    return probs
end

function _prob(n, data, tol, T)
    @unpack Es, p, ρBs, povm_outcomes = data
    prob = zero(T)
    for (y, outcome) in enumerate(povm_outcomes)
        for x in eachindex(p, ρBs)
            if findfirst(==(x), outcome) == n
                prob += p[x] * tr(Es[y] * ρBs[x])
            end
        end
    end
    if !isreal(prob)
        abs(imag(prob)) > tol && @warn "Imaginary probability" n prob
    end
    prob = real(prob)
    if prob < 0
        prob < -tol && @warn "Probability negative " n prob
        prob = zero(T)
    end
    return prob
end


using CurveProximityQueries
using ConvexBodyProximityQueries
using Test
using LinearAlgebra
using StaticArrays
using IntervalArithmetic
using Random: seed!
using Unitful: nm, μm, mm, cm

import CurveProximityQueries: differentiate, integrate

@testset "CurveProximityQueries" begin
    @testset "Constructors" begin
        # should return SVector when constructed without SVectors
        @test Bernstein([[0.0, 0.0], [1.0, 1.0]])(0.5) === @SVector([0.5, 0.5])

        # non-floats arguments should become floats to satisfy assumptions
        # elsewhere in the code (e.g. use of `eps(T)`)
        @test eltype(Bernstein([[0, 0], [1, 1]])) <: Float64

        # none of the following should error; some should be inferrable
        @test Bernstein([
                @SVector([0.0, 0.0]),
                @SVector([1.0, 1.0])
            ])(0.5) === @SVector([0.5, 0.5])
        @test @inferred(Bernstein(@SVector([
                @SVector([0.0, 0.0]),
                @SVector([1.0, 1.0])
            ])))(0.5) === @SVector([0.5, 0.5])
        @test @inferred(Bernstein(
                @SVector([0.0, 0.0]),
                @SVector([1.0, 1.0])
            ))(0.5) === @SVector([0.5, 0.5])

        # similar reasoning below, but allow for any StaticVector
        struct Point{T} <: StaticArrays.FieldVector{2, T}
            x::T
            y::T
        end
        StaticArrays.similar_type(::Type{P}, ::Type{T},
            ::StaticArrays.Size{(2,)}) where {P <: Point, T} = Point{T}

        @test Bernstein([Point(0.0, 0.0), Point(1.0, 1.0)])(0.5) ===
            Point(0.5, 0.5)
        @test @inferred(
            Bernstein(@SVector([Point(0.0, 0.0), Point(1.0, 1.0)])))(0.5) ===
                Point(0.5, 0.5)
        @test @inferred(Bernstein(Point(0.0, 0.0), Point(1.0, 1.0)))(0.5) ===
            Point(0.5, 0.5)
    end
    @testset "Bernstein Polynomials" begin
        B2 = rand(Bernstein{2,8})
        B3 = rand(Bernstein{3,5})
        @test typeof(B2) <: Bernstein{2,8}
        @test typeof(B3) <: Bernstein{3,5}
        @test sum(norm.((differentiate(integrate(B2.f)) - B2.f).control_points)) ≤ 1e-10
        @test sum(norm.((differentiate(integrate(B3.f)) - B3.f).control_points)) ≤ 1e-10
    end
    @testset "Curve - Polygon" begin
        seed!(1)
        obs = @point zeros(2)
        c = rand(Bernstein{2,7})
        @test abs(minimum_distance(obs, c) - 0.3683522584741768) ≤ 1e-5
        @test tolerance_verification(obs, c, 0.3) == true
        @test tolerance_verification(obs, c, 0.4) == false
        @test collision_detection(obs, c) == false
        @test abs(minimum_distance(c, obs) - 0.3683522584741768) ≤ 1e-5
        @test tolerance_verification(c, obs, 0.3) == true
        @test tolerance_verification(c, obs, 0.4) == false
        @test collision_detection(c, obs) == false

        seed!(1)
        obs = @point zeros(3)
        c = rand(Bernstein{3,11})
        @test abs(minimum_distance(obs, c) - 0.511627527056288) ≤ 1e-5
        @test tolerance_verification(obs, c, 0.5) == true
        @test tolerance_verification(obs, c, 0.6) == false
        @test collision_detection(obs, c) == false
        @test abs(minimum_distance(c, obs) - 0.511627527056288) ≤ 1e-5
        @test tolerance_verification(c, obs, 0.5) == true
        @test tolerance_verification(c, obs, 0.6) == false
        @test collision_detection(c, obs) == false
    end
    @testset "Curve - Curve" begin
        seed!(1)
        c = rand(Bernstein{2,7})
        d = rand(Bernstein{2,3})
        @test abs(minimum_distance(c, d)) ≤ 1e-5
        @test tolerance_verification(c, d, 0.3) == false
        @test collision_detection(c, d) == true

        seed!(1)
        obs = @point zeros(3)
        c = rand(Bernstein{3,11})
        d = rand(Bernstein{3,5})
        @test abs(minimum_distance(c, d) - 0.137873546917387) ≤ 1e-5
        @test tolerance_verification(c, d, 0.1) == true
        @test tolerance_verification(c, d, 0.3) == false
        @test collision_detection(c, d) == false
    end
    @testset "Dimensionful curves" begin
        b2 = Bernstein([[0.0, 1.0]μm, [1.0, 1.0]μm, [1.0, 0.0]μm])
        @test b2(0.0) == [0, 1000]nm
        @test b2(1.0) == [0.001, 0]mm
        @test arclength(b2) == arclength(Bernstein([[0,1], [1,1], [1,0]])) * μm
        @test length(CurveProximityQueries.cvxhull(b2, Interval(0.0, 1.0))) == 8

        b3 = Bernstein([[0.0, 0.0, 0.0]mm, [1.0, 0.0, 0.0]mm, [1.0, 0.0, 1.0]mm])
        c3 = Bernstein([[2.0, 1.0, 1.0]cm, [2.0, 3.0, 2.0]cm, [3.0, 1.0, 1.0]cm])
        @test length(CurveProximityQueries.cvxhull(b3, Interval(0.0, 1.0))) == 16
        @test abs(minimum_distance(b3, c3, atol=1e-8mm) - 2.328cm) ≤ 1e-4cm
        @test tolerance_verification(b3, c3, 2cm, atol=1e-8mm) == true
        @test collision_detection(b3, c3, atol=1e-8mm) == false
    end
    @testset "Support" begin
        @test CurveProximityQueries.support(Point(1.0, 2.0), Point(1.0, 1.0)) ===
            Point(1.0, 2.0)
        @test CurveProximityQueries.support(@SVector(
            [Point(0.0, 1.0), Point(1.0, 2.0)]), @SVector([1.0, 1.0])) ===
                Point(1.0, 2.0)
        @test CurveProximityQueries.support(@SVector(
            [@SVector([0.0, 1.0]), @SVector([1.0, 2.0])]),
                @SVector([1.0, 1.0])) === @SVector([1.0, 2.0])
        @test CurveProximityQueries.support(@SVector(
            [Point(0.0, 1.0), Point(1.0, 2.0), Point(3.0, 4.0)]), Point(1.0, 1.0)) ===
                Point(3.0, 4.0)

    end
end


module Assemblers

using SparseArrays: sparse
using LinearAlgebra
using ..LocalAssemblers: LocalMatrixAssembler, init!

"""
    AbstractSysmatAssembler

Abstract type of system-matrix assembler.
"""
abstract type AbstractSysmatAssembler end;

"""
    SysmatAssemblerSparse{T<:Number} <: AbstractSysmatAssembler

Type for assembling a sparse global matrix from individual entries.
"""
mutable struct SysmatAssemblerSparse{T<:Number} <: AbstractSysmatAssembler
    row::Vector{Int64}
    col::Vector{Int64}
    val::Vector{T}
    nrow::Int64; ncol::Int64;
end

"""
    SysmatAssemblerSparse(zero::T=0.0) where {T<:Number}

Construct blank system matrix assembler. The matrix entries are of type `T`.

# Example

This is how a sparse matrix is assembled from two rectangular dense matrices.
```
a = SysmatAssemblerSparse(0.0)
start!(a, 7, 7)
m = [0.24406   0.599773    0.833404  0.0420141
0.786024  0.00206713  0.995379  0.780298
0.845816  0.198459    0.355149  0.224996]
gi = [1 7 5]
gj = [5 2 1 4]
for j in 1:size(m, 2), i in 1:size(m, 1)
    assemble!(a, gi[i], gj[j], m[i, j])
end
m = [0.146618  0.53471   0.614342    0.737833
0.479719  0.41354   0.00760941  0.836455
0.254868  0.476189  0.460794    0.00919633
0.159064  0.261821  0.317078    0.77646
0.643538  0.429817  0.59788     0.958909]
gi =  [2 3 1 7 5]
gj = [6 7 3 4]
for j in 1:size(m, 2), i in 1:size(m, 1)
    assemble!(a, gi[i], gj[j], m[i, j])
end
A = finish!(a)
```
"""
function SysmatAssemblerSparse(zero::T=0.0) where {T<:Number}
    return SysmatAssemblerSparse{T}(fill(0, 0), fill(0, 0), fill(zero, 0), 0, 0)
end

"""
    start!(self::SysmatAssemblerSparse{T}, nrow, ncol) where {T<:Number}

Start the assembly of a global matrix.
"""
function start!(self::SysmatAssemblerSparse{T}, nrow, ncol) where {T<:Number}
    nexpected = Int64(round(0.00001 * max(nrow, ncol, 10000)^2))
    self.row = fill(0, 0)
    sizehint!(self.row, nexpected)
    self.col = fill(0, 0)
    sizehint!(self.col, nexpected)
    self.val = fill(zero(T), 0)
    sizehint!(self.val, nexpected)
    self.nrow = nrow;
    self.ncol = ncol;
    return self
end

"""
    assemble!(self::SysmatAssemblerSparse{T}, r, c, v::T) where {T<:Number}

Assemble a single entry of a rectangular matrix.
"""
function assemble!(self::SysmatAssemblerSparse{T}, r, c, v::T) where {T<:Number}
    push!(self.row, r)
    push!(self.col, c)
    push!(self.val, v)
    return self
end

"""
    assemble!(self::SysmatAssemblerSparse{T}, lma::LocalMatrixAssembler{IT, T}) where {IT<:Integer, T<:Number}

Assemble the row numbers, column numbers, and values from a local assembler.
"""
function assemble!(self::SysmatAssemblerSparse{T}, lma::LocalMatrixAssembler{IT, T}) where {IT<:Integer, T<:Number}
    append!(self.row, lma.row)
    append!(self.col, lma.col)
    append!(self.val, lma.M)
    return self
end

"""
    assemble!(self::SysmatAssemblerSparse{T}, lma::Transpose{T,LocalMatrixAssembler{IT,T}}) where {IT<:Integer, T<:Number}

Assemble the row numbers, column numbers, and values from a local assembler.
"""
function assemble!(self::SysmatAssemblerSparse{T}, lma::Transpose{T,LocalMatrixAssembler{IT,T}}) where {IT<:Integer, T<:Number}
    append!(self.row, transpose(lma.parent.col))
    append!(self.col, transpose(lma.parent.row))
    append!(self.val, transpose(lma.parent.M))
    return self
end

"""
    finish!(self::SysmatAssemblerSparse)

Make a sparse matrix.
"""
function finish!(self::SysmatAssemblerSparse)
    return sparse(self.row, self.col, self.val, self.nrow, self.ncol);
end

"""
    AbstractSysvecAssembler

Abstract type of system vector assembler.
"""
abstract type AbstractSysvecAssembler end;

"""
    start!(self::SV,  nrow) where {SV<:AbstractSysvecAssembler, T<:Number}

Start assembly.

The method makes the buffer for the vector assembly. It must be called before
the first call to the method assemble.

`nrow`= Total number of degrees of freedom.
"""
function start!(self::SV,  nrow) where {SV<:AbstractSysvecAssembler, T<:Number}
end

"""
    assemble!(self::SV, i, val::T) where {SV<:AbstractSysvecAssembler, T<:Number}

Assemble an elementwise vector.

The method assembles a column element vector using the vector of degree of
freedom numbers for the rows.
"""
function assemble!(self::SV, i, val::T) where {SV<:AbstractSysvecAssembler, T<:Number}
end

"""
    makevector!(self::SysvecAssembler)

Make the global vector.
"""
function finish!(self::SV) where {SV<:AbstractSysvecAssembler}
end

"""
    SysvecAssembler

Assembler for the system vector.
"""
mutable struct SysvecAssembler{T<:Number} <: AbstractSysvecAssembler
    val::Vector{T};
    ndofs::Int64
end

"""
    SysvecAssembler(zero::T=0.0) where {T<:Number}

Construct blank system vector assembler. The vector entries are of type `T`.
"""
function SysvecAssembler(zero::T=0.0) where {T<:Number}
    return SysvecAssembler([zero], 0)
end

"""
    start!(self::SysvecAssembler{T},  nrow::Int64) where {T<:Number}

Start assembly.

The method makes the buffer for the vector assembly. It must be called before
the first call to the method assemble.

`nrow`= Total number of degrees of freedom.
"""
function start!(self::SysvecAssembler{T},  nrow::Int64) where {T<:Number}
    self.ndofs = nrow
    self.val = zeros(T, self.ndofs);
    return self
end

"""
    assemble!(self::SysvecAssembler{T}, i, val::T) where {T<:Number}

Assemble a single value into the row `i`.
"""
function assemble!(self::SysvecAssembler{T}, i, val::T) where {T<:Number}
    self.val[i] = self.val[i] + val;
    return self
end

"""
    assemble!(self::SysvecAssembler{T}, lva) where {T<:Number}

Assemble local vector assembler.
"""
function assemble!(self::SysvecAssembler{T}, lva) where {T<:Number}
    for i in 1:length(lva.row)
        gi = lva.row[i]
        self.val[gi] += lva.V[i];
    end
    return self
end

"""
    finish!(self::SysvecAssembler)

Make the global vector.
"""
function finish!(self::SysvecAssembler)
  return deepcopy(self.val);
end

end


#DA algorythm

#recursive version
function call_match(m, n, m_prefs, f_prefs)
  males = 1:m
  females = 1:n

  pointer_males = ones(m)
  temp_matched_males = zeros(n)

  da_match(m, n, m_prefs, f_prefs, pointer_males, temp_matched_males)
end

function da_match(m, n, m_prefs, f_prefs, pointer_males, temp_matched_males)#m offers to f
  candidate_males = Array{Int32}(m, n)
  for i in 1:m
    push!(candidate_males[m_prefs[i, pointer_males[i]]], i)
  end

  choose_best_male!(n, f_prefs, candidate_males, temp_matched_males)
  if all(pointer_males .== n)
    return temp_matched_males
  else
    da_match(m, n, m_prefs, f_prefs, pointer_males, temp_matched_males)
  end
end

function choose_best_male(n, f_prefs, candidate_males, temp_matched_males)
  for j in 1:n
    for i in f_prefs[j]
      if i in candidate_males[:, j]
        temp_matched_males[j] = i
        break
      end
    end
    temp_matched_males[j] = 0
  end
end

call_match(2, 2, [1 2; 2 1], [1 2; 2 1])


using ThreeBodyProblem
using DifferentialEquations
using Plots
gr()

sys = jupiter_europa()

rv₀ = [1.000045941880842, -0.000083429669677, 0.002323145484474, -0.093200789960739, 0.099227116703640, -0.033175567182715]

tspan = (0.,-22.)
prob = ODEProblem(CR3BPdynamics!,rv₀,tspan,sys)
sol = solve(prob,reltol=1e-6)

c = [1-sys.μ, 0, 0]
xyz_hemisphere = deserno_hemisphere(100, c)

plot(sol,vars=(1,2),title="Lunar return trajectory in rotating frame")
plot!(sys)