"Heteroscedastic PPCA Technique (HePPCAT) module. Provides probabilistic PCA for data with samples of heterogeneous quality (heteroscedastic noise)."
module HePPCAT

# Imports
using Base: IdentityUnitRange   # todo: swap with IdentityRanges.IdentityRange if working (https://github.com/JuliaArrays/IdentityRanges.jl/issues/12)
using IdentityRanges: IdentityRange
using IntervalArithmetic: interval, mid
using IntervalRootFinding: Newton, roots
using LinearAlgebra: Diagonal, I, opnorm, norm, qr, svd, /
using Logging: @warn
import PolynomialRoots
using Roots: find_zero
using Statistics: mean

# Exports
export HePPCATModel, heppcat, loglikelihood

# More convenient form of argmax not yet available
# + findmax from https://github.com/cmcaine/julia/blob/argmax-2-arg-harder/base/reduce.jl#L704-L705
# + argmax from https://github.com/cmcaine/julia/blob/argmax-2-arg-harder/base/reduce.jl#L830
# + part of pull request https://github.com/JuliaLang/julia/pull/35316
_findmax(f, domain) = mapfoldl(x -> (f(x), x), _rf_findmax, domain)
_rf_findmax((fm, m), (fx, x)) = isless(fm, fx) ? (fx, x) : (fm, m)
_argmax(f, domain) = _findmax(f, domain)[2]

include("types.jl")
include("extern.jl")
include("updatev.jl")
include("updateF.jl")
include("updateU.jl")
include("updatelambda.jl")

end


function mehrmann(L,z;maxiter=10,tol=0., relax=1., x0=[],v=[], output=true)
  if output
      println("Launching Mehrmann...")
  end
  if x0==[]
    x0=ones(ComplexF64,size(L(0))[1])
  end
  if v==[]
    v=ones(ComplexF64,size(L(0))[1])
  end
  #normalize
  x0/=v'*x0

  active=L.active #TODO: better integrate activity into householder

  z0=complex(Inf)
  n=0
  flag=1
  try
    while abs(z-z0)>tol && n<maxiter
      if output;  println(n,"\t\t",abs(z-z0),"\t", z );flush(stdout); end
      z0=z
      u=L(z,0)\(L(z,1)*x0)
      #println(size(v)," ",size(x0)," ",size(u))
      z=z0-(v'*x0)/(v'*u)
      x0=u/(v'*u)
      #println("###",v'x0)
      n+=1
    end
  catch excp
    if output
        println("Error occured:")
        println(excp)
        println(typeof(excp))
        println("...aborted Mehrmann!")
    end

    flag=-2
  end

  #convergence checks
  if flag==1
    if output;  println(n,"\t\t",abs(z-z0),"\t", z );flush(stdout); end
    if n>=maxiter
      flag=-1
      if output; println("Warning: Maximum number of iterations has been reached!");end
    elseif abs(z-z0)<=tol
      flag=1
      if output; println("Solution has converged!"); end
    else
      if output; println("Warning: This should not be possible....\n If you can read this contact GAM!");end
      flag=-3
      println(z)
    end
  end

  # if flag==1
  #   #compute left eigenvector
  #   #TODO: degeneracy
  #   lam,y = Arpack.eigs(L(z)',nev=1, sigma = 0,v0=x0)
  #   #normalize
  #   println("size:::$size(y)")
  #   x0./sqrt(x0'*x0)
  #   y./=conj(y'*L(z,1)*x0)
  # else
  #   y=[]
  # end
  y=[]

  return  Solution(L.params,x0,y,L.eigval), n , flag

end
##
function lancaster(L,z;maxiter=10,tol=0., relax=1., x0=[],y0=[], output=true)
  if output
      println("Launching Lancaster...")
  end
  if x0==[]
    x0=ones(ComplexF64,size(L(0))[1])
  end
  if y0==[]
    y0=ones(ComplexF64,size(L(0))[1])
  end

  z0=complex(Inf)
  n=0
  flag=1
  try
    while abs(z-z0)>tol && n<maxiter
            if output;  println(n,"\t\t",abs(z-z0),"\t", z );flush(stdout); end
      z0=z
      #solve
      ξ=L(z)\x0
      η=L(z)'\y0
      z=z0-(η'*L(z,0)*ξ)/(η'*L(z,1)*ξ)
      n+=1
    end
  catch excp
    if output
        println("Error occured:")
        println(excp)
        println(typeof(excp))
        println("...aborted Lancaster!")
    end

    flag=-2
  end

  #convergence checks
  if flag==1
    if output;  println(n,"\t\t",abs(z-z0),"\t", z );flush(stdout); end
    if n>=maxiter
      flag=-1
      if output; println("Warning: Maximum number of iterations has been reached!");end
    elseif abs(z-z0)<=tol
      flag=1
      if output; println("Solution has converged!"); end
    else
      if output; println("Warning: This should not be possible....\n If you can read this contact GAM!");end
      flag=-3
      println(z)
    end
  end

  return  Solution(L.params,zeros(ComplexF64,length(x0)),[],L.eigval), n , flag
end

#%%
import LinearAlgebra
function juniper(L,z;maxiter=10,tol=0.,relax=1.,output=true)
    if output
        println("Launching Juniper...")
    end
  z0=complex(Inf)
  n=0
  flag=1
  try
    while abs(z-z0)>tol && n<maxiter
      if output;  println(n,"\t\t",abs(z-z0),"\t", z ); end
      z0=z
      #dz=-1/LinearAlgebra.tr(Array(L(z))\Array(L(z,1))) #TODO: better implementation for sparse types
      #QR=SparseArrays.qr(L(z))
      LU=SparseArrays.lu(L(z),check=false)
      L1=L(z,1)
      dz=0.0+0.0im
      for idx=1:size(LU)[1]
        dz+=(LU\Array(L1[:,idx]))[idx]
      end
      dz=-1/dz
      z=z0+relax*dz
      n+=1
    end
  catch excp
    if output
        println("Error occured:")
        println(excp)
        println(typeof(excp))
        println("...aborted Juniper!")
    end
    flag=-2
  end

  #convergence checks
  if flag==1
    if output;  println(n,"\t\t",abs(z-z0),"\t", z );flush(stdout); end
    if n>=maxiter
      flag=-1
      if output; println("Warning: Maximum number of iterations has been reached!");end
    elseif abs(z-z0)<=tol
      flag=1
      if output; println("Solution has converged!"); end
    else
      if output; println("Warning: This should not be possible....\n If you can read this contact GAM!");end
      flag=-3
      println(z)
    end
  end

    return Solution(L.params,[],[],L.eigval), n, flag
end

##

function guettel(L,z;maxiter=10,tol=0., relax=1., x0=[],y0=[], output=true)
  if output
      println("Launching Guettel...")
  end
  if x0==[]
    x0=zeros(ComplexF64,size(L(0))[1])
    x0[1]=1
  end
  if y0==[]
    y0=zeros(ComplexF64,size(L(0))[1])
    y0[1]=1
  end
  #normalize
  x0/=sqrt(x0'*x0)
  y0/=sqrt(y0'*y0)
  z0=complex(Inf)
  n=0
  flag=1
  try
    while abs(z-z0)>tol && n<maxiter
      if output;  println(n,"\t\t",abs(z-z0),"\t", z ); end
      z0=z
      LU=SparseArrays.lu(L(z),check=false)
      x0=LU\(L(z,1)*x0)
      y0=LU'\(L(z,1)'*y0)
      #normalize
      x0/=sqrt(x0'*x0)
      y0/=sqrt(y0'*y0)
      #inner iteration
      idx=0
      z00=complex(Inf)
      while abs(z-z00)>tol && idx<10
        z00=z
        z=z-(y0'*L(z)*x0)/(y0'*L(z,1)*x0)
        idx+=1
      end
      n+=1
    end
  catch excp
    if output
        println("Error occured:")
        println(excp)
        println(typeof(excp))
        println("...aborted Guettel!")
    end
    flag=-2
  end
  #convergence checks
  if flag==1
    if output;  println(n,"\t\t",abs(z-z0),"\t", z );flush(stdout); end
    if n>=maxiter
      flag=-1
      if output; println("Warning: Maximum number of iterations has been reached!");end
    elseif abs(z-z0)<=tol
      flag=1
      if output; println("Solution has converged!"); end
    elseif isnan(z)
        flag=-5
        if output; println("Warning: computer arithmetics problem. Eigenvalue is NaN"); end
    else
      if output; println("Warning: This should not be possible....\n If you can read this contact GAM!");end
      flag=-3
      println(z)
    end
  end
      return Solution(L.params,x0,y0,L.eigval), n, flag
end

## count poles and zeros

function count_eigvals(L,Γ,N;output=false)
  function integrand(z)
    z,w=z
    LU=SparseArrays.lu(L(z),check=false)
    L1=L(z,1)
    dz=0.0+0.0im
    for idx=1:size(LU)[1]
      dz+=(LU\Array(L1[:,idx]))[idx]
    end
    return dz*w
  end

  return gauss(0.0+0.0im,integrand,Γ,N,output)/2/pi/1.0im
end


# Julia native arrays are ones based
x = [1,2,3];
println( "x[1] = ", x[1] );

# Can use 2n and skip "*" in some cases
a=7;
println( "2a = ", 2a );

# Can use unicode, including for some operators
b=4
println( "√b = ",√b);

# Not every unicode as operator works
# in second println cᵈ is seen as a variable
c=2
d=2
println( "cᵈ = ",c^d);
#E println( "cᵈ = ",cᵈ);

# Dot used for element wise operations
# Otherwise interpreted as whole object function
println( "x.^3 = ",x.^3);
#E println( "x^3 = ",x^3);

#### ERROR: LoadError: MethodError:  .....
#    is a very common message from Julia runtime error

# Assignment copies values for scalars, pointers for arrays
a1=3
a2=a1
a1=4
println("a1, a2 = ",a1,", ", a2);
a1=[3]
a2=a1
a1[1]=4
println("a1, a2 = ",a1,", ", a2);

# Dot is used to "broadcast" arrays in assingment (which copies values)
a1=[3]
a2.=a1
a1[1]=4
println("a1, a2 = ",a1,", ", a2);


struct DiscreteMeasure{X<:AbstractVector,P<:AbstractVector}
    xs::X
    ps::P

    function DiscreteMeasure{X,P}(xs::X, ps::P) where {X,P}
        length(xs) == length(ps) ||
            error("length of support `xs` and probabilities `ps` must be equal")
        return new{X,P}(xs, ps)
    end
end

"""
    DiscreteMeasure(xs::AbstractVector, ps::AbstractVector)

Construct a discrete measure with support `xs` and corresponding weights `ps`.
"""
function DiscreteMeasure(xs::AbstractVector, ps::AbstractVector)
    return DiscreteMeasure{typeof(xs),typeof(ps)}(xs, ps)
end

Base.length(μ::DiscreteMeasure) = length(μ.xs)

# initial gradient step (unregularized subgradient)
function gradient_step(c, τ, ν::DiscreteMeasure, x, ::Nothing)
    tmp = @. c((x,), ν.xs)
    i = argmin(tmp)
    z = τ .* ν.ps
    z[i] -= τ
    return z
end

# initial gradient step (regularized gradient)
function gradient_step(c, τ, ν::DiscreteMeasure, x, ε::Real)
    tmp = @. -c((x,), ν.xs) / ε
    LogExpFunctions.softmax!(tmp)
    z = @. τ * (ν.ps - tmp)
    return z
end

# gradient step (unregularized subgradient)
function gradient_step!(
    c,
    z::AbstractVector,
    τ,
    v::AbstractVector,
    ν::DiscreteMeasure,
    x,
    tmp::AbstractVector,
    ::Nothing;
    reset::Bool=false,
)
    @. tmp = c((x,), ν.xs) - v
    if reset
        @. z = τ * ν.ps
    else
        @. z += τ * ν.ps
    end
    z[argmin(tmp)] -= τ
    return z
end

# gradient step (regularized gradient)
function gradient_step!(
    c,
    z::AbstractVector,
    τ,
    v::AbstractVector,
    ν::DiscreteMeasure,
    x,
    tmp::AbstractVector,
    ε::Real;
    reset::Bool=false,
)
    @. tmp = (v - c((x,), ν.xs)) / ε
    LogExpFunctions.softmax!(tmp)
    if reset
        @. z = τ * (ν.ps - tmp)
    else
        @. z += τ * (ν.ps - tmp)
    end
    return z
end

@doc raw"""
    ctransform(c, v, x, ν, ε)

Compute the c-transform
```math
v^{c,ε}(x) = \begin{cases}
- ε \log\bigg(\int \exp{\Big(\frac{v_y - c(x, y)}{ε}\Big)} \, ν(\mathrm{d}y)\bigg) & \text{if } ε > 0,\\
\min_{y} c(x, y) - v_y & \text{otherwise}.
\end{cases}
```
"""
function ctransform(c, v::AbstractVector, x, ν::DiscreteMeasure, ::Nothing)
    return minimum(c(x, yᵢ) - vᵢ for (vᵢ, yᵢ) in zip(v, ν.xs))
end
function ctransform(c, v::AbstractVector, x, ν::DiscreteMeasure, ε::Real)
    t = LogExpFunctions.logsumexp(
        (vᵢ - c(x, yᵢ)) / ε + log(νᵢ) for (vᵢ, yᵢ, νᵢ) in zip(v, ν.xs, ν.ps)
    )
    return -ε * (t + 1)
end


@testset "Vortex Sheets" begin
    @testset "Sheet construction" begin
        N = 100
        zs = rand(ComplexF64, N)
        Γs = accumulate(+, rand(N))
        δ = rand()

        sheet = Vortex.Sheet(zs, Γs, δ)

        @test circulation(sheet) ≈ circulation(sheet.blobs)
    end

    @testset "Sheet induced velocities" begin
        N = 100
        zs = rand(ComplexF64, N)
        Γs = accumulate(+, rand(N))
        δ = rand()
        sheet = Vortex.Sheet(zs, Γs, δ)

        ws = allocate_velocity(sheet)
        wb = allocate_velocity(sheet.blobs)

        self_induce_velocity!(ws, sheet, 0)
        for (i, t) in enumerate(sheet.blobs), s in sheet.blobs
            wb[i] += induce_velocity(t, s, 0)
        end
        @test wb ≈ ws
        @test wb ≈ induce_velocity(sheet, sheet, 0)

        sheet₊ = Vortex.Sheet(Vector{Vortex.Blob}(undef, 10), rand(10), 0.0)
        Δt = 1e-2
        advect!(sheet₊, sheet, ws, Δt)
        advect!(sheet, sheet, ws, Δt)

        @test sheet₊.Ss ==  sheet.Ss
        @test !(sheet₊.Ss === sheet.Ss)

        @test sheet₊.blobs == sheet.blobs
        @test !(sheet₊.blobs === sheet.blobs)

        @test sheet₊.δ  ==  sheet.δ
    end

    @testset "Sheet surgery" begin
        N = 100
        zs = rand(ComplexF64, N)
        Γs = accumulate(+, rand(N))
        δ = 0.05
        sheet = Vortex.Sheet(zs, Γs, δ)

        truncate_n = 50
        ΣΓ = circulation(sheet)
        ΔΓ₀ = Γs[truncate_n] - Γs[1]
        ΔΓ  = Sheets.truncate!(sheet, truncate_n)
        @test ΔΓ₀ ≈ ΔΓ

        @test Sheets.compute_trapezoidal_weights(sheet.Ss) ≈ circulation.(sheet.blobs)
        @test ΔΓ + circulation(sheet) ≈ ΣΓ

        Γ = circulation(sheet)
        sheet₊ = Sheets.split!(sheet, rand(3:length(sheet)-3))
        @test Γ ≈ sum(circulation, (sheet, sheet₊))

        @test Sheets.compute_trapezoidal_weights(sheet.Ss) ≈ circulation.(sheet.blobs)
        @test Sheets.compute_trapezoidal_weights(sheet₊.Ss) ≈ circulation.(sheet₊.blobs)

        N = 200
        θ = range(π, 0, length = N)
        zs = complex.(cos.(θ))
        Γs = sin.(θ)
        Sheets.redistribute_points!(sheet, zs, Γs)
        fill!(Γs, 0.0)
        @test sheet.Ss != Γs
        @test Sheets.compute_trapezoidal_weights(sheet.Ss) ≈ getfield.(sheet.blobs, :S)
        @test Elements.position.(sheet.blobs) == getfield.(sheet.blobs, :z)
        @test sheet.zs == Elements.position.(sheet.blobs)

        Sheets.remesh!(sheet, 0.005)
        @test length(sheet) == 400
        @test sheet.zs ≈ range(-1, 1, length = 400)
        @test norm(sheet.Ss .- sqrt.(1 .- range(-1,1, length=400).^2)) ≤ 1e-3

        θ = range(π, 0, length = N)
        zs = complex.(cos.(θ))
        Γs = sin.(θ)
        Sheets.redistribute_points!(sheet, zs, Γs)
        sheet₁ = deepcopy(sheet)
        sheet₂ = deepcopy(sheet)

        @test_logs (:warn, "Cannot remesh, sheet length smaller than nominal spacing") (:warn, "Filter not applied, total sheet length smaller than nominal spacing") Sheets.filter!(sheet₁, 3.0, 0.03)
        @test sheet₁.zs == sheet₂.zs
        @test sheet₁.Ss == sheet₂.Ss
        @test sheet₁.blobs == sheet₂.blobs

        @test_logs (:warn, "Cannot remesh, sheet length smaller than nominal spacing") Sheets.remesh!(sheet₂, 3.0)
        @test sheet₁.zs == sheet₂.zs
        @test sheet₁.Ss == sheet₂.Ss
        @test sheet₁.blobs == sheet₂.blobs

        Sheets.remesh!(sheet₂, 0.01)
        Sheets.filter_position!(sheet₂, 0.03)

        property = copy(sheet₁.Ss)
        Sheets.filter!(sheet₁, 0.01, 0.03, (property,))

        @test sheet₁.zs == sheet₂.zs
        @test sheet₁.Ss == sheet₂.Ss
        @test sheet₁.Ss == property
    end
end


"""
Functions that help the automation of tests
"""

using Test
using Catlab.CategoricalAlgebra, Catlab.Present
using Catlab.CategoricalAlgebra.CSetDataStructures: struct_acset

# Functions to make schemas programmatically
############################################
xs(x::Int)::Symbol = Symbol("x$x")
xs(xx::AbstractVector{Int})::Vector{Symbol} = [Symbol("x$x") for x in xx]
es(x::Int)::Symbol = Symbol("e$x")
es(xx::AbstractVector{Int})::Vector{Symbol} = [Symbol("e$x") for x in xx]

"""
Create a CSet type specified by a graph
Vertices are x₁,x₂,..., edges are e₁, e₂,...
all edges are indexed
"""
function graph_to_cset(grph::StructACSet, name::Symbol)::Pair{Presentation,StructACSet}
  pres = Presentation(FreeSchema)
  xobs = [Ob(FreeSchema,xs(i)) for i in 1:nv(grph)]
  for x in xobs
    add_generator!(pres, x)
  end
  for (i,(src, tgt)) in enumerate(zip(grph[:src], grph[:tgt]))
    add_generator!(pres, Hom(es(i), xobs[src], xobs[tgt]))
  end

  expr = struct_acset(name, StructACSet, pres, index=es(1:ne(grph)))
  eval(expr)
  csettype = eval(name)
  return pres => Base.invokelatest(csettype)
end

  """Create n copies of a CSet based on a graph schema"""
function init_graphs(name::Symbol, schema::StructACSet, consts::Vector{Int},
             n::Int=2)::Pair{Presentation, Vector{StructACSet}}
  pres, cset = graph_to_cset(schema, name)
  for (i, con) in enumerate(consts)
    add_parts!(cset, Symbol("x$i"), con)
  end
  return pres => [deepcopy(cset) for _ in 1:n]
end

# Testing nauty.jl and CSetAutomorphisms on the same input
##########################################################

"""Confirm canonical hash tracks with whether two ACSets are iso"""
function test_iso(a::StructACSet,b::StructACSet, pres::Union{Nothing,Presentation}=nothing)::Test.Pass
  @test a != b  # confirm they're not literally equal
  eq = is_isomorphic(a,b)
  tst = (x,y) -> eq ? x==y : x!=y # hashes should be equal iff they're iso
  if !isnothing(pres)
    @test tst(canonical_hash(a,pres=pres), canonical_hash(b; pres=pres))
  end
  @test tst(canonical_hash(a), canonical_hash(b))
end

# Particular Schemas
####################
@present TheoryDecGraph(FreeSchema) begin
  E::Ob
  V::Ob
  src::Hom(E,V)
  tgt::Hom(E,V)

  X::AttrType
  dec::Attr(E,X)
end

@acset_type Labeled(TheoryDecGraph)


using Test
using TestSetExtensions
using LinearAlgebra
using Qaintessent
using Random
using Qaintessent.MaxKColSubgraphQAOA

graphs = [
    Graph(3, [(1,2), (2,3)]),
    Graph(5, [(1,2), (2,3), (3,4), (4,5), (5,1), (1,3)]),
    Graph(4, [(1,2), (1,3), (1,4), (2,3), (2,4), (3,4)])
]

# Utility function to count the properly colored edges in a graph coloring (i.e. endpoints have different colors)
function properly_colored_edges(graph::Graph, coloring::Vector{Int})
    length(coloring) == graph.n || throw(ArgumentError("Length of coloring must equal the number of vertices."))
    return count(coloring[a] != coloring[b] for (a, b) ∈ graph.edges)
end

# Create a state ψ that corresponds to a given coloring
function ψ_from_coloring(n::Int, κ::Int, colors::Vector{Int})::Vector{ComplexF64}
    (n > 0 && κ > 0) || throw(DomainError("Parameters n and κ must be positive integers."))
    colors ⊆ 1:κ || throw(ArgumentError("Parameter `colors` may only contain colors in the range 1:$(κ)."))

    # Create ψ with |1> entries in the indices corresponding to the given colors, |0> elsewhere
    ψ = kron((color == colors[vertex] ? [0.0im, 1] : [1, 0.0im] for vertex ∈ 1:n for color ∈ 1:κ)...)
    ψ
end

# Utility function to compute the probabilities of the outcomes represented by a wavefunction ψ, sorted by descending probability
function wavefunction_distribution(ψ::Vector{ComplexF64}; as_bitstrings::Bool = true,
        include_zero = false)::Union{Vector{Tuple{Int, Float64}}, Vector{Tuple{Vector{Int}, Float64}}}
    distribution = [(i-1, abs(amplitude)^2) for (i, amplitude) ∈ enumerate(ψ) if abs(amplitude) > 0 || include_zero]
    if as_bitstrings
        N = Int(log2(length(ψ)))
        distribution = [(digits(i, base=2, pad=N) |> reverse, p) for (i, p) ∈ distribution]
    end

    return sort(distribution, by=(t -> -t[2]))
end

@testset ExtendedTestSet "max-k-col. subgraph QAOA Hamiltonians and gate matrices" begin
    @testset ExtendedTestSet "max-k-col. subgraph QAOA Hamiltonians" begin
        # Test that the phase separation Hamiltonian used in MaxKColSubgraphPhaseSeparationGate
        # implements the objective function correctly.
        @testset "Hamiltonian MaxKColSubgraphPhaseSeparationGate" begin
            κs = [4, 2, 3]

            hamiltonians = max_k_col_subgraph_phase_separation_hamiltonian.(graphs, κs)
            colorings = [rand(1:κ, graph.n) for (κ, graph) ∈ zip(κs, graphs)]

            for (H, graph, κ, coloring) ∈ zip(hamiltonians, graphs, κs, colorings)
                f_val = properly_colored_edges(graph, coloring)
                ψ_coloring = ψ_from_coloring(graph.n, κ, coloring)
                scaling = κ * length(graph.edges) - 4 * f_val # ψ_coloring should be scaled by this factor
                @test H * ψ_coloring ≈ ψ_coloring * scaling
            end
        end

        # Utility functions to perform circular shift on Vectors (that represent bitstrings)
        left_circ_shift(l::Vector) = [l[2:length(l)]; l[1:1]]
        right_circ_shift(l::Vector) = [l[length(l):length(l)]; l[1:length(l)-1]]

        # Shorthand for repeated function application
        # Taken from https://stackoverflow.com/questions/39895672/apply-function-repeatedly-a-specific-number-of-times
        (repeatApply)(f::Function, i::Int) = i==1 ? f : x->(repeatApply(f, i-1))(f(x))

        # Test that the Hamiltonian which the RNearbyValuesMixerGate is based on
        # maps coloring states to a superposition of their neighboring states
        @testset "Hamiltonian RNearbyValuesMixerGate" begin
            ds = [2, 5, 9, 10] # different numbers of colors
            rs = [1, 2, 5]
            nums_samples = [2, 5, 5, 5]

            for (d, num_samples) ∈ zip(ds, nums_samples)
                for r ∈ rs[rs .< d] # only test valid rs
                    # compute the Hamiltonian
                    gate = RNearbyValuesMixerGate(0., r, d)
                    hamiltonian = r_nearby_values_hamiltonian_onehot(gate)

                    # for each color (i.e. |1000..00>, |0100..00>, ..., |0000..01>):
                    for color ∈ randperm(d)[1:num_samples]
                        # compute superposition state bitstrings after applying Hamiltonian
                        ψ = ψ_from_coloring(1, d, [color])
                        ψ_bits = wavefunction_distribution(ψ)[1][1]
                        out_vec = hamiltonian * ψ # not a wavefunction (we apply a Hamiltonian, not a gate)
                        superposition_states = (wavefunction_distribution(out_vec)
                            |> (distr -> filter(d -> !(d[2] ≈ 0), distr))
                            .|> (d -> d[1]))

                        # compute bitstrings of nearby colors by shifting
                        near_values = [
                            [repeatApply(left_circ_shift, k)(ψ_bits) for k ∈ 1:r];
                            [repeatApply(right_circ_shift, k)(ψ_bits) for k ∈ 1:r]
                        ] # repeatApply is repeated function application defined above

                        # r-NV Mixer should map state to superposition of its neighbor colors
                        @test Set(superposition_states) == Set(near_values)
                    end
                end
            end
        end
    end

    # Utility function that returns the Hamming weights of all basis states of ψ
    hamming_weights(ψ) = (wavefunction_distribution(ψ)
    |> (distr -> filter(d -> !(d[2] ≈ 0), distr))
    .|> (d -> sum(d[1])))

    @testset ExtendedTestSet "max-k-col. subgraph QAOA mixer Hamming weights" begin
        # Test that the RNearbyValuesMixerGate preserves the Hamming weight of coloring
        # states (i.e. maps states of Hamming weight one to a superposition of states
        # with Hamming weight one)
        @testset "Hamming weight RNearbyValuesMixerGate" begin
            ds = [2, 5, 9, 10] # different numbers of colors
            rs = [1, 2, 5]
            nums_samples = [2, 5, 5, 5]

            for (d, num_samples) ∈ zip(ds, nums_samples)
                for r ∈ rs[rs .< d] # only test valid rs
                    gate = RNearbyValuesMixerGate(rand() * 10 + 0.1, r, d)
                    U = Qaintessent.matrix(gate)

                    for color ∈ randperm(d)[1:num_samples]
                        ψ = ψ_from_coloring(1, d, [color])
                        ψ_out = U * ψ

                        # All basis states should have Hamming weight 1
                        out_hamming_weights = hamming_weights(ψ_out)
                        @test hamming_weights(ψ_out) == ones(Int64, length(out_hamming_weights))
                    end
                end
            end
        end

        # Test that the ParityRingMixerGate preserves the Hamming weight of coloring
        # states (i.e. maps states of Hamming weight one to a superposition of states
        # with Hamming weight one), if d == 2
        @testset "Hamming weight ParityRingMixerGate (d = 2)" begin
            ds = [2] # different numbers of colors

            for d ∈ ds
                gate = ParityRingMixerGate(rand() * 10 + 0.1, d)
                U = Qaintessent.matrix(gate)

                for color ∈ 1:d
                    ψ = ψ_from_coloring(1, d, [color])
                    ψ_out = U * ψ

                    # All basis states should have Hamming weight 1
                    out_hamming_weights = hamming_weights(ψ_out)
                    @test hamming_weights(ψ_out) == ones(Int64, length(out_hamming_weights))
                end
            end
        end

        # Test that the ParityRingMixerGate preserves the Hamming weight of coloring
        # states (i.e. maps states of Hamming weight one to a superposition of states
        # with Hamming weight one), if d ≥ 3
        @testset "Hamming weight ParityRingMixerGate (d ≥ 3)" begin
            ds = [5, 8, 9] # different numbers of colors
            num_samples = 5

            for d ∈ ds
                gate = ParityRingMixerGate(rand() * 10 + 0.1, d)
                U = Qaintessent.matrix(gate)

                for color ∈ randperm(d)[1:num_samples]
                    ψ = ψ_from_coloring(1, d, [color])
                    ψ_out = U * ψ

                    # All basis states should have Hamming weight 1
                    out_hamming_weights = hamming_weights(ψ_out)
                    @test hamming_weights(ψ_out) == ones(Int64, length(out_hamming_weights))
                end
            end
        end

        # Test that the PartitionMixerGate preserves the Hamming weight of coloring
        # states (i.e. maps states of Hamming weight one to a superposition of states
        # with Hamming weight one)
        @testset "Hamming weight PartitionMixerGate" begin
            ds = [5, 8, 9] # different numbers of colors
            num_samples = 5
            nums_partition_layers = [3, 2, 7]

            for (d, num_partition_layers) ∈ zip(ds, nums_partition_layers)
                # create num_partition_layers partition parts, each of
                # random length and with random entries
                partition = [
                    # create h x 2 matrix with random height 2 from a random permutation of the indices...
                    (reshape(randperm(d)[1:(2*rand(1:div(d, 2)))], (:, 2))
                    # ... then map the rows to tuples
                        |> eachrow .|> Tuple)
                    for _ ∈ 1:num_partition_layers
                ]

                gate = PartitionMixerGate(rand() * 10 + 0.1, d, partition)
                U = Qaintessent.matrix(gate)

                for color ∈ randperm(d)[1:num_samples]
                    ψ = ψ_from_coloring(1, d, [color])
                    ψ_out = U * ψ

                    # All basis states should have Hamming weight 1
                    out_hamming_weights = hamming_weights(ψ_out)
                    @test hamming_weights(ψ_out) == ones(Int64, length(out_hamming_weights))
                end
            end
        end
    end

    # Test that the parity ring mixer is a special case of the partition mixer
    @testset "PartitionMixerGate generalizes ParityRingMixerGate" begin
        ds = [2, 5, 9, 10]

        for d ∈ ds
            β = rand() * 10 + 0.1

            # construct the parity mixer
            parity_mixer = ParityRingMixerGate(β, d)

            # construct the equivalent partition mixer
            odd_part = [(a, a+1) for a ∈ 1:2:(d-1)]
            even_part = [(a, (a % d) + 1) for a ∈ 2:2:d]
            last_part = isodd(d) ? [(d, 1)] : nothing
            partition = isodd(d) ? [odd_part, even_part, last_part] : [odd_part, even_part]
            partition_mixer = PartitionMixerGate(β, d, partition)

            @test Qaintessent.matrix(parity_mixer) ≈ Qaintessent.matrix(partition_mixer)
        end
    end

    @testset ExtendedTestSet "max-k-col. subgraph QAOA adjoints" begin
        # Test that `MaxKColSubgraphPhaseSeparationGate` has the correct adjoint.
        @testset "adjoint MaxKColSubgraphPhaseSeparationGate" begin
            γs = rand(length(graphs)) * 2π
            κs = [3, 2, 2]

            gates = MaxKColSubgraphPhaseSeparationGate.(γs, κs, graphs)
            for i ∈ 1:length(gates)
                @test Qaintessent.matrix(gates[i]) * Qaintessent.matrix(Base.adjoint(gates[i])) ≈ I
            end
        end

        # Test that `ParityRingMixerGate` has the correct adjoint.
        @testset "adjoint ParityRingMixerGate" begin
            N = 3 # number of test cases
            βs = rand(N) * 2π
            ds = rand(2:9, N)

            gates = ParityRingMixerGate.(βs, ds)
            for i ∈ 1:length(gates)
                @test Qaintessent.matrix(gates[i]) * Qaintessent.matrix(Base.adjoint(gates[i])) ≈ I
            end
        end

        # Test that `RNearbyValuesMixerGate` has the correct adjoint.
        @testset "adjoint RNearbyValuesMixerGate" begin
            βs = rand(4) * 2π
            ds = [2, 3, 6, 9]
            rs = [1, 2, 1, 5]

            gates = RNearbyValuesMixerGate.(βs, rs, ds)
            for i ∈ 1:length(gates)
                @test Qaintessent.matrix(gates[i]) * Qaintessent.matrix(Base.adjoint(gates[i])) ≈ I
            end
        end

        # Test that `PartitionMixerGate` has the correct adjoint.
        @testset "adjoint PartitionMixerGate" begin
            ds = [2, 3, 6, 9]
            nums_partition_layers = [2, 3, 4, 5]

            for (d, num_partition_layers) ∈ zip(ds, nums_partition_layers)
                # create num_partition_layers partition parts, each of
                # random length and with random entries
                partition = [
                    # create h x 2 matrix with random height 2 from a random permutation of the indices...
                    (reshape(randperm(d)[1:(2*rand(1:div(d, 2)))], (:, 2))
                    # ... then map the rows to tuples
                        |> eachrow .|> Tuple)
                    for _ ∈ 1:num_partition_layers
                ]

                β = rand() * 10 + 0.1
                gate = PartitionMixerGate(β, d, partition)
                @test Qaintessent.matrix(gate) * Qaintessent.matrix(Base.adjoint(gate)) ≈ I
            end
        end
    end

    @testset ExtendedTestSet "max-k-col. subgraph QAOA num wires" begin
        @testset "num wires MaxKColSubgraphPhaseSeparationGate" begin
            γs = rand(length(graphs)) * 2π
            κs = [3, 2, 2]

            gates = MaxKColSubgraphPhaseSeparationGate.(γs, κs, graphs)
            for (gate, graph, κ) ∈ zip(gates, graphs, κs)
                @test Qaintessent.num_wires(gate) == graph.n * κ
            end
        end

        @testset "num wires ParityRingMixerGate" begin
            N = 3 # number of test cases
            βs = rand(N) * 2π
            ds = rand(2:9, N)

            gates = ParityRingMixerGate.(βs, ds)
            for (gate, d) ∈ zip(gates, ds)
                @test Qaintessent.num_wires(gate) == d
            end
        end

        @testset "num wires RNearbyValuesMixerGate" begin
            βs = rand(4) * 2π
            ds = [2, 3, 6, 9]
            rs = [1, 2, 1, 5]

            gates = RNearbyValuesMixerGate.(βs, rs, ds)
            for (gate, d) ∈ zip(gates, ds)
                @test Qaintessent.num_wires(gate) == d
            end
        end

        @testset "num wires PartitionMixerGate" begin
            ds = [2, 3, 6, 9]
            nums_partition_layers = [2, 3, 4, 5]

            for (d, num_partition_layers) ∈ zip(ds, nums_partition_layers)
                # create num_partition_layers partition parts, each of
                # random length and with random entries
                partition = [
                    # create h x 2 matrix with random height 2 from a random permutation of the indices...
                    (reshape(randperm(d)[1:(2*rand(1:div(d, 2)))], (:, 2))
                    # ... then map the rows to tuples
                        |> eachrow .|> Tuple)
                    for _ ∈ 1:num_partition_layers
                ]

                β = rand() * 10 + 0.1
                gate = PartitionMixerGate(β, d, partition)
                @test Qaintessent.num_wires(gate) == d
            end
        end
    end
end


using Bridge, StaticArrays, Bridge.Models
const R = ℝ
Random.seed!(2)
Random.seed!(3)
πH = 2000. # prior

t = 1.0
T = 5.00
n = 50001 # total imputed length
m = 100 # number of segments
M = div(n-1,m)
skippoints = 2
dt = (T-t)/(n-1)
tt = t:dt:T
si = 3.
# 10, 20, 8/3 srand(2)
θ² = 1.
γ = 1.
P = Bridge.Models.Pendulum(θ², γ)
Psm = Bridge.Models.Pendulum(θ², 0.0)
x0 = Models.x0(P)

W = sample(tt, Wiener())
X = SamplePath(tt, zeros(ℝ{2}, length(tt)))
Bridge.solve!(Euler(), X, x0, W, P)
W = sample(tt, Wiener())
Xsm = SamplePath(tt, zeros(ℝ{2}, length(tt)))
Bridge.solve!(Euler(), Xsm, x0, W, Psm)

Xtrue = copy(X)

# Observation scheme and subsample
_pairs(collection) = Base.Generator(=>, keys(collection), values(collection))
SV = ℝ{2}
SM = typeof(one(Bridge.outer(zero(SV))))
if isdefined(:partial)
    if !partial
        L = I
        Σ = SDiagonal(1., 1.)
        lΣ = chol(Σ)'
        RV = SV
        RM = SM

        V = SamplePath(collect(_pairs(Xtrue))[1:M:end])
        map!(y -> L*y + lΣ*randn(RV), V.yy, V.yy)
    else
        L = @SMatrix [1.0 0.0]
        Σ = isdefined(:Σ_) ? Σ_ : 1.00
        @show Σ
        lΣ = chol(Σ)'
        RV = Float64
        RM = Float64 #typeof(one(Bridge.outer(zero(RV))))

        V_ = SamplePath(collect(_pairs(Xtrue))[1:M:end])
        V = SamplePath(V_.tt, map(y -> (L*y)[1] + lΣ*randn(RV), V_.yy))
    end
else
    println("Info:  No observation scheme defined. Set `partial in [true, false]`")
end


# Aiyagari Model Test
# Michele Fornino
#---------------------------------------------------#
#					Aiyagari MODEL 					#
#---------------------------------------------------#

"""
	Aiyagari_params{F<:Real, I<:Integer}

	Composite type containing all parameters of the standard Aiyagari Model.

"""
@with_kw struct Aiyagari_params{F<:Real, I<:Integer} @deftype F

	# Parameters
	σ = 1.0
	w = 1.0
	r = log(1 + 0.01)
	ρ = r
	σz = 0.2
	θz = 0.1
	zd = 1.0

	# Derived Parameters
	γ = 1.0 / σ

	# Grids Parameters
	Na::I = 50
	Nz::I = 20
	zGrid_lower_quantile = 1e-2
	zGrid_upper_quantile = 1 - 1e-2
	aGrid_lower = 0.
	aGrid_upper = 20.

	# Assertions
	@assert σ > 0
	@assert r > 0
	@assert w > 0
	@assert ρ > 0
	@assert θz > 0
	@assert σz > 0
	@assert zd > 0

	# Check values for grids make sense
	@assert Na > 1
	@assert Nz > 0
	@assert zGrid_upper_quantile > zGrid_lower_quantile
	@assert zGrid_lower_quantile > 0.0
	@assert zGrid_upper_quantile < 1.0
	@assert aGrid_upper > aGrid_lower

	# Enforce definitions (prevents accidental creation of object with wrong derived parameters)
	@assert γ == 1 / σ

end

"""
	function HamiltonJacobiBellmanEquations.HJBEquation(params::Aiyagari_params)

Overloaded constructor for input of type Aiyagari_params.

"""
function HamiltonJacobiBellmanEquations.HJBEquation(params::Aiyagari_params)

	# Use UnPack.jl to localize all members of params
	@unpack_Aiyagari_params params

	# Utility Function and derivatives.
	function u(c::T)::T where T<:Real
		if σ == 1.0
			return log(c)
		else
			return (c ^ (1 - γ) - 1) / (1 - γ)
		end
	end

	function u′(c::T)::T where T<:Real
		return c ^ -γ
	end

	function u′⁻¹(x::T)::T where T<:Real
		return x ^ (- 1 / γ)
	end

	# function u(c::T)::T where T<:Real
	# 	return - 1 / γ * exp(-γ * c)
	# end

	# function u′(c::T)::T where T<:Real
	# 	return exp( - γ * c)
	# end

	# function u′⁻¹(x::T)::T where T<:Real
	# 	return - 1 / γ * log(x)
	# end

	# Returns Function
	function flow_returns(state::AbstractVector{T}, c::Union{T, AbstractVector{T}})::T where T<:Real
		return u(c[1])
	end

	# Policy Function in term of marginal value function
	function optimal_controls(state::AbstractVector{T}, ∂V∂a::Union{T, AbstractVector{T}})::T where T<:Real
		return u′⁻¹(∂V∂a[1])
	end

	# State equations
	function state_equations(state::AbstractVector{T}, c::Union{T, AbstractVector{T}})::T where T<:Real
		return w * state[2] + r * state[1] - c[1]
	end

	# State constraint
	function state_constraints(state::AbstractVector{T})::T where T<:Real
		return u′(w * state[2] + r * state[1])
	end

	# Tuple of endogenous state grids
	aGrid = Vector(range(aGrid_lower, aGrid_upper, length = Na))
	state_grids_endog = (aGrid, )

	# Exogenous shock
	exog_shock = ContinuousTimeMarkovChain(params)

	# Allocate hjb object
	return HJBEquation(flow_returns,
					   state_equations,
					   optimal_controls,
					   state_constraints,
					   state_grids_endog,
					   exog_shock,
					   ρ)

end

"""
	function MarkovChains.ContinuousTimeMarkovChain(params::Aiyagari_params)

Overloaded constructor for input of type Aiyagari_params.

"""
function MarkovChains.ContinuousTimeMarkovChain(params::Aiyagari_params)

	@unpack_Aiyagari_params params

	# Construct exog_shock ContinuousTimeMarkovChain object (if problem is stochastic).
	if Nz > 1

		# Obtain discretize diffusion process object as a MarkovChains.ContinuousTimeMarkovChain.
		μ = z -> -θz .* (z .- zd)
		σ = z -> sqrt(2 * θz * σz^2 / zd) .* sqrt.(z)
		dp = ItoDiffusionProcess(μ, σ)

		# Theoretical benchmark for CIR stationary distribution, used for zGrid.
		shape = zd^2 / σz^2
		rate = zd / shape
		Fz_theory = Gamma(shape, rate)
		zGrid_lower = quantile(Fz_theory, zGrid_lower_quantile)
		zGrid_upper = quantile(Fz_theory, zGrid_upper_quantile)

		# Allocate Grid
		zGrid = Vector(range(zGrid_lower, zGrid_upper, length = Nz))

		# Exogenous Shock: ContinuousTimeMarkovChain object
		exog_shock = ContinuousTimeMarkovChain(dp, zGrid)

	elseif Nz == 1

		exog_shock = ContinuousTimeMarkovChain(reshape([0], 1, 1), [zd])

	end

	return exog_shock

end


using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:2] <= 1)

                     @variable(m, -4.847398418942631 <= q <= 3.5461678827819525)

                     add_NL_constraint(m, :(sigmoid(-0.19745953189550702 + -0.698583212677272*sigmoid(-0.2024495957518173 + -0.1501654471281335*sigmoid(-0.06497322351233636 + -0.4876409540129627*$(x[1]) + -0.026006048535915927*$(x[2])) + 0.7965619538518913*sigmoid(-0.4201092923480956 + 0.6085672502473995*$(x[1]) + -0.465188513706837*$(x[2])) + 0.8656395810315112*sigmoid(-0.9600965095983032 + 0.022609326108050443*$(x[1]) + -0.12508836215483532*$(x[2]))) + 0.07925279857906231*sigmoid(-0.4774715210476401 + -0.18966514690787006*sigmoid(-0.06497322351233636 + -0.4876409540129627*$(x[1]) + -0.026006048535915927*$(x[2])) + 0.1687755600264067*sigmoid(-0.4201092923480956 + 0.6085672502473995*$(x[1]) + -0.465188513706837*$(x[2])) + 0.9882336248352352*sigmoid(-0.9600965095983032 + 0.022609326108050443*$(x[1]) + -0.12508836215483532*$(x[2]))) + -0.8658353105986021*sigmoid(-0.5967554189497073 + 0.7632066186026765*sigmoid(-0.06497322351233636 + -0.4876409540129627*$(x[1]) + -0.026006048535915927*$(x[2])) + -0.767873838532454*sigmoid(-0.4201092923480956 + 0.6085672502473995*$(x[1]) + -0.465188513706837*$(x[2])) + 0.02410968604346264*sigmoid(-0.9600965095983032 + 0.022609326108050443*$(x[1]) + -0.12508836215483532*$(x[2])))) + sigmoid(0.9047173820312868 + 0.09375520961098927*sigmoid(-0.2024495957518173 + -0.1501654471281335*sigmoid(-0.06497322351233636 + -0.4876409540129627*$(x[1]) + -0.026006048535915927*$(x[2])) + 0.7965619538518913*sigmoid(-0.4201092923480956 + 0.6085672502473995*$(x[1]) + -0.465188513706837*$(x[2])) + 0.8656395810315112*sigmoid(-0.9600965095983032 + 0.022609326108050443*$(x[1]) + -0.12508836215483532*$(x[2]))) + -0.9282166334997988*sigmoid(-0.4774715210476401 + -0.18966514690787006*sigmoid(-0.06497322351233636 + -0.4876409540129627*$(x[1]) + -0.026006048535915927*$(x[2])) + 0.1687755600264067*sigmoid(-0.4201092923480956 + 0.6085672502473995*$(x[1]) + -0.465188513706837*$(x[2])) + 0.9882336248352352*sigmoid(-0.9600965095983032 + 0.022609326108050443*$(x[1]) + -0.12508836215483532*$(x[2]))) + 0.3928411969784107*sigmoid(-0.5967554189497073 + 0.7632066186026765*sigmoid(-0.06497322351233636 + -0.4876409540129627*$(x[1]) + -0.026006048535915927*$(x[2])) + -0.767873838532454*sigmoid(-0.4201092923480956 + 0.6085672502473995*$(x[1]) + -0.465188513706837*$(x[2])) + 0.02410968604346264*sigmoid(-0.9600965095983032 + 0.022609326108050443*$(x[1]) + -0.12508836215483532*$(x[2])))) + sigmoid(0.16193188153327842 + 0.4964335766578394*sigmoid(-0.2024495957518173 + -0.1501654471281335*sigmoid(-0.06497322351233636 + -0.4876409540129627*$(x[1]) + -0.026006048535915927*$(x[2])) + 0.7965619538518913*sigmoid(-0.4201092923480956 + 0.6085672502473995*$(x[1]) + -0.465188513706837*$(x[2])) + 0.8656395810315112*sigmoid(-0.9600965095983032 + 0.022609326108050443*$(x[1]) + -0.12508836215483532*$(x[2]))) + -0.6165233713312808*sigmoid(-0.4774715210476401 + -0.18966514690787006*sigmoid(-0.06497322351233636 + -0.4876409540129627*$(x[1]) + -0.026006048535915927*$(x[2])) + 0.1687755600264067*sigmoid(-0.4201092923480956 + 0.6085672502473995*$(x[1]) + -0.465188513706837*$(x[2])) + 0.9882336248352352*sigmoid(-0.9600965095983032 + 0.022609326108050443*$(x[1]) + -0.12508836215483532*$(x[2]))) + -0.4911670927879337*sigmoid(-0.5967554189497073 + 0.7632066186026765*sigmoid(-0.06497322351233636 + -0.4876409540129627*$(x[1]) + -0.026006048535915927*$(x[2])) + -0.767873838532454*sigmoid(-0.4201092923480956 + 0.6085672502473995*$(x[1]) + -0.465188513706837*$(x[2])) + 0.02410968604346264*sigmoid(-0.9600965095983032 + 0.022609326108050443*$(x[1]) + -0.12508836215483532*$(x[2])))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m


"""
    EGPpowermix(σ, ξ, κ₁, κ₂, p)

*EGPpowermix* corresponds to the second extended GP model of Naveau et al. (2016), which is a mixture of power laws with G(v) = pv^κ₁ + (1-p)v^κ₂.

It is a five parameters family: κ₁ controls the lower tail behavior, κ₂ modifies the shape of the density in its central part, σ is a scale parameter, and ξ is a shape parameter.

```julia
EGPpowermix(σ, ξ, κ₁, κ₂, p)   # EGP of Naveau et al. (2016) (type 2) with parameters σ, ξ, κ₁, κ₂, p.

params(d)           # Get the parameters, i.e. (σ, ξ, κ₁, κ₂, p)
```

Reference :

* Naveau, P., Huser, R., Ribereau, P., and Hannart, A. (2016), Modeling jointly low, moderate, and heavy rainfall intensities without a threshold selection, Water Resour. Res., 52, 2753– 2769, doi:10.1002/2015WR018552.
"""
struct EGPpowermix{T<:Real} <: ContinuousUnivariateDistribution
    σ::T    # scale parameter
    ξ::T    # rate of upper tail decay
    κ₁::T    # shape of the lower tail
    κ₂::T   # shape of the density in the central part
    p::T    # probability

    function EGPpowermix{T}(σ::T, ξ::T, κ₁::T,κ₂::T, p::T) where {T <: Real}
        new{T}(σ, ξ, κ₁, κ₂, p)
    end
end

function EGPpowermix(σ::T, ξ::T, κ₁::T, κ₂::T, p::T; check_args=true) where {T <: Real}
    check_args && @check_args(EGPpowermix, σ > zero(σ) && κ₁ > zero(κ₁) && κ₂ > zero(κ₂))  # manque test pour p
    return EGPpowermix{T}(σ, ξ, κ₁, κ₂, p)
end

EGPpowermix(σ::Real, ξ::Real, κ₁::Real, κ₂::Real, p::Real) = EGPpowermix(promote(σ, ξ, κ₁, κ₂, p)...)
EGPpowermix(σ::Integer, ξ::Integer, κ₁::Integer, κ₂::Integer, p::Integer) = EGPpowermix(float(σ), float(ξ), float(κ₁), float(κ₂), float(p))

minimum(d::EGPpowermix) = 0.0
maximum(d::EGPpowermix) = Inf * (d.ξ >= 0) - ( d.σ / d.ξ ) * (d.ξ < 0)
insupport(d::EGPpowermix, x::Real) = minimum(d) <= x <= maximum(d)

#### Parameters

scale(d::EGPpowermix) = d.σ
decay(d::EGPpowermix) = d.ξ    # noms à vérifier...
shape_tail(d::EGPpowermix) = d.κ₁
shape_central(d::EGPpowermix) = d.κ₂
prob(d::EGPpowermix) = d.p

params(d::EGPpowermix) = (d.σ, d.ξ, d.κ₁, d.κ₂, d.p)
partype(::EGPpowermix{T}) where {T} = T

#### Evaluation

function logpdf(d::EGPpowermix{T}, x::Real) where T<:Real
    μ = 0
    (σ, ξ, κ₁, κ₂, p) = params(d)

    pd = GeneralizedPareto(μ, 1, ξ)

    g(v::Real, p::Real, κ₁::Real, κ₂::Real) = p*κ₁*v^(κ₁-1) + (1-p)*κ₂*v^(κ₂-1)
    lg(v::Real, p::Real, κ₁::Real, κ₂::Real) = log(g(v,p,κ₁,κ₂))

    lf = -log(σ) + Distributions.logpdf(pd, x/σ) + lg(Distributions.cdf(pd, x/σ),p,κ₁,κ₂)

    return lf
end

pdf(d::EGPpowermix, x::Real) = exp(logpdf(d, x))

function logcdf(d::EGPpowermix{T}, x::Real) where T<:Real
    μ = 0
    (σ, ξ, κ₁, κ₂, p) = params(d)

    pd = GeneralizedPareto(μ, 1, ξ)

    G(v::Real, p::Real, κ₁::Real, κ₂::Real) = p*v^κ₁ + (1-p)*v^κ₂

    lf = log(G(Distributions.cdf(pd, x/σ), p, κ₁, κ₂))

    return lf
end

cdf(d::EGPpowermix, x::Real) = exp(logcdf(d, x))

function quantile(d::EGPpowermix{T}, p::Real) where T<:Real
    @assert zero(p)<p<one(p) "the quantile level should be between 0 and 1."

    (σ, ξ, κ₁, κ₂, m) = params(d)

    G(v::Real, m::Real, κ₁::Real, κ₂::Real) = m*v^κ₁ + (1-m)*v^κ₂

    function invG(p, m, κ₁, κ₂)
        a = 0
        b = 1
        δ = 0.00001
        crit = 1
        X = NaN
        while crit > 2*δ
            X = (a + b)/2
            F = G(X, m, κ₁, κ₂)
            if F <= p
                a = X
            else
                b = X
            end
            crit = b - a
        end
        return X
    end

    x = (σ/ξ)*(((1 - invG(p,m,κ₁, κ₂))^(-ξ)) - 1)

    return x
end

#### Sampling

function rand(rng::AbstractRNG, d::EGPpowermix)
    # Generate a Float64 random number uniformly in (0,1].
    u = 1 - rand(rng)

    G(v::Real, m::Real, κ₁::Real, κ₂::Real) = m*v^κ₁ + (1-m)*v^κ₂

    function invG(p, m, κ₁, κ₂)
        a = 0
        b = 1
        δ = 0.00001
        crit = 1
        X = NaN
        while crit > 2*δ
            X = (a + b)/2
            F = G(X, m, κ₁, κ₂)
            if F <= p
                a = X
            else
                b = X
            end
            crit = b - a
        end
        return X
    end

    return (d.σ/d.ξ)*(((1 - invG(u,d.p,d.κ₁,d.κ₂))^(-d.ξ)) - 1)
end

sampler(d::EGPpowermix) = d


# ---
# title: 35. Search Insert Position
# id: problem35
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array, Binary Search
# link: <https://leetcode.com/problems/search-insert-position/description/>
# hidden: true
# ---
#
# Given a sorted array of distinct integers and a target value, return the index
# if the target is found. If not, return the index where it would be if it were
# inserted in order.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,3,5,6], target = 5
#     Output: 2
#
#
# **Example 2:**
#
#
#
#     Input: nums = [1,3,5,6], target = 2
#     Output: 1
#
#
# **Example 3:**
#
#
#
#     Input: nums = [1,3,5,6], target = 7
#     Output: 4
#
#
# **Example 4:**
#
#
#
#     Input: nums = [1,3,5,6], target = 0
#     Output: 0
#
#
# **Example 5:**
#
#
#
#     Input: nums = [1], target = 0
#     Output: 0
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 104`
#   * `-104 <= nums[i] <= 104`
#   * `nums` contains **distinct** values sorted in **ascending** order.
#   * `-104 <= target <= 104`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using Plots
include_ddp()

# Model
include_model("quadrotor")

function fd(model::Quadrotor{Midpoint, FixedTime}, x, u, w, h, t)
	return view(x, 1:model.n) + h * f(model, view(x, 1:model.n) + 0.5 * h * f(model, view(x, 1:model.n), view(u, 1:model.m), w), view(u, 1:model.m), w)
end

n = model.n
m = model.m

# Time
T = 101
h = 0.05

# Initial conditions, controls, disturbances
x1 = zeros(model.n)
x1[3] = 1.0

xT = copy(x1)
# xT[1] = 5.0
# xT[2] = 0.0
xT[3] = 5.0

u_ref = -1.0 * model.mass * model.g[3] / 4.0 * ones(model.m)
ū = [u_ref + 0.001 * randn(model.m) for t = 1:T-1]
w = [zeros(model.d) for t = 1:T-1]

# Rollout
x̄ = rollout(model, x1, ū, w, h, T)
# x̄ = linear_interpolation(x1, xT, T)
# plot(hcat(x̄...)')

# Objective
Q = h * [(t < T ? Diagonal([1.0 * ones(6); 1.0 * ones(6)])
        : Diagonal([0.0 * ones(6); 0.0 * ones(6)])) for t = 1:T]
q = h * [-2.0 * Q[t] * xT for t = 1:T]

R = h * [Diagonal(1.0 * ones(model.m)) for t = 1:T-1]
r = h * [-2.0 * R[t] * u_ref * 0.0 for t = 1:T-1]

obj = StageCosts([QuadraticCost(Q[t], q[t],
	t < T ? R[t] : nothing, t < T ? r[t] : nothing) for t = 1:T], T)

function g(obj::StageCosts, x, u, t)
	T = obj.T
    if t < T
		Q = obj.cost[t].Q
		q = obj.cost[t].q
	    R = obj.cost[t].R
		r = obj.cost[t].r
        return x' * Q * x + q' * x + u' * R * u + r' * u
    elseif t == T
		Q = obj.cost[T].Q
		q = obj.cost[T].q
        return x' * Q * x + q' * x
    else
        return 0.0
    end
end

# Constraints
p = [t < T ? 2 * m : n for t = 1:T]
info_t = Dict(:ul => zeros(model.m), :uu => 10.0 * ones(model.m), :inequality => (1:2 * m))
info_T = Dict(:xT => xT)
con_set = [StageConstraint(p[t], t < T ? info_t : info_T) for t = 1:T]

function c!(c, cons::StageConstraints, x, u, t)
	T = cons.T
	p = cons.con[t].p

	if t < T
		ul = cons.con[t].info[:ul]
		uu = cons.con[t].info[:uu]
		c .= [ul - u; u - uu]
	elseif t == T
		xT = cons.con[T].info[:xT]
		c .= x - xT
	else
		nothing
	end
end

prob = problem_data(model, obj, con_set, copy(x̄), copy(ū), w, h, T)

# Solve
@time constrained_ddp_solve!(prob,
    max_iter = 1000, max_al_iter = 10,
	con_tol = 1.0e-3,
	ρ_init = 1.0, ρ_scale = 10.0)

x, u = current_trajectory(prob)
x̄, ū = nominal_trajectory(prob)

# Trajectories
plot(hcat(u...)', linetype = :steppost)
plot(hcat(x...)[1:3, :]', linetype = :steppost)

# Visualize
# include(joinpath(pwd(), "models/visualize.jl"))
# vis = Visualizer()
# render(vis)
# # open(vis)
# visualize!(vis, model, x̄, Δt = h)


using Test
using QuantumOpticsBase

@testset "subspace" begin

b = FockBasis(3)

u = Ket[fockstate(b, 1), fockstate(b, 2)]
v = Ket[fockstate(b, 2), fockstate(b, 1)]

bu = SubspaceBasis(u)
bv = SubspaceBasis(v)

T1 = projector(bu, b)
T2 = projector(bv, b)
T12 = projector(bu, bv)

state = fockstate(b, 2)
state_u = Ket(bu, [0, 1])
state_v = Ket(bv, [1., 0])

@test T1*state == state_u
@test T2*state == state_v

state_v = Ket(bv, [1, -1])
state_u = Ket(bu, [-1, 1])

@test T12*state_v == state_u

u2 = Ket[1.5*fockstate(b, 1), fockstate(b, 1) + fockstate(b, 2)]
bu2_orth = QuantumOpticsBase.orthonormalize(SubspaceBasis(u))
@test bu2_orth.basisstates == bu.basisstates

# Test sparse version
T1 = sparseprojector(bu, b)
T2 = sparseprojector(bv, b)
T12 = sparseprojector(bu, bv)

@test isa(T1, SparseOpType)

state_u = Ket(bu, [0, 1])
state_v = Ket(bv, [1., 0])

@test T1*state == state_u
@test T2*state == state_v

state_v = Ket(bv, [1, -1])
state_u = Ket(bu, [-1, 1])

@test T12*state_v == state_u

# Test errors
b2 = FockBasis(4)
@test_throws ArgumentError SubspaceBasis(b2, u)
@test_throws ArgumentError projector(bu, b2)
@test_throws ArgumentError projector(b2, bu)
b2_sub = SubspaceBasis(b2, [fockstate(b2, 1)])
@test_throws ArgumentError projector(bu, b2_sub)

end # testset


@block SebastianM_C ["3d", "diffeq"] begin
    @cell "DifferentialEquations path animation" [lines, meshscatter, animated, record, diffeq] begin
        using OrdinaryDiffEq
        using Statistics

        function lorenz(du,u,p,t)
            du[1] = 10.0*(u[2]-u[1])
            du[2] = u[1]*(28.0-u[3]) - u[2]
            du[3] = u[1]*u[2] - (8/3)*u[3]
        end

        u0 = [1.0; 0.0; 0.0]
        tspan = (0.0, 10.0)
        prob = ODEProblem(lorenz, u0, tspan)
        sol = solve(prob, Tsit5())

        t = Node(tspan[1])
        trajectory = lift(t; init = [Point3f0(sol(t[], idxs = [1,2,3]))]) do t
            push!(trajectory[], Point3f0(sol(t, idxs=[1,2,3])))
        end
        endpoint = lift(t->[Point3f0(sol(t, idxs = [1,2,3]))], t)

        xm, xM = extrema(sol[1,:])
        ym, yM = extrema(sol[2,:])
        zm, zM = extrema(sol[3,:])

        limits = FRect3D((xm,ym,zm), (xM-xm,yM-ym,zM-zm))

        scene = lines(trajectory, limits = limits, markersize = 0.7)
        meshscatter!(scene, endpoint, limits = limits, markersize = 0.5, color = :red)
        eyepos = Vec3f0(100, 5.0, 25)
        N = 1000
        record(scene, @replace_with_a_path(mp4), LinRange(tspan..., N)) do tᵢ
            global eyepos
            t[] = tᵢ
            eyepos = Vec3f0(eyepos[1] - (N / 10000), eyepos[2], eyepos[3])
            update_cam!(scene, eyepos, mean(trajectory[]))
            sleep(0.001)
        end
    end

end

@block AnshulSinghvi ["interactive", "diffeq"]  begin

    @cell "Interactive Differential Equation" [lines, slider, animated, record, diffeq, interactive] begin
        using OrdinaryDiffEq, ParameterizedFunctions
        import AbstractPlotting: textslider
        lorenz = @ode_def Lorenz begin           # define the system
            dx = σ * (y - x)
            dy = x * (ρ - z) - y
            dz = x * y - β*z
        end σ ρ β

        u0 = [1.0,0.0,0.0]                       # initial conditions
        tspan0 = (0.0,100.0)                      # initial timespan
        p0 = [10.0,28.0,8/3]                      # initial parameters
        prob = ODEProblem(lorenz, u0, tspan0, p0)  # define the problem

        ## setup sliders and plotting

        "The order of magnitude to range between."
        OME = 8

        sσ, oσ = textslider(exp10.(-OME:0.001:OME), "σ", start = p0[1]);

        sρ, oρ = textslider(exp10.(-OME:0.001:OME), "ρ", start = p0[2]);

        sβ, oβ = textslider(exp10.(-OME:0.001:OME), "β", start = p0[3]);

        st, ot = textslider(exp10.(-OME:0.001:OME), "tₘₐₓ", start = tspan0[end]);

        sr, or = textslider(100:10000, "resolution", start = 2000);

        trange = lift(ot, or) do tmax, resolution
            LinRange(0.0, tmax, resolution)
        end

        data = lift(oσ, oρ, oβ, trange) do σ, ρ, β, ts
            Point3f0.(
                solve(
                    remake(
                        prob;
                        p = [σ, ρ, β],
                        tspan = (ts[1], ts[end])
                    ),
                    Tsit5(),
                )(ts).u
            )  # change to fit the dimensionality - maybe even return 2 arrays, or a set of `Point2`s...
        end
        parent = Scene(resolution = (1000, 500))

        three = lines(
            data, linewidth = 2,
            transparency = true, color = ("#fe4a49", 0.4),
            show_axis = false
        )

        scatter!(
            three, data, markersize = 0.3, color = (:white, 0.3),
            strokecolor = :black, strokewidth = 1
        )
        on(data) do x
            # center camera etc
            update!(three)
        end
        scene = vbox(hbox(sσ, sρ, sβ, st, sr), three, parent = parent)

        # Do not execute beyond this point!

        RecordEvents(scene, @replace_with_a_path)

    end

end


#=

    distributions.jl

    Contains code to create alternatively parameterised distributions, including copulae.

    Author:     Alexander P. Browning
                ======================
                School of Mathematical Sciences
                Queensland University of Technology
                ======================
                ap.browning@icloud.com
                alexbrowning.me

=#

import Base: rand, minimum, maximum
import Distributions: pdf, cdf
import Statistics: mean, var, std, quantile
import StatsBase: skewness, sample

abstract type Copula end

##############################################################
## Alternative parameterisation of the gamma distribution
##############################################################
struct GammaAlt{T<:Real} <: ContinuousUnivariateDistribution
    μ::T
    σ::T
    ω::T
    d::ContinuousUnivariateDistribution  # Underlying (truncated) Gamma distribution
    GammaAlt{T}(μ,σ,ω,d) where {T} = new{T}(μ,σ,ω,d)
end
struct GammaAltNegative{T<:Real} <: ContinuousUnivariateDistribution
    μ::T
    σ::T
    ω::T
    d::ContinuousUnivariateDistribution  # Underlying (truncated) Gamma distribution
    GammaAltNegative{T}(μ,σ,ω,d) where {T} = new{T}(μ,σ,ω,d)
end

"""
    GammaAlt(μ,σ,ω)

Construct a truncated (x > 0) Gamma distribution `d` where the mean, standard deviation
and skewness of the untruncated distribution are given by μ, σ and ω.

"""
function GammaAlt(μ::T,σ::T,ω::T) where {T <: Real}
    μ > 0.0 || error("Mean must be positive.")
    σ > 0.0 || error("Standard deviation must be positive.")
    if ω < 0
        α = 4/ω^2
        θ = -σ * ω / 2
        d = Truncated(Gamma(α, θ) - α*θ - μ,-Inf,0.0)
        GammaAltNegative{T}(μ,σ,ω,d)
    else
        α = 4/ω^2
        θ = σ * ω / 2
        d = Truncated(Gamma(α, θ) - α*θ + μ,0.0,Inf)
        GammaAlt{T}(μ,σ,ω,d)
    end
end

#### Evaluation
rand(rng::AbstractRNG, d::GammaAltNegative) = -rand(rng,d.d)
pdf(d::GammaAltNegative,x::Real) = pdf(d.d,-x)
logpdf(d::GammaAltNegative,x::Real) = logpdf(d.d,-x)
cdf(d::GammaAltNegative,x::Real) = 1 .- cdf(d.d,-x)
quantile(d::GammaAltNegative,p::AbstractArray) = -quantile(d.d,1.0 .- p)
quantile(d::GammaAltNegative,p::Number) = -quantile(d.d,1.0 .- p)
minimum(d::GammaAltNegative) = 0.0
maximum(d::GammaAltNegative) = Inf

rand(rng::AbstractRNG, d::GammaAlt) = rand(rng,d.d)
pdf(d::GammaAlt,x::Real) = pdf(d.d,x)
logpdf(d::GammaAlt,x::Real) = logpdf(d.d,x)
cdf(d::GammaAlt,x::Real) = cdf(d.d,x)
quantile(d::GammaAlt,p::AbstractArray) = quantile(d.d,p)
quantile(d::GammaAlt,p::Number) = quantile(d.d,p)
minimum(d::GammaAlt) = 0.0
maximum(d::GammaAlt) = Inf

##############################################################
## Alternative parameterisation of the Log Normal distribution
##############################################################
"""
    LogNormalAlt(μ,μ₁,σ₁)

Construct a truncated (x > 0) Log-Normal distribution `d` where the mean of the
untruncated distribution is `μ`. `μ₁` and `σ₁` are the standard Log Normal parameters
in the unshifted distribution.
"""
LogNormalAlt(μ,μ₁,σ₁) = Truncated(LogNormal(μ₁,σ₁) - exp(μ₁ + σ₁^2/2) + μ,0,Inf)

##############################################################
## Gaussian copula
##############################################################
"""
    GaussianCopula(P)

Construct a Gaussian copula with correlation matrix `P`.
"""
struct GaussianCopula{N} <: Copula
    P::Matrix           # Covariance matrix
    L::LowerTriangular  # Cholesky decomposition of `P`
end
GaussianCopula(P::Matrix) = GaussianCopula{size(P,1)}(P,cholesky(P).U')
GaussianCopula(p) = (P = vec_to_cor(p); GaussianCopula{size(P,1)}(P,cholesky(P).U'))

rand(rng::AbstractRNG,C::GaussianCopula{N}) where {N} = normcdf.(C.L * randn(rng,N))
rand(rng::AbstractRNG,C::GaussianCopula{N},n::Int) where {N} = normcdf.(C.L * randn(rng,N,n))
rand(C::GaussianCopula{N},n::Int) where {N} = normcdf.(C.L * randn(N,n))
sample(C::GaussianCopula{N},n::Int=1) where {N} = normcdf.(C.L * randn(N,n))

pdf(C::GaussianCopula{N},u::Vector) where {N} = (x = norminvcdf.(u); 1 / sqrt(det(C.P)) * exp.(-0.5 * x' * (inv(C.P) - I) * x))

marginalize(C::GaussianCopula{N};dims=[1,2]) where {N} = GaussianCopula(C.P[dims,dims])

##############################################################
## t copula
##############################################################
"""
    TCopula(P,ν=1)

Construct a t-copula with correlation matrix `P` and degrees of freedom `ν`.

WARNING: CURRENTLY UNTESTED

"""
struct TCopula{N} <: Copula
    P::Matrix           # Covariance matrix
    ν::Real             # df parameter
    L::LowerTriangular  # Cholesky decomposition of `P`
end
TCopula(P::Matrix,ν::Real=1) = TCopula{size(P,1)}(P,ν,cholesky(P).U')
TCopula(p,ν::Real=1) = (P = vec_to_cor(p); TCopula{size(P,1)}(P,ν,cholesky(P).U'))

rand(rng::AbstractRNG,C::TCopula{N}) where {N} = tdistcdf.(C.ν,C.L * randn(rng,N) / sqrt.(rand(Chisq(C.ν))' / C.ν))
rand(rng::AbstractRNG,C::TCopula{N},n::Int) where {N} = tdistcdf.(C.ν,C.L * randn(rng,N,n) ./ sqrt.(rand(Chisq(C.ν),n)' / C.ν))
rand(C::TCopula{N},n::Int) where {N} = tdistcdf.(C.ν,C.L * randn(N,n) ./ sqrt.(rand(Chisq(C.ν),n)' / C.ν))
sample(C::TCopula{N},n::Int=1) where {N} = tdistcdf.(C.ν,C.L * randn(N,n) ./ sqrt.(rand(Chisq(C.ν),n)' / C.ν))

Γ = gamma
function pdf(C::TCopula{N},u::Vector) where {N}
    x = tdistinvcdf.(C.ν,u);
    Γ((C.ν + N)/2) /
    (Γ(C.ν/2) * C.ν^(N/2) * π^(N/2) * sqrt(det(C.P))) *
    (1 + x' * inv(C.P) * x / C.ν)^(-(C.ν + N) / 2) /
    prod(tdistpdf.(C.ν,x))
end

##############################################################
## Copula-MVDistribution
##############################################################
"""
    MvDependent(C,v)

Construct dependent multivariate distribution where the marginals are described by
elements of `v` (a vector), and the dependence structure is described a copula `C`.
"""
struct MvDependent{N}
    C::Copula
    v::Vector
    q::Vector{Bool}
end
MvDependent(C::Copula,v::Vector;q=falses(length(v))) = MvDependent{length(v)}(C,v,q)

function sample(d::MvDependent{N},n=1) where {N}
    U = rand(d.C,n)
    for i = 1:N
        U[i,:] = d.q[i] ? quantile_interpolated(d.v[i],U[i,:]) : quantile(d.v[i],U[i,:])
    end
    return U
end
rand(d::MvDependent,n::Int=1) = sample(d,n)

marginalize(d::MvDependent;dims=[1,2]) = MvDependent(
    marginalize(d.C;dims),
    d.v[dims],
    q=d.q[dims]
)

pdf(d::MvDependent,x::Vector) = pdf(d.C,max.(0.0,cdf.(d.v,x))) * prod(pdf.(d.v,x))

##############################################################
## Helpful functions
##############################################################
"""
    vec_to_cor(p)

Convert number of vector `p` to a correlation matrix.

Example:

    vec_to_cor([0.1,0.2,0.3])
    3×3 Matrix{Float64}:
     1.0  0.1  0.2
     0.1  1.0  0.3
     0.2  0.3  1.0
"""
function vec_to_cor(p::Vector)
    n = 0.5 * (1 + sqrt(1 + 8length(p)))
    isinteger(n) || error("Check length of input.")
    n = Int(n)
    P = zeros(n,n)
    P[(triu(ones(n,n)) - I) .== 1.0] = p
    P += P' + I
end
vec_to_cor(p::Real) = vec_to_cor([p])

"""
    quantile_interpolated(d,u)

Approximate quantile function using interpolation. Output is comparable to
`quantile(d,u)`, with a considerable speed-up if `length(u)` is large and
the quantile function of `d` is expensive.
"""
function quantile_interpolated(d::Distribution,u::Vector;pts::Int=100)
    length(u) > 2 || "`u` must contain more than 2 elements."
    x = range(quantile(d,[extrema(u)...])...,length=pts)
    y = cdf(d,x)
    LinearInterpolation(y,x,extrapolation_bc=Line()).(u)
end

Distributions.cdf(x::Vector) = invperm(sortperm(x)) / length(x) .- 1 / 2length(x)


using bayesian_nlde
using DataFrames
using Plots
using Turing
# using ReverseDiff
# using DiffEqSensitivity

Turing.setadbackend(:forwarddiff)

# load data and format for inference
dat = import_data(download_data = false)
grouped_dat = [DataFrame(df) for df in groupby(dat, :Country_Region)];
grouped_sir_data = [Matrix((df[:, [:susc_prop, :case_prop, :recov_prop]]))' for df in groupby(dat, :Country_Region)];
country_pops = [df.pop[1] for df in grouped_dat]
country_names = [df.Country_Region[1] for df in grouped_dat]

# look at countries with more than a months worth of reporting
sub_idx = [i for (i,x) in enumerate(grouped_sir_data) if size(x, 2) > 28]
sub_idx = sub_idx[1:10]
sub_grouped_sir_data = grouped_sir_data[sub_idx]
sub_country_pops = country_pops[sub_idx]

model = hierarchical_sir_model(grouped_sir = sub_grouped_sir_data,
    m_countries = length(sub_grouped_sir_data),
    country_pops = sub_country_pops)

varinfo = Turing.VarInfo(hierarchical_sir_model)
spl = Turing.SampleFromPrior()
@code_warntype model.f(varinfo, spl, Turing.DefaultContext(), model)

# infer paramters
@time chains = sample(hierarchical_sir_model(grouped_sir = sub_grouped_sir_data,
    m_countries = length(sub_grouped_sir_data), country_pops = sub_country_pops),
    NUTS(100, 0.8), 500; progress = true, verbose = true)
write(string(dirname(@__DIR__)) * "/data/chains/posterior_samples.jls", chains)

#plotlyjs()

# plot distribution over trajectories
for i in sub_idx
    plt = plot_country_posterior(chains, grouped_dat, i; subsample_size = 200)
    plt_path = string(dirname(@__DIR__)) * "/data/plots/" * string(country_names[i]) * ".svg"
    savefig(plt, plt_path)
end


using StaticArrays: SVector
using UniqueVectors: UniqueVector
using Combinatorics: combinations

Cell(points::AbstractVector{Point{N}}, K::Int) where N = Cell{N}(Cell{N}[],
    Dict{Cell{N}, Bool}(), Vector(points), K)

Cell(s::Simplex{N, K}) where {N, K} = Cell(s.points, K)

import Base: show
show(io::IO, c::Cell{N}) where N = print(io,
    "Cell{$N,$(c.K)}(children: $(length(c.children)), " *
    "parents: $(length(c.parents)), points: $(c.points))")

CellComplex{N,K}() where {N, K} = CellComplex{N,K}(
    SVector{K}([UniqueVector(Cell{N}[]) for _ in 1:K]))

CellComplex{N,K}(cells::AbstractVector{Cell{N}}) where {N, K} =
    append!(CellComplex{N,K}(), cells)

CellComplex(simplices::AbstractVector{Simplex{N, K}}) where {N, K} =
    TriangulatedComplex(simplices).complex

"""
    CellComplex(cells::AbstractVector{Cell{N}}) where N

Create a cell complex consisting of `cells` and all descendants of `cells`.
"""
function CellComplex(cells::AbstractVector{Cell{N}}) where N
    K = maximum(c.K for c in cells)
    comp = CellComplex{N, K}(cells)
    for k in reverse(1:K)
        for c in comp.cells[k]
            append!(comp, c.children)
        end
    end
    return comp
end

show(io::IO, comp::CellComplex{N,K}) where {N, K} = print(
     io, "CellComplex{$N,$K}$(tuple(map(length, comp.cells)...))")

export parent!
"""
    parent!(c::Cell{N}, p::Cell{N}, o::Bool) where N

Make `p` a parent of `c` with relative orientation `o` and make `c` a child of `p`.
Return c.
"""
function parent!(c::Cell{N}, p::Cell{N}, o::Bool) where N
    @assert p.K == c.K+1
    c.parents[p] = o
    push!(p.children, c)
    return c
end

import Base: push!
"""
    push!(comp::CellComplex{N,K}, c::Cell{N}) where {N, K}

Add a Cell to a CellComplex unless it is already in the CellComplex. Return comp.
"""
function push!(comp::CellComplex{N,K}, c::Cell{N}) where {N, K}
    @assert c.K <= K
    if !(c in comp.cells[c.K])
        push!(comp.cells[c.K], c)
    end
    return comp
end

import Base: append!
"""
    append!(comp::CellComplex{N}, cells::AbstractVector{Cell{N}}) where N
    append!(comp1::CellComplex{N, K}, comp2::CellComplex{N, J}) where {N, K, J}

`push!` all cells in an array or a cell complex into another cell complex. Return the first
cell complex.
"""
function append!(comp::CellComplex{N}, cells::AbstractVector{Cell{N}}) where N
    for c in cells
        push!(comp, c)
    end
    return comp
end

function append!(comp1::CellComplex{N, K}, comp2::CellComplex{N, J}) where {N, K, J}
    @assert 1 <= J <= K
    for cells in comp2.cells
        append!(comp1, cells)
    end
    return comp1
end

export subcomplex
"""
    subcomplex(comp::CellComplex{N}, simplices::AbstractVector{Simplex{N, K}}) where {N, K}

Find the subcomplex of the given cell complex that includes the given simplices.
"""
function subcomplex(comp::CellComplex{N},
    simplices::AbstractVector{Simplex{N, K}}) where {N, K}
    simplex_sets = [Set(s.points) for s in simplices]
    return CellComplex(filter(c -> Set(c.points) in simplex_sets, comp.cells[K]))
end

export simplicial
"""
    simplicial(comp::CellComplex{N, K}) where {N, K}

Return `true` if every cell of `comp` is a simplex, else `false`.
"""
function simplicial(comp::CellComplex{N, K}) where {N, K}
    for k in 1:K
        for cell in comp.cells[k]
            if !(length(cell.points) == cell.K == k)
                return false
            end
        end
    end
    return true
end

export pairwise_delaunay
"""
    pairwise_delaunay(m::Metric{N}, comp::CellComplex{N, K}, ϵ::Real=0.0) where {N, K}

Return `true` if every pair of cells sharing a face are within `ϵ` of pairwise delaunay,
else `false`.
"""
function pairwise_delaunay(m::Metric{N}, comp::CellComplex{N, K},
                           ϵ::Real=0.0) where {N, K}
    for k in 1:K
        for cell in comp.cells[k]
            for (c1, c2) in combinations(collect(keys(cell.parents)), 2)
                s1, s2 = Simplex(c1), Simplex(c2)
                if !(pairwise_delaunay(m, s1, s2) > ϵ &&
                     pairwise_delaunay(m, s2, s1) > ϵ)
                    return false
                end
            end
        end
    end
    return true
end

export one_sided
"""
    one_sided(m::Metric{N}, comp::CellComplex{N, K}, ϵ::Real=0.0) where {N, K}

Return `true` if every boundary cell has its circumcenter inside the cell complex up to `ϵ`.
"""
function one_sided(m::Metric{N}, comp::CellComplex{N, K},
    ϵ::Real=0.0) where {N, K}
    for k in 1:K
        for cell in comp.cells[k]
            if length(cell.parents) == 1 # boundary parent
                s = Simplex(collect(keys(cell.parents))[1])
                center, _ = circumsphere_barycentric(m, s)
                p = first_setdiff(s.points, cell.points)
                i = findfirst(isequal(p), s.points)
                if center.coords[i] <= ϵ
                    return false
                end
            end
        end
    end
    return true
end

export pairwise_noncocyclic
"""
    pairwise_noncocyclic(m::Metric{N}, comp::CellComplex{N, K}, ϵ::Real=0.0) where {N, K}

Return `true` if every pair of cells sharing a face are more than `ϵ` away from sharing
a circumsphere, else `false`.
"""
function pairwise_noncocyclic(m::Metric{N}, comp::CellComplex{N, K},
    ϵ::Real=0.0) where {N, K}
    for k in 1:K
        for cell in comp.cells[k]
            for (c1, c2) in combinations(collect(keys(cell.parents)), 2)
                s1, s2 = Simplex(c1), Simplex(c2)
                if (ϵ >= pairwise_delaunay(m, s1, s2) >= -ϵ) ||
                    (ϵ >= pairwise_delaunay(m, s2, s1) >= -ϵ)
                    return c1, c2
                end
            end
        end
    end
    return true
end

export well_centered
"""
    well_centered(m::Metric{N}, comp::CellComplex{N, K}, ϵ::Real=0.0) where {N, K}

Return `true` if every cell contains is circumcenter, within `ϵ`, else `false`.
"""
function well_centered(m::Metric{N}, comp::CellComplex{N, K},
    ϵ::Real=0.0) where {N, K}
    for k in 1:K
        for cell in comp.cells[k]
            s = Simplex(cell)
            center, _ = circumsphere_barycentric(m, s)
            if any(center.coords .<= ϵ)
                return false
            end
        end
    end
    return true
end


"""

Returns prototype fitlering matrix for BF-OFDM  with specified by overlapping factor K, PPN size nFBMC, filter type. Additional parameters are the CP size of the precoding stage (set to optimal value by default) and filter structure for specific parameter overset.
# --- Syntax
(p,pF) = getBFOFDMFilter(K,nFBMC,filter;GI,BT,filterStopBand,fS)>
# --- Input parameters
- K	  : Overlapping factor [Int]
- nFBMC : PPN size (Number of FBMC carriers) [Int]
- filter: Filter type: "Gaussian","phydyas","FS","DC": [String]
- BT	  : Bandwidth time product for custom gaussian filter (used with "gaussian") [Float64]
- filterStopBand : Attenuation with custom DC window (used with "DC") [Float64]
- fS	  : Frequency sampling custom  coefficients (Array{Float64,K}) (used with "fs")
# --- Output parameters
- p		: Filter impulse response [Array{Float64,K*nFBMC}].
- pF		: Filter coefficient in frequency domain (For Frequency sampling implementation)
---
Filter types
All filter type can be called with 2 ways. "filter" and "filter_opt". When "filter_opt" is used, the filter coefficient are obtained through numerical optimisation (SIR optimisation). Otherwise, it uses its associated parameter to extract appropriate coefficients. See [1] and [2] for filter optimisation principles. See [4] for special cases about 5G-NR compatibilty.
The supported filter/windows are:
- Gaussian	: Gaussian filter shape, specified by BT. Optimized Gaussian shape can be used [1]
- phydyas		: Classic FBMC pulse shape defined in frequency domain [3]
- FS			: Coefficient defined in frequency domain. Filter coefficient are based on intrinsic SIR optimisation [1]
- DC			: Dolph Chebyshev window. Can be based on window optimisation (SIR optimisation) or defined by filterStopBand parameter.
# ---
# References
- [1]	Demmer, D and Gerzaguet, R and Doré, J-B and Le Ruyet, D. and Kténas, D, "Filter Design for 5G BF-OFDM Waveform", 2017.
- [2] A. Sahin, I. Guvenc, and H. Arslan, “A Survey on Multicarrier Communications: Prototype Filters, Lattice Structures, and Implementation Aspects", 2014
- [3] Phydyas project, "Deliverable D5.1 : Prototype filter and structure optimization", 2009
- [4]	Demmer, D and Rostom, Z and Gerzaguet, R and Doré, J-B and Le Ruyet, D. "Study of OFDM Precoded Filter-Bank Waveforms", 2018
# Examples
Get a classic FBMC filter with phydyas and an overlapping factor of 4, with a input PPN of size 64. The precodinng stage is a OFDM of size 64 and a GI of size 4.

	# --- phydyas filter
	( hphydyas,hFphydyas )	 = getBFOFDMFilter(4,64,"phydyas");

	# --- Gaussian shape
	( hGaussian,hFGaussian )	 = getBFOFDMFilter(4,64,"Gaussian_opt",GI=4);

	# --- Custom Gaussian window
	( hGaussianC,hFGaussianC ) = getBFOFDMFilter(4,64,"Gaussian",GI=4,BT=0.5)

# ---
# v 1.0 - Robin Gerzaguet.
"""
function getBFOFDMFilter(nFBMC,K,filter;GI=-1,BT=-1,filterStopBand =-1,fS=[])
	# --- Avoid type issue with lowercase casting
	filterSwitch = lowercase(filter);
	# ----------------------------------------------------
	# --- Filter switch
	# ----------------------------------------------------
	if filterSwitch == "phydyas"
		# ----------------------------------------------------
		# --- phydyas prototype filter [3]
		# ----------------------------------------------------
		# --- Getting frequency coefficients
		if K == 2
			p1			= sqrt(2)/2;
			P			= p1;
		elseif K == 3
			p1			= 0.91143783;
			P			= [p1;sqrt(1-p1^2)];
		elseif K == 4
			p1			= 0.97195983;
			p2			= 1/sqrt(2);
			p3          = sqrt(1-p1^2);
			P			= [p1;p2;p3];
		elseif K == 6
			p1          = 0.99722723;
			p2          = 0.94136732;
			P           = [ p1;p2 sqrt(2)/2;sqrt(1-p2^2);sqrt(1-p1^2)];
		elseif  K== 8
			p1          = 0.99988389;
			p2          = 0.99315513;
			p3          = 0.92708081;
			P          = [ p1;p2;p3;sqrt(2)/2;sqrt(1-p3^2);sqrt(1-p2^2);sqrt(1-p1^2)];
		else
			error("For Phydyas prototype filter, supported overlapping factor are 2, 3, 4, 6 or 8");
		end
		# --- Getting time impulse response
		# Time indexes
		t	      = collect(1:K*nFBMC);
		# Frequency vector
		P_K	      = [P[end:-1:1];1;P];
		# Init filter time impulse reponse
		p         = zeros(Float64,K*nFBMC);
		for k     = K+1 : 2*length(P)+1
			# --- Iterative generation (see [3])
			p[ t ]  = p[ t ] .+  2*P_K[ k ]*(-1)^(k-K) .* cos.(2*pi/(K*nFBMC)*(k-K) .*(t .- 1));
		end
		p         =p .+ P_K[ K ];
		# --- Final impulse response
		imp_resp  = (-1)^K * p;
		# --- Final frequency response
		freq_coef = P_K;
		# --- Output
		return (imp_resp,freq_coef)
	elseif filterSwitch == "gaussian_opt"
		#FIXME: Optimize value for any (K,GI) tuple
		# ----------------------------------------------------
		# --- Gaussian optimized window [1]
		# ----------------------------------------------------
		# To use an optimized Gaussian function as it depends on GI size, we have to use a double switch (on K and on GI)
		if K == 2
			if GI ==  2
				a = sqrt(log(2)/2)/(0.40);
			elseif GI ==  1
				a = sqrt(log(2)/2)/(0.45);
			elseif GI == 9
				# 5G-NR value [3]
				a = sqrt(log(2)/2)/(1.08);
			else
				# --- Of other value (>), it corresponds to the optimized Gaussian shape without ISI (2K-1)
				a = sqrt(log(2)/2)/(0.38);
			end
		elseif K == 3
			if GI ==  9 || GI == 18 || GI == 27 GI == 36
				# --- 5G NR with 0.25 compression rate
				a = sqrt(log(2)/2)/(0.3);
			elseif GI ==  5
				a = sqrt(log(2)/2)/(0.32);
			elseif GI ==  4
				a = sqrt(log(2)/2)/(0.33);
			elseif GI ==  2
				a = sqrt(log(2)/2)/(0.40);
			elseif GI ==  1
				a = sqrt(log(2)/2)/(0.40);
			end
		elseif K == 4
			if GI ==  7
				a = sqrt(log(2)/2)/(0.28);
			elseif GI ==  4
				a = sqrt(log(2)/2)/(0.33);
			elseif GI ==  3
				a = sqrt(log(2)/2)/(0.357);
			elseif GI ==  2
				a = sqrt(log(2)/2)/(0.40);
			else
				# --- Raise an error for other values
				error("unknown optimized filter for GI (gaussian)");
			end
		elseif K == 6
			a = sqrt(log(2)/2)/(0.327);
		else
			error("Unknown optimized filter for K (gaussian)");
		end
		# --- Time index
		t = range(-K/2,stop=K/2,length=K*nFBMC);
		# --- Gaussian filter shape
		p = sqrt(pi)/a*exp.(-(pi/a .* t).^2);
		# --- Output
		imp_resp  = p;
		freq_coef = [];
		return(imp_resp,freq_coef)
	elseif filterSwitch == "gaussian"
		# ----------------------------------------------------
		# --- Manual Gaussian filter
		# ----------------------------------------------------
		if BT == -1
			# --- Corresponds to the default value for BT.
			# Raise an error, as BT should be defined for a specific Gaussian filter
			error("For custom Gaussian filter, BT should be defined. For an optimized Gaussian shape, use gaussian_opt instead");
		end
		a = sqrt(log(2)/2)/(BT);
		# --- Time index
		t = range(-K/2,stop=K/2,length=K*nFBMC);
		# --- Gaussian filter shape
		p = sqrt(pi)/a*exp.(-(pi/a .* t).^2);
		# --- Output
		imp_resp  = p;
		freq_coef = [];
		return(imp_resp,freq_coef)
	elseif filterSwitch == "fs_opt"
		# ----------------------------------------------------
		# --- Frequency sampling optimized filter
		# ----------------------------------------------------
		# --- Get frequency coefficients
		if K == 4
			if GI ==  7
				P = [0.72071;0.25452;0.03380];
			elseif GI ==  4
				P = [0.792;0.375;0.082];
			elseif GI ==  2
				P = [0.870;0.581;0.211];
			else
				error("unknown optimized filter for GI (fs_opt)");
			end
		elseif K == 3
			if GI == 5
				P = [0.633;0.133];
			elseif GI ==  4
				P = [0.656;0.154];
			elseif GI ==  2
				P = [0.806;0.301];
			else
				error("unknown optimized filter for GI (fs_opt)");
			end
		elseif K == 2
			if GI == 3
				P = 0.475;
			elseif GI ==  2
				P = 0.485;
			elseif GI ==  1
				P = 0.524;
			else
				error("unknown optimized filter for GI (fs_opt)");
			end
		else
			error("unknown optimized filter for K (fs_opt)");
		end
		# --- Deduce filter taps
		t	  = collect(1:K*nFBMC);
		P_K	  = [P[ end:-1:1 ];1;P];
		p = zeros(Float64,K*nFBMC);
		for k = K+1 : 2*length(P)+1
			p[ t ] = p[ t ] +2*P_K[ k ]*(-1)^(k-K) .* cos.(2*pi/(K*nFBMC)*(k-K) .*(t .- 1));
		end
		p =p .+ P_K(K);
		# --- Final impulse response
		imp_resp  = (-1)^K * p;
		# --- Final frequency response
		freq_coef = P_K;
		# --- Output
		return (imp_resp,freq_coef)
	elseif filterSwitch == "fs"
		# ----------------------------------------------------
		# --- Customized Frequency sampling filter
		# ----------------------------------------------------
		if isempty(fS)
			# --- Corresponds to the default value for freuency points.
			# Raise an error, as the vector should be defined properly.
			error("For customFrequency sampling filter, fS should be defined. For an optimized FS system, use fs_opt instead");
		end
		t	  = collect(1:K*nFBMC);
		P_K	  = fS;
		p = zeros(Float64,K*nFBMC);
		for k = K+1 : 2*length(P)+1
			p[ t ] = p[ t ] +2*P_K[ k ]*(-1)^(k-K) .* cos.(2*pi/(K*nFBMC)*(k-K) .*(t .- 1));
		end
		p =p .+ P_K(K);
		# --- Final impulse response
		imp_resp  = (-1)^K * p;
		# --- Final frequency response
		freq_coef = P_K;
		# --- Output
		return (imp_resp,freq_coef)
	elseif filterSwitch == "dc_opt"
		# ----------------------------------------------------
		# --- Dolph Chebyshev optimized window
		# ----------------------------------------------------
		# --- Getting optimized stop band attenuation
		if K == 6
			if G == 7
				filterStopBand = 159;
			elseif GI ==  4
				filterStopBand = 340;
			elseif GI ==  2
				filterStopBand = 317;
			else
				error("unknown optimized filter for GI (DC)");
			end
		elseif K == 4
			if GI ==  7
				filterStopBand = 159;
			elseif GI ==  4
				filterStopBand = 220;
			elseif GI ==  2
				filterStopBand = 317;
			else
				error("unknown optimized filter for GI (DC)");
			end
		elseif K == 3
			if GI == 5
				filterStopBand = 119.30;
			elseif GI ==  4
				filterStopBand = 129.19;
			elseif GI ==  2
				filterStopBand = 182.00;
			else
				error("unknown optimized filter for GI (DC)");
			end
		elseif K ==2
			if GI ==  3
				filterStopBand = 80.04;
			elseif GI ==  2
				filterStopBand = 86.31;
			elseif GI ==  1
				filterStopBand = 106.18;
			else
				error("unknown optimized filter for GI (DC)");
			end
		else
			error("unknown optimized filter for K (DC)");
		end
		# --- Final impulse response
		imp_resp  = dolphChebyshev(K*nFBMC,Float64(filterStopBand));
		# --- Final frequency response
		freq_coef = [];
		# --- Output
		return (imp_resp,freq_coef)
	elseif filterSwitch == "dc"
		# ----------------------------------------------------
		# --- Customized Dolph Chebyshev window
		# ----------------------------------------------------
		if filterStopBand == -1
			# --- Corresponds to the default value for BT.
			# Raise an error, as BT should be defined for a specific Gaussian filter
			error("For custom Dolph Chebyshev window, filterStopBand should be defined. For an optimized DC window, use dc_opt instead");
		end
		# --- Final impulse response
		imp_resp  = dolphChebyshev(K*nFBMC,Float64(filterStopBand));
		# --- Final frequency response
		freq_coef = [];
		# --- Output
		return (imp_resp,freq_coef)
	end
end


using Test, ApproxFun, SingularIntegralEquations
import ApproxFunBase: ∞, testbandedoperator, testfunctional, testblockbandedoperator, testraggedbelowoperator
import ApproxFunOrthogonalPolynomials: JacobiZ
import SingularIntegralEquations: testsies, testsieeval, stieltjesmoment, Directed, _₂F₁, ⁺, ⁻

@testset "Log Kernel" begin
    @testset "Segment" begin
        testsieeval(Jacobi(0,0))

        a=1.0;b=2.0+im
        d=Segment(a,b)
        z=Fun(d)
        f=real(exp(z)/(sqrt(z-a)*sqrt(b-z)))
        S=space(f)
        testsies(S)

        a=1.0;b=2.0+im
        d=Segment(a,b)
        z=Fun(d)
        f=real(exp(z)*(sqrt(z-a)*sqrt(b-z)))
        S=space(f)
        testsies(S)

        a=1.0;b=2.0+im
        d=Segment(a,b)
        z=Fun(d)
        f=real(exp(z)/(sqrt(z-a)*sqrt(b-z)))
        S=JacobiWeight(-0.5,-0.5,ChebyshevDirichlet{1,1}(d))
        # TODO: move to testsies
        H=OffSingularIntegral(S,Chebyshev(3..4),0)
        @test (H*f)(3.5) ≈ logkernel(f,3.5)

        H=OffSingularIntegral(S,Chebyshev(Segment(3,4.0+im)),0)
        @test (H*f)(3.5+0.5im) ≈ logkernel(f,3.5+0.5im)
    end

    @testset "Circle" begin
        d=Circle(0.2,3.0)
        S=Fourier(d)
        ζ=Fun(d)
        f=real(ζ+1/(ζ-0.1))
        #z=0.1+0.1im;@test linesum(logabs(ζ-z)*f) ≈ logkernel(f,z)*π
        #z=5.0+0.1im;@test linesum(logabs(ζ-z)*f) ≈ logkernel(f,z)*π

        d=Circle(0.2,3.0)
        S=Fourier(d)
        H=Hilbert(S,0)
        ζ=Fun(d)
        f=real(ζ+1/(ζ-0.1))
        z=0.2+3im;@test (H*f)(z) ≈ logkernel(f,z)
    end

    @testset "Legendre and Jacobi on interval" begin
        x=Fun()
        f=exp(x)

        @test isa(logkernel(f,2.0+im),Real)
        @test logkernel(f,2.0+im) ≈ sum(f*logabs(2.0+im-x))/π

        f=sqrt(1-x^2)*exp(x)
        @test logkernel(f,2.0+im) ≈ sum(f*logabs(2.0+im-x))/π

        x=Fun()
        f=(1-x)^0.1
        sp=space(f)
        @test logjacobimoment(sp.α,sp.β,2.0) ≈ sum((1-x)^sp.α*(1+x)^sp.β*logabs(2.0-x))

        @test logkernel(f,2.0) ≈ sum(f*logabs(2.0-x)/π)

        f=(1-x)^0.1*exp(x)

        @test stieltjes(f,2.0) ≈ sum(f/(2.0-x))

        @test logkernel(f,2.0) ≈ sum(f*logabs(2.0-x)/π)

        @test isa(logkernel(f,2.0+im),Real)
        @test logkernel(f,2.0+im) ≈ sum(f*logabs(2.0+im-x)/π)

        f=(1-x^2)^0.1*exp(x)
        sp=space(f)
        @test logjacobimoment(sp.α,sp.β,2.0) ≈ sum((1-x^2)^0.1*logabs(2.0-x))
        @test logkernel(f,2.0+im) ≈ sum(f*logabs(2.0+im-x)/π)

        f=(1-x)^(-0.1)*(1+x)^(-0.2)*exp(x)
        @test logkernel(f,2.0+im) ≈ sum(f*logabs(2.0+im-x)/π)
        @test isa(logkernel(f,2.0+im),Real)
    end

    @testset "other intervals" begin
        x = Fun()
        f = exp(x)
        z = 2+im
        @test logkernel(f,z) ≈ linesum(f*logabs(x-z))/π
        @test logkernel(f, -10.0) ≈ linesum(f*logabs(x+10))/π

        x = Fun()
        f = real(sqrt(1-x)*exp(x))
        z = 2+im
        @test logkernel(f,z) ≈ linesum(f*logabs(x-z))/π
        @test logkernel(f, -10.0) ≈ linesum(f*logabs(x+10))/π
        @test_broken logkernel(f, 0.1) ≈ logkernel(f, 0.1+eps()im)

        x = Fun(0.1..1)
        f = real(sqrt(1-x)*exp(x))
        z = 2+im
        @test logkernel(f,z) ≈ linesum(f*logabs(x-z))/π
        @test logkernel(f, -10.0) ≈ linesum(f*logabs(x+10))/π
        @test_broken logkernel(f, 0.2) ≈ logkernel(f, 0.2+eps()im)
    end

    @testset "Arc" begin
        a=Arc(0.,1.,0.,π/2)
        testsieeval(Legendre(a);posdirection=(-1-im))
    end

    @testset "Chebyshev singularities" begin
        for d in (ChebyshevInterval(), Interval(0,1), Segment(1+im, 2+3im))
            for k = 1:10
                u = Fun(JacobiWeight.(-0.5,-0.5,Ref(ChebyshevDirichlet{1,1}(d))), [zeros(4);1])
                x = Fun(d)
                for z in (20.0+im, 3.0)
                    ex = linesum(u*logabs(z-x))/π
                    @test logkernel(u, z) ≈ ex
                end
            end
        end

        for d in (ChebyshevInterval(), Interval(0,1))
            for k = 1:10
                u = Fun(JacobiWeight.(-0.5,-0.5,Ref(ChebyshevDirichlet{1,1}(d))), [zeros(4);1])
                x = Fun(d)
                for z in (20.0+im, 3.0)
                    ex_s = sum(u*log(z-x))
                    @test stieltjesintegral(u, z) ≈ ex_s
                end
            end
        end
    end
end


struct FNN <: AbstractApproximator
    n::Int
    m::Int
    NN::Flux.Chain
    function FNN(n::Int, m::Int, h_array::Vector{Int}, act)
        node_array = [n+m, h_array..., 1]
        new(n, m, construct_layer_array(node_array, act))
    end
end

"""
    (nn::FNN)(x, u)

Infer using two input arguments `x` and `u`.
# Notes
size(x) = (n, d)
size(u) = (m, d)
"""
function (nn::FNN)(x, u)
    res = nn.NN(vcat(x, u))
end

Flux.params(approximator::FNN) = Flux.params(approximator.NN)


"""
Bag-of-random-embedding-projections (BOREP) structure for
document embedding using word vectors.

# References
  * [Wieting, Kiela ICLR 2019, "No training required: Exploring random encoders
     for sentence classification"](https://arxiv.org/abs/1901.10444)
"""
struct BOREPEmbedder{T,S} <: WordVectorsEmbedder{T,S}
    embeddings::EmbeddingsLibrary{T,S}
    R::Matrix{T}
    fpool::Function
    config::NamedTuple
end

function BOREPEmbedder(embeddings::EmbeddingsLibrary{T,S},
                       config;
                       dim=DEFAULT_BOREP_DIMENSION,
                       pooling_function=DEFAULT_BOREP_POOLING_FUNCTION,
                       initialization::Symbol=:heuristic,
                       kwargs...
                      ) where {T<:AbstractFloat, S<:AbstractString}
    # Check initialization option and generate random matrix
    d = size(embeddings)[1]  # number of vector components
    if initialization == :heuristic
        R = rand(T[-1/sqrt(d), 1/sqrt(d)], dim, d)
    elseif initialization == :uniform
        R = rand(T[-0.1, 0.1], dim, d)
    elseif initialization == :normal
        R = randn(T, dim, d)
    end

    # Check pooling function option anf generate pooling function
    if pooling_function == :sum
        fpool = x->vec(sum(x, dims=2))
    elseif pooling_function == :max
        fpool = x->vec(maximum(x, dims=2))
    end

    return BOREPEmbedder(embeddings, R, fpool, config)
end

# Dimensionality function
function dimensionality(embedder::BOREPEmbedder)
    # The final number of dimensions is given by
    # the dimension of the random space (i.e. number
    # of rows of the projection matrix)
    return size(embedder.R, 1)
end

# Sentence embedding function
function sentences2vec(embedder::BOREPEmbedder,
                       document_embedding::Vector{Matrix{T}};
                       kwargs...) where {T,S}
    n = length(document_embedding)
    X = zeros(T, dimensionality(embedder), n)
    @inbounds @simd for i in 1:n
        X[:,i] = embedder.fpool(embedder.R * document_embedding[i])
    end
    return X
end


using Test: @test, @testset, @inferred
using Statistics: mean, std, var, cor
using Random: randn!, bitrand
using LogExpFunctions: softmax
using RestrictedBoltzmannMachines: BinaryRBM, energy, free_energy, metropolis, metropolis!
import RestrictedBoltzmannMachines as RBMs

@testset "metropolis β=$β" for β in [0.5, 1.0, 2.0]
    N = 5
    M = 2
    rbm = BinaryRBM(randn(N), randn(M), randn(N, M) / √N)
    T = 10000
    B = 100
    v = bitrand(N, B, T)
    for t in 2:T
        v[:,:,t] .= metropolis(rbm, v[:,:,t - 1]; β)
    end
    counts = Dict{BitVector, Int}()
    for t in 1000:T, n in 1:B
        counts[v[:,n,t]] = get(counts, v[:,n,t], 0) + 1
    end
    freqs = Dict(v => c / sum(values(counts)) for (v,c) in counts)
    𝒱 = [BitVector(digits(Bool, x; base=2, pad=N)) for x in 0:(2^N - 1)]
    @test cor([get(freqs, v, 0.0) for v in 𝒱], softmax(-β * free_energy.(Ref(rbm), 𝒱))) > 0.99
end

@testset "metropolis! β=$β" for β in [0.5, 1.0, 2.0]
    N = 5
    M = 2
    rbm = BinaryRBM(randn(N), randn(M), randn(N, M) / √N)
    T = 10000
    B = 100
    v = bitrand(N, B, T)
    metropolis!(v, rbm; β)
    counts = Dict{BitVector, Int}()
    for t in 1000:T, n in 1:B
        counts[v[:,n,t]] = get(counts, v[:,n,t], 0) + 1
    end
    freqs = Dict(v => c / sum(values(counts)) for (v,c) in counts)
    𝒱 = [BitVector(digits(Bool, x; base=2, pad=N)) for x in 0:(2^N - 1)]
    @test cor([get(freqs, v, 0.0) for v in 𝒱], softmax(-β * free_energy.(Ref(rbm), 𝒱))) > 0.99
end


module hex_ip_constr_examples
using FinEtools
using FinEtools.MeshExportModule.CSV: savecsv
using FinEtoolsDeforLinear
using FinEtoolsDeforLinear.AlgoDeforLinearModule
using FinEtools.AlgoBaseModule: evalconvergencestudy
using FinEtoolsDeforLinear.AlgoDeforLinearModule: linearstatics, exportstresselementwise, exportstress
using Statistics: mean
using LinearAlgebra: Symmetric, cholesky, eigen, svd
using PGFPlotsX

mesh() = (FinEtools.FENodeSetModule.FENodeSet([
	0.0 0.0 0.0;
	1.0 0.0 0.0;
	1.0 1.0 0.0;
	0.0 1.0 0.0;
	0.0 0.0 1.0;
	1.0 0.0 1.0;
	1.0 1.0 1.0
	0.0 1.0 1.0;
	]), FinEtools.FESetModule.FESetH8(reshape([1, 2, 3, 4, 5, 6, 7, 8], 1, 8)))

function hex_ip_constr_par()
	fens,fes = mesh()

	# integration_rule = TrapezoidalRule(3)
	integration_rule = GaussRule(3, 2)
	pc = integration_rule.param_coords;
	w  =  integration_rule.weights
	npts = integration_rule.npts;

	Q = fill(0.0, npts, 24)
	for j in 1:npts
		gradNpar = bfundpar(fes, vec(pc[j,:]))
		k = 1
		for i in 1:size(gradNpar, 1)
			Q[j, k:(k+2)] .= gradNpar[i, :]
			k = k + 3
		end
	end

	decomp = svd(Q)
	@show decomp.S

		true
end # hex_ip_constr_par

function hex_ip_constr_xyz()
	xyzperturbation = [
	0.0767656  -0.983206    -0.14343
	0.45767     0.981479     0.450997
	-0.295854    0.542922     0.321333
	-0.85204    -0.97824     -0.772874
	-0.0539756   0.994907     0.822798
	0.447173    0.528742     0.0445352
	-0.468417    0.00673427   0.953151
	-0.898513   -0.915871     0.933237  ] ./ 5.0
	fens,fes = mesh()
	fens.xyz .+= xyzperturbation

	# integration_rule = TrapezoidalRule(3)
	integration_rule = GaussRule(3, 2)
	pc = integration_rule.param_coords;
	w  =  integration_rule.weights
	npts = integration_rule.npts;

	Q = fill(0.0, npts, 24)
	for j in 1:npts
		gradNpar = bfundpar(fes, vec(pc[j,:]))
		J = (transpose(fens.xyz) * gradNpar)
		gradN = gradNpar / J
		k = 1
		for i in 1:size(gradN, 1)
			Q[j, k:(k+2)] .= gradN[i, :]
			k = k + 3
		end
	end

	decomp = svd(Q)
	@show decomp.S

	true
end # hex_ip_constr_xyz

function allrun()
	println("#####################################################")
	println("# hex_ip_constr_par ")
	hex_ip_constr_par()
	println("#####################################################")
	println("# hex_ip_constr_xyz ")
	hex_ip_constr_xyz()
	return true
end # function allrun

end # module hex_ip_constr_examples


module Quadratures

    using ..Utils

    export Quadrature
    export RiemannQuadratureLeft, RiemannQuadratureRight
    export GaussLegendreQuadrature, LobattoLegendreQuadrature,
           GaussChebyshevQuadrature, LobattoChebyshevQuadrature,
           ClenshawCurtisQuadrature
    export quadrature, weights

    include("quadratures/quadrature.jl")

end


###############################################################################
#Package of functions related to the Butterfly relaxation
#
#Citation:
#Dussault, J. P., Haddou, M., Kadrani, A., & Migot, T. (2017).
#How to Compute a M-stationary point of the MPCC.
###############################################################################

"""
psi(x) : relaxation function
"""
function psi(x :: Union{Float64,Vector},
             r :: Float64,
             s :: Float64,
             t :: Float64)

 return s .+ t.*theta(x .- s, r)
end

function invpsi(x :: Union{Float64,Vector},
                r :: Float64,
                s :: Float64,
                t :: Float64)

 #inverse not defined for t = 0.
 return t != 0 ? s .+ invtheta((x.-s)./t, r) : -Inf
end

function dpsi(x :: Union{Float64,Vector},
              r :: Float64,
              s :: Float64,
              t :: Float64)

 return t .* dtheta(x .- s, r)
end

function ddpsi(x :: Union{Float64,Vector},
               r :: Float64,
               s :: Float64,
               t :: Float64)

 return t .* ddtheta(x .- s,r)
end

"""
phi(x,r,s,t) : évalue la fonction de relaxation
de la contrainte de complémentarité en (yg,yh) in (ncc,ncc)

output : Array of size 1 x ncc
"""
function phi(yg :: Union{Float64,Vector},
             yh :: Union{Float64,Vector},
             r  :: Float64,
             s  :: Float64,
             t  :: Float64)

 return (yg - psi(yh, r, s, t)) .* (yh - psi(yg, r, s, t))
end
"""
phi(x,r,s,t) : évalue la fonction de relaxation
 de la contrainte de complémentarité en x in n+2ncc

output : Array of size 1 x ncc
"""
function phi(x   :: Vector,
             ncc :: Int64,
             r   :: Float64,
             s   :: Float64,
             t   :: Float64)

 n = length(x)-2*ncc

 yg = x[n+1:n+ncc]
 yh = x[n+ncc+1:n+2*ncc]

 return phi(yg, yh, r, s, t)
end

"""
dphi(x,r,s,t) : évalue la fonction de relaxation de la contrainte de complémentarité en (yg,yh) in (ncc,ncc)

output : Array of size 1 x 2ncc
"""
function dphi(yg :: Union{Float64,Vector},
              yh :: Union{Float64,Vector},
              r  :: Float64,
              s  :: Float64,
              t  :: Float64)

 dg = (yh - psi(yg,r,s,t)) - dpsi(yg,r,s,t).*(yg - psi(yh,r,s,t))
 dh = (yg - psi(yh,r,s,t)) - dpsi(yh,r,s,t).*(yh - psi(yg,r,s,t))

 return [dg;dh]
end

"""
dphi(yg,yh,r,s,t) : évalue la fonction de relaxation de la contrainte de complémentarité en x in n+2ncc

output : Array of size 1 x (n +2ncc)
"""
function dphi(x   :: Vector,
              ncc :: Int64,
              r   :: Float64,
              s   :: Float64,
              t   :: Float64)

 n = length(x)-2*ncc

 yg = x[n+1:n+ncc]
 yh = x[n+ncc+1:n+2*ncc]

 return [zeros(n);dphi(yg,yh,r,s,t)]
end

"""
ddphi(x,r,s,t) : évalue la fonction de relaxation de la contrainte de complémentarité en (yg,yh) in (ncc,ncc)

output : Array of size 2ncc x 2ncc
"""
function ddphi(yg :: Union{Float64,Vector},
               yh :: Union{Float64,Vector},
               r  :: Float64,
               s  :: Float64,
               t  :: Float64)

 dgg = -2*dpsi(yg,r,s,t) - ddpsi(yg,r,s,t).*(yg - psi(yh,r,s,t))
 dhh = -2*dpsi(yh,r,s,t) - ddpsi(yh,r,s,t).*(yh - psi(yg,r,s,t))
 dgh = 1 .+ dpsi(yg,r,s,t).*dpsi(yh,r,s,t)

 return dgg, dgh, dhh
end

"""
alpha_max : see ThetaFct.alpha_theta_max

output : Array of size 1 x 2
"""
function alpha_max(x  :: Union{Float64,Vector},
                   dx :: Union{Float64,Vector},
                   y  :: Union{Float64,Vector},
                   dy :: Union{Float64,Vector},
                   r  :: Float64,
                   s  :: Float64,
                   t  :: Float64)

 return alpha_theta_max(x, dx, y, dy, r, s, t)
end


import Clustering
import Distances
import LinearAlgebra: norm

using MLJBase
using MLJClusteringInterface
using Random:seed!
using Test

const Dist = Distances

seed!(132442)
X, y = @load_crabs

####
#### KMEANS
####

@testset "Kmeans" begin
    barekm = KMeans()
    fitresult, cache, report = fit(barekm, 1, X)
    R = matrix(transform(barekm, fitresult, X))
    X_array = matrix(X)
    # distance from first point to second center
    @test R[1, 2] ≈ norm(view(X_array, 1, :) .- view(fitresult[1], :, 2))^2
    @test R[10, 3] ≈ norm(view(X_array, 10, :) .- view(fitresult[1], :, 3))^2

    p = predict(barekm, fitresult, X)
    @test argmin(R[1, :]) == p[1]
    @test argmin(R[10, :]) == p[10]

end

####
#### KMEDOIDS
####

@testset "Kmedoids" begin
    barekm = KMedoids()
    fitresult, cache, report = fit(barekm, 1, X)
    X_array = matrix(X)
    R = matrix(transform(barekm, fitresult, X))
    @test R[1, 2] ≈ Dist.evaluate(
        barekm.metric, view(X_array, 1, :), view(fitresult[1], :, 2)
    )
    @test R[10, 3] ≈ Dist.evaluate(
        barekm.metric, view(X_array, 10, :), view(fitresult[1], :, 3)
    )
    p = predict(barekm, fitresult, X)
    @test all(report.assignments .== p)
end

####
#### HierarchicalClustering
####

@testset "HierarchicalClustering" begin
    barehc = HierarchicalClustering(linkage = :complete)
    hc = transform(barehc, X)
    @test mean(int(y) .== hc(k = 2)) > .5 # that's a somewhat silly test...
end