## Mesh convergence tests
using AeroMDAO
using Plots
using DataFrames

## Wing
wing = Wing(foils     = fill(naca4((0,0,1,2)), 2),
            chords    = [1.0, 0.6],
            twists    = [0.0, 0.0],
            spans     = [5.0],
            dihedrals = [5.],
            sweeps      = [5.]);
print_info(wing, "Wing")
S, b, c = projected_area(wing), span(wing), mean_aerodynamic_chord(wing);

## Assembly
ρ       = 1.225
ref     = [0.25, 0., 0.]
V, α, β = 1.0, 5.0, 5.0
Ω       = [0.0, 0.0, 0.0]
fs      = Freestream(V, α, β, Ω)

## Panelling error closure
panel_case(x, y) = solve_case(wing, fs;
                              rho_ref   = ρ,
                              r_ref     = ref,
                              area_ref  = S,
                              span_ref  = b,
                              chord_ref = c,
                              span_num  = y,
                              chord_num = x,
                              viscous   = false,
                              x_tr      = [0.3, 0.3],
                              spacing   = Uniform()
                             );

## Distributions
xs = 10:2:50
ys = xs .* 2# fill(maximum(xs) * 4, length(xs))

num = length(xs)

coeff_data = zeros(num, 9 + 2)
for (i, (x, y)) in enumerate(zip(xs, ys))
    @time nf_coeffs, ff_coeffs, CFs, CMs, horseshoe_panels, camber_panels, horseshoes, Γs =
    panel_case(x, y)

    println("Distribution: nx = $x, ny = $y")
    print_coefficients(nf_coeffs, ff_coeffs)

    coeff_data[i,1]     = x
    coeff_data[i,2]     = y
    coeff_data[i,3:5]   = ff_coeffs
    coeff_data[i,6:end] = nf_coeffs
end

##
data = DataFrame(coeff_data, :auto)
rename!(data, [:chord_num, :span_num, :CD_ff, :CY_ff, :CL_ff, :CD_nf, :CY_nf, :CL_nf, :Cl, :Cm, :Cn])

##
errors = @. abs((data[3:end,3:end] - data[2:end-1,3:end]) / data[2:end-1,3:end])

##
plot(xlabel = "Chordwise Panels", ylabel = "Error", yscale = :log10)

# CDs Errors
plot!(data[2:end-1,"chord_num"], errors[!,"CD_nf"], label = "CD Nearfield")
plot!(data[2:end-1,"chord_num"], errors[!,"CD_ff"], label = "CD Farfield")

# CLs Errors
plot!(data[2:end-1,"chord_num"], errors[!,"CL_nf"], label = "CL Nearfield")
plot!(data[2:end-1,"chord_num"], errors[!,"CL_ff"], label = "CL Farfield")

# CYs Errors
plot!(data[2:end-1,"chord_num"], errors[!,"CY_nf"], label = "CY Nearfield")
plot!(data[2:end-1,"chord_num"], errors[!,"CY_ff"], label = "CY Farfield")

##
CL = plot(data[2:end,"chord_num"], data[2:end,"CL_nf"], marker = :dot, label = "Nearfield", ylabel = "CL")
plot!(data[2:end,"chord_num"], data[2:end,"CL_ff"], marker = :dot, label = "Farfield")

CD = plot(data[2:end,"chord_num"], data[2:end,"CD_nf"], marker = :dot, label = "Nearfield", ylabel = "CD")
plot!(data[2:end,"chord_num"], data[2:end,"CD_ff"], marker = :dot, label = "Farfield")

CY = plot(data[2:end,"chord_num"], data[2:end,"CY_nf"], marker = :dot, label = "Nearfield", ylabel = "CY")
plot!(data[2:end,"chord_num"], data[2:end,"CY_ff"], marker = :dot, label = "Farfield")

Cl = plot(data[2:end,"chord_num"], data[2:end,"Cl"], marker = :dot, label = "Nearfield", ylabel = "Cl")
Cm = plot(data[2:end,"chord_num"], data[2:end,"Cm"], marker = :dot, label = "Nearfield", ylabel = "Cm")
Cn = plot(data[2:end,"chord_num"], data[2:end,"Cn"], marker = :dot, label = "Nearfield", ylabel = "Cn")

plot(CL, CD, CY, Cl, Cm, Cn, layout = (2,3), size = (1280, 720))


# MemoryFunction.jl
# References:
# [1] R. P. Feynman, R. W. Hellwarth, C. K. Iddings, and P. M. Platzman, Mobility of slow electrons in a polar crystal, PhysicalReview127, 1004 (1962) https://doi.org/10.1103/PhysRev.127.1004.
# [2] Devreese, J. De Sitter, and M. Goovaerts, “Optical absorption of polarons in thefeynman-hellwarth-iddings-platzman approximation,”Phys. Rev. B, vol. 5, pp. 2367–2381, Mar 1972.
# [3] F. Peeters and J. Devreese, “Theory of polaron mobility,” inSolid State Physics,pp. 81–133, Elsevier, 1984.

"""
    polaron_memory_function(Ω, β, α, v, w; ω = 1.0, rtol = 1e-3)

Calculate the memory function χ(Ω) of the polaron at finite temperatures (equation (35) in FHIP 1962) for a given frequency Ω.

Includes conditionals for zero-temperature and DC limits.

# Arguments
- `Ω::Float64`: electric field frequency (2π THz).
- `β::Union{Float64, Vector{Float64}` reduced thermodynamic temperatures (unitless). Either a single value for one phonon frequency or a vector of values for multiple phonon frequencies.
- `α::Union{Float64, Vector{Float64}}`: Frohlich alpha coupling parameter (unitless). Either a single value for one phonon frequency or a vector of values for multiple phonon frequencies.
- `v::Union{Float64, Vector{Float64}}`: variational 'v' parameters that minimise the polaron free energy (unitless).
- `w::Union{Float64, Vector{Float64}}`: variational 'w' parameters that minimise the polaron free energy (unitless).
- `ω::Union{Float64, Vector{Float64}}`: phonon mode frequencies (2π THz). Predefined as `ω = 1.0` for a single mode in polaron units.
- `rtol = 1e-3`: specifies the relative error tolerance for the QuadGK integral.

See R. P. Feynman, R. W. Hellwarth, C. K. Iddings, and P. M. Platzman, Mobility of slow electrons in a polar crystal, PhysicalReview127, 1004 (1962). https://doi.org/10.1103/PhysRev.127.1004.

See also [`polaron_memory_function_thermal`](@ref), [`polaron_memory_function_athermal`](@ref), [`polaron_memory_function_dc`](@ref).
"""
function polaron_memory_function(Ω, β, α, v, w; ω = 1.0, rtol = 1e-3)

    # Zero temperature and frequency is just zero.
	if Ω == 0 && any(x -> x == Inf64, β)
		return 0.0 + im * 0.0

    # DC zero frequency limit at finite temperatures.
	elseif Ω == 0 && any(x -> x != Inf64, β)
		return polaron_memory_function_dc(β, α, v, w; ω = ω, rtol = rtol)

    # Zero temperature limit at AC finite frequencies.
	elseif Ω != 0 && any(x -> x == Inf64, β)
		return polaron_memory_function_athermal(Ω, α, v, w; ω = ω, rtol = rtol)

    # Finite temperatures and frequencies away from zero limits.
	elseif Ω != 0 && any(x -> x != Inf64, β)
		return polaron_memory_function_thermal(Ω, β, α, v, w; ω = ω, rtol = rtol)

    # Any other frequencies or temperatures (e.g. negative or complex) prints error message.
	else
		println("Photon frequency Ω and thermodynamic temperature β must be ≥ 0.0.")
	end
end

"""
----------------------------------------------------------------------
Multiple Branch Polaron Memory Function and Complex Conductivity
----------------------------------------------------------------------

This section of the code is dedicated to calculating the polaron memory function and complex conductivity, generalised from FHIP's expression to the case where multiple phonon modes are present in the material.
"""

"""
    function polaron_memory_function_thermal(Ω, β, α, v, w; ω = 1.0, rtol = 1e-3)

Calculate polaron complex memory function at finite temperature inclusive of multiple phonon branches j, each with angular frequency ω[j] (2π THz).

# Arguments
- `Ω::Float64`: is the frequency (THz) of applied electric field.
- `β::Vector{Float64}`: is an array of reduced thermodynamic betas, one for each phonon frequency ω[j].
- `α::Vector{Float64}`: is an array of decomposed Frohlich alphas, one for each phonon frequency ω[j].
- `v::Union{Float64, Vector{Float64}}`: is a vector or single value of the 'v' variational parameter(s).
- `w::Union{Float64, Vector{Float64}}`: is a vector or single value of the 'w' variational parameter(s).
- `m_eff::Float64`: is the bare-band mass of the particle (typically electron / hole, in units of electron mass m_e).
- `ω::Union{Float64, Vector{Float64}}`: phonon mode frequencies (2π THz). Predefined as `ω = 1.0` for a single mode in polaron units.
- `rtol = 1e-3`: specifies the relative error tolerance for the QuadGK integral.

See also [`polaron_memory_function`](@ref).
"""
function polaron_memory_function_thermal(Ω, β, α, v, w; ω = 1.0, rtol = 1e-3)

    # FHIP1962, page 1009, eqn (36).
    S(t, β) = (exp(1im * t) + exp(-1im * t - β)) / (1 - exp(-β)) / D_j(-1im * t, β, v, w)^(3 / 2)

    # FHIP1962, page 1009, eqn (35a).
    integrand(t, β, Ω) = (1 - exp(1im * Ω * 2π * t)) * imag(S(t, β))

    integral = map((x, y) -> quadgk(t -> integrand(t, x, Ω / y), 0.0, Inf64, rtol = rtol)[1], β, ω)

     memory = sum(2 .* α .* ω .^2 .* integral ./ (3 * √π * Ω * 2π))

    return memory
end

"""
    function polaron_memory_function_athermal(Ω, α, v, w; ω = 1.0, rtol = 1e-3)

Calculate polaron complex memory function at zero-temperature inclusive of multiple phonon branches j, each with angular frequency ω[j] (2π THz).

# Arguments
- `Ω::Float64`: is the frequency (THz) of applied electric field.
- `α::Vector{Float64}`: is an array of decomposed Frohlich alphas, one for each phonon frequency ω[j].
- `v::Union{Float64, Vector{Float64}}`: is a vector or single value of the 'v' variational parameter(s).
- `w::Union{Float64, Vector{Float64}}`: is a vector or single value of the 'w' variational parameter(s).
- `m_eff::Float64`: is the bare-band mass of the particle (typically electron / hole, in units of electron mass m_e).
- `ω::Union{Float64, Vector{Float64}}`: phonon mode frequencies (2π THz). Predefined as `ω = 1.0` for a single mode in polaron units.
- `rtol = 1e-3`: specifies the relative error tolerance for the QuadGK integral.

See also [`polaron_memory_function`](@ref).
"""
function polaron_memory_function_athermal(Ω, α, v, w; ω = 1.0, rtol = 1e-3)
    # FHIP1962, page 1009, eqn (36).
    S(t) = exp(1im * t) / D_j(-1im * t, v, w)^(3 / 2)

    # FHIP1962, page 1009, eqn (35a).
    integrand(t, Ω) = (1 - exp(1im * 2π * Ω * t)) * imag(S(t))

    memory = 0.0

    for j in 1:length(ω)

        integral = quadgk(t -> integrand(t, Ω / ω[j]), 0.0, 1 / rtol, rtol = rtol)[1]
        memory += 2 * α[j] * ω[j]^2 * integral / (3 * √π * Ω * 2π)

    end

    return memory
end

"""
    function polaron_memory_function_dc(β, α, v, w; ω = 1.0, rtol = 1e-3)

Calculate zero-frequency polaron complex memory function at finite temperature inclusive of multiple phonon branches j, each with angular frequency ω[j] (2π THz).

# Arguments
- `β::Vector{Float64}`: is an array of reduced thermodynamic betas, one for each phonon frequency ω[j].
- `α::Vector{Float64}`: is an array of decomposed Frohlich alphas, one for each phonon frequency ω[j].
- `v::Union{Float64, Vector{Float64}}`: is a vector or single value of the 'v' variational parameter(s).
- `w::Union{Float64, Vector{Float64}}`: is a vector or single value of the 'w' variational parameter(s).
- `m_eff::Float64`: is the bare-band mass of the particle (typically electron / hole, in units of electron mass m_e).
- `ω::Union{Float64, Vector{Float64}}`: phonon mode frequencies (2π THz). Predefined as `ω = 1.0` for a single mode in polaron units.
- `rtol = 1e-3`: specifies the relative error tolerance for the QuadGK integral.

See also [`polaron_memory_function`](@ref).
"""
function polaron_memory_function_dc(β, α, v, w; ω = 1.0, rtol = 1e-3)
    # FHIP1962, page 1009, eqn (36).
    S(t, β) = (exp(1im * t) + exp(-1im * t - β)) / (1 - exp(-β)) / D_j(-1im * t, β, v, w)^(3 / 2)

    # FHIP1962, page 1009, eqn (35a).
    integrand(t, β) = -im * t * imag(S(t, β))

    integral = map(x -> quadgk(t -> integrand(t, x), 0.0, 1 / rtol, rtol = rtol)[1], β)

    memory = sum(2 .* α .* ω .* integral ./ (3 * √π))

    return memory
end


struct Scordis
    dis
    fm
    Sinv::Matrix{Float64}
    cutoff::Real
    nlv::Int64
end

struct Odis
    dis
    fm
    cutoff::Real
    nlv::Int64
end

"""
    scordis(object::Union{Pca, Plsr}; nlv = nothing, rob = true, alpha = .01)
Compute the score distances (SDs) from a PCA or PLS model

* `object` : The fitted model.
* `nlv` : Nb. components (PCs or LVs) to consider. If nothing,
    it is the maximum nb. of components.
* `rob` : If true, the moment estimation of the distance cutoff is robustified.
    This may be relevant after robust PCA or PLS on small data sets
        containing extreme values.
* `alpha` : Risk-I level used for computing the cutoff detecting extreme values.

SDs are the Mahalanobis distances of the projections of row observations on the
score plan to the center of the score space.

They are computed for the training `Xtrain` and the eventual `X`.

A cutoff is computed from the training, using a moment estimation of the parameters of
a Chi-squared distrbution for SD^2 (see e.g. Pomerantzev 2008).
In the output, column dstand is a standardized distance, defined as SD / cutoff.
A value dstand > 1 may be considered as extreme.

The Winisi "GH" is also provided (usually considered as extreme if GH > 3).

## References

M. Hubert, P. J. Rousseeuw, K. Vanden Branden (2005). ROBPCA: a new approach to robust
principal components analysis. Technometrics, 47, 64-79.

Pomerantsev, A.L., 2008. Acceptance areas for multivariate classification derived by
projection methods. Journal of Chemometrics 22, 601-609. https://doi.org/10.1002/cem.1147
"""
function scordis(object::Union{Pca, Plsr};
    nlv = nothing, rob = true, alpha = .01)
    a = size(object.T, 2)
    isnothing(nlv) ? nlv = a : nlv = min(nlv, a)
    T = @view(object.T[:, 1:nlv])
    S = Statistics.cov(T, corrected = false)
    LinearAlgebra.inv!(cholesky!(S))   # ==> S := Sinv
    d2 = mahsq(T, zeros(nlv)', S)
    d = sqrt.(d2)
    if !rob
        mu = mean(d2)
        s2 = var(d2)
    else
        mu = median(d2)
        s2 = mad(d2)^2
    end
    nu = 2 * mu^2 / s2
    cutoff = sqrt(mu / nu * quantile.(Distributions.Chisq(nu), 1 - alpha))
    dstand = d / cutoff
    dis = DataFrame((d = d, dstand = dstand, gh = d2 / nlv))
    Scordis(dis, object, S, cutoff, nlv)
end

function predict(object::Scordis, X)
    nlv = object.nlv
    T = transform(object.fm, X; nlv = nlv)
    d2 = mahsq(T, zeros(nlv)', object.Sinv)
    d = sqrt.(d2)
    dstand = d / object.cutoff
    dis = DataFrame((d = d, dstand = dstand, gh = d2 / nlv))
    (dis = dis,)
end

"""
    odis(object::Union{Pca, Plsr}, X; nlv = nothing, rob = true, alpha = .01)
Compute the orthogonal distances (ODs) from a PCA or PLS model

* `object` : The fitted model.
* `X` : X-data that were used to compute the model (training).
* `nlv` : Nb. components (PCs or LVs) to consider. If nothing, it is the maximum
    nb. of components.
* `rob` : If true, the moment estimation of the distance cutoff is robustified.
    This may be relevant after robust PCA or PLS on small data sets
        containing extreme values.
* `alpha` : Risk-I level used for computing the cutoff detecting extreme values.

ODs ("X-residuals") are the Euclidean distances of row observations to their projections
to the score plan (see e.g. Hubert et al. 2005, Van Branden & Hubert 2005, p. 66;
Varmuza & Filzmoser, 2009, p. 79).

They are computed for the training `Xtrain` and the eventual `X`.

A cutoff is computed from the training, using a moment estimation of the parameters of a
Chi-squared distrbution for OD^2 (see Nomikos & MacGregor 1995, and Pomerantzev 2008).
In the output, column dstand is a standardized distance, defined as OD / cutoff.
A value dstand > 1 may be considered as extreme.

The cutoff for detecting extreme OD values is computed using a moment estimation of
a Chi-squared distrbution for the squared distance.

## References

M. Hubert, P. J. Rousseeuw, K. Vanden Branden (2005). ROBPCA: a new approach to robust principal
components analysis. Technometrics, 47, 64-79.

Nomikos, P., MacGregor, J.F., 1995. Multivariate SPC Charts for Monitoring Batch Processes.
null 37, 41-59. https://doi.org/10.1080/00401706.1995.10485888

Pomerantsev, A.L., 2008. Acceptance areas for multivariate classification derived by projection methods.
Journal of Chemometrics 22, 601â-609. https://doi.org/10.1002/cem.1147

K. Vanden Branden, M. Hubert (2005). Robuts classification in high dimension based on the SIMCA method.
Chem. Lab. Int. Syst, 79, 10-21.

K. Varmuza, P. Filzmoser (2009). Introduction to multivariate statistical analysis in chemometrics.
CRC Press, Boca Raton.
"""
function odis(object::Union{Pca, Plsr}, X;
        nlv = nothing, rob = true, alpha = .01)
    a = size(object.T, 2)
    isnothing(nlv) ? nlv = a : nlv = min(nlv, a)
    E = xresid(object, X; nlv = nlv)
    d = sqrt.(sum(E .* E, dims = 2))
    d2 = d.^2
    if(!rob)
        mu = mean(d2)
        s2 = var(d2)
    else
        mu = median(d2)
        s2 = mad(d2)^2
    end
    nu = 2 * mu^2 / s2
    cutoff = sqrt(mu / nu * quantile.(Distributions.Chisq(nu), 1 - alpha))
    dstand = d / cutoff
    dis = DataFrame((d = d, dstand = dstand))
    Odis(dis, object, cutoff, nlv)
end

function predict(object::Odis, X)
    nlv = object.nlv
    E = xresid(object.fm, X; nlv = nlv)
    d = sqrt.(sum(E .* E, dims = 2))
    dstand = d / object.cutoff
    dis = DataFrame((d = d, dstand = dstand))
    (dis = dis,)
end


function hilbert2int(coords::Vector{Int64})
  nD = length(coords)
  coord_chunks = unpack_coords(coords)
  nChunks = length(coord_chunks)
  mask = 2^nD - 1
  start, finish = initial_start_end(nChunks, nD)
  #Int here?
  index_chunks = zeros(Int64,nChunks)
  for j = 0:(nChunks-1)
    i = gray_decode_travel(start, finish, mask, coord_chunks[j+1])
    index_chunks[j+1] = i
    start, finish = child_start_end(start, finish, mask, i)
  end
  return pack_index(index_chunks, nD)
end

function unpack_coords(coords::Vector{Int64})
  nD = length(coords)
  #biggest = reduce(max, coords)  # the max of all coords
  #nChunks = max(1, ceil(Int64,log(2, biggest + 1))) # max # of bits
  nChunks = max(1, ceil(Int64,log(2, maximum(coords) + 1))) # max # of bits
  return transpose_bits(coords, nChunks)
end

function pack_indexSlow(chunks::Vector{Int64}, nD::Int64)
  p = 2^nD  # Turn digits mod 2**nD back into a single number:
  #return reduce( lambda n, chunk: n * p + chunk, chunks )
  return reduce((n,chunk)->n*p+chunk, chunks)
end

function pack_index(chunks::Vector{Int64}, nD::Int64)
  p = 2^nD  # Turn digits mod 2**nD back into a single number:
  #return reduce( lambda n, chunk: n * p + chunk, chunks )
  sum = 0::Int64
  @simd for i = 1:length(chunks)
    sum *= p
    sum += chunks[i]
  end
  return sum
end

function initial_start_end(nChunks::Int64, nD::Int64)
  # This orients the largest cube so that
  # its start is the origin (0 corner), and
  # the first step is along the x axis, regardless of nD and nChunks:
  return 0, 2^mod(-nChunks-1,nD)  # in Python 0 <=  a % b  < b.
end

function gray_encode( bn::Int64)
    assert(bn >= 0)
    return bn $ fld(bn,2)
end

function gray_decode(n::Int64)
  sh = 1
  while true
    div = n >> sh
    n $= div
    if div <= 1
      return n
    end
    sh = sh << 1
  end
end

function gray_decode_travel(start::Int64, finish::Int64, mask::Int64, g::Int64)
  travel_bit = start $ finish
  modulus = mask + 1          # == 2**nBits
  rg = (g $ start)*fld(modulus ,travel_bit*2)
  return gray_decode((rg | fld(rg,modulus)) & mask)
end

## gray_encode_travel -- gray_encode given start and end using bit rotation.
#    Modified Gray code.  mask is 2**nbits - 1, the highest i value, so
#        gray_encode_travel( start, end, mask, 0 )    == start
#        gray_encode_travel( start, end, mask, mask ) == end
#        with a Gray-code-like walk in between.
#    This method takes the canonical Gray code, rotates the output word bits,
#    then xors ("^" in Python) with the start value.
#
function gray_encode_travel(start::Int64, finish::Int64, mask::Int64, i::Int64)
    travel_bit = start $ finish
    modulus = mask + 1          # == 2**nBits
    # travel_bit = 2**p, the bit we want to travel.
    # Canonical Gray code travels the top bit, 2**(nBits-1).
    # So we need to rotate by ( p - (nBits-1) ) == (p + 1) mod nBits.
    # We rotate by multiplying and dividing by powers of two:
    g = gray_encode( i ) * ( travel_bit * 2 )
    return ((g | fld(g, modulus)) & mask) $ start
end

## child_start_end( parent_start, parent_end, mask, i ) -- Get start & end for child.
#    i is the parent's step number, between 0 and mask.
#    Say that parent( i ) =
#           gray_encode_travel( parent_start, parent_end, mask, i ).
#    And child_start(i) and child_end(i) are what child_start_end()
#    should return -- the corners the child should travel between
#    while the parent is in this quadrant or child cube.
#      o  child_start( 0 ) == parent( 0 )       (start in a corner)
#      o  child_end( mask ) == parent( mask )   (end in a corner)
#      o  child_end(i) - child_start(i+1) == parent(i+1) - parent(i)
#         (when parent bit flips, same bit of child flips the opposite way)
#    Those constraints still leave choices when nD (# of bits in mask) > 2.
#    Here is how we resolve them when nD == 3 (mask == 111 binary),
#    for parent_start = 000 and parent_end = 100 (canonical Gray code):
#         i   parent(i)    child_
#         0     000        000   start(0)    = parent(0)
#                          001   end(0)                   = parent(1)
#                 ^ (flip)   v
#         1     001        000   start(1)    = parent(0)
#                          010   end(1)                   = parent(3)
#                ^          v
#         2     011        000   start(2)    = parent(0)
#                          010   end(2)                   = parent(3)
#                 v          ^
#         3     010        011   start(3)    = parent(2)
#                          111   end(3)                   = parent(5)
#               ^          v
#         4     110        011   start(4)    = parent(2)
#                          111   end(4)                   = parent(5)
#                 ^          v
#         5     111        110   start(5)    = parent(4)
#                          100   end(5)                   = parent(7)
#                v          ^
#         6     101        110   start(6)    = parent(4)
#                          100   end(6)                   = parent(7)
#                 v          ^
#         7     100        101   start(7)    = parent(6)
#                          100   end(7)                   = parent(7)
#    This pattern relies on the fact that gray_encode_travel()
#    always flips the same bit on the first, third, fifth, ... and last flip.
#    The pattern works for any nD >= 1.
#
function child_start_end(parent_start::Int64, parent_end::Int64, mask::Int64, i::Int64)
  start_i = max(0,    (i - 1) & ~1)  # next lower even number, or 0
  end_i =   min(mask, (i + 1) |  1)  # next higher odd number, or mask
  child_start = gray_encode_travel(parent_start, parent_end, mask, start_i)
  child_end   = gray_encode_travel(parent_start, parent_end, mask, end_i)
  return child_start, child_end
end

## transpose_bits --
#    Given nSrcs source ints each nDests bits long,
#    return nDests ints each nSrcs bits long.
#    Like a matrix transpose where ints are rows and bits are columns.
#    Earlier srcs become higher bits in dests;
#    earlier dests come from higher bits of srcs.

@inbounds function transpose_bits(srcs::Vector{Int64}, nDests::Int64)
  #srcs = copy(srcs)  # Make a copy we can modify safely.
  nSrcs = length(srcs)
  dests = Array{Int64}(nDests)#zeros(Int64,nDests)
  # Break srcs down least-significant bit first, shifting down:
  for j = (nDests - 1):-1:0
    # Put dests together most-significant first, shifting up:
    dest = 0
    for k = 0:(nSrcs-1)
      dest = dest*2 + mod(srcs[k+1],2)
      srcs[k+1] = fld(srcs[k+1],2)
    end
    dests[j+1] = dest
  end
  return dests
end

#     Overwrites srcs! and writes data to dests,
#apparently slower!!!!
@inbounds function transpose_bits!(srcs::Vector{Int64}, nDests::Int64, dests::Vector{Int64})
  nSrcs = length(srcs)
  # Break srcs down least-significant bit first, shifting down:
  for j = (nDests - 1):-1:0
    # Put dests together most-significant first, shifting up:
    dest = 0
    for k = 0:(nSrcs-1)
      dest = dest*2 + mod(srcs[k+1],2)
      srcs[k+1] = fld(srcs[k+1],2)
    end
    dests[j+1] = dest
  end
  return dests
end


#jl #! format: off
# # PTDF with [PowerSimulations.jl](https://github.com/NREL-SIIP/PowerSimulations.jl)

# **Originally Contributed by**: Sourabh Dalvi

# ## Introduction

# PowerSimulations.jl supports linear PTDF optimal power flow formulation. This example shows a
# single multi-period optimization of economic dispatch with a linearized DC-OPF representation of
# using PTDF power flow and how to extract duals values or locational marginal prices for energy.

# ## Dependencies
using SIIPExamples
using PowerSystems
using PowerSimulations
using PowerSystemCaseBuilder
using DataFrames

# Since we'll be retrieving duals, we need a solver that returns duals values
# here we use Ipopt.
using Ipopt
solver = optimizer_with_attributes(Ipopt.Optimizer)

# We can use the same RTS data and some of the initialization as in
# [OperationsProblem example](https://nbviewer.jupyter.org/github/NREL-SIIP/SIIPExamples.jl/blob/master/notebook/3_PowerSimulations_examples/01_operations_problems.ipynb)
sys = build_system(PSITestSystems, "modified_RTS_GMLC_DA_sys")

# Here, we want do define an economic dispatch (linear generation decisions) with
# linear DC-OPF using PTDF network representation.
# So, starting with the network, we can select from _almost_ any of the endpoints on this
# tree:
print_tree(PowerSimulations.PM.AbstractPowerModel)

# For now, let's just choose a standard PTDF formulation.
ed_template = template_economic_dispatch(network = StandardPTDFModel)

# Currently  energy budget data isn't stored in the RTS-GMLC dataset.
set_device_model!(ed_template, HydroEnergyReservoir, HydroDispatchRunOfRiver)

# Calculate the PTDF matrix.
PTDF_matrix = PTDF(sys)

# Now we can build a 4-hour economic dispatch / OPF problem with the RTS data.
# Here, we have to pass the keyword argument `constraint_duals` to OperationsProblem
# with the name of the constraint for which duals are required for them to be returned in the results.
problem = OperationsProblem(
    EconomicDispatchProblem,
    ed_template,
    sys,
    horizon = 1,
    optimizer = solver,
    balance_slack_variables = true,
    constraint_duals = [
        :CopperPlateBalance,
        :network_flow__Line,
        :network_flow__TapTransformer,
    ],
    PTDF = PTDF_matrix,
)
build!(problem, output_dir = mktempdir())

# And solve the problem and collect the results
solve!(problem)

# Here we collect the dual values from the results for the `:CopperPlateBalance` and `:network_flow`
# constraints. In the case of PTDF network formulation we need to compute the final LMP for each bus in the system by
# subtracting the duals (μ) of `:network_flow` constraints multiplied by the PTDF matrix
# from the  dual (λ) of `:CopperPlateBalance` constraint.
res = ProblemResults(problem)
duals = get_duals(res)
λ = convert(Array, duals[:CopperPlateBalance][:, :CopperPlateBalance])
flow_duals = outerjoin(
    [duals[k] for k in [:network_flow__Line, :network_flow__TapTransformer]]...,
    on = :DateTime,
)
μ = Matrix(flow_duals[:, PTDF_matrix.axes[1]])

# Here we create Dict to store the calculate congestion component of the LMP which is a product of μ and the PTDF matrix.
buses = get_components(Bus, sys)
congestion_lmp = Dict()
for bus in buses
    congestion_lmp[get_name(bus)] = μ * PTDF_matrix[:, get_number(bus)]
end
congestion_lmp = DataFrame(congestion_lmp)

# Finally here we get the LMP for each node in a lossless DC-OPF using the PTDF formulation.
LMP = λ .- congestion_lmp


using Jackknife
using Base.Test

import Jackknife: bias, estimate, leaveoneout, variance

@test bias(mean, 1:10) == 0.0
@test estimate(mean, 1:10) == mean(1:10)
@test_throws ArgumentError leaveoneout(mean, Int[])
@test_throws ArgumentError leaveoneout(mean, [1])

@test bias(sum, 1:10) == -49.5
@test estimate(sum, 1:10) == 104.5
@test variance(sum, 1:10) == 74.25
@test leaveoneout(sum, 1:10) == collect(54:-1:45)


@testset "Mass matrices and inverses" begin
    R = FiniteDifferences(20,0.2)
    @testset "Mass matrix" begin
        @test R'R == step(R)*I
    end
    @testset "Inverses" begin
        R⁻¹ = pinv(R)
        @test R⁻¹*R == I
        @test R*R⁻¹ == I

        cu = rand(size(R,2))
        cv = rand(size(R,2))
        cuv = [cu cv]

        u = R*cu
        v = R*cv
        uv = R*cuv

        @test R⁻¹*u === cu
        @test R⁻¹*v === cv
        @test R⁻¹*uv === cuv

        ut = u'
        # Fails with: ERROR: MethodError: no method matching axes(::UniformScaling{Float64}, ::Int64)
        # @test ut*R⁻¹' === ut.args[1]
    end
end

@testset "Scalar operators" begin
    B = FiniteDifferences(5,1.0)
    x = axes(B,1)
    x² = B'QuasiDiagonal(x.^2)*B
    @test x² isa Diagonal

    v = ones(5)
    @test x²*v == locs(B).^2

    y = Inclusion(0.0..5.0)
    @test_throws DimensionMismatch B'QuasiDiagonal(y.^2)*B

    @testset "Restricted bases" begin
        N = 10
        ρ = 1.0

        R = FiniteDifferences(N, ρ)
        r = axes(R, 1)

        sel = 3:6
        sel2 = 8:10

        R̃ = R[:, sel]
        R′ = R[:, sel2]

        x = QuasiDiagonal(r)

        A = R'x*R
        @test A isa Diagonal

        a = R̃'x*R
        @test a isa BandedMatrix
        @test a == Matrix(A)[sel,:]
        @test bandwidths(a) == (-2,2)

        a = R'x*R̃
        @test a isa BandedMatrix
        @test a == Matrix(A)[:,sel]
        @test bandwidths(a) == (2,-2)

        a = R̃'x*R̃
        @test a isa BandedMatrix
        @test isdiag(a)
        @test a == Matrix(A)[sel,sel]

        # Non-overlapping restrictions
        a = R̃'x*R′
        @test a isa BandedMatrix
        @test size(a) == (length(sel), length(sel2))
        @test iszero(a)
        @test bandwidths(a) == (3,-5)
    end
end


# # An unstable Bickley jet in Shallow Water model
#
# This example uses Oceananigans.jl's `ShallowWaterModel` to simulate
# the evolution of an unstable, geostrophically balanced, Bickley jet
# The example is periodic in ``x`` with flat bathymetry and
# uses the conservative formulation of the shallow water equations.
# The initial conditions superpose the Bickley jet with small-amplitude perturbations.
# See ["The nonlinear evolution of barotropically unstable jets," J. Phys. Oceanogr. (2003)](https://doi.org/10.1175/1520-0485(2003)033<2173:TNEOBU>2.0.CO;2)
# for more details on this problem.
#
# The mass transport ``(uh, vh)`` is the prognostic momentum variable
# in the conservative formulation of the shallow water equations,
# where ``(u, v)`` are the horizontal velocity components and ``h``
# is the layer height.
#
# ## Install dependencies
#
# First we make sure that we have all of the packages that are required to
# run the simulation.
#
# ```julia
# using Pkg
# pkg"add Oceananigans, NCDatasets, Plots, Printf, Polynomials"
# ```

using Oceananigans
using Oceananigans.Models: ShallowWaterModel

# ## Two-dimensional domain
#
# The shallow water model is a two-dimensional model and thus the number of vertical
# points `Nz` must be set to one.  Note that ``L_z`` is the mean depth of the fluid.

Lx, Ly, Lz = 2π, 20, 1
Nx, Ny = 128, 128

grid = RectilinearGrid(size = (Nx, Ny),
                       x = (0, Lx), y = (-Ly/2, Ly/2),
                       topology = (Periodic, Bounded, Flat))

# ## Physical parameters
#
# We choose non-dimensional parameters

const U = 1.0         # Maximum jet velocity

f = 1           # Coriolis parameter
g = 9.8         # Gravitational acceleration
Δη = f * U / g  # Maximum free-surface deformation as dictated by geostrophy

# ## Building a `ShallowWaterModel`
#
# We build a `ShallowWaterModel` with the `WENO5` advection scheme and
# 3rd-order Runge-Kutta time-stepping,

model = ShallowWaterModel(
                          timestepper = :RungeKutta3,
                          advection = WENO5(),
                          grid = grid,
                          gravitational_acceleration = g,
                          coriolis = FPlane(f=f))

# Use `architecture = GPU()` to run this problem on a GPU.

# ## Background state and perturbation
#
# The background velocity ``ū`` and free-surface ``η̄`` correspond to a
# geostrophically balanced Bickely jet with maximum speed of ``U`` and maximum
# free-surface deformation of ``Δη``,

h̄(x, y, z) = Lz - Δη * tanh(y)
ū(x, y, z) = U * sech(y)^2

# The total height of the fluid is ``h = L_z + \eta``. Linear stability theory predicts that
# for the parameters we consider here, the growth rate for the most unstable mode that fits
# our domain is approximately ``0.139``.

# The vorticity of the background state is

ω̄(x, y, z) = 2 * U * sech(y)^2 * tanh(y)

# The initial conditions include a small-amplitude perturbation that decays away from the
# center of the jet.

small_amplitude = 1e-4

 uⁱ(x, y, z) = ū(x, y, z) + small_amplitude * exp(-y^2) * randn()
uhⁱ(x, y, z) = uⁱ(x, y, z) * h̄(x, y, z)

# We first set a "clean" initial condition without noise for the purpose of discretely
# calculating the initial 'mean' vorticity,

ū̄h(x, y, z) = ū(x, y, z) * h̄(x, y, z)

set!(model, uh = ū̄h, h = h̄)

# We next compute the initial vorticity and perturbation vorticity,

uh, vh, h = model.solution

## Build velocities
u = uh / h
v = vh / h

## Build and compute mean vorticity discretely
ω = ComputedField(∂x(v) - ∂y(u))
compute!(ω)

## Copy mean vorticity to a new field
ωⁱ = Field(Face, Face, Nothing, model.architecture, model.grid)
ωⁱ .= ω

## Use this new field to compute the perturbation vorticity
ω′ = ComputedField(ω - ωⁱ)

# and finally set the "true" initial condition with noise,

set!(model, uh = uhⁱ)

# ## Running a `Simulation`
#
# We pick the time-step so that we make sure we resolve the surface gravity waves, which
# propagate with speed of the order ``\sqrt{g L_z}``. That is, with `Δt = 1e-2` we ensure
# that `` \sqrt{g L_z} Δt / Δx,  \sqrt{g L_z} Δt / Δy < 0.7``.

simulation = Simulation(model, Δt = 1e-2, stop_time = 150)

# ## Prepare output files
#
# Define a function to compute the norm of the perturbation on the cross channel velocity.
# We obtain the `norm` function from `LinearAlgebra`.

using LinearAlgebra: norm

perturbation_norm(args...) = norm(v)

# Build the `output_writer` for the two-dimensional fields to be output.
# Output every `t = 1.0`.

simulation.output_writers[:fields] = NetCDFOutputWriter(model, (; ω, ω′),
                                                        filepath = joinpath(@__DIR__, "shallow_water_Bickley_jet_fields.nc"),
                                                        schedule = TimeInterval(1),
                                                        mode = "c")

# Build the `output_writer` for the growth rate, which is a scalar field.
# Output every time step.

simulation.output_writers[:growth] = NetCDFOutputWriter(model, (; perturbation_norm),
                                                        filepath = joinpath(@__DIR__, "shallow_water_Bickley_jet_perturbation_norm.nc"),
                                                        schedule = IterationInterval(1),
                                                        dimensions = (; perturbation_norm = ()),
                                                        mode = "c")

# And finally run the simulation.

run!(simulation)

# ## Visualize the results

# Load required packages to read output and plot.

using NCDatasets, Plots, Printf
nothing # hide

# Define the coordinates for plotting.

x, y = xnodes(ω), ynodes(ω)
nothing # hide

# Define keyword arguments for plotting the contours.

kwargs = (
         xlabel = "x",
         ylabel = "y",
         aspect = 1,
           fill = true,
         levels = 20,
      linewidth = 0,
          color = :balance,
       colorbar = true,
           ylim = (-Ly/2, Ly/2),
           xlim = (0, Lx)
)
nothing # hide

# Read in the `output_writer` for the two-dimensional fields and then create an animation
# showing both the total and perturbation vorticities.

ds = NCDataset(simulation.output_writers[:fields].filepath, "r")

iterations = keys(ds["time"])

anim = @animate for (iter, t) in enumerate(ds["time"])
    ω = ds["ω"][:, :, 1, iter]
    ω′ = ds["ω′"][:, :, 1, iter]

    ω′_max = maximum(abs, ω′)

    plot_ω = contour(x, y, ω',
                     clim = (-1, 1),
                     title = @sprintf("Total vorticity, ω, at t = %.1f", t); kwargs...)

    plot_ω′ = contour(x, y, ω′',
                      clim = (-ω′_max, ω′_max),
                      title = @sprintf("Perturbation vorticity, ω - ω̄, at t = %.1f", t); kwargs...)

    plot(plot_ω, plot_ω′, layout = (1, 2), size = (800, 440))
end

close(ds)

mp4(anim, "shallow_water_Bickley_jet.mp4", fps=15)

# Read in the `output_writer` for the scalar field (the norm of ``v``-velocity).

ds2 = NCDataset(simulation.output_writers[:growth].filepath, "r")

iterations = keys(ds2["time"])

     t = ds2["time"][:]
norm_v = ds2["perturbation_norm"][:]

close(ds2)
nothing # hide

# We import the `fit` function from `Polynomials.jl` to compute the best-fit slope of the
# perturbation norm on a logarithmic plot. This slope corresponds to the growth rate.

using Polynomials: fit

I = 6000:7000

degree = 1
linear_fit_polynomial = fit(t[I], log.(norm_v[I]), degree, var = :t)

# We can get the coefficient of the ``n``-th power from the fitted polynomial by using `n`
# as an index, e.g.,

constant, slope = linear_fit_polynomial[0], linear_fit_polynomial[1]

# We then use the computed linear fit coefficients to construct the best fit and plot it
# together with the time-series for the perturbation norm for comparison.

best_fit = @. exp(constant + slope * t)

plot(t, norm_v,
        yaxis = :log,
        ylims = (1e-3, 30),
           lw = 4,
        label = "norm(v)",
       xlabel = "time",
       ylabel = "norm(v)",
        title = "growth of perturbation norm",
       legend = :bottomright)

plot!(t[I], 2 * best_fit[I], # factor 2 offsets fit from curve for better visualization
           lw = 4,
        label = "best fit")

# The slope of the best-fit curve on a logarithmic scale approximates the rate at which instability
# grows in the simulation. Let's see how this compares with the theoretical growth rate.

println("Numerical growth rate is approximated to be ", round(slope, digits=3), ",\n",
        "which is very close to the theoretical value of 0.139.")


using Reduce
using ReduceLinAlg
using Test

# write your own tests here
@test mat_jacobian((:x,),(:x,)) == mat_jacobian([:x],[:x])
@test hessian(:x,(:x,)) == hessian(:x,[:x])
@test jordan_block(:x,1) == jordan_block(Reduce.RExpr("x"),Reduce.RExpr("1")) |> parse


using Turing, ReverseDiff, Memoization
using DifferentialEquations
using Plots
using Statistics
using StatsPlots
using DataFrames
using CSV
using NLsolve
Turing.setadbackend(:reversediff)
Turing.setrdcache(true)
Turing.turnprogress(true)
using Distributions
using LinearAlgebra

# here we are simulating the signal of two coupled oscillators
# with amplitudes A1 and A2 of kBT/k and kBT/(k+c) where c is the coupling constant
# D = kBT/gamma
# in terms of the standard parameters of a OU process:
# sigma = sqrt(2*D)
# Theta = D/A

function corr_osc(c,N,delta_t)
    # function that returns two time series that are correlated through a coupling coefficient c
    μ = 0.0 # mean is zero
    σ = sqrt(2) # D=1
    Θ1 = 1.0
    Θ2 = 1.0+abs(c)
    total_time = N*delta_t
    W1 = OrnsteinUhlenbeckProcess(Θ1,μ,σ,0.0,1.0)
    W2 = OrnsteinUhlenbeckProcess(Θ2,μ,σ,0.0,1.0)
    prob1 = NoiseProblem(W1,(0.0,total_time))
    prob2 = NoiseProblem(W2,(0.0,total_time))
    sol1 = solve(prob1;dt=delta_t)
    sol2 = solve(prob2;dt=delta_t)

    # creating the two correlated
    x1 = (sol1.u .+ sol2.u)/2
    if c>0
        x2 = (sol1.u .- sol2.u)/2
    else
        x2 = (sol2.u .- sol1.u)/2
    end
    return x1,x2
end

function calc_fundstats(x)
    return x[1]^2+x[end]^2,sum(x[2:end-1].^2),sum(x[1:end-2].*x[2:end-1])
end

function phi_deriv(F,x,a1ep,a1ss,a1c,a2ep,a2ss,a2c,delta_t,N)
    # x[1] = A1, x[2] = A2, x[3]=D
    A1 = x[1]
    A2 = x[2]
    D = x[3]
    b1 = b(D,A1,delta_t)
    b2 = b(D,A2,delta_t)
    Q1 = q(a1ep,a1ss,a1c,b1)
    Q2 = q(a2ep,a2ss,a2c,b2)
    dQ1 = dqdB(a1ep,a1ss,a1c,b1)
    dQ2 = dqdB(a2ep,a2ss,a2c,b2)
    F[1] = -N*A1^2/2 + A1*Q1/2 + b1*D*delta_t*(A1*b1*(N-1)/(1-b1^2)-dQ1/2)
    F[2] = -N*A2^2/2 + A2*Q2/2 + b2*D*delta_t*(A2*b2*(N-1)/(1-b2^2)-dQ2/2)
    F[3] = (b1*(N-1)/(1-b1^2)-dQ1/A1/2)*b1/A1 + (b2*(N-1)/(1-b2^2)-dQ2/A2/2)*b2/A2
end

@model ou_corr(rn1,rn2,T,delta_t) = begin
    ampl1 ~ Uniform(0.0,5.0)
    ampl2 ~ Uniform(0.0,5.0)
    d ~ Uniform(0.0,5.0)
    b1 = exp(-delta_t*d/ampl1)
    b2 = exp(-delta_t*d/ampl2)

    rn1[1] ~ Normal(0,sqrt(ampl1))
    rn2[1] ~ Normal(0,sqrt(ampl2))

    for i=2:T
        rn1[i] ~ Normal(rn1[i-1]*b1,sqrt(ampl1*(1-b1^2)))
        rn2[i] ~ Normal(rn2[i-1]*b2,sqrt(ampl2*(1-b2^2)))
    end
end

# here we would like to run MCMC on several artificial data sets with varying correlation coefficients
# we should test rho from -0.9 to 0.9

results = DataFrame(rho = Float64[],
                    coupling = Float64[],
                    pearson = Float64[],
                    a1 = Float64[],
                    da1 = Float64[],
                    a2 = Float64[],
                    da2 = Float64[],
                    d = Float64[],
                    dd = Float64[],
                    da1da2 = Float64[],
                    da1dd = Float64[],
                    dasdd = Float64[],
                    a1ep = Float64[],
                    a1ss = Float64[],
                    a1c = Float64[],
                    a2ep = Float64[],
                    a2ss = Float64[],
                    a2c = Float64[],
                    c = Float64[],
                    dc = Float64[])

tries = 500
rho = 0.5
delta_t = 0.3
N = 1000
coupling = 2*abs(rho)/(1-abs(rho))*sign(rho)
println("coupling: ",coupling," rho: ",rho)
# do rho tries times
for i in 1:tries
     x1, x2 = corr_osc(coupling,N,delta_t)
    pearson = Statistics.cor(x1,x2)
    println("i = ",i,"Pearson: ",pearson)

    # lets see whether we can estimate c from the data
    y1 = x1 .+ x2
    y2 = x1 .- x2

    a1ep,a1ss,a1c = calc_fundstats(y1)
    a2ep,a2ss,a2c = calc_fundstats(y2)

    # Ornstein-Uhlenbeck process of two coupled oscillators
    model_fct = ou_corr(y1,y2,length(y1),delta_t)
    chn = sample(model_fct, NUTS(0.65), 5000)
#        Turing.emptyrdcache()

    print(describe(chn))
    println(chn)
    ampl1 = Array(chn[:ampl1])
    ampl2 = Array(chn[:ampl2])
    d = Array(chn[:d])

    a1 = mean(ampl1)
    a2 = mean(ampl2)
    da1 = std(ampl1)
    da2 = std(ampl2)
    da1da2 = cov(ampl1,ampl2)
    da1dd = cov(ampl1,d)
    da2dd = cov(ampl2,d)
    d_mean = mean(d)
    dd = std(d)

    if a1>a2
        c = (ampl1 .- ampl2)./ampl2
    else
        c = (ampl1 .- ampl2)./ampl1
    end
    c_mean = mean(c)
    dc = std(c)

    println("A1,A2: ",a1,",",a2)
    println("c estimate: ",c_mean,"std: ",dc)
    println("cross corr: ",da1da2,da1dd,da2dd)

    push!(results,[rho,coupling,pearson,a1,da1,a2,da2,d_mean,dd,da1da2[1],da1dd[1],da2dd[1],a1ep,a1ss,a1c,a2ep,a2ss,a2c,c_mean,dc])
end

println(results)
CSV.write("correlations1k05.csv",results)


# coding: utf-8

# In[1]:

#details of insertion sort can be found in the following link
# https://github.com/je-suis-tm/search-and-sort/blob/master/bubble%2C%20selection%20and%20insertion%20sort.jl
function insertion_sort(arr)

    for i in 2:length(arr)

        target=arr[i]

        j=i

        while (j!=1) && (target<arr[j-1])

            arr[j]=arr[j-1]
            j-=1
        end

        arr[j]=target

    end

    return arr

end

# In[2]:

#shell sort is a variation of insertion sort
#slicing method [i:group:length(arr)] is the key
#we gotta cut the original list into a few small groups

#assume we have a list of n elements
#the first step, we get a few sublists by [i:group:length(arr)]
#we apply insertion sort on those sublists and concatenate
#next, we do group/default
#and we obtain a few new sublists by [i:group:length(arr)]
#we apply insertion sort on new sublists and concatenate
#we keep using /default method and do insertion sort and concatenate
#til group reaches zero
#we concatenate sublists and do a final insertion sort
#we shall end up with a sorted list
#the time complexity is quite difficult to calculate
#it falls somewhere between linear and quadratic

# In[3]:

function shell_sort(arr,default=3)

    #the first step is to initialize default value
    #we will use this variable to divide the list and do slicing
    #in this case, default equals to 3
    #you can change the default number to any digit
    #bear in mind that this variable keeps the size of each sublist reasonable
    #we keep group divided by 3 until it reaches zero
    group=Int32(floor(length(arr)/default))

    while group>0

        for i in 1:group

            #equivalent to [i::group] in python
            subset=arr[i:group:length(arr)]

            #instead of reinventing the wheel
            #i directly call it in shell sort
            subset=insertion_sort(subset)

            arr[i:group:length(arr)]=subset

        end

        group=Int32(floor(group/default))

    end

    return insertion_sort(arr)

end

# In[4]:

for _ in 1:100

    test_arr=rand(1000)

    if !(shell_sort(test_arr)==sort(test_arr))
        printstyled("Erreur",color=:red)
    end

end


using Pkg, DrWatson
#@quickactivate "StatisticalRethinkingStan"
using StanSample
using StatisticalRethinking

function sim_train_test(N::Int, K::Int, rho = [0.15, -0.4])

    n_dim = 1 + length(rho)
    n_dim = n_dim < K ? K : n_dim

    Rho = Matrix{Float64}(I, n_dim, n_dim)
    for i in 1:length(rho)
        Rho[i+1, 1] = Rho[1, i + 1] = rho[i]
    end

    x_train = Matrix(rand(MvNormal(zeros(n_dim), Rho), N)')
    x_test = Matrix(rand(MvNormal(zeros(n_dim), Rho), N)')
    y = x_train[:, 1]
    x_train = x_train[:, 2:n_dim]
    (y, x_train, x_test[:, 2:n_dim])
end

stan_01 = "
data {
    int<lower=1> K;
    int<lower=0> N;
    matrix[N, K] x;
    vector[N] y;

    int<lower=0> N_new;
    matrix[N_new, K] x_new;
}
parameters {
    real a;
    vector[K] b;
    real<lower=0> sigma;

    vector[N_new] y_new;                  // predictions
}
transformed parameters {
    vector[N] mu;
    vector[N_new] mu_new;

    mu = a + x * b;
    mu_new = a + x_new * b;
}
model {
    y ~ normal(mu, sigma);          // observed model
    y_new ~ normal(mu_new, sigma);  // prediction model
}
generated quantities {
    vector[N] log_lik;
    vector[N_new] log_lik_new;

    for (i in 1:N)
        log_lik[i] = normal_lpdf(y[i] | mu[i], sigma);

    for (j in 1:N_new)
        log_lik_new[j] = normal_lpdf(y_new[j] | mu_new[j], sigma);

}
";

sm = SampleModel("stan_01", stan_01)
N = 100
K = 6
rho = [0.15, -0.4]
y, x_train, x_test = sim_train_test(N, K, rho)
println(length(y))
println(size(x_train))
println(size(x_test))
k = K < 2 ? K : K - 1
data = (N = size(x_train, 1), K = size(x_train, 2), y = y,
    x = x_train, N_new = size(x_test, 1), x_new = x_test)
rc = stan_sample(sm; data)
if success(rc)
    sm_df = read_samples(sm, :dataframe)
    precis(sm_df[:, vcat([:a], [Symbol("b.$i") for i in 1:k])])
end

nt = read_samples(sm)
log_lik = nt.log_lik'
log_lik_new = nt.log_lik_new'

loo, loos, pk = psisloo(log_lik)
loo_new, loos_new, pk_new = psisloo(log_lik_new)

[K, loo, loo_new] |> display

#plot(y)
#plot!(nt.y_new')


using Test
using MathOptInterface
const MOI = MathOptInterface

using MathOptSetDistances
const MOD = MathOptSetDistances

using FiniteDifferences
using LinearAlgebra

const bfdm = FiniteDifferences.backward_fdm(5, 1)
const ffdm = FiniteDifferences.forward_fdm(5, 1)
const cfdm = FiniteDifferences.central_fdm(5,1)

import ChainRulesCore
const CRC = ChainRulesCore

"""
A multivariate Gaussian generator without points too close to 0
"""
function safe_randn(n)
    v = 2 * randn(n)
    for i in eachindex(v)
        while v[i] ≈ 0
            v[i] = 2 * randn()
        end
    end
    return v
end

@testset "Test gradients with finite differences" begin
    Ntrials = 10
    @testset "Dimension $n" for n in (1, 3, 10)
        vector_sets = (
            MOI.Zeros(n),
            MOI.Reals(n),
            MOI.Nonnegatives(n),
            MOI.Nonpositives(n),
        )
        vs = [safe_randn(n) for _ in 1:Ntrials]
        for _ in 1:Ntrials
            push!(vs, rand(-5:2:5, n))
        end
        @testset "Vector set $s" for s in vector_sets
            @testset "Vector $v" for v in vs
                dΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, s)
                grad_fdm1 = FiniteDifferences.jacobian(bfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                grad_fdm2 = FiniteDifferences.jacobian(ffdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                @test size(grad_fdm1) == size(grad_fdm2) == size(dΠ)
                # finite diff can mess up for low v
                @test all(eachindex(v)) do idx
                    dΠ[idx,idx] ≈ grad_fdm1[idx,idx] || dΠ[idx,idx] ≈ grad_fdm2[idx,idx]
                end
            end
        end
        @testset "PSD cone" begin
            s = MOI.PositiveSemidefiniteConeTriangle(n)
            for _ in 1:Ntrials
                L = 3 * tril(rand(n, n))
                M = L * L'
                @testset "Positive definite" begin
                    v = MOD.vec_symm(M)
                    dΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, s)
                    grad_fdm1 = FiniteDifferences.jacobian(ffdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                    grad_fdm2 = FiniteDifferences.jacobian(bfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                    @test size(grad_fdm1) == size(grad_fdm2) == size(dΠ)
                    @test dΠ ≈ I
                end
                @testset "Negative definite" begin
                    v = MOD.vec_symm(-M)
                    dΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, s)
                    grad_fdm1 = FiniteDifferences.jacobian(ffdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                    grad_fdm2 = FiniteDifferences.jacobian(bfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                    @test size(grad_fdm1) == size(grad_fdm2) == size(dΠ)
                    if !isapprox(det(M), 0, atol=10e-6)
                        @test all(dΠ .≈ 0)
                    end
                end
            end
        end
        @testset "SOC" begin
            s = MOI.SecondOrderCone(n+1)
            for _ in 1:Ntrials
                x = safe_randn(n)
                @testset "SOC interior and negative bound" begin
                    t = LinearAlgebra.norm2(x) + 2 * rand()
                    v = vcat(t, x)
                    dΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, s)
                    grad_fdm1 = FiniteDifferences.jacobian(ffdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                    grad_fdm2 = FiniteDifferences.jacobian(bfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                    @test size(grad_fdm1) == size(grad_fdm2) == size(dΠ)
                    @test dΠ ≈ grad_fdm1 || dΠ ≈ grad_fdm2
                    v = vcat(-t, x)
                    dΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, s)
                    grad_fdm1 = FiniteDifferences.jacobian(ffdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                    grad_fdm2 = FiniteDifferences.jacobian(bfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                    @test size(grad_fdm1) == size(grad_fdm2) == size(dΠ)
                    @test dΠ ≈ grad_fdm1 || dΠ ≈ grad_fdm2
                end
                @testset "Out of cone point" begin
                    for tscale in (0.1, 0.5, 0.9)
                        t = tscale * LinearAlgebra.norm2(x)
                        v = vcat(t, x)
                        dΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, s)
                        grad_fdm1 = FiniteDifferences.jacobian(ffdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                        grad_fdm2 = FiniteDifferences.jacobian(bfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                        @test size(grad_fdm1) == size(grad_fdm2) == size(dΠ)
                        @test dΠ ≈ grad_fdm1 || dΠ ≈ grad_fdm2
                        t = tscale
                        xr = x / norm(x)
                        v = vcat(t, xr)
                        dΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, s)
                        # theoretical expression for unit vector
                        @test 2dΠ ≈ [
                            1 xr'
                            xr ((t + 1) * I - t * xr * xr')
                        ]
                    end
                end
            end
        end
    end
    @testset "Indefinite matrix" begin
        s = MOI.PositiveSemidefiniteConeTriangle(2)
        Q = [
            1 0
            0 -1
        ]
        Qi = Q'
        B = [
            0 1/2
            1/2 1
        ]
        for _ in 1:5
            # scale factor
            f = 20 * rand() + 5
            A = [
                -f 0
                0 f
            ]
            Λ = Diagonal([-f, f])
            Λp = Diagonal([0, f])
            @test A ≈ Q * Λ * Qi
            v = MOD.vec_symm(A)
            Πv = MOD.projection_on_set(MOD.DefaultDistance(), v, s)
            Π = MOD.unvec_symm(Πv, 2)
            @test Π ≈ Q * Λp * Qi
            DΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, s)
            # directional derivative
            for _ in 1:Ntrials
                Xd = randn(2,2)
                xd = MOD.vec_symm(Xd)
                @test DΠ * xd ≈ MOD.vec_symm(Q * (B .* (Q' * Xd * Q)) * Q')
            end
            grad_fdm1 = FiniteDifferences.jacobian(ffdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
            grad_fdm2 = FiniteDifferences.jacobian(bfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
            @test grad_fdm1 ≈ DΠ || grad_fdm2 ≈ DΠ
        end
    end
    @testset "Scalar $ST" for ST in (MOI.LessThan, MOI.GreaterThan, MOI.EqualTo)
        s = ST(10 * randn())
        for v in 1:Ntrials
            v = randn()
            while v ≈ MOI.constant(s)
                v += 2 * randn()
            end
            dfor = ffdm(x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)
            dback = bfdm(x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)
            dΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, s)
            @test ≈(dfor, dΠ, atol=1e-5) || ≈(dback, dΠ, atol=1e-5)
        end
    end

    @testset "Exp Cone" begin
        function det_case_exp_cone(v; dual=false)
            v = dual ? -v : v
            if MOD._in_exp_cone(v; dual=false)
                return 1
            elseif MOD._in_exp_cone(-v; dual=true)
                return 2
            elseif v[1] <= 0 && v[2] <= 0 #TODO: threshold here??
                return 3
            else
                return 4
            end
        end

        Random.seed!(0)
        s = MOI.ExponentialCone()
        sd = MOI.DualExponentialCone()
        case_p = zeros(4)
        case_d = zeros(4)
        # Adjust tolerance down because a 1-2 errors when projection ends up
        #   very close to the z axis
        # For intuition, see Fig 5.1 https://docs.mosek.com/modeling-cookbook/expo.html
        #   Note that their order is reversed: (x, y, z) = (x3, x2, x1) [theirs]
        tol = 1e-6
        for ii in 1:100
            v = 5*randn(3)
            @testset "Primal Cone" begin
                case_p[det_case_exp_cone(v; dual=false)] += 1
                dΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, s)
                grad_fdm1 = FiniteDifferences.jacobian(ffdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                grad_fdm2 = FiniteDifferences.jacobian(bfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                @test size(grad_fdm1) == size(grad_fdm2) == size(dΠ)
                @test ≈(dΠ, grad_fdm1,atol=tol) || ≈(dΠ, grad_fdm2, atol=tol)
            end

            @testset "Dual Cone" begin
                case_d[det_case_exp_cone(v; dual=true)] += 1
                dΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, sd)
                grad_fdm1 = FiniteDifferences.jacobian(ffdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, sd), v)[1]'
                grad_fdm2 = FiniteDifferences.jacobian(bfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, sd), v)[1]'
                @test size(grad_fdm1) == size(grad_fdm2) == size(dΠ)
                @test ≈(dΠ, grad_fdm1,atol=tol) || ≈(dΠ, grad_fdm2, atol=tol)
            end
        end
        @test all(case_p .> 0) && all(case_d .> 0)
    end

    @testset "Power Cone" begin
        function det_case_pow_cone(x, α; dual=false)
            v = dual ? -x : x
            s = MOI.PowerCone(α)
            if MOD._in_pow_cone(v, s)
                return 1
            elseif MOD._in_pow_cone(-v, MOI.dual_set(s))
                return 2
            elseif abs(v[3]) <= 1e-8
                return 3
            else
                return 4
            end
        end

        case_p = zeros(4)
        case_d = zeros(4)
        Random.seed!(0)
        tol = 1e-5
        for ii in 1:100
            v = 5*randn(3)
            for α in [0.5; rand(0.05:0.05:0.95)]
                if ii % 10 == 1
                    v[3] = 0.0
                end
                s = MOI.PowerCone(α)
                sd = MOI.dual_set(s)
                @testset "Primal Cone" begin
                    case = det_case_pow_cone(v, α; dual=false)
                    case_p[case] += 1
                    dΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, s)
                    grad_fdm1 = FiniteDifferences.jacobian(ffdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                    grad_fdm2 = FiniteDifferences.jacobian(bfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                    grad_fdm3 = FiniteDifferences.jacobian(cfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, s), v)[1]'
                    @test size(grad_fdm1) == size(grad_fdm2) == size(dΠ)
                    @test ≈(dΠ, grad_fdm1,atol=tol) || ≈(dΠ, grad_fdm2, atol=tol) || ≈(dΠ, grad_fdm3, atol=tol)
                    if !(≈(dΠ, grad_fdm1,atol=tol) || ≈(dΠ, grad_fdm2, atol=tol) || ≈(dΠ, grad_fdm3, atol=tol))
                        error("v=$v\ndΠ = $dΠ\ncase=$case\nFD=$grad_fdm3")
                    end
                end

                @testset "Dual Cone" begin
                    case = det_case_pow_cone(v, α; dual=true)
                    case_d[case] += 1
                    dΠ = MOD.projection_gradient_on_set(MOD.DefaultDistance(), v, sd)
                    grad_fdm1 = FiniteDifferences.jacobian(ffdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, sd), v)[1]'
                    grad_fdm2 = FiniteDifferences.jacobian(bfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, sd), v)[1]'
                    grad_fdm3 = FiniteDifferences.jacobian(cfdm, x -> MOD.projection_on_set(MOD.DefaultDistance(), x, sd), v)[1]'
                    @test size(grad_fdm1) == size(grad_fdm2) == size(dΠ)
                    @test ≈(dΠ, grad_fdm1,atol=tol) || ≈(dΠ, grad_fdm2, atol=tol) || ≈(dΠ, grad_fdm3, atol=tol)
                    if !(≈(dΠ, grad_fdm1,atol=tol) || ≈(dΠ, grad_fdm2, atol=tol) || ≈(dΠ, grad_fdm3, atol=tol))
                        error("v=$v\ndΠ = $dΠ\ncase=$case\nFD=$grad_fdm3")
                    end
                end
            end
        end
        @test all(case_p .> 0) && all(case_d .> 0)
    end
end


using Random
using Plots
include("../src/sample/direct.jl")
include("../src/vas.jl")

function sir_vas(β, c, γ)
    take = [
        1 0;
        1 1;
        0 0;
    ]
    give = [
        0 0;
        2 0;
        0 1;
    ]
    rates = [(state) -> Exponential(β*c*state[2]/sum(state)*state[1]), (state) -> Exponential(state[2] * γ)]
    (take, give, rates)
end

p = [0.05,10.0,0.25]; # β,c,γ
u0 = [990,10,0]; # S,I,R

take, give, rates = sir_vas(p...)

rng = MersenneTwister()

vas = VectorAdditionSystem(take, give, rates)
state = zero_state(vas)

vam = VectorAdditionModel(vas, state, 0.0)
fsm = VectorAdditionFSM(vam, DirectCall(Int))

initial_state = vas_initial(vas, u0)

out = Matrix{Float64}(undef, u0[2] + 2*u0[1] + 1, 4)

when, next_transition = simstep!(fsm, initial_state, rng)

event_cnt = 0
tnow = 0.0
out[event_cnt + 1, 2:4] = u0
out[event_cnt + 1, 1] = tnow

while next_transition !== nothing
    tnow += when
    out[event_cnt + 1, 1] = tnow
    out[event_cnt + 1, 2:4] = vam.state
    token = vas_delta(vas, next_transition)
    when, next_transition = simstep!(fsm, token, rng)
    event_cnt += 1
end

plot(
    out[1:event_cnt, 1],
    out[1:event_cnt, 2:4],
    label=["S" "I" "R"],
    xlabel="Time",
    ylabel="Number"
)


 n=rand(1:10)
 @test diag(full(Toeplitz([-n:n]))) == zeros(n+1)


#Write a program that prints a multiplication table for numbers up to 12.

using Printf

println(" X |   1   2   3   4   5   6   7   8   9  10  11  12")
println("---+------------------------------------------------")

for i=1:12, j=0:12
    if j == 0
        @printf("%2d | ", i)
    elseif i <= j
        @printf("%3d%c", i * j, j == 12 ? '\n' : ' ')
    else
        print("    ")
    end
end


#=
@author : Spencer Lyon

@date: 07/09/2014
=#

using PyPlot

const r = 0.05
const beta = 1.0 / (1.0 + r)
const T = 20  # Time horizon
const S = 5   # Impulse date
const sigma1 = 0.15
const sigma2 = 0.15

function time_path(permanent=false)
    w1 = zeros(T+1)
    w2 = zeros(T+1)
    b = zeros(T+1)
    c = zeros(T+1)

    if permanent === false
        w2[S+2] = 1.0
    else
        w1[S+2] = 1.0
    end

    for t=2:T
        b[t+1] = b[t] - sigma2 * w2[t]
        c[t+1] = c[t] + sigma1 * w1[t+1] + (1 - beta) * sigma2 * w2[t+1]
    end

    return b, c
end

function main()
    fix, axes = subplots(2, 1)
    plt.subplots_adjust(hspace=0.5)
    p_args = {:lw=> 2, :alpha => 0.7}

    L = 0.175

    for ax in axes
        ax[:grid](alpha=0.5)
        ax[:set_xlabel]("Time")
        ax[:set_ylim](-L, L)
        ax[:plot]((S, S), (-L, L), "k-", lw=0.5)
    end

    ax = axes[1]
    b, c = time_path(false)
    ax[:set_title]("impulse-response, transitory income shock")
    ax[:plot](0:T, c, "g-", label="consumption"; p_args...)
    ax[:plot](0:T, b, "b-", label="debt"; p_args...)
    ax[:legend](loc="upper right")

    ax = axes[2]
    b, c = time_path(true)
    ax[:set_title]("impulse-response, permanent income shock")
    ax[:plot](0:T, c, "g-", label="consumption"; p_args...)
    ax[:plot](0:T, b, "b-", label="debt"; p_args...)
    ax[:legend](loc="lower right")

    return Void
end


n = rand(3:10)
@test matrixdepot("grcar", n, 3) == matrixdepot("grcar", n)

A = matrixdepot("grcar", n, n-1)
@test istril(A - ones(n,n))
@test istriu(A + diagm(ones(n-1), -1))
println("'grcar' passed test...")


include("basic_setup.jl")
using Statistics

function estimator_variance(num_responses, integration_nodes, chain_length)
    algorithm = TIEstimate(1024, integration_nodes, chain_length)
    result = Trajectories.marginal_entropy(gen, algorithm=algorithm; num_responses=num_responses, duration=100.0)
    estimator_var = var(result["marginal_entropy"].Sample) / (num_responses - 1)
    time = sum(result["marginal_entropy"].TimeElapsed)
    (estimator_var, time)
end

integration_nodes = 4

chains = map(x -> round(Int, x), 2 .^ (8:0.25:11))

results = estimator_variance.(100, integration_nodes, chains)
times = map(x->x[2], results)
variance = map(x->x[1], results)

using DataFrames, GLM
data = DataFrame(N=chains, Time=times, Variance=variance)

ols = lm(@formula(Time ~ N), data)
print(ols)

fixtime = coef(ols)[1] / 100
tau_s = coef(ols)[2] / 100

nr(ns) = round(Int, 10 * 60 / (fixtime + ns * tau_s))

chains2 = map(x -> round(Int, x), 2 .^ (9:0.125:12))
nr.(chains2)

results2 = estimator_variance.(nr.(chains2), integration_nodes, chains2)

using Plots
plot(data.N, [data.Time, predict(ols)])

data2 = DataFrame(N=chains2, Time=map(x->x[2], results2), Variance=map(x->x[1], results2))

plot(data2.N, data2.Variance, xlim=(0, :auto), ylim=(0, :auto))

using DrWatson
using CSV

mkpath(projectdir("data", "estimator_variance"))
CSV.write(projectdir("data", "estimator_variance", "integration_nodes=4.csv"), data2)
CSV.write(projectdir("data", "estimator_variance", "timing_integration_nodes=4.csv"), data)


proclist = workers()         # Get list of workers

# test that can use tasks on all workers
for proc in proclist
  println("Julia says welcome from proc ",proc," w/ hostname ",@fetchfrom(proc,gethostname() ))
end
#Libc.flush_cstdio()

# Functions for demo calculation
function calc_pi(N::Integer)
  piapprox = @sync @parallel (+) for i in 1:N
    x, y = rand(), rand()
    (x*x+y*y <= 1.0) ? 1 : 0
  end
  piapprox *= 4.0/N
end

# Perform demo calculation
srand(42)        # make sure different processors have different seeds
calc_pi(10)
@time piapprox = calc_pi(1000000)
@printf("pi is estimated to be %14.12f\n", piapprox)


#=
Popis
  Funkcia vytvori histogram zaznamenavajuci pocetnost stupna vrcholov v grafe.

Parametre
  g : datova struktura grafu :: GenericGraph

Navratova hodnota
  histogram : slovnik s mapovanim kluc=stupen_vrchola => hodnota=pocetnost_vrcholov :: Dict{Int64, Int64}()

Poznamky
=#

function degree_histogram(g::GenericGraph)

  histogram = [out_degree(v, g) => 0 for v = vertices(g)]

  for v in vertices(g)
    histogram[out_degree(v, g)] += 1
  end

  return histogram

end


# Type
"""
	(type) Segment{T<:AnyComplex} in the complex plane

Each `Segment` type is parameterized according to the common type of its complex input arguments.
"""
struct Segment{T<:AnyComplex} <: AbstractCurve
	za::T
	zb::T
	function Segment{T}(a,b) where T<:AnyComplex
		@assert isfinite(a) && isfinite(b)
		new(a,b)
	end
end

# Untyped constructor
"""
	Segment(a,b)

Consruct a segment that starts at value `a` and ends at `b`.
"""
function Segment(a::Number,b::Number)
	a,b = promote(complex(float(a)),b)
	Segment{typeof(a)}(a,b)
end

# Complex type converters
for ctype in [:Spherical,:Polar,:Complex]
	@eval begin
		function $ctype(S::Segment{T}) where T<:AnyComplex
			Segment($ctype(S.za),$ctype(S.zb))
		end
	end
end

# Required methods
arclength(S::Segment) = abs(S.zb-S.za)
point(S::Segment,t::Real) = (1-t)*S.za + t*S.zb
(C::Segment)(t::Real) = point(C,t)
"""
	arg(S::Segment,z)

Find the parameter argument `t` such that `S(t)==z` is true.

This gives undefined results if `z` is not actually on the segment.
"""
arg(S::Segment,z::Number) = (real(z) - real(S.za)) / (real(S.zb) - real(S.za))
tangent(S::Segment,t::Real) = S.zb - S.za
unittangent(S::Segment,t::Real=0) = sign(S.zb-S.za)

+(S::Segment,z::Number) = Segment(S.za+z,S.zb+z)
-(S::Segment) = Segment(-S.za,-S.zb)
*(S::Segment,z::Number) = Segment(S.za*z,S.zb*z)

"""
	inv(S)
Invert the segment `S` through the origin. In general the inverse is an `Arc`, though the result is a `Segment` if `S` would pass through the origin when extended.
"""
function inv(S::Segment)
	w = 1 ./ point(S,[0,0.5,1])
	Arc(w...)
end

# Other methods
isfinite(::Segment) = true
conj(S::Segment) = Segment(conj(S.za),conj(S.zb))
reverse(S::Segment) = Segment(S.zb,S.za)

sign(S::Segment) = unittangent(S)

"""
	isapprox(S1::Segment,S2::Segment; tol=<default>)
	S1 ≈ S2

Determine if `S1` and `S2` represent the same segment, irrespective of the type or values of its parameters. Identity is determined by agreement within `tol`, which is interpreted as the weaker of absolute and relative differences.
"""
function isapprox(S1::Segment,S2::Segment;tol=DEFAULT[:tol])
	return isapprox(S1.za,S2.za,rtol=tol,atol=tol) &&
		isapprox(S1.zb,S2.zb,rtol=tol,atol=tol)
end

"""
	isleft(z,S::Segment)

Determine whether the number `z` lies "to the left" of segment `S`. This means that the angle it makes with `tangent(S)` is in the interval (0,π).

Note that `isleft` and `isright` are *not* logical opposites; a point on the (extended) segment should give `false` in both cases.
"""
function isleft(z::Number,S::Segment)
	a,b = S.za,S.zb
	(real(b)-real(a)) * (imag(z)-imag(a)) > (real(z)-real(a)) * (imag(b)-imag(a))
end
"""
	isright(z,S::Segment)

Determine whether the number `z` lies "to the right" of segment `S`. This means that the angle it makes with `tangent(S)` is in the interval (-π,0).

Note that `isleft` and `isright` are *not* logical opposites; a point on the (extended) segment should give `false` in both cases.
"""
function isright(z::Number,S::Segment)
	a,b = S.za,S.zb
	(real(b)-real(a)) * (imag(z)-imag(a)) < (real(z)-real(a)) * (imag(b)-imag(a))
end

"""
	dist(z,S::Segment)

Compute the distance from number `z` to the segment `S`.
"""
dist(z::Number,S::Segment) = abs(z - closest(z,S))
"""
	closest(z,S::Segment)

Find the point on segment `S` that lies closest to `z`.
"""
function closest(z::Number,S::Segment)
	# translate and rotate segment to positive Re axis
	d = S.zb-S.za
	s = sign(d)
	ζ = (z-S.za)/s
	S.za + s*min( max(real(ζ),0), abs(d) )
end
"""
	reflect(z,S::Segment)

Reflect the value `z` across the extension of segment `S` to a line. (For reflection of a segment through a point, use translation and negation.)
"""
reflect(z::Number,S::Segment) = reflect(z,Line(S.za,S.zb))

# Display methods
function show(io::IO,S::Segment{T}) where {T}
	print(IOContext(io,:compact=>true),"Segment(",point(S,0),",",point(S,1),")")
end

function show(io::IO,::MIME"text/plain",S::Segment{T}) where {T}
	print(io,"Segment{$T} in the complex plane:\n   from (",point(S,0),") to (",point(S,1),")")
end

plotdata(S::Segment{T}) where T<:Union{Complex,Polar} = [S.za,S.zb]


# coding: utf-8

# In[1]:

#gnome sort is another variation from bubble sort
#go the link below to check out the details of bubble sort
# https://github.com/je-suis-tm/search-and-sort/blob/master/bubble,%20selection%20and%20insertion%20sort.jl
#similar to bubble sort, it does comparison and swap from one side to the other
#the different part is, once the swap occurs, it will reduce one step
#to make sure the newly swapped element is in the right spot of this partially sorted left part
#the concept is kinda like depth first search in graph theory

# In[2]:

function gnome_sort(arr)

    i=2

    maxval=length(arr)

    while i<=maxval

        if arr[i]>arr[i-1]
            i+=1
        else
            arr[i],arr[i-1]=arr[i-1],arr[i]

            #this step is to ensure the index to stay in the range
            #this is the crucial step
            #repeatedly move the smallest to the front
            if i>2
                i-=1
            else
                i+=1
            end
        end

    end

    return arr

end

# In[3]:

for _ in 1:100

    test_arr=rand(1000)

    if !(gnome_sort(test_arr)==sort(test_arr))
        printstyled("Erreur",color=:red)
    end

end


using Arrowhead
using DoubleDouble
using Polynomials
using MatrixDepot
using Base.Test

# There are four tests, a random matrix test and three tests from the arrowhead paper (see README for details)
println("There are four tests for arrowhead matrices, a random matrix test and three")
println("tests from the arrowhead paper [1] (see README for details)","\n")

@show tols = [1e2,1e2,1e2,1e2,1e2] # set the tolerances for eig

println("\n","1st test - Random SymArrow matrix","\n")

@show A = Arrowhead.GenSymArrow( 10, 10 )

@show U, Lambda = eig( full(A) ) # the standard eig command
@show Ua, Lambdaa, Sind, Kb, Kz, Knu, Krho, Qout = eig( A, tols ) # arrowhead solver
@show Residual = full(A) * Ua - Ua * diagm(Lambdaa) # compute the residual
@show normR = norm( Residual )
@test normR < 100.0 * eps() # norm of residual should be small

println("\n","2nd test - Example 1, p. 15 from [1]","\n")

@show A = SymArrow( [ 2e-3, 1e-7, 0.0, -1e-7, -2e-3 ], [ 1e7, 1e7, 1.0, 1e7, 1e7 ], 1e20, 6 )
@show Ua, Lambdaa, Sind, Kb, Kz, Knu, Krho, Qout = eig( A, tols )
@show Lambda=[ 1.0e20, 0.001999001249000113, 4.987562099722814e-9, -9.99999999998e-21, -2.0049855621017174e-6, -0.002001001251000111 ] # accurate eigenvalues from [1]
@test maximum( abs( Lambdaa - Lambda ) ./ abs( Lambda ) ) < 3.0 * eps() # relative errors in eigenvalues should be small
@show v4 = [ 4.999999999985e-11, 9.999999999969e-7, 0.9999999999989999, -9.999999999970999e-7, -4.999999999985e-11, -9.99999999997e-21] # v_4 from [1]
@test maximum( abs( Ua[:,4] - v4 ) ./ abs( v4 ) ) < 3.0 * eps() # component-wise relative errors in eigenvector(s) should be small

println("\n","3rd test - Example 2, p. 16 from [1]","\n")

@show A = SymArrow( [ 1+4eps(), 1+3eps(), 1+2eps(), 1+eps() ], [ 1.0, 2.0, 3.0, 4.0 ], 0.0, 5 )
@show Ua, Lambdaa, Sind, Kb, Kz, Knu, Krho, Qout = eig( A, tols )
@show Lambda = [ 6.000000000000001, 1+4eps(), 1+3eps(), 1+2eps(), -4.999999999999999 ] # accurate eigenvalues from [1]
@test maximum( abs( Lambdaa - Lambda ) ./ abs( Lambda ) ) < 3.0 * eps() # relative errors in eigenvalues should be small

println("\n","4th test - Example 3, p. 17 from [1]","\n")

@show A = SymArrow( [ 1e10+1.0/3.0, 4.0, 3.0, 2.0, 1.0 ], [ 1e10 - 1.0/3.0, 1.0, 1.0, 1.0, 1.0 ], 1e10, 6 )
@show Ua, Lambdaa, Sind, Kb, Kz, Knu, Krho, Qout = eig( A, tols )
@show Lambda = [ 1.999999999983333e10, 4.174722501468362, 3.188318635336404 , 2.223251566590035, 1.261850509234366, -0.3481422590562395 ] # accurate eigenvalues from [1]
@test maximum( abs( Lambdaa - Lambda ) ./ abs( Lambda ) ) < 3.0 * eps() # relative errors in eigenvalues should be small

# There are four tests for DPR1, a random matrix test and three tests from the arrowhead paper (see README for details)
println("There are four tests for DPR1 matrices, a random matrix test and three")
println("tests from the DPR1 paper [2] (see README for details)","\n")

@show tols = [1e3,1e3,1e3,1e3,1e3] # set the tolerances for eig

println("\n","1st test - Random SymDPR1 matrix","\n")

@show A = Arrowhead.GenSymDPR1( 10 )

@show U, Lambda = eig( full(A) ) # the standard eig command
@show Ua, Lambdaa, Sind, Kb, Kz, Knu, Krho, Qout = eig( A, tols ) # arrowhead solver
@show Residual = full(A) * Ua - Ua * diagm(Lambdaa) # compute the residual
@show normR = norm( Residual )
@test normR < 100.0 * eps() # norm of residual should be small

println("\n","2nd test - Example 1, p. 9 from [2]","\n")

@show A = SymDPR1( [ 1e10, 5.0, 4e-3, 0.0, -4e-3,-5.0 ], [ 1e10, 1.0, 1.0, 1e-7, 1.0,1.0 ], 1.0 )
@show Ua, Lambdaa, Sind, Kb, Kz, Knu, Krho, Qout = eig( A, tols )
@show Lambda=[ 1.000000000100000e20, 5.000000000100000, 4.000000100000001e-3, 9.999999998999997e-25, -3.999999900000001e-3, -4.999999999900000] # accurate eigenvalues from [1]
@test maximum( abs( Lambdaa - Lambda ) ./ abs( Lambda ) ) < 3.0 * eps() # relative errors in eigenvalues should be small
@show v4 = [ 9.999999998999996e-18, 1.999999999800000e-18, 2.499999999749999e-15, -1.000000000000000, -2.499999999749999e-15, -1.999999999800000e-18] # v_4 from [1]
@test maximum( abs( Ua[:,4] - v4 ) ./ abs( v4 ) ) < 3.0 * eps() # component-wise relative errors in eigenvector(s) should be small

println("\n","3rd test - Example 2, p. 10 from [2]","\n")

@show A = SymDPR1( [ 1+40eps(), 1+30eps(), 1+20eps(), 1+10eps() ], [ 1.0, 2.0, 2.0, 1.0 ], 1.0 )
@show Ua, Lambdaa, Sind, Kb, Kz, Knu, Krho, Qout = eig( A, tols )
@show Lambda = [ 11.0+32eps(), 1.0+39eps(), 1.0+25eps(), 1.0+11eps()] # accurate eigenvalues from [1]
@test maximum( abs( Lambdaa - Lambda ) ./ abs( Lambda ) ) < 3.0 * eps() # relative errors in eigenvalues should be small

println("\n","4th test - Example 3, p. 11 from [2], see also Ming Gu's paper","\n")

@show A = SymDPR1( [ 10.0/3.0, 2.0+1e-7, 2.0-1e-7, 1.0 ], [ 2.0, 1e-7, 1e-7, 2.0], 1.0 )
@show Ua, Lambdaa, Sind, Kb, Kz, Knu, Krho, Qout = eig( A, tols )
@test norm(Ua'*Ua -eye(4)) < 3.0 * eps() # the eigenvectors are orthogonal (AND componentwise accurate)

println("\n There is one test for SVD of HalfArrow matrices")

println("\n","HalfArrow with entries varying in magnitude","\n")
@show A=HalfArrow(sort(exp(20*(rand(8)-0.5)),rev=true),(exp(20*(rand(8)-0.5))))
@show Ua, Lambdaa, Va, Sind, Kb, Kz, Knu, Krho, Qout = svd( A, tols )
@show N1=norm(Ua'*Ua-eye(8))
@show N2=norm(Va'*Va-eye(8))
@test N1 < 300.0*eps() && N2 < 300.0*eps() # Just test the orthogonality of the eigenvectors

println("\n We test tridiagonal divide and conquer on the Wlikinson's matrix W21")

@show W=SymTridiagonal(matrixdepot("wilkinson",21))
@show U,E=tdc(W)
@show Etrue= [10.746194182903393,
 10.746194182903322,
 9.210678647361332,
 9.210678647304919,
 8.038941122829023,
 8.038941115814273,
 7.003952209528675,
 7.003951798616375,
 6.000234031584167,
 6.000217522257098,
 5.000244425001913,
 4.999782477742902,
 4.004354023440857,
 3.9960482013836245,
 3.0430992925788236,
 2.9610588841857264,
 2.1302092193625057,
 1.7893213526950813,
 0.9475343675292932,
 0.25380581709667804,
 -1.1254415221199847]
@test norm(sort(E)-sort(Etrue))<5.0*eps()

println("\n There are three tests for roots of polynomials")
println("\n Example 1 from [3]")
println("\n The Wilkinson's polynomial p18")

@show p18=Poly([   6402373705728000,
 -22376988058521600,
  34012249593822720,
 -30321254007719424,
  17950712280921504,
  -7551527592063024,
   2353125040549984,
   -557921681547048,
    102417740732658,
    -14710753408923,
      1661573386473,
      -147560703732,
        10246937272,
         -549789282,
           22323822,
            -662796,
              13566,
               -171,
                1])
D=roots(polyder(p18))
@show R,Qout=rootsah(p18,D)
@test norm(R-collect(18.0:-1:1.0))<5*eps()

println("\n Example 2 from [3]")
e5=[-618970019642690000010608640,
4181389724724490601097907890741292883247104,
-6277101735386680066937501969125693243111159424202737451008,
713623846352979940529142984724747568191373312,
-20282409603651670423947251286016, 1]
@show p5=Poly(map(Float64,e5))
pd=polyder(p5)
pd=pd/pd[end]
r=roots(pd)
D=rootsWDK(pd,r,1)
@show R,Qout=rootsah(p5,D)
@show Rtrue=  [2.028240960365167e31,
 1.759218623050247e13,
 1.7592185858329531e13,
 4.440892098500623e-16,
               2.2204460492503136e-16]

@test norm(sort(R)-sort(Rtrue))<5.0*eps()

println("\n Example 3 from [3]")
println("\n The Chebyshev's polynomial T[374]")

# Generate first 374 Chebyshev polynomials in BigFloat
n=374
T=Array(Any,n+1)
T[1]=Poly([BigFloat(1)])
T[2]=Poly([BigFloat(0),1])
for i=3:n+1
    T[i]=2*T[2]*T[i-1]-T[i-2]
end

# True zeros
Rtrue=map(Float64,[cos((2*k-1)*pi/(2*n)) for k=1:n])
# Zeros of T[n]
D=map(Float64,[cos((2*k-1)*pi/(2*(n-1))) for k=1:n-1])

# Computed zeros
R,qout=rootsah(T[n+1],D);

@show mean(abs(sort(R)-sort(Rtrue))./abs(sort(Rtrue))), maximum(abs(sort(R)-sort(Rtrue))./abs(sort(Rtrue)))
@test norm(sort(R)-sort(Rtrue))<10.0*eps()


struct GridData
    voxelNums::Vector{Int}
    startVox::Vector{Float64}
    voxelWidths::Vector{Float64}
    endVox::Vector{Float64}
end

struct ROIData
    Mins::Vector{Float64}
    Maxs::Vector{Float64}
end

function createGrid( minPoint, maxPoint, voxelNums )

    @assert sum( minPoint .<= maxPoint ) == 3

    ### MODIFY USING MULTIPLE DISPATCH ###
    if length( voxelNums ) == 1
        voxelNums = voxelNums * ones(3)
    end
    # ----------------------------------

    voxelWidths   = ( maxPoint - minPoint ) ./ voxelNums
    startVox      =   minPoint + voxelWidths * 0.5
    endVox        =   maxPoint - voxelWidths * 0.5

    return GridData( voxelNums, startVox, voxelWidths, endVox )
end


include("../examples/nonconvex_smooth_noncompact.jl")
function lyap_exp(w, η)
	N = 20000
	le = 0.0
	for n = 1:N
		le += log(abs(dnext(w, η)))/N
		w = next(w, η)
	end
	return le
end
function run(w, η, n, s)
	for i = 1:n
		w = next(w, η, s)
	end
	return w
end


@testset "Objective" begin

    # Test GameObjective
    v = [1,2,3]
    inds = [1,3,5]
    n = 5
    V = [1,0,2,0,3]
    @test V == Algames.expand_vector(v,inds,n)

    # Test Game Objective
    T = Float64
    N = 10
    n = 12
    m = 6
    p = 3
    model = UnicycleGame(p=p)
    Q = [Diagonal(rand(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(rand(SVector{model.mi[i],T})) for i=1:p]
    xf = [i*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [2i*ones(SVector{model.mi[i],T}) for i=1:p]

    game_obj = GameObjective(Q,R,xf,uf,N,model)

    X = [1.0*SVector{model.n,T}([1,2,3,1,2,3,1,2,3,1,2,3]) for k=1:N]
    U = [1.0*SVector{model.m,T}([2,4,6,2,4,6]) for k=1:N-1]
    dt = 0.1
    Dt = dt*ones(N-1)
    traj = Traj(X, U,Dt)

    @test abs(cost(game_obj.obj[1][1], traj)) <= 1e-10
    @test abs(cost(game_obj.obj[2][1], traj)) <= 1e-10
    @test abs(cost(game_obj.obj[3][1], traj)) <= 1e-10

    # Test cost_gradient! and cost_hessian!
    probsize = ProblemSize(N, model)
    pdtraj = PrimalDualTraj(probsize, dt, f=rand, amplitude=1e1)
    for i = 1:p
        for k = 1:N
            game_obj.E[i][1].cost[k].Q += rand(n,n)
            game_obj.E[i][1].cost[k].R += rand(m,m)
            game_obj.E[i][1].cost[k].H += rand(m,n)
            game_obj.E[i][1].cost[k].q += rand(n)
            game_obj.E[i][1].cost[k].r += rand(m)
            game_obj.E[i][1].cost[k].c += rand()
        end
    end
    cost_gradient!(game_obj, pdtraj)
    cost_hessian!(game_obj, pdtraj)

    # Gradient
    # Stage cost
    @test norm(game_obj.E[1][1].cost[1].q - (game_obj.obj[1][1].cost[1].Q*state(pdtraj.pr[1]) + game_obj.obj[1][1].cost[1].q)*dt, 1) < 1e-10
    @test norm(game_obj.E[1][1].cost[1].r - (game_obj.obj[1][1].cost[1].R*control(pdtraj.pr[1]) + game_obj.obj[1][1].cost[1].r)*dt, 1) < 1e-10
    # Terminal cost
    @test norm(game_obj.E[1][1].cost[end].q - (game_obj.obj[1][1].cost[end].Q*state(pdtraj.pr[end]) + game_obj.obj[1][1].cost[end].q), 1) < 1e-10
    @test norm(game_obj.E[1][1].cost[end].r, 1) < 1e-10

    # Hessian
    # Stage cost
    @test norm(game_obj.E[1][1].cost[1].Q - game_obj.obj[1][1].cost[1].Q*dt, 1) < 1e-10
    @test norm(game_obj.E[1][1].cost[1].R - game_obj.obj[1][1].cost[1].R*dt, 1) < 1e-10
    # Terminal cost
    @test norm(game_obj.E[1][1].cost[end].Q - game_obj.obj[1][1].cost[end].Q, 1) < 1e-10
    @test norm(game_obj.E[1][1].cost[end].R, 1) < 1e-10

    # Test add_collision_cost!
    T = Float64
    dt = 0.1
    N = 10
    p = 3
    model = UnicycleGame(p=p)
    probsize = ProblemSize(N,model)
    n = model.n
    m = model.m
    Q = [Diagonal(rand(SVector{model.ni[i],T})) for i=1:p]
    R = [Diagonal(rand(SVector{model.mi[i],T})) for i=1:p]
    xf = [i*ones(SVector{model.ni[i],T}) for i=1:p]
    uf = [2i*ones(SVector{model.mi[i],T}) for i=1:p]

    game_obj = GameObjective(Q,R,xf,uf,N,model)

    radius = [1.0, 2.0, 3.0]
    μ = [10.0, 20.0, 30.0]
    i = 1
    j = 2
    c = CollisionCost{n,m,T,length(model.px[i])}(μ[i], radius[i], model.px[i], model.px[j])
    add_collision_cost!(game_obj, radius, μ)

    length(game_obj.obj[1]) == 1 + p-1
    length(game_obj.obj[2]) == 1 + p-1
    length(game_obj.obj[3]) == 1 + p-1
    add_collision_cost!(game_obj, radius, μ)
    length(game_obj.obj[1]) == 1 + 2*(p-1)
    length(game_obj.obj[2]) == 1 + 2*(p-1)
    length(game_obj.obj[3]) == 1 + 2*(p-1)

    typeof(game_obj.obj[1][2].cost) <: Vector{<:CollisionCost}
    game_obj.obj[1][2].cost[1].μ == 10.0
    game_obj.obj[2][2].cost[1].μ == 20.0
    game_obj.obj[3][2].cost[1].μ == 30.0

    game_obj.obj[1][2].cost[1].r == 1.0
    game_obj.obj[2][2].cost[1].r == 2.0
    game_obj.obj[3][2].cost[1].r == 3.0

    # Test CollisionCost
    T = Float64
    dt = 0.1
    model = DoubleIntegratorGame(p=2)
    μ = 10.0
    r = 0.2
    pxi = SVector{length(model.px[1]),Int}(model.px[1])
    pxj = SVector{length(model.px[2]),Int}(model.px[2])
    cs = CollisionCost{model.n,model.m,T,length(model.px[1])}(μ,r,pxi,pxj)

    cs0 = copy(cs)
    @test cs0.μ == cs.μ
    @test cs0.r == cs.r
    @test cs0.terminal == cs.terminal
    @test all(cs0.pxi .== cs.pxi)
    @test all(cs0.pxj .== cs.pxj)
    @test state_dim(cs) == model.n
    @test control_dim(cs) == model.m

    @test cs.μ == μ
    @test cs.r == r
    @test cs.pxi == pxi
    @test cs.pxj == pxj
    @test cs.terminal == false
    @test typeof(cs) <: CollisionCost{model.n,model.m,T,length(model.px[1])}
    @test typeof(cs) <: TrajectoryOptimization.CostFunction

    x = @SVector [1.0, 1.1, 2.0, 2.0, 0.0, 0.0, 0.0, 0.0]
    u = rand(SVector{model.m,T})
    z = KnotPoint(x,u,dt)
    @test typeof(z) <: TrajectoryOptimization.AbstractKnotPoint

    # Test Stage Cost
    @test abs(TrajectoryOptimization.stage_cost(cs, x) - 0.05) < 1e-10
    @test abs(TrajectoryOptimization.stage_cost(cs, state(z)) - 0.05) < 1e-10
    @test abs(TrajectoryOptimization.stage_cost(cs, x, u) - 0.05) < 1e-10
    @test abs(TrajectoryOptimization.stage_cost(cs, state(z), control(z)) - 0.05) < 1e-10
    @test abs(TrajectoryOptimization.stage_cost(cs, x) -  0.05) < 1e-10
    @test abs(TrajectoryOptimization.stage_cost(cs, x, u) - 0.05) < 1e-10
    @test abs(TrajectoryOptimization.stage_cost(cs, z) - 0.005) < 1e-10

    @test (@ballocated TrajectoryOptimization.stage_cost($cs, $x)) == 0
    @test (@ballocated TrajectoryOptimization.stage_cost($cs, $x, $u)) == 0
    @test (@ballocated TrajectoryOptimization.stage_cost($cs, $z)) == 0

    # Test Gradient
    function easy_gradient(cs, z)
        function local_eval(x)
            z_ = KnotPoint(x, control(z), z.dt)
            return TrajectoryOptimization.stage_cost(cs,state(z_))
        end
        return ForwardDiff.gradient(local_eval, state(z))
    end

    zr = KnotPoint(rand(SVector{model.n,T}), rand(SVector{model.m,T}), dt)
    cs_active = CollisionCost{model.n,model.m,T,length(model.px[1])}(μ,1e3,pxi,pxj)
    cs_notactive = CollisionCost{model.n,model.m,T,length(model.px[1])}(μ,1e-3,pxi,pxj)
    E = TrajectoryOptimization.QuadraticCost{T}(model.n,model.m)

    TrajectoryOptimization.gradient!(E,cs_active,state(zr))
    grad_x = easy_gradient(cs_active, zr)
    @test norm(E.q - grad_x, 1)/norm(E.q) < 1e-7
    @test E.r == zeros(model.m)

    TrajectoryOptimization.gradient!(E,cs_notactive,state(zr))
    grad_x = easy_gradient(cs_notactive, zr)
    @test norm(E.q - grad_x, 1) < 1e-7
    @test E.r == zeros(model.m)

    @test (@ballocated TrajectoryOptimization.gradient!($E, $cs, $x)) == 0
    @test (@ballocated TrajectoryOptimization.gradient!($E, $cs, $x, $u)) == 0

    # Test Hessian
    function easy_hessian(cs, z)
        function local_eval(x)
            z_ = KnotPoint(x, control(z), z.dt)
            return TrajectoryOptimization.stage_cost(cs,state(z_))
        end
        return ForwardDiff.hessian(local_eval, state(z))
    end

    TrajectoryOptimization.hessian!(E,cs_active,state(zr))
    hess_x = easy_hessian(cs_active, zr)
    E.Q
    hess_x
    @test norm(E.Q - hess_x, 1) < 1e-2
    @test E.r == zeros(model.m)

    TrajectoryOptimization.hessian!(E,cs_notactive,state(zr))
    hess_x = easy_hessian(cs_notactive, zr)
    E.Q
    @test norm(E.Q - hess_x, 1) < 1e-2
    @test E.r == zeros(model.m)

    @test (@ballocated TrajectoryOptimization.hessian!($E, $cs, $x)) == 0
    @test (@ballocated TrajectoryOptimization.hessian!($E, $cs, $x, $u)) == 0

end


using HPAT
using MPI

@acc @inline function cndf2( In::Array{Float64,1} )
    Out = 0.5 .+ 0.5 .* erf(0.707106781 .* In);
    return Out;
end

@acc hpat function blackscholes(iterations)

    sptprice   = Float64[ 42.0 for i=1:iterations ]
    strike     = Float64[ 40.0 + (i / iterations) for i=1:iterations ]
    rate       = Float64[ 0.5 for i=1:iterations ]
    volatility = Float64[ 0.2 for i=1:iterations ]
    time       = Float64[ 0.5 for i=1:iterations ]

    logterm = log10(sptprice ./ strike);
    powterm = .5 .* volatility .* volatility;
    den = volatility .* sqrt(time);
    d1 = (((rate .+ powterm) .* time) .+ logterm) ./ den;
    d2 = d1 .- den;
    NofXd1 = cndf2(d1);
    NofXd2 = cndf2(d2);
    futureValue = strike .* exp(- rate .* time);
    c1 = futureValue .* NofXd2
    Call = sptprice .* NofXd1 .- c1;
    Put  = Call .- futureValue .+ sptprice;
    sum(Put)
end

n = 10^8

if (length(ARGS) > 0)
	n = parse(Int, ARGS[1])
end

rank = MPI.Comm_rank(MPI.COMM_WORLD)
pes = MPI.Comm_size(MPI.COMM_WORLD)

checksum = blackscholes(10000)

MPI.Barrier(MPI.COMM_WORLD)
t1 = time_ns()
checksum = blackscholes(n)
t2 = time_ns()

if rank==0
	println("nodes: ", pes, "\ntime: ", (t2-t1)/1.0e9, "\nchecksum: ", checksum)
end

MPI.Barrier(MPI.COMM_WORLD)


# load requried packages
using IncrementalInference
using Test

## parameters

lm_prior_noise = 0.01
meas_noise = 0.25
odom_noise = 0.1
n_samples = 200

# initialize mean landmark locations
l0 = 0.0
l1 = 10.0
l2 = 40.0

# "Ground-truth" robot poses
x0 = 0.0
x1 = 10.0
x2 = 20.0
x3 = 40.0

@testset "2door basic binary multihypothesis tests..." begin

## Initialize empty factor graph
fg = initfg()
fg.solverParams.N = n_samples
# Place strong prior on locations of three "doors"
addVariable!(fg, Symbol("l0"), ContinuousScalar, N=n_samples)
addFactor!(fg, [:l0], Prior(Normal(l0, lm_prior_noise)))

addVariable!(fg, Symbol("l1"), ContinuousScalar, N=n_samples)
addFactor!(fg, [:l1], Prior(Normal(l1, lm_prior_noise)))

# Add first pose
addVariable!(fg, :x0, ContinuousScalar, N=n_samples)

# Make first "door" measurement
# addFactor!(fg, [:x0; :l0], LinearConditional(Normal(0, meas_noise)))
addFactor!(fg, [:x0; :l0; :l1], LinearConditional(Normal(0, meas_noise)), multihypo=[1.0; 1.0/2.0; 1.0/2.0])

# Add second pose
addVariable!(fg, :x1, ContinuousScalar, N=n_samples)

# Gaussian transition model
addFactor!(fg, [:x0; :x1], LinearConditional(Normal(x1-x0, odom_noise)))

# Make second "door" measurement
# addFactor!(fg, [:x1; :l1], LinearConditional(Normal(0, meas_noise)) )
addFactor!(fg, [:x1; :l0; :l1], LinearConditional(Normal(0, meas_noise)), multihypo=[1.0; 1.0/2.0; 1.0/2.0])

## Add one more pose/odometry to invoke issue #236

# Add third pose
addVariable!(fg, :x2, ContinuousScalar, N=n_samples)
addFactor!(fg, [:x1; :x2], LinearConditional(Normal(x2-x1, odom_noise)))

## Solve graph
tree, smt, hist = solveTree!(fg)
# tree = batchSolve!(fg, drawpdf=false, show=false, recursive=false)
# drawTree(tree, show=true)

@test abs(getKDEMean(getKDE(fg, :x0))[1]) < 1.0
@test abs(getKDEMean(getKDE(fg, :x1))[1]-10) < 1.0
@test abs(getKDEMean(getKDE(fg, :x2))[1]-20) < 1.0

@test abs(getKDEMean(getKDE(fg, :l0))[1]) < 2.0
@test abs(getKDEMean(getKDE(fg, :l1))[1]-10) < 2.0

# using RoMEPlotting
# plotKDE(fg, [:l0;:l1])

end

#


module Test_dlyap

using LinearAlgebra
using MatrixEquations
using Test

@testset "Testing discrete Lyapunov equation solvers" begin

n = 100
Ty = Float64

@testset "Discrete Lyapunov equations" begin

reltol = eps(float(100))
a = -2+im; b = 2; @time x = lyapd(a,b)
@test abs(a*x*a'-x+b) < reltol

reltol = eps(float(100f0))
a = 1f0-2f0im; b = 2f0; @time x = lyapd(a,b)
@test abs(a*x*a'-x+b)  < reltol

for Ty in (Float64, Float32)

ar = rand(Ty,n,n)
ac = rand(Ty,n,n)+im*rand(Ty,n,n)
c = rand(Ty,n,n)+im*rand(Ty,n,n)
qc = c'*c
Qr = real(qc)
Ty == Float64 ? reltol = eps(float(100)) : reltol = eps(100*n*one(Ty))

@time x = lyapd(ac,qc);
@test norm(ac*x*ac'-x+qc)/norm(x)/max(1.,norm(ac)^2) < reltol

α = 3+im; # α = 1; #SingularException
@time x = lyapd(α*I,qc);
@test norm(α*x*α'-x+qc)/norm(x) < reltol

try
   α = 1; #SingularException
   @time x = lyapd(α,qc);
catch
   @test true
end

α = 3. +im;
@time x = lyapd(α,qc);
@test norm(α*x*α'-x+qc)/norm(x) < reltol

@time x = lyapd(ac',qc);
@test norm(ac'*x*ac-x+qc)/norm(x)/max(1.,norm(ac)^2) < reltol

@time x = lyapd(ar,Qr)
@test norm(ar*x*ar'-x+Qr)/norm(x)/max(1.,norm(ar)^2) < reltol

@time x = lyapd(ar',Qr)
@test norm(ar'*x*ar-x+Qr)/norm(x)/max(1.,norm(ar)^2) < reltol

@time x = lyapd(ac,Qr);
@test norm(ac*x*ac'-x+Qr)/norm(x)/max(1.,norm(ac)^2)  < reltol

@time x = lyapd(ac',Qr);
@test norm(ac'*x*ac-x+Qr)/norm(x)/max(1.,norm(ac)^2)  < reltol

@time x = lyapd(ar,qc)
@test norm(ar*x*ar'-x+qc)/norm(x)/max(1.,norm(ar)^2) < reltol

@time x = lyapd(ar',qc)
@test norm(ar'*x*ar-x+qc)/norm(x)/max(1.,norm(ar)^2) < reltol
end
end

@testset "Discrete generalized Lyapunov equations" begin

reltol = eps(float(100))
ac = -2+im; ec = 4+im; b = 2; @time x = lyapd(ac,ec,b)
@test abs(ac*x*ac'-ec*x*ec'+b) < reltol

for Ty in (Float64, Float32)

ar = rand(Ty,n,n)
ac = rand(Ty,n,n)+im*rand(Ty,n,n)
er = rand(Ty,n,n)
ec = er+im*rand(Ty,n,n)

c = rand(Ty,n,n)+im*rand(Ty,n,n)
qc = c'*c
Qr = real(qc)
Ty == Float64 ? reltol = eps(float(100)) : reltol = eps(100*n*one(Ty))

@time x = lyapd(ac,ec,qc);
@test norm(ac*x*ac'-ec*x*ec'+qc)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

β = 3
@time x = lyapd(ac,β*I,qc);
@test norm(ac*x*ac'-β*x*β'+qc)/norm(x)/norm(ac) < reltol

β = (1+im);
@time x = lyapd(ac,β,qc);
@test norm(ac*x*ac'-β*x*β'+qc)/norm(x)/norm(ac) < reltol

α = 2+3im
@time x = lyapd(α*I,ec,qc);
@test norm(α*x*α'-ec*x*ec'+qc)/norm(x)/norm(ec) < reltol

α = 2+3im
@time x = lyapd(α,ec,qc);
@test norm(α*x*α'-ec*x*ec'+qc)/norm(x)/norm(ec) < reltol

try
   @time x = lyapd(0*ac,0*ec,qc);
   @test false
catch
   @test true
end

@test norm(α*x*α'-ec*x*ec'+qc)/norm(x)/norm(ec) < reltol

α = 2+3im; β = (1+im);
@time x = lyapd(α*I,β*I,qc);
@test norm(α*x*α'-β*x*β'+qc)/norm(x)/norm(ac)/norm(ec) < reltol

try
   α = 1; β = 1; #SingularException
   @time x = lyapd(α,β,qc);
   @test false
catch
   @test true
end

α = 2+3im; β = (1+im);
@time x = lyapd(α,β,qc);
@test norm(α*x*α'-β*x*β'+qc)/norm(x)/norm(ac)/norm(ec) < reltol

@time x = lyapd(ac',ec',qc);
@test norm(ac'*x*ac-ec'*x*ec+qc)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

β = (1+im);
@time x = lyapd(ac',β*I,qc);
@test norm(ac'*x*ac-β*x*β'+qc)/norm(x)/norm(ac) < reltol

α = 2+3im
@time x = lyapd(α*I,ec',qc);
@test norm(α*x*α'-ec'*x*ec+qc)/norm(x)/norm(ec) < reltol

@time x = lyapd(ac',ec,qc);
@test norm(ac'*x*ac-ec*x*ec'+qc)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time x = lyapd(ac,ec',qc);
@test norm(ac*x*ac'-ec'*x*ec+qc)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time x = lyapd(ar,er,Qr);
@test norm(ar*x*ar'-er*x*er'+Qr)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar',er',Qr);
@test norm(ar'*x*ar-er'*x*er+Qr)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar',er,Qr);
@test norm(ar'*x*ar-er*x*er'+Qr)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar,er',Qr);
@test norm(ar*x*ar'-er'*x*er+Qr)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ac,ec,Qr);
@test norm(ac*x*ac'-ec*x*ec'+Qr)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time x = lyapd(ac',ec',Qr);
@test norm(ac'*x*ac-ec'*x*ec+Qr)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time x = lyapd(ac',ec,Qr);
@test norm(ac'*x*ac-ec*x*ec'+Qr)/norm(x)/max(norm(ac)^2,norm(ec)^2) < reltol

@time x = lyapd(ac,ec',Qr);
@test norm(ac*x*ac'-ec'*x*ec+Qr)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar,er,qc);
@test norm(ar*x*ar'-er*x*er'+qc)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar',er',qc);
@test norm(ar'*x*ar-er'*x*er+qc)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar',er,qc);
@test norm(ar'*x*ar-er*x*er'+qc)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol

@time x = lyapd(ar,er',qc);
@test norm(ar*x*ar'-er'*x*er+qc)/norm(x)/max(norm(ar)^2,norm(er)^2) < reltol
end
end

@testset "Discrete Lyapunov equations - Schur form" begin

for Ty in (Float64, Float32)

ar = rand(Ty,n,n)
ac = rand(Ty,n,n)+im*rand(Ty,n,n)
er = rand(Ty,n,n)
ec = er+im*rand(Ty,n,n)
as, es = schur(ar,er)
acs, ecs = schur(ac,ec)

c = rand(Ty,n,n)+im*rand(Ty,n,n)
qc = c'*c
Qr = real(qc)
Ty == Float64 ? reltol = eps(float(100)) : reltol = eps(100*n*one(Ty))

x = copy(Qr)
@time lyapds!(as,x);
@test norm(as*x*as'+Qr-x)/norm(x)/max(norm(as)^2,1.) < reltol

x = copy(qc)
@time lyapds!(as,x);
@test norm(as*x*as'+qc-x)/norm(x)/max(norm(as)^2,1.) < reltol

x = copy(Qr)
@time lyapds!(as,x,adj=true);
@test norm(as'*x*as+Qr-x)/norm(x)/max(norm(as)^2,1.) < reltol

x = copy(qc)
@time lyapds!(as,x,adj=true);
@test norm(as'*x*as+qc-x)/norm(x)/max(norm(as)^2,1.) < reltol

x = copy(qc)
@time lyapds!(acs,ecs,x);
@test norm(acs*x*acs'+qc-ecs*x*ecs')/norm(x)/max(norm(acs)^2,norm(ecs)^2) < reltol

x = copy(qc)
@time lyapds!(acs,I,x);
@test norm(acs*x*acs'+qc-x)/norm(x)/max(norm(acs)^2,norm(ecs)^2) < reltol

x = copy(qc)
@time lyapds!(acs,ecs,x,adj=true);
@test norm(acs'*x*acs+qc-ecs'*x*ecs)/norm(x)/max(norm(acs)^2,norm(ecs)^2) < reltol

x = copy(Qr)
@time lyapds!(as,es,x);
@test norm(as*x*as'+Qr-es*x*es')/norm(x)/max(norm(as)^2,norm(es)^2) < reltol

x = copy(Qr)
@time lyapds!(as,es,x,adj=true);
@test norm(as'*x*as+Qr-es'*x*es)/norm(x)/max(norm(as)^2,norm(es)^2) < reltol

x = copy(qc)
@time lyapds!(as,es,x);
@test norm(as*x*as'+qc-es*x*es')/norm(x)/max(norm(as)^2,norm(es)^2) < reltol

x = copy(qc)
@time lyapds!(as,es,x,adj=true);
@test norm(as'*x*as+qc-es'*x*es)/norm(x)/max(norm(as)^2,norm(es)^2) < reltol

end
end

end

end


xr = 56:76
yr = -162:-134
# xr = 20:30
# yr = -10:-5

function test(sx, sy, n)
  y_end = (2 * sy - n + 1) * n / 2
  if (n <= sx)
      x_end = (2 * sx - n + 1) * n / 2
  else
      x_end = (sx + 1) * sx / 2
  end
  if ((x_end in xr) && (y_end in yr))
    return true
  end
  return false
end

# Coarse but safe ranges for x/y velocities
max_sx = maximum(xr)
max_sy = max(abs(minimum(yr)), abs(maximum(yr)))
max_n = max_sy * 2 + 1

answers = 0
for sy in (-max_sy):max_sy
  for sx in 0:max_sx
    for n in 0:max_n
      if (test(sx, sy, n))
        # println(sx, ",", sy)
        global answers = answers + 1
        break
      end
    end
  end
end

println("answers found ", answers)


using DelaySSAToolkit, Catalyst

rn = @reaction_network begin
    ρ, S+I --> E+I
    r, I --> R
end ρ r

jumpsys = convert(JumpSystem, rn, combinatoric_ratelaws=false)
states(jumpsys)

u0 = [999,1,0,0]
de_chan0 = [[]]
tf = 400.
tspan = (0,tf)

ps = [1e-4, 1e-2]
τ = 20.
dprob = DiscreteProblem(jumpsys,u0,tspan, ps)
delay_trigger_affect! = function (integrator, rng)
    append!(integrator.de_chan[1], τ)
end
delay_trigger = Dict(1=>delay_trigger_affect!)
# delay_trigger = Dict(1=>[1=>τ])
delay_complete = Dict(1=>[2=>1, 3=>-1])
delay_interrupt = Dict()
# algo = DelayDirect()
# algo = DelayDirectCR()
algo = DelayRejection()
delayjumpset = DelayJumpSet(delay_trigger, delay_complete, delay_interrupt)
jprob = DelayJumpProblem(jumpsys, dprob, algo, delayjumpset, de_chan0, save_positions=(true,true))
@time sol = solve(jprob, SSAStepper(), seed = 1234)
using Plots; theme(:vibrant)
fig = plot(sol, label = ["S" "I" "E" "R"], linewidth = 3, legend = :top, ylabel = "# of individuals", xlabel = "Time", fmt=:png)
savefig(fig,"docs/src/assets/seir.png")