using DiscreteFiltering
using LinearAlgebra
using Plots
using ProgressLogging
using OrdinaryDiffEq: RK4

## Parameters

# Domain
a = 0.0
b = 2π
domain = PeriodicIntervalDomain(a, b)

# Exact solutions (advection equation)
u(x, t) = sin(x - t) + 3 / 5 * cos(5(x - t)) + 7 / 25 * sin(20(x - 1 - t))
u_int(x, t) = -cos(x - t) + 3 / 25 * sin(5(x - t)) - 7 / 25 / 20 * cos(20(x - 1 - t))

## Time
T = 1.0

# Ode solver tolerances
# tols = (;)
# tols = (; abstol = 1e-6, reltol = 1e-4)
# tols = (; abstol = 1e-7, reltol = 1e-5)
tols = (; abstol = 1e-10, reltol = 1e-8)
degmax = 1000
subspacedim = 1000
λ = 1e-8

# Amplitudes
c = [
    [1],
    [1, 3],
    [1, 3 / 5, 1 / 25],
    rand(5),
]

# Frequencies
ω = [
    [0],
    [2, 3],
    [1, 5, 20],
    (rand(1:10, 5)),
]

# Phase-shifts
ϕ = [
    [π / 2],
    [0, 0],
    [0, 0, 20],
    2π * rand(5),
]

# Number of time steps for fitting C̄
nₜ = 500
t = LinRange(0, T, nₜ)
UU = sum_of_sines.([domain], c, ω, ϕ)
u₀_list = [UU[1] for UU ∈ UU]
U₀_list = [UU[2] for UU ∈ UU]

# Number of mesh points
nrefine = 10
MM = floor.(Int, 10 .^ LinRange(1, 4, nrefine))
NN = MM
# NN = floor.(Int, 5 // 4 .* MM)
# λλ = [0, 0, 0, 0, 1e-1]

# Errors
err = zeros(nrefine)
err_bar = zeros(nrefine)
err_allbar = zeros(nrefine)

## Solve
@time for i = 1:nrefine
    # λ = λλ[i]
    λ = 1e-8
    λ_ridge = 1e-8
    # M = 100
    M = MM[i]
    # N = NN[i]
    N = 1000 # NN[i]
    # M = N

    println("Solving for M = $M and N = $N")
    # Discretization
    x = discretize(domain, M)
    ξ = discretize(domain, N)
    Δx = (b - a) / M

    # Filter
    h₀ = 1.0Δx
    # h₀ = (b - a) / 200
    # h₀ = Δx / 2
    h(x) = h₀ * (1 - 1 / 2 * cos(x))
    # h(x) = 0.49999Δx
    filter = TopHatFilter(h)
    # filter = ConvolutionalFilter(x -> 2h(x), x -> abs(x) ≤ h(x) ? 1/ 2h(x) : zero(x))
    # h(x) = 2h₀
    # σ = h₀/2
    # filter = GaussianFilter(h, σ)

    C̄ = fit_Cbar(domain, filter, u₀_list, U₀_list, M, t; λ, method = :ridge)

    # Equations
    equation = AdvectionEquation(domain, IdentityFilter())
    equation_filtered = AdvectionEquation(domain, filter)

    # Solve discretized problem
    sol = solve(
        equation,
        ξ -> u(ξ, 0.0),
        (0.0, T),
        M,
        M;
        method = "discretizefirst",
        subspacedim,
        tols...,
        degmax,
        λ,
    )

    # Solve filtered-then-discretized problem
    # sol_bar = solve(
    #     equation_filtered,
    #     ξ -> u(ξ, 0.0),
    #     (0.0, T),
    #     M,
    #     M;
    #     method = "filterfirst",
    #     subspacedim,
    #     tols...,
    #     degmax,
    #     λ,
    # )

    # Solve discretized-then-filtered problem
    sol_allbar = solve(
        equation_filtered,
        ξ -> u(ξ, 0.0),
        (0.0, T),
        M,
        N;
        # method = "discretizefirst",
        # method = "discretizefirst-without-R",
        method = "discretizefirst-fit-Cbar",
        # solver = RK4(),
        subspacedim,
        tols...,
        degmax,
        λ,
        λ_ridge,
        C̄,
    )

    # Exact filtered solution
    ū = apply_filter_int(ξ -> u_int(ξ, T), filter, domain)
    # ū = apply_filter(ξ -> u(ξ, T), filter, domain)

    # Relative error
    # u_exact = u.(ξ, T)
    u_exact = u.(x, T)
    ū_exact = ū.(x)
    # err[i] = norm(sol(T) - u_exact) / norm(u_exact)
    err[i] = norm(sol(T) - ū_exact) / norm(ū_exact)
    # err_bar[i] = norm(sol_bar(T) - ū.(ξ)) / norm(ū.(ξ))
    # err_bar[i] = norm(sol_bar(T) - ū.(x)) / norm(ū.(x))
    err_allbar[i] = norm(sol_allbar(T) - ū_exact) / norm(ū_exact)
    # println("λ = $λ, err = $(err_allbar[i])")
    println("λ = $λ, err_bar = $(err_bar[i]), err_allbar = $(err_allbar[i])")
end

## Set PGFPlotsX plotting backend to export to tikz
pgfplotsx()

## In terminal plots
unicodeplots()

## Default backend
gr()

##
p = plot(
    xaxis = :log10,
    yaxis = :log10,
    minorgrid = true,
    # size = (400, 300),
    legend = :bottomleft,
    # xlims = (10, 10^4),
    ylims = (1e-6, 1e-0),
    xticks = 10 .^ (1:4),
    xlabel = "M",
    # title = "N = 1000"
);
plot!(p, MM, err_allbar; label = "Data driven Cbar, HF", marker = :d, color = 2);
# plot!(p, MM, err_allbar_lf; label = "Data driven Cbar, LF", marker = :d, color = 1);
plot!(p, MM, err; label = "Without filter, HF", marker = :c, linestyle = :dash, color = 2);
# plot!(p, MM, err_lf; label = "Without filter, LF", marker = :c, linestyle = :dash, color = 1);
# display(p)
# output = "output/advection/Cbar.tikz"
output = "output/advection/Cbar.pdf"
savefig(p, output)


# SteinKernel

abstract type SteinKernel end

## METHODS TO IMPLEMENT ##
function k(kernel::SteinKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    error("Must implement the method k")
end

function gradxk(kernel::SteinKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    error("Must implement the method gradxk")
end

function gradxyk(kernel::SteinKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    error("Must implement the method gradxyk")
end

# Computes \grad_y k(x, y) = \grad_x k(y, x)
function gradyk(kernel::SteinKernel, x::Array{Float64, 1}, y::Array{Float64, 1})
    gradxk(kernel, y, x)
end

# This is the Chris Oates Langevin Stein kernel
function k0(kernel::SteinKernel,
            x::Array{Float64, 1},
            y::Array{Float64, 1},
            gradlogpx::Array{Float64, 1},
            gradlogpy::Array{Float64, 1})
    k(kernel, x, y) * dot(gradlogpx, gradlogpy) +
        dot(gradlogpx, gradyk(kernel, x, y)) +
        dot(gradlogpy, gradxk(kernel, x, y)) +
        gradxyk(kernel, x, y)
end


# groupby aggregators

#
# Aggregators are structs with fit! and value methods
#
# fit!(aggregator, row) updates aggregator with row
# value(aggregator) returns current value
#

export fit!,
    value,
    Sum,
    Count,
    Avg

abstract type AbstractAggregator end

mutable struct Sum{F} <: AbstractAggregator
    valfunc::F
    val::Float64
    Sum(valfunc::F, val = 0.0) where {F <: Function} = new{F}(valfunc, val)
end

Base.eltype(::Type{Sum{F}}) where {F} = Float64

"""
    Sum(column::Symbol, val = 0.0)

Aggregator for sum of `column`.
"""
function Sum(column::Symbol, val = 0.0)
    valfunc = s -> s[column]
    Sum(valfunc, val)
end

function fit!(s::Sum, x)
    s.val += s.valfunc(x)
end

value(s::Sum) = s.val

"""
    Count()

Aggregator for count.
"""
mutable struct Count <: AbstractAggregator
    val::Int64
    Count() = new(0)
end

Base.eltype(::Type{Count}) = Int64

function fit!(c::Count, x)
    c.val += 1
end

value(c::Count) = c.val

mutable struct Avg{F} <: AbstractAggregator
    keyfunc::F
    total::Float64
    n::Int64
    Avg(keyfunc::F) where {F <: Function} = new{F}(keyfunc, 0.0, 0)
end

Base.eltype(::Type{Avg{F}}) where {F} = Float64

"""
    Avg(column::Symbol)

Aggregator for average of `column`.
"""
function Avg(column::Symbol)
    valfunc = s -> s[column]
    Avg(valfunc)
end

function fit!(a::Avg, x)
    a.total += a.keyfunc(x)
    a.n += 1
end

value(a::Avg) = a.total / a.n

# d is collection of key -> Aggregator pairs
# fit!(d, row) fit!s all aggregators with row
function fit!(d, x)
    for (k, v) in d
        fit!(v, x)
    end
end


#### Teacup GIF
## Inspired by an example of Lorez attractor: http://docs.juliaplots.org/latest/
## FFmpeg has to be installed prior to running the code

## Preparation

# Activate project environment
import Pkg; Pkg.activate(".")

# Load required packages
using DifferentialEquations
using Plots

# Load Hastings & Powell's model
include("HPmodel.jl")

## Produce the GIF

# define the model
mutable struct TeaCup
    dt; a1; a2; b1; b2; d1; d2; x; y; z
end

function step!(t::TeaCup)
    dx = t.x*(1-t.x) - t.a1*t.x*t.y/(1 + t.b1*t.x)     ; t.x += t.dt * dx
    dy = t.a1*t.x*t.y/(1 + t.b1*t.x) - t.a2*t.y*t.z/(1 + t.b2*t.y) - t.d1*t.y ; t.y += t.dt * dy
    dz = t.a2*t.y*t.z/(1 + t.b2*t.y) - t.d2*t.z   ; t.z += t.dt * dz
end

# solution from initial values
TeaCup_sol = TeaCup((dt = 0.05, a1 = 5.0, a2 = 0.1, b1 = 3.0, b2 = 2.0, d1 = 0.4, d2 = 0.01, x = 0.76, y = 0.16, z = 9.9)...)

# initialize a 3D plot with 1 empty series
# requires same trick to flip axis Y as for figure2D
plt = plot3d(1, xlabel = "x", xlim = (0.0, 1.0), xticks = 0.0:0.5:1.0, xrotation = 30.0,
                ylabel = "y", ylim = (0.0, 0.5), yticks = (0.0:0.25:0.5, string.(0.5:-0.25:0.0)), yrotation = -10.0,
                zlabel = "z", zlim = (7.5, 10.5), zticks = 7.5:1.5:10.5, zrotation = 0.0,
                linewidth = 0.1,
                marker = :none, legend = :none, size = (600, 600), c = :black,
                right_margin = 5Plots.mm, left_margin = 2Plots.mm,
                )

# build an animated gif by pushing new points to the plot, saving every 20th frame
anim = @animate for i = 1:30000
    step!(TeaCup_sol)
    push!(plt, TeaCup_sol.x, (0.5 - TeaCup_sol.y), TeaCup_sol.z)
end every 20
gif(anim, joinpath("..", "article", "figures", "figure2D.gif"), fps = 40)


using Plexifications
using Base.Test: @test, @test_throws

@test_throws ErrorException begin
    inv(Prplex(1, 1))
end

@test begin
    iszerodivisor(Prplex(1,1))
end

@test begin
    !iszerodivisor(Prplex(0,1))
end

@test begin
    length(unreal(1)) == 0
end

@test begin
    a = Plexifications.randomBigFloat()
    isreal(Prplex(a))
end

@test begin
    a = 1
    !isreal(Prplex(a, a))
end

@test begin
    z = random(Prplex{Int})
    z == +(z)
end

@test begin
    a = Plexifications.randomBigFloat()
    real(Prplex(a)) == a
end

@test begin
    io = IOBuffer()
    show(io, Prplex(1, 2))
    l = String(take!(io))
    r = "[1: 1, R: 2]"
    l == r
end

@test begin
    x = random(Prplex{BigFloat})
    y = random(Prplex{BigFloat})
    x != y
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(Prplex{BigInt})
    l = a + z
    r = z + a
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    z = random(Prplex{BigInt})
    l = a - z
    r = -(z - a)
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    l = x + y
    r = y + x
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    l = x * y
    r = y * x
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    l = x - y
    r = -(y - x)
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    z = random(Prplex{BigInt})
    l = (x + y) + z
    r = x + (y + z)
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    z = random(Prplex{BigInt})
    l = (x * y) * z
    r = x * (y * z)
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    barR(barR(x)) == x
end

@test begin
    x = random(Prplex{Rational{BigInt}})
    inv(inv(x)) == x
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    l = barR(x * y)
    r = barR(y) * barR(x)
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    l = barR(x + y)
    r = barR(x) + barR(y)
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    l = barR(x - y)
    r = barR(x) - barR(y)
    l == r
end

@test begin
    x = random(Prplex{Rational{BigInt}})
    y = random(Prplex{Rational{BigInt}})
    l = inv(x * y)
    r = inv(y) * inv(x)
    l == r
end

@test begin
    x = random(Prplex{Rational{BigInt}})
    y = random(Prplex{Rational{BigInt}})
    l = inv(x / y)
    r = y / x
    l == r
end

@test begin
    x = random(Prplex{Rational{BigInt}})
    y = random(Prplex{Rational{BigInt}})
    l = inv(x \ y)
    r = y \ x
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    z = random(Prplex{BigInt})
    l = x * (y + z)
    r = (x * y) + (x * z)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    l = a * (x + y)
    r = (a * x) + (a * y)
    l == r
end

@test begin
    a = Plexifications.randomBigInt()
    x = random(Prplex{BigInt})
    l = a * x
    r = x * a
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    z = random(Prplex{BigInt})
    l = (x + y) * z
    r = (x * z) + (y * z)
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    z = random(Prplex{BigInt})
    l = x * (y - z)
    r = (x * y) - (x * z)
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    z = random(Prplex{BigInt})
    l = (x - y) * z
    r = (x * z) - (y * z)
    l == r
end

@test begin
    x = random(Prplex{BigInt})
    abs2(x) > big(0) || abs2(x) < big(0) || abs2(x) == big(0)
end

@test begin
    x = random(Prplex{BigInt})
    y = random(Prplex{BigInt})
    l = abs2(x * y)
    r = abs2(x) * abs2(y)
    l == r
end

@test_throws ErrorException begin
    inv(Prplex(0, 0.0))
end

@test_throws ErrorException begin
    random(Prplex{Int}) / 0
end

@test_throws ErrorException begin
    0 \ random(Prplex{Int})
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(Prplex{Rational{BigInt}})
    l = z / a
    r = a \ z
    l == r
end

@test begin
    a = Plexifications.randomBigRational()
    z = random(Prplex{Rational{BigInt}})
    l = a / z
    r = z \ a
    l == r
end

@test begin
    z = random(Prplex{BigInt})
    z == z + zero(Prplex{BigInt})
end

@test begin
    z = random(Prplex{BigInt})
    z == z * one(Prplex{BigInt})
end

@test begin
    z = random(Prplex{BigInt})
    z == z + zero(z)
end

@test begin
    z = random(Prplex{BigInt})
    z == z * one(z)
end


# This file is pat of RLinearAlgebra.jl
# 1. Specifies type
# 2. Implements sample function
# 3. Exports Type

# using Random

"""
    LinSysVecRowDetermCyclic <: LinSysVecRowSelect

An immutable structure without any fields. Specifies deterministic cycling through the
equations of a linear system.
"""
struct LinSysVecRowDetermCyclic <: LinSysVecRowSelect end

# Common sample interface for linear systems
function sample(
    type::LinSysVecRowDetermCyclic,
    A::AbstractArray,
    b::AbstractVector,
    x::AbstractVector,
    iter::Int64
)

    eqn_ind = mod(iter, 1:length(b))
    return A[eqn_ind, :], b[eqn_ind]
end

#export LinSysVecRowDetermCyclic


using MultipleViewGeometry, Base.Test
using MultipleViewGeometry.ModuleCostFunction
using MultipleViewGeometry.ModuleTypes
using BenchmarkTools, Compat
using StaticArrays

𝒳 = [Point3DH(x,y,z,1.0)
                        for x=-1000:5:1000 for y=-1000:5:1000 for z=1:-5:-1000]

# Intrinsic and extrinsic parameters of camera one.
𝐊₁ = eye(3)
𝐑₁ = eye(3)
𝐭₁ = [0.0, 0.0, 0.0]

# Intrinsic and extrinsic parameters of camera two.
𝐊₂ = eye(3)
𝐑₂ = eye(3)
𝐭₂ = [100.0, 2.0, -100.0]

# Camera projection matrices.
𝐏₁ = construct(ProjectionMatrix(),𝐊₁,𝐑₁,𝐭₁)
𝐏₂ = construct(ProjectionMatrix(),𝐊₂,𝐑₂,𝐭₂)

# Set of corresponding points.
ℳ = project(Pinhole(),𝐏₁,𝒳)
ℳʹ = project(Pinhole(),𝐏₂,𝒳)

@time ℳ = project(Pinhole(),𝐏₁,𝒳)
@time ℳʹ = project(Pinhole(),𝐏₂,𝒳)

#@btime  project(Pinhole(),𝐏₂,𝒳)

#@time ℳ = project(Pinhole(),SMatrix{3,4}(𝐏₁),𝒳)
#@time ℳʹ = project(Pinhole(),SMatrix{3,4}(𝐏₂),𝒳)

#@btime project(Pinhole(),SMatrix{3,4}(𝐏₂),𝒳)


using JLD, Images

function main(dim)
    points = load("./out/20k.jld", "points")
    img = zeros(RGB,dim,dim)
    for c in points
        z = c
        while abs(z) < 2
            x = trunc(Int,(real(z) + 2)*dim/3); y = trunc(Int,(imag(z) + 1.5)*dim/3)
            0<x<dim>y>0 && (img[x,y] += RGB(0.0005,0,0))
            z = z^2 + c
        end
    end
    return map(clamp01nan, img)
end

@time save("./images/t1.png", main(1500))


using LinearAlgebra

struct ValueOne; end
ValueOne()

# Compute X <- a X + b I.
function matfun_axpby!(X,a,b,Y::UniformScaling)
    m,n=size(X)
    if ~(a isa ValueOne)
        rmul!(X,a)
    end
    @inbounds for i=1:n
        X[i,i]+=(b isa ValueOne) ? 1 : b
    end
end

@inline function exp_ps_m3(A)
    T=promote_type(eltype(A),Float64)
    A_copy=similar(A,T); A_copy .= A;
    return exp_ps_m3!(A_copy)
end

@inline function exp_ps_m3!(A)
    T=promote_type(eltype(A),Float64) # Make it work for many 'bigger' types (matrices and scalars)
    # max_memslots=5
    n=size(A,1)
    # The first slots are precomputed nodes [:A]
    memslots2 = similar(A,T)
    memslots3 = similar(A,T)
    memslots4 = similar(A,T)
    memslots5 = similar(A,T)
    # Assign precomputed nodes memslots
    memslots1=A # overwrite A
    # Uniform scaling is exploited.
    # No matrix I explicitly allocated.
    value_one=ValueOne()
    # Computation order: A2 A3 P1 C0 P0
    # Computing A2 with operation: mult
    mul!(memslots2,memslots1,memslots1)
    # Computing A3 with operation: mult
    mul!(memslots3,memslots1,memslots2)
    # Computing P1 = x*I+x*A+x*A2+x*A3
    coeff1=0.16666666666666666
    coeff2=0.041666666666666664
    coeff3=0.008333333333333333
    coeff4=0.001388888888888889
    memslots4 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots3
    mul!(memslots4,true,I*coeff1,true,true)
    # Computing C0 with operation: mult
    mul!(memslots5,memslots4,memslots3)
    # Deallocating P1 in slot 4
    # Deallocating A3 in slot 3
    # Computing P0 = x*I+x*A+x*A2+x*C0
    coeff1=1.0
    coeff2=1.0
    coeff3=0.5
    coeff4=1.0
    # Smart lincomb recycle A
    memslots1 .= coeff2.*memslots1 .+ coeff3.*memslots2 .+ coeff4.*memslots5
    mul!(memslots1,true,I*coeff1,true,true)
    # Deallocating A2 in slot 2
    # Deallocating C0 in slot 5
    return memslots1 # Returning P0
end


"""
    refine_triangulation(triang_vertices, triang_simplex_indices, split_indices, k)

Refine a triangulation with explicitly providing the splitting rules. In this function,
these are computed internally, using a splitting factor of k.
"""
function refine_triangulation(triang_vertices, triang_simplex_indices, split_indices, k)

    if length(split_indices) == 0
        untouched_indices = collect(1:size(triang_simplex_indices, 1))
        return triang_vertices, triang_simplex_indices, untouched_indices
    end

    # The number of simplices to split
    n_split_simplices = length(split_indices)

    # The dimension of the space
    E = size(triang_vertices, 2)

    # Generic rules for splitting a simplex
    splitting_rules = simplicial_subdivision(k, E)

    # Rules for forming the strictly new vertices of the subtriangulation
    rules = splitting_rules[1]

    # Array where each row represents one of the new simplices in the splitted simplex
    subtriangulation = splitting_rules[2]

    # How many new vertices are created each split?
    n_newvertices_eachsplit = size(rules, 1)

    # We need an array that can accomodate all of them. Each row in this
    # array will be a new vertex. Stacks of n_newvertices_eachsplit * E arrays.
    # We have as many as we have simplices to split.
    new_vertices = zeros(n_newvertices_eachsplit * n_split_simplices, E)

    # Fill the array by looping over all simplices that we want to split
    for i = 1:n_split_simplices
        # Figure out what the row indices corresponding to the ith simplex
        # must be. Marks the beginning of each of the simplex stacks in new_vertices
        ind = n_newvertices_eachsplit * (i -1)

        # Index of the simplex we need to split
        simplex_idx = split_indices[i]

        # Get the vertices of the simplex currently being splitted. Each of the
        # n_newvertices_eachsplit new vertices will be a linear combination
        # of these vertices. Each row is a vertex.
        vertices = triang_vertices[triang_simplex_indices[simplex_idx, :], :]

        # Generate the strictly new vertices for each sub
        for j = 1:n_newvertices_eachsplit
            # The index for a particular new vertex of the ith new simplex
            ind_newvertex = ind + j

            # Compute the jth vertex of the subtriangulation of the ith simplex

            # Go the jth new subsimplex of the ith splitted simplex. The entries of this
            # vector
            jth_subsimplex = rules[j, :]

            # Pick the corresponding original vertices with indices contained in rules[j, :]
            original_vertices = vertices[rules[j, :], :]

            new_vertices[ind_newvertex, :] = sum(original_vertices, dims=1) ./ k

        end
    end

    # Find the unique new vertices
    new_vertices_noreps = unique(new_vertices, dims=1)
    Ind = Vector{Int}(undef, size(new_vertices, 1))
    for i = 1:size(new_vertices_noreps, 1)
        for j = 1:size(new_vertices, 1)
            if new_vertices_noreps[i, :] == new_vertices[j, :]
                Ind[j] = i
            end
        end
    end

    # Combine old and newly introduced vertices
    num_vertices_beforesplit = size(triang_vertices, 1)
    triang_vertices = vcat(triang_vertices, new_vertices_noreps)

    # Update the Ind array, so that we start at the new vertices
    Ind = Ind .+ num_vertices_beforesplit
    num_simplices_each_split = size(subtriangulation, 1)

    # The subsimplices formed by the splitting. Each row contains E + 1 indices referencing
    # the vertices furnishing that particular subsimplex (now found in the updated
    # triang_vertices array).
    newtriangulation = Array{Float64, 2}(undef, num_simplices_each_split * n_split_simplices, E + 1)

    # For each simplex that we need to split,
    for i = 1:n_split_simplices
        # The beginning of the stack we need to fill.
        index = num_simplices_each_split * (i - 1)

        # Figure out what the row indices corresponding to the ith simplex
        # must be. Marks the beginning of each of the simplex stacks in new_vertices
        ind = n_newvertices_eachsplit * (i -1)

        # Index of the simplex we need to split
        simplex_idx = split_indices[i]

        # Pick the indices of the original vertices. Should be a column vector.
        inds_original_vertices = triang_simplex_indices[simplex_idx, :]

        # Indices of the new vertices. Should be a column vector
        inds_new_vertices = Ind[(ind + 1):(ind + n_newvertices_eachsplit)]
        inds_all_vertices = vcat(inds_original_vertices, inds_new_vertices)

        # Go through each of the subsimplices formed by splitting that simplex
        for j = 1:num_simplices_each_split
            newtriangulation[index + j, :] = inds_all_vertices[subtriangulation[j, :]]
        end
    end

    # Indices of the simplices that are not split
    untouched_indices = setdiff(1:size(triang_simplex_indices, 1), split_indices)

    triang_simplex_indices = round.(Int, vcat(triang_simplex_indices[untouched_indices, :],
                                  newtriangulation))

    return triang_vertices, triang_simplex_indices, untouched_indices
end


function generate_bloobs(n_features = 2, n_centers = 3, n_samples = 100)
    centers = @parallel (hcat) for i = 1:n_features
        rand(Uniform(-10, 10), n_centers)
    end

    cluster_std = @parallel (hcat) for i = 1:n_features
      rand(Uniform(0.5, 4), n_centers)
    end

    n_samples_per_center = ones(Integer, n_centers) * int(n_samples / n_centers)

    for i = 1:(n_centers % n_samples)
        n_samples_per_center[i] += 1
    end

    X = @parallel (hcat) for i = 1:n_centers
      rand(MvNormal(ones(n_features) .* centers[i], eye(n_features) .* cluster_std[i]), n_samples_per_center[i])
    end

    X = X';

    Y = @parallel (vcat) for i = 1:n_centers
      y = ones(n_samples_per_center[i]) * i
    end

    ids = [1:size(X)[1]]
    shuffle!(ids)

    X = X[ids,:]
    Y = Y[ids];

    return (X, Y)
end

function generateBarsDataset(img_size::Int, noise_level::Float64, num_per_mixture::Array{Float64}, num_group::Int, num_data::Int)

  numdim   = img_size^2;
  numbars  = img_size*2

  mix_theta = zeros(numdim, numbars)

  # add horizontal bars
  for i in 1:img_size
    img = ones(img_size, img_size) .* noise_level / img_size^2;
    img[:,i] = img[:,i] + 1/img_size;
    mix_theta[:,i] = img[:] / sum(img[:]);
  end

  # add vertical bars
  for i = 1:img_size
    img = ones(img_size, img_size) .* noise_level / img_size^2;
    img[i,:] = img[i,:] + 1/img_size;
    mix_theta[:,img_size+i] = img[:]/sum(img[:]);
  end

  # generate samples
  # Note: strange behavior of cumsum!!!
  cumbar = cumsum(num_per_mixture, 2)

  samples = Vector{Vector{Int}}(num_group)

  for g in 1:num_group

      # random number of mixture components
      n_bars = 1 + sum(rand() .> cumbar)
      k = randperm(numbars)[1:n_bars]

      # get weigths
      theta = mean(mix_theta[:,k], 2)

      # get samples (e.g. words)
      samples[g] = vec( 1 + sum(repmat(rand(num_data)', numdim, 1) .> repmat(cumsum(theta, 1), 1, num_data), 1) )
  end

  # return samples and bars
  return (samples, mix_theta)

end


#=
   Car buyer - Neapolital

   From the Netica model Car_Buyer Neapolital.dne
   """
   This decision network is from Neapolitan90 (p.373), which is a simpler version of the car
   buyer example of Howard62.  Eliminating the rationale given there of how the numbers are
   arrived at, we have the following story:

   Joe is going to buy a used car, which could be good with probability 0.8 or a lemon with
   probability 0.2.  After accounting for repairs, Joe's profit will be $60 if the car is
   good, and $-100 if it is bad.  Before buying the car he has the option of having one test
   or two tests done on it.  The first test costs $9, and both together cost $13.  The first
   test has a 90% chance of returning positive if the car is good, and a 40% chance if it's
   a lemon.  If the first test returns positive, then the second test has a 88.89% chance of
   returning positive if the car is good, and a 33.33% chance if it's a lemon.  If the first
   test returns negative, then the second test has a 100% chance of returning positive if
   the car is good, and a 44.44% chance if it's a lemon.

   Joe has 2 decisions to make: whether to do the tests, and whether to buy the car.  These
   are represented by the "Do Tests?" and "Buy It?" decision nodes.  The outcome of the tests
   are given by nodes "First Test" and "Second Test".  The costs of the tests are represented
   by utility node U, and the profits after repairs (not including test costs) by utility
   node V.

   When Joe decides whether to do the tests, he doesn't know the value of any of these
   variables, so there are no links entering the "Do Tests?" node.  When he decides whether
   to buy, he will know the outcome of both tests (the outcomes may be "not done"), and so
   there are links from those two nodes to "Buy It?".  He will also know the value of
   "Do Tests?" since he has already made that decision, so you could put a link from that
   node to "Buy It?", but it is not necessary since it is a no-forgetting link and there is
   already a directed path from "Do Tests?" to "Buy It?".
   """

   Distributions of variable condition
   lemon      =>    5813  (0.581300)
   good       =>    4187  (0.418700)

   Distributions of variable do_tests
   both       =>   10000  (1.000000)

   Distributions of variable first_test
   negative   =>   10000  (1.000000)

   Distributions of variable second_test
   notDone    =>    4187  (0.418700)
   negative   =>    3276  (0.327600)
   positive   =>    2537  (0.253700)

   Distributions of variable buy_it
   dont_buy   =>    5813  (0.581300)
   buy        =>    4187  (0.418700)

   Distributions of variable v (num:0)
   0.00000 =>    5813  (0.581300)
   60.00000 =>    4187  (0.418700)

   Distributions of variable u (num:0)
  -13.00000 =>   10000  (1.000000)

   do_tests, first_test, second_test, buy_it, v, u, v+y:
   Distributions of variable (num:0)
   [3, 3, 1, 1, 60, -13, 47]	=>	4187 (0.4187)
   [3, 3, 3, 2, 0, -13, -13]	=>	3276 (0.3276)
   [3, 3, 2, 2, 0, -13, -13]	=>	2537 (0.2537)

  Cf ~/blog/car_buyer.blog
     ~/webppl/car_buyer.wppl

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function car_buyer()
    good = 1
    lemon = 2
    # conditions = [good, lemon]

    none = 1
    first = 2
    both = 3
    # doTests = [none", first", both"]

    notDone = 1
    positive = 2
    negative = 3
    # tests = [notDone, positive, negative]

    buy = 1
    dont_buy = 2

    condition ~ Categorical(simplex([80,20])) # [good, lemon]
    do_tests ~ Categorical(simplex([28,26.2,22.7333])) # [none,first,both]

    # [notDone,positive,negative]
    first_test ~
        (condition==good && do_tests==none)  ? Categorical(simplex([100,0,0])) :
        (condition==good && do_tests==first) ? Categorical(simplex([0,90,10])) :
        (condition==good && do_tests==both)  ? Categorical(simplex([0,90,10])) :
        (condition==lemon&& do_tests==none)  ? Categorical(simplex([100,0,0])) :
        (condition==lemon&& do_tests==first) ? Categorical(simplex([0,40,60])) :
        (condition==lemon&& do_tests==both)  ? Categorical(simplex([0,40,60])) : Dirac(0)

    second_test ~ (first_test==notDone && condition==good && do_tests==none)    ? Categorical(simplex([100, 0, 0])) :
        (first_test==notDone && condition==lemon && do_tests==none)   ? Categorical(simplex([100, 0, 0])) :

        (first_test==positive && condition==good && do_tests==first)  ? Categorical(simplex([100, 0, 0])) :
        (first_test==positive && condition==good && do_tests==both)   ? Categorical(simplex([0, 88.889, 11.111])) :

        (first_test==positive && condition==lemon && do_tests==first) ? Categorical(simplex([100, 0, 0])) :
        (first_test==positive && condition==lemon && do_tests==both)  ? Categorical(simplex([100, 33.3333, 66.6667])) :

        (first_test==negative && condition==good && do_tests==first)  ? Categorical(simplex([100, 0, 0])) :
        (first_test==negative && condition==good && do_tests==both)   ? Categorical(simplex([100, 0, 0])) :

        (first_test==negative && condition==lemon && do_tests==first) ? Categorical(simplex([100, 0, 0])) :
        (first_test==negative && condition==lemon && do_tests==both)  ? Categorical(simplex([0, 44.4444, 55.5556])) : Dirac(0)

    buy_it ~ Dirac(
            (do_tests == first && second_test == notDone && first_test == negative)
            ||
            (do_tests == both && second_test == positive && first_test == negative)
            ||
            (do_tests == both && second_test == negative && first_test == positive)
            ||
            (do_tests == both && second_test == negative && first_test == negative) ? dont_buy : buy)

    v ~ Dirac(
        (condition==good  && buy_it==buy)      ? 60   :
        (condition==good  && buy_it==dont_buy) ? 0    :
        (condition==lemon && buy_it==buy)      ? -100 :
        (condition==lemon && buy_it==dont_buy) ?  0   : -1
        )

    u ~ Dirac(
        do_tests == none  ?  0  :
        do_tests == first ? -9  :
        do_tests == both  ? -13 : -1
    )

    true ~ Dirac(do_tests == both)
    true ~ Dirac(first_test == negative)
    # true ~ Dirac(second_test == negative)

    # true ~ Dirac(do_tests == first)
    # true ~ Dirac(first_test == negative)
    # true ~ Dirac(buy_it==buy)
    # true ~ Dirac(v>0)
    return [do_tests,first_test,second_test,buy_it,v,u,v+u]
end

model = car_buyer()

num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
# chns = sample(model, PG(15), 10_000)
chns = sample(model, SMC(10_000), 10_000)
# chns = sample(model, IS(), 10_000)

# chns = sample(model, NUTS(), 10_000)
# chns = sample(model, HMC(0.1,10), 10_000)

display(chns)
# display(plot(chns))

println("\ngood = 1 lemon = 2")
show_var_dist_pct(chns, :condition,["good","lemon"])

println("\nnone = 1 first = 2  both = 3")
show_var_dist_pct(chns, :do_tests,["none","first","both"])

println("\nnotDone = 1 positive = 2 negative = 3")
show_var_dist_pct(chns, :first_test,["notDone","positive","negative"])

show_var_dist_pct(chns, :second_test,["notDone","positive","negative"])

println("\nbuy = 1 dont_buy = 2\n")
show_var_dist_pct(chns, :buy_it,["buy","dont_buy"])

println()
show_var_dist_pct(chns, :v)

println()
show_var_dist_pct(chns, :u)

println()
chains_params = Turing.MCMCChains.get_sections(chns, :parameters)
genq = generated_quantities(model, chains_params)
println("do_tests, first_test, second_test, buy_it, v, u, v+y:")
show_var_dist_pct(genq)


export TransposeOperator

mutable struct TransposeOperator{T<:Number,B<:Operator} <: Operator{T}
    op::B
end

TransposeOperator(B::Operator{T}) where {T<:Number}=TransposeOperator{T,typeof(B)}(B)

convert(::Type{Operator{T}},A::TransposeOperator) where {T}=TransposeOperator(convert(Operator{T},A.op))

domainspace(P::TransposeOperator)=rangespace(P.op)
rangespace(P::TransposeOperator)=domainspace(P.op)

domain(P::TransposeOperator)=domain(P.op)

bandwidths(P::TransposeOperator) = reverse(bandwidths(P.op))

getindex(P::TransposeOperator,k::Integer,j::Integer) =
    P.op[j,k]

function BandedMatrix(S::SubOperator{T,TO}) where {T,TO<:TransposeOperator}
    kr,jr=parentindices(S)
    transpose(BandedMatrix(view(parent(S).op,jr,kr)))
end

Base.transpose(A::Operator)=TransposeOperator(A)
adjoint(A::Operator{T}) where {T<:Real}=TransposeOperator(A)


using Combinatorics
using DataStructures
using Random

include("Thresholds.jl")

"""
    S_matrix(ts, window)

Returns S matrix containing all segments of length 'window' in 'ts'.
Done with a sliding window.
"""
function S_matrix(ts, window)
    S = Array{typeof(ts[1]),2}(undef, length(ts) - window, window)
    for index in 1:length(ts) - window
        S[index,:] = ts[index:index+window-1]
    end
    return S
end

"""
Returns all possible masks of length d among the possible w positions.
positive masks: the indexes represent the columns that are KEPT.
    w : motif length
    t : length of masked motifs
"""
function get_masks(w, t)
    if t > w
        error("mask length must be smaller than w")
    end
    return collect(combinations(collect(1:w), t))
end

"""
    exclusion_mask!(c_matrix, exclusion_zone)

Applies exclusion zone to collision matrix (in place).
"""
function exclusion_mask!(c_matrix, exclusion_zone)
    for column in 1:size(c_matrix)[2] - exclusion_zone
        for row in 1:column + exclusion_zone
            c_matrix[row, column] = -Inf
        end
    end
    for i in 0:exclusion_zone
        c_matrix[:, end - i] .= -Inf
    end
end

"""
A version of `findall` adapted for loops to go faster.
"""
function findall_fast(f, a::Array{T, N}, stocker) where {T, N}
    j = 1
    @inbounds for i in eachindex(a)
        @inbounds if f(a[i])
        @inbounds    stocker[j] = i
            j += 1
        end
    end
    return stocker[1:j-1]
end

"""
    update_c_matrix!(collision_matrix, masked_S)

Updates the collision matrix by adding 1 in the row where repetitions of a motif are found.
The row are the index of the repetitions, the column represent the first found motif of this type.
An exclusion zone is applied to tackle trivial neighbours motifs.
"""
function update_c_matrix!(collision_matrix, masked_S)
    masked_S_list = [masked_S[index,:] for index in 1:size(masked_S)[1]]
    total_count = counter(masked_S_list)
    # sub_count = Dict(key => total_count[key] for key in keys(total_count) if total_count[key] > 1)
    stocker = Vector{Int}(undef, length(masked_S_list))
    for motif in keys(total_count)
        if total_count[motif] > 1
            p = findall_fast(x -> x == motif, masked_S_list, stocker)
            # stocker .= 0
            for index in 1:length(p)-1
                @inbounds if abs(p[index+1] - p[index]) >= length(masked_S[1,:])
                @inbounds collision_matrix[p[index+1], p[1]] += 1
                end
            end
        end
    end
end

"""
    collision_matrix(ts, w, d)

Constructs and returns the collision matrix of time-series 'ts'.
inputs (Int):
    ts : input time-series
    w : motif size (window size)
    d : # of allowed errors between motifs
    e : exclusion zone to get rid of trivial matches.
    t : length of projections after applying mask. Defaults to w - d.
    iters : the number of cycles used to construct the collision matrix.
returns (Int):
    C : collision matrix
"""
function collision_matrix(ts, w, d, t = w - d, e = div(w,2); iters = 1000)
    S = S_matrix(ts, w)
    # pre-alocate memory
    C = zeros(length(ts) - w, length(ts) - w)
    exclusion_mask!(C, e) #apply exclusion zone to collision matrix
    masks = get_masks(w, t)
    for i in 1:iters
        mask = rand(masks)
        update_c_matrix!(C, S[:, mask]) #selecting masked version of S
    end
    return C
end


@testset "Diffusion Simulation" begin

gx = complete_graph(5)

for g in testgraphs(gx)  # this makes graphs of different eltypes
    # Most basic
    @test @inferred(diffusion_rate(g, 1.0, 4)) == [1, 5, 5, 5]
end

for i in 1:5
  add_vertex!(gx)
end

for g in testgraphs(gx)  # this makes graphs of different eltypes

    ######
    # Check on fully connected, prob = 1
    ######

    # Add disconnected for more dynamics

    # Basic test. Watch connected vertices
    @test @inferred(diffusion_rate(g, 1.0, 4,
                                   watch=collect(1:5),
                                   initial_infections=[2]
                                   )) == [1, 5, 5, 5]

    # Watching unconnected vertices
    @test @inferred(diffusion_rate(g, 1.0, 4,
                                   watch=collect(6:10),
                                   initial_infections=[2]
                                   )) == [0, 0, 0, 0]

    # Watch subset
    @test @inferred(diffusion_rate(g, 1.0, 4,
                                   watch=collect(1:2),
                                   initial_infections=[2]
                                   )) == [1, 2, 2, 2]

    @test @inferred(diffusion_rate(g, 1.0, 4,
                                   watch=collect(1:5),
                                   initial_infections=[10]
                                   )) == [0, 0, 0, 0]

end

######
# Check along path graph
######

gx = path_graph(5)

for g in testgraphs(gx)  # this makes graphs of different eltypes

    @test @inferred(diffusion_rate(g, 1.0, 4,
                                   watch=collect(1:5),
                                   initial_infections=[1]
                                   )) == [1, 2, 3, 4]

    @test @inferred(diffusion_rate(g, 1.0, 4,
                                   watch=collect(1:5),
                                   initial_infections=[3]
                                   )) == [1, 3, 5, 5]
end

gx = path_graph(30)
for g in testgraphs(gx)
    # Check normalize
    @test @inferred(diffusion_rate(g,
                                  1.0,
                                  6,
                                  initial_infections=[15],
                                  normalize=false
                                  )) == [1, 3, 5, 7, 9, 11]

    @test @inferred(diffusion_rate(g, 2.0, 6,
                                   initial_infections=[15],
                                   normalize=true)
                                   ) == [1, 3, 5, 7, 9, 11]

    # Test probability accurate
    # In a Path network,
    # number of nodes infected (minus 1)
    # is equal to number of successes of a
    # Burnoulli process.
    # So if p = 0.2, in 5 steps (seed +
    # 4 trials) expected value is 0.8,
    # with standard deviation 0.8.

    means = Dict(0.2 => 0.8, 0.4 => 1.6)
    stds = Dict(0.2 => 0.8, 0.4 => 0.98)
    runs = 20
    for p in [0.2, 0.4]
        final_value = 0.0

        for i in 1:20
            result = @inferred(diffusion_rate(g, p, 5,
                                     initial_infections=[1]))
            final_value += result[5]
        end

        # Pretty loose bounds so don't get lots of failed tests.
        # Just want some safeguard.
        # Note 5 steps = 4 Bernoullis + initial infection.
        # False rate less than 1 in 1000
        # Subtract 1 for initial infection
        avg = final_value / runs - 1
        @test avg < means[p] + stds[p] / sqrt(runs) * 3.5
        @test avg > means[p] - stds[p] / sqrt(runs) * 3.5
    end
end

gx = path_digraph(10)

for g in testdigraphs(gx)

    ######
    # Check on digraphs
    ######

    @test @inferred(diffusion_rate(g, 1.0, 9,
                                  initial_infections=[1]
                                  )) == collect(1:9)

    @test @inferred(diffusion_rate(g, 1.0, 9,
                                  initial_infections=[10]
                                  )) == ones(Int, 9)

    # Check probabilities.
    # See note in analogous tests above for undirected tests.
    runs = 20
    means = Dict(0.2 => 2, 0.4 => 4)
    stds = Dict(0.2 => 1.2649110640673518,
                0.4 => 1.5491933384829668)

    for p in [0.2, 0.4]
        final_value = 0.0

        for i in 1:20
            result = @inferred(diffusion_rate(g, p, 11,
                                     initial_infections=[1]))
            final_value += result[11]

        end

        # Pretty loose bounds so don't get lots of failed tests.
        # Just want some safeguard.
        # False rate less than 1 in 1000
        # Subtract 1 for initial infection
        avg = final_value / runs - 1
        @test avg < means[p] + stds[p] / sqrt(runs) * 3.5
        @test avg > means[p] - stds[p] / sqrt(runs) * 3.5
    end
end
end


using RiemannTheta
using Test, Random
using LinearAlgebra

######### some testing values

Ω1 = [ 1.690983006 + 0.9510565162im 1.5+0.363271264im ;
      1.5+0.363271264im 1.309016994+0.9510565162im ]
T1 = Matrix(cholesky(imag.(Ω1)).U)

Ω2 = -1/(2π * im) * [ 111.207 96.616 ; 96.616 83.943 ]
T2 = Matrix(cholesky(imag.(Ω2)).U)

Random.seed!(0)
tmp = 5*rand(10,10) .- 2.5
Ω3 = 5*rand(10,10) .- 2.5 + (tmp * tmp') * im
T3 = Matrix(cholesky(imag.(Ω3)).U)

derivs1 = [ rand(ComplexF64, 2) for i in 1:1 ]
derivs2 = [ rand(ComplexF64, 10) for i in 1:5 ]

ϵ1, ϵ2 = 1e-3, 1e-8

##########

include("internals.jl")
include("mainfuncs.jl")


using DataFrames, GLM

function construct_variables(cause_cnt, data_cnt)
    exact = zeros(Float64, cause_cnt * data_cnt)
    weight = zeros(Float64, cause_cnt * data_cnt)
    prevalence = zeros(Float64, cause_cnt * data_cnt)
    total = zeros(Float64, cause_cnt * data_cnt)
    bc = zeros(Float64, cause_cnt * data_cnt)
    bc2 = zeros(Float64, cause_cnt * data_cnt)
    w = rand(cause_cnt)
    for i in 1:data_cnt
        p = rand(cause_cnt)
        e = exact_burden(w, p)
        assign = ((i - 1) * cause_cnt + 1):(i * cause_cnt)
        exact[assign] = e
        weight[assign] = w
        prevalence[assign] = p
        bc[assign] = (w .* p) / sum(w .* p)
        bc2[assign] = (w .* p) / sum(e)
        total[assign] .= sum(e)
    end
    weight, prevalence, exact, total, bc, bc2
end

w, p, e, t, bc, bc2 = construct_variables(20, 20)
data = DataFrame(E = e, W = w, P = p, BC = bc, BC2 = bc2, T = t)
ols = lm(@formula(E ~ BC + W) , data)
scatter(e, predict(ols))

w1 = 0.3 * rand(10)
p1 = 0.6 * rand(10)
b1 = exact_burden(w1, p1)

w2 = 0.5 * rand(10)
p2 = 0.2 * rand(10)
b2 = exact_burden(w2, p2)

b12 = vcat(b1, b2)
b3 = exact_burden(vcat(w1, w2), vcat(p1, p2))
sum(b3)
1 - (1 - sum(b1)) * (1 - sum(b2))

# This is an oddly successful way to combine the two smaller exact
# calculations into a larger calculation. It's off by a constant.
b1p = b1 * (1 - sum(b2))
b2p = b2 * (1 - sum(b1))
scatter(b3[1:10], b1p, markercolor = :green)
scatter!(b3[11:20], b2p, markercolor = :blue)
xlabel!("Exact Burden")
ylabel!("Estimated Burden")

guess = vcat(b1p, b2p)
guess = guess * sum(b3) / sum(guess)
rel = (guess - b3) ./ b3
# Max relative error is 0.0363.
maximum(abs.(rel))


using BenchmarkTools, Plots, Plots.PlotMeasures, LaTeXStrings
using BenchmarkTools: minimum, median

SUITE = BenchmarkGroup()
SUITE["LaubLoomis"] = BenchmarkGroup()

# ==============================================================================
# Jet-based approach using Taylor Models
# ==============================================================================
include("laubloomis.jl")

# ----------------------------------------
# --- Case 1: smaller initial states ---
# ----------------------------------------
𝑃, 𝑂 = laubloomis(W=0.01, property=(t,x)->x[4] < 4.5)

𝑂₁ = Options(:abs_tol=>1e-10, :orderT=>7, :orderQ=>1, :max_steps=>1000)

# first run
sol_case_1 = solve(𝑃, 𝑂, op=TMJets(𝑂₁))

# verify that specification holds
v4 = [0.0, 1.0] # the flowpipe has been projected so we check for the second component which is x4
@assert all([ρ(v4, sol_case_1.Xk[i].X) < 4.5 for i in eachindex(sol_case_1.Xk)])

# print width of final box
H = sol_case_1.Xk[end].X
println("width of final box, case W = 0.01 : $(high(H)[2] - low(H)[2])")

# benchmark
SUITE["LaubLoomis"]["W=0.01"] = @benchmarkable solve($𝑃, $𝑂, op=TMJets($𝑂₁))

# ----------------------------------------
# --- Case 2: intermediate initial states ---
# ----------------------------------------
𝑃, 𝑂 = laubloomis(W=0.05, property=(t,x)->x[4] < 4.5)

𝑂₂ = Options(:abs_tol=>1e-10, :orderT=>7, :orderQ=>1, :max_steps=>1000)

# first run
sol_case_2 = solve(𝑃, 𝑂, op=TMJets(𝑂₂))

# verify that specification holds
@assert all([ρ(v4, sol_case_2.Xk[i].X) < 4.5 for i in eachindex(sol_case_2.Xk)])

# print width of final box
H = sol_case_2.Xk[end].X
println("width of final box, case W = 0.05 : $(high(H)[2] - low(H)[2])")

# benchmark
SUITE["LaubLoomis"]["W=0.05"] = @benchmarkable solve($𝑃, $𝑂, op=TMJets($𝑂₂))

# ----------------------------------------
# --- Case 3: larger initial states ---
# ----------------------------------------
𝑃, 𝑂 = laubloomis(W=0.1, property=(t,x)->x[4] < 5.0)

𝑂₃ = Options(:abs_tol=>1e-10, :orderT=>7, :orderQ=>1, :max_steps=>1000)

# first run
sol_case_3 = solve(𝑃, 𝑂, op=TMJets(𝑂₃))

# verify that specification holds
@assert all([ρ(v4, sol_case_3.Xk[i].X) < 5.0 for i in eachindex(sol_case_3.Xk)])

# print width of final box
H = sol_case_3.Xk[end].X
println("width of final box, case W = 0.1 : $(high(H)[2] - low(H)[2])")

# benchmark
SUITE["LaubLoomis"]["W=0.1"] = @benchmarkable solve($𝑃, $𝑂, op=TMJets($𝑂₃))

# ==============================================================================
# Execute benchmarks and save benchmark results
# ==============================================================================

# tune parameters
tune!(SUITE)

# run the benchmarks
results = run(SUITE, verbose=true)

# return the sample with the smallest time value in each test
println("minimum time for each benchmark:\n", minimum(results))

# return the median for each test
println("median time for each benchmark:\n", median(results))

# ==============================================================================
# Execute benchmarks and save benchmark results
# ==============================================================================

plot(sol_case_1,
     tickfont=font(30, "Times"), guidefontsize=45,
     xlab=L"t\raisebox{2.0mm}{\textcolor{white}{.}}",
     ylab=L"x_{4}\raisebox{1.2mm}{\textcolor{white}{.}}",
     xtick=[0., 2., 4., 6., 8., 10., 12., 14., 16., 18., 20.],
     ytick=[2, 2.5, 3, 3.5, 4, 4.5],
     xlims=(0., 20.), ylims=(1.5, 4.5),
     bottom_margin=6mm, left_margin=8mm, right_margin=4mm, top_margin=3mm,
     size=(1000, 1000), linecolor="blue")

plot!(x->x, x->4.5, 0., 20., line=2, color="red", linestyle=:dash, legend=nothing)
savefig(@relpath "laubloomis_case_1.png")

plot(sol_case_2,
     tickfont=font(30, "Times"), guidefontsize=45,
     xlab=L"t\raisebox{2.0mm}{\textcolor{white}{.}}",
     ylab=L"x_{4}\raisebox{1.2mm}{\textcolor{white}{.}}",
     xtick=[0., 2., 4., 6., 8., 10., 12., 14., 16., 18., 20.],
     ytick=[2, 2.5, 3, 3.5, 4, 4.5, 5.0],
     xlims=(0., 20.), ylims=(1.5, 5.0),
     bottom_margin=6mm, left_margin=8mm, right_margin=4mm, top_margin=3mm,
     size=(1000, 1000), linecolor="blue")

plot!(x->x, x->5.0, 0., 20., line=2, color="red", linestyle=:dash, legend=nothing)
savefig(@relpath "laubloomis_case_2.png")

plot(sol_case_3,
     tickfont=font(30, "Times"), guidefontsize=45,
     xlab=L"t\raisebox{2.0mm}{\textcolor{white}{.}}",
     ylab=L"x_{4}\raisebox{1.2mm}{\textcolor{white}{.}}",
     xtick=[0., 2., 4., 6., 8., 10., 12., 14., 16., 18., 20.],
     ytick=[2, 2.5, 3, 3.5, 4, 4.5, 5.0],
     xlims=(0., 20.), ylims=(1.5, 5.0),
     bottom_margin=6mm, left_margin=8mm, right_margin=4mm, top_margin=3mm,
     size=(1000, 1000), linecolor="blue")

plot!(x->x, x->5.0, 0., 20., line=2, color="red", linestyle=:dash, legend=nothing)
savefig(@relpath "laubloomis_case_3.png")

plot(sol_case_1, color="red")

plot!(sol_case_2, alpha=0.6, color="green")

plot!(sol_case_3, alpha=0.2,
     tickfont=font(30, "Times"), guidefontsize=45,
     xlab=L"t\raisebox{2.0mm}{\textcolor{white}{.}}",
     ylab=L"x_{4}\raisebox{1.2mm}{\textcolor{white}{.}}",
     xtick=[0., 2., 4., 6., 8., 10., 12., 14., 16., 18., 20.],
     ytick=[2, 2.5, 3, 3.5, 4, 4.5, 5.0],
     xlims=(0., 20.), ylims=(1.5, 5.0),
     bottom_margin=6mm, left_margin=8mm, right_margin=4mm, top_margin=3mm,
     size=(1000, 1000), color="blue")

plot!(x->x, x->5.0, 0., 20., line=2, color="red", linestyle=:dash, legend=nothing)
plot!(x->x, x->4.5, 0., 20., line=2, color="red", linestyle=:dash, legend=nothing)
savefig(@relpath "laubloomis_case_all.png")


module Radial

export R_vals, set_d1!, set_d2!, filter!

"""
Computes array of compactified radial points

R_vals(
   nx::Int64,
   dr::Float64
   )::Vector{Float64}
"""
function R_vals(nx::Int64, dr::Float64)::Vector{Float64}
   return [dr*i for i=1:nx]
end

"""
Takes first derivative in radial direction.

set_d1!(
   dv,
   v,
   dr::Float64,
   )
"""
function set_d1!(
      dv,
      v,
      dr::Float64,
     )
   inv_dr   = 1.0/dr

   nx, ny = size(dv)

   for j=1:ny
      for i=3:nx-2
         dv[i,j] = (-
                    (1.0/12.0)*v[i+2,j]
                    +
                    (2.0/3.0 )*v[i+1,j]
                    -
                    (2.0/3.0 )*v[i-1,j]
                    +
                    (1.0/12.0)*v[i-2,j]
                   )*inv_dr
      end
      dv[1,j] = (-
                 3.0 *v[5,j]
                 +
                 16.0*v[4,j]
                 -
                 36.0*v[3,j]
                 +
                 48.0*v[2,j]
                 -
                 25.0*v[1,j]
                )*(1.0/12.0)*inv_dr
      dv[2,j] = (+
                 1.0 *v[5,j]
                 -
                 6.0 *v[4,j]
                 +
                 18.0*v[3,j]
                 -
                 10.0*v[2,j]
                 -
                 3.0 *v[1,j]
                )*(1.0/12.0)*inv_dr

      dv[nx,j] = -(-
                   3.0 *v[nx-4,j]
                   +
                   16.0*v[nx-3,j]
                   -
                   36.0*v[nx-2,j]
                   +
                   48.0*v[nx-1,j]
                   -
                   25.0*v[nx,j]
                  )*(1.0/12.0)*inv_dr
      dv[nx-1,j] = -(+
                     1.0 *v[nx-4,j]
                     -
                     6.0 *v[nx-3,j]
                     +
                     18.0*v[nx-2,j]
                     -
                     10.0*v[nx-1,j]
                     -
                     3.0 *v[nx,j]
                    )*(1.0/12.0)*inv_dr
   end
   return nothing
end
"""
Takes second derivative in radial direction

set_d2!(
   dv,
   v,
   dr::Float64
   )
"""
function set_d2!(
      dv,
      v,
      dr::Float64
     )
   inv_dr2 = 1.0/dr^2

   nx, ny = size(dv)

   for j=1:ny
      for i=3:nx-2
         dv[i,j] = (-
                    1.0 *v[i+2,j]
                    +
                    16.0*v[i+1,j]
                    -
                    30.0*v[i  ,j]
                    +
                    16.0*v[i-1,j]
                    -
                    1.0 *v[i-2,j]
                   )*(1.0/12.0)*inv_dr2
      end
      dv[1,j] = (-
                 10.0 *v[6,j]
                 +
                 61.0 *v[5,j]
                 -
                 156.0*v[4,j]
                 +
                 214.0*v[3,j]
                 -
                 154.0*v[2,j]
                 +
                 45.0 *v[1,j]
                )*(1.0/12.0)*inv_dr2
      dv[2,j] = (+
                 1.0 *v[6,j]
                 -
                 6.0 *v[5,j]
                 +
                 14.0*v[4,j]
                 -
                 4.0 *v[3,j]
                 -
                 15.0*v[2,j]
                 +
                 10.0*v[1,j]
                )*(1.0/12.0)*inv_dr2
      dv[nx,j] = (-
                  10.0 *v[nx-5,j]
                  +
                  61.0 *v[nx-4,j]
                  -
                  156.0*v[nx-3,j]
                  +
                  214.0*v[nx-2,j]
                  -
                  154.0*v[nx-1,j]
                  +
                  45.0 *v[nx  ,j]
                 )*(1.0/12.0)*inv_dr2
      dv[nx-1,j] = (+
                    1.0 *v[nx-5,j]
                    -
                    6.0 *v[nx-4,j]
                    +
                    14.0*v[nx-3,j]
                    -
                    4.0 *v[nx-2,j]
                    -
                    15.0*v[nx-1,j]
                    +
                    10.0*v[nx  ,j]
                   )*(1.0/12.0)*inv_dr2
   end
   return nothing
end

"""
6th order Kreiss-Oliger dissipation
"""
function filter!(
      v,
      tmp,
      eps_KO::Float64
     )
   nx, ny = size(v)

   for j=1:ny
      for i=1:nx
         tmp[i,j] = v[i,j]
      end
   end

   for j=1:ny
      for i=4:nx-3
         v[i,j] = (
                         (eps_KO/64.) *tmp[i-3,j]
           +        (-6.0*eps_KO/64.) *tmp[i-2,j]
           +        (15.0*eps_KO/64.) *tmp[i-1,j]
           + (1.0 - (20.0*eps_KO/64.))*tmp[i  ,j]
           +        (15.0*eps_KO/64.) *tmp[i+1,j]
           +        (-6.0*eps_KO/64.) *tmp[i+2,j]
           +             (eps_KO/64.) *tmp[i+3,j]
          )
      end
   end

   return nothing
end

end


module PlayGround

    import Base: product
    using StaticArrays

    abstract type Grid{d} end

    struct UCGrid{d} <: Grid{d}
        a::SVector{d, Float64}
        b::SVector{d, Float64}
        n::SVector{d, Int64}
        nodes:: Vector{SVector{d, Float64}}
    end

    ndims(g::Grid{d}) where d = d

    import Base
    Base.iterate(ucg::UCGrid{d}, args...) where d = iterate(ucg.nodes, args...)
    Base.length(ucg::UCGrid{d}) where d = length(ucg.nodes)
    Base.eltype(ucg::UCGrid{d}) where d = SVector{d, Float64}

    function UCGrid{d}(a::SVector{d, Float64},
                         b::SVector{d, Float64},
                         n::SVector{d, Int64}
    ) where d
        intervals = tuple( [range(a[i], b[i]; length=n[i]) for i=1:d]... )
        # print( [product(intervals)...])
        # nn = [(e) for e in product(intervals...)]
        # return nn
        cnodes = [SVector(e...) for e in product(intervals...)]
        # TODO: can we avoid the cnodes[:] allocation
        return UCGrid{d}(a,b,n,cnodes[:])
    end

    function UCGrid(
        a::AbstractVector{Float64},
        b::AbstractVector{Float64},
        n::AbstractVector{Int64})

        d = length(n)
        @assert (length(a)==d) & (length(b)==d)

        aa = SVector(a...)
        bb = SVector(b...)
        nn = SVector(n...)

        return UCGrid{d}(aa, bb, nn)
    end

    # Concatenate SVectors
end

import Main.PlayGround

grid_1 = Main.PlayGround.UCGrid([0.0, 0.0], [1.0, 1.0], [2, 2])
grid_2 = Main.PlayGround.UCGrid([0.0, 0.0], [1.0, 1.0], [2, 2])

pg = Main.PlayGround.PGrid(grid_1, grid_2)

show(pg)

for k in Main.PlayGround.iter(pg)
    println(k)
end

for k in Main.PlayGround.enumerate_product(pg)
    println(k)
end


"""
    Cell(;has_mine=false, flagged=false, revealed=false, mine_count=0, idx=(0,0))

Cell represents the state of a cell in Minesweeper
- `has_mine`: true if cell contains a mine
- `flagged`: true if cell has been flagged as having a mine
- `revealed`: true if the cell has been revealed
- `mine_count`: number of mines in adjecent cells
- `idx`: Cartesian indices of cell
"""
mutable struct Cell
    has_mine::Bool
    flagged::Bool
    revealed::Bool
    mine_count::Int
    idx::CartesianIndex
end

function Cell(;has_mine=false, flagged=false, revealed=false, mine_count=0, idx=(0,0))
    coords = CartesianIndex(idx)
    return Cell(has_mine, flagged, revealed, mine_count, coords)
end

"""
    Game(;dims=(12,12), n_mines=40, mines_flagged=0, mine_detonated=false, trials=0)

Minesweeper game object

- `cells`: an array of cells
- `dims`: a Tuple indicating dimensions of cell array
- `n_mines`: number of mines in the game
- `mines_flagged`: number of mines flagged
- `mine_detonated`: indicates whether a mine has been detonated
- `score`: score for game, which includes hits, misses, false alarms and correct rejections
- `trials`: the number of trials or moves
"""
mutable struct Game{T}
    cells::Array{Cell,2}
    dims::Tuple{Int,Int}
    n_mines::Int
    mines_flagged::Int
    mine_detonated::Bool
    score::T
    trials::Int
end

function Game(;dims=(12,12), n_mines=40, mines_flagged=0,
    mine_detonated=false, trials=0)
    score = (hits=0.0,false_alarms=0.0,misses=0.0,correct_rejections=0.0)
    cells = initilize_cells(dims)
    add_mines!(cells, n_mines)
    mine_count!(cells)
    return Game(cells, dims, n_mines, mines_flagged, mine_detonated,
        score, trials)
end

function initilize_cells(dims)
    return [Cell(idx=(r,c)) for r in 1:dims[1], c in 1:dims[2]]
end

function add_mines!(cells, n_mines)
    mines = sample(cells, n_mines, replace=false)
    for m in mines
        m.has_mine = true
    end
    return nothing
end

"""
    mine_count!(cells)

For each cell, omputes the number of neighboring cells containing a mine.
"""
function mine_count!(cells)
    for cell in cells
        neighbors = get_neighbors(cells, cell.idx)
        cell.mine_count = count(x->x.has_mine, neighbors)
    end
    return nothing
end

get_neighbors(game::Game, idx) = get_neighbors(game.cells, idx.I...)
get_neighbors(game::Game, cell::Cell) = get_neighbors(game.cells, cell.idx.I...)
get_neighbors(game::Game, x, y) = get_neighbors(game.cells, x, y)
get_neighbors(cells, idx) = get_neighbors(cells, idx.I...)

function get_neighbors(cells, xr, yr)
    v = [-1,0,1]
    X = xr .+ v; Y = yr .+ v
    neighbors = Vector{Cell}()
    for y in Y, x in X
        (x == xr) && (y == yr) ? (continue) : nothing
        in_bounds(cells, x, y) ? push!(neighbors, cells[x,y]) : nothing
    end
    return neighbors
end

get_neighbor_indices(game::Game, idx) = get_neighbor_indices(game.cells, idx.I...)
get_neighbor_indices(game::Game, x, y) = get_neighbor_indices(game.cells, x, y)
get_neighbor_indices(cells, idx) = get_neighbor_indices(cells, idx.I...)

function get_neighbor_indices(cells, xr, yr)
    v = [-1,0,1]
    X = xr .+ v; Y = yr .+ v
    neighbors = Vector{CartesianIndex}()
    for y in Y, x in X
        (x == xr) && (y == yr) ? (continue) : nothing
        in_bounds(cells, x, y) ? push!(neighbors, CartesianIndex(x,y)) : nothing
    end
    return neighbors
end

function in_bounds(cells, x, y)
    if (x < 1) || (y < 1)
        return false
    elseif (x > size(cells, 1)) || (y > size(cells, 2))
        return false
    else
        return true
    end
end

reveal_zeros!(game::Game, x, y) = reveal_zeros!(game.cells, CartesianIndex(x, y))
reveal_zeros!(game::Game, idx) = reveal_zeros!(game.cells, idx)

function reveal_zeros!(cells, idx)
    cells[idx].mine_count ≠ 0 ? (return nothing) : nothing
    indices = get_neighbor_indices(cells, idx)
    for i in indices
        c = cells[i]
        if !c.has_mine && !c.revealed
            c.revealed = true
            c.mine_count == 0 ? reveal_zeros!(cells, i) : nothing
        end
    end
    return nothing
end

reveal(game::Game) = reveal(game.cells::Array{Cell,2})

"""
    reveal(cells::Array{Cell,2})

Reveals game state in REPL
"""
function reveal(cells::Array{Cell,2})
    s = ""
    for x in 1:size(cells,1)
        for y in 1:size(cells,2)
            c = cells[x,y]
            if c.has_mine
                s *= string(" ", "💣", " ")
            else
                s *= string(" ", c.mine_count, " ")
            end
        end
        s *= "\r\n"
    end
    println(s)
end

Base.show(io::IO, game::Game) = Base.show(io::IO, game.cells)

function Base.show(io::IO, cells::Array{Cell,2})
    s = ""
    for x in 1:size(cells,1)
        for y in 1:size(cells,2)
            c = cells[x,y]
            if c.flagged
                s *= string(" ", "🚩", " ")
                continue
            end
            if c.revealed
                if c.has_mine
                    s *= string(" ", "💣", " ")
                    continue
                end
                if c.mine_count == 0
                    s *= string(" ", "∘", " ")
                else
                    s *= string(" ", c.mine_count, " ")
                end
            else
                s *= string(" ", "■", " ")
            end
        end
        s *= "\r\n"
    end
    println(s)
end

"""
    game_over(game)

Terminates game if mine is detonated or all cells have been revealed or flagged.
"""
function game_over(game)
    if game.mine_detonated
        return true
    end
    return all(x->x.revealed || x.flagged, game.cells)
end

"""
    compute_score!(game)

Computes hits, false alarms, misses and correct rejections
"""
function compute_score!(game)
    cells = game.cells
    n_mines = game.n_mines
    no_mines = prod(game.dims) - n_mines
    hits = count(x->x.flagged && x.has_mine, cells)
    misses = n_mines - hits
    false_alarms = count(x->x.flagged && !x.has_mine, cells)
    cr = no_mines - false_alarms
    game.score = (hits=hits,false_alarms=false_alarms,misses=misses,correct_rejections=cr)
    return nothing
end

"""
    select_cell!(game, x, y)

Select cell given a game and row and column indices.
"""
select_cell!(game, x, y) = select_cell!(game, CartesianIndex(x, y))

"""
    select_cell!(game, x, y)

Select cell given a game and cell object.
"""
select_cell!(game, cell::Cell) = select_cell!(game, cell.idx)

"""
    select_cell!(game, idx)

Select cell given a game and Cartesian index.
"""
function select_cell!(game, idx)
    cell = game.cells[idx]
    cell.revealed = true
    reveal_zeros!(game, idx)
    game.mine_detonated = cell.has_mine
    return nothing
end

function flag_cell!(cell, game, gui::Nothing)
    cell.flagged = true
    game.mines_flagged += 1
    return nothing
end

function setup()
    println("Select Difficulty")
    waiting = true
    levels = [20,50,70]
    i = 0
    while waiting
        d = input("1: easy 2: medium, 3: difficult")
        i = parse(Int, d)
        if i ∉ 1:3
            println("Please select 1, 2, or 3")
        else
            waiting = false
        end
    end
    return Game(;dims=(15,15), n_mines=levels[i])
end

function play()
    game = setup()
    play(game)
end

"""
    play(game)

Play Minesweeper in REPL

Run play() to use with default easy, medium and difficult settings
Pass game object to play custom game
Enter exit to quit
Enter row_index column_index to select cell (e.g. 1 2)
"""
function play(game)
    playing = true
    coords = fill(0, 2)
    println(game)
    while playing
        str = input("select coordinates")
        str == "exit" ? (return) : nothing
        try
            coords = parse.(Int, split(str, " "))
        catch
            error("Please give x and y coordinates e.g. 3 2")
        end
        select_cell!(game, coords...)
        if game.mine_detonated
            reveal(game)
            println("Game Over")
            println(" ")
            d = input("Press y to play again")
            d == "y" ? (game = setup()) : (playing = false)
        else
            println(game)
        end
    end
end

function input(prompt)
    println(prompt)
    return readline()
end

"""
    update!(game, gui::Nothing)

`update!` is used to display game state during simulations.
`update!` can be used with REPL or Gtk GUI if imported via import_gui()
"""
function update!(game, gui::Nothing)
    println(game)
end

"""
imports Gtk related code for the GUI
"""
function import_gui()
    path = @__DIR__
    include(path*"/Gtk_Gui.jl")
end


using Colors
using CoordinateTransformations
using GeometryBasics
using MeshCat
using Rotations
using LinearAlgebra

# second-order cone
function κ_so(z)
    z1 = z[1:end-1]
    z2 = z[end]

    z_proj = zero(z)

    if norm(z1) <= z2
        z_proj = copy(z)
    elseif norm(z1) <= -z2
        z_proj = zero(z)
    else
        a = 0.5 * (1.0 + z2 / norm(z1))
        z_proj[1:end-1] = a * z1
        z_proj[end] = a * norm(z1)
    end

    return z_proj
end

# visualizer
function second_order_cone_viz(; h = 1.0, trans = 0.1)
    # visualizer
    vis = Visualizer()

    # set empty background
    setvisible!(vis["/Background"], true)
    setprop!(vis["/Background"], "top_color", RGBA(1.0, 1.0, 1.0, 1.0))
    setprop!(vis["/Background"], "bottom_color", RGBA(1.0, 1.0, 1.0, 1.0))
    setvisible!(vis["/Axes"], false)

    # generate cone
    l = h * 1.0
    w = h * 1.45
    pyramid = Pyramid(Point3(0.0, 0.0, 0.0), l, w)

    n = 500
    for i = 1:n
        setobject!(vis["pyramid$i"], pyramid,
            MeshPhongMaterial(color = RGBA(1.0, 0.0, 0.0, trans)))
        settransform!(vis["pyramid$i"],
            compose(Translation(0.0, 0.0, l),
                LinearMap(RotX(π) * RotZ(π * i / n))))
    end

    # generate point
    setobject!(vis["point"], Sphere(Point3f0(0),
        convert(Float32, 0.05)),
        MeshPhongMaterial(color = RGBA(0.0, 0.0, 0.0, 1.0)))

    settransform!(vis["point"], Translation(0.5, 0.5, 0.5))

    return vis
end

function set_point!(vis, p)
    settransform!(vis["point"], Translation(p...))
end

function project_point!(vis, p)
    p_proj = κ_so(p)
    settransform!(vis["point"], Translation(p_proj...))
    return p_proj
end

vis = second_order_cone_viz(h = 2.0)
render(vis)

p = [1.0, 0.0, 1.0]

set_point!(vis, p)
p_proj = project_point!(vis, p)


"""
    Categorical(p)
A *Categorical distribution* is parameterized by a probability vector `p` (of length `K`).

```math
P(X = k) = p[k]  \\quad \\text{for } k = 1, 2, \\ldots, K.
```

```julia
Categorical(p)   # Categorical distribution with probability vector p
params(d)        # Get the parameters, i.e. (p,)
probs(d)         # Get the probability vector, i.e. p
ncategories(d)   # Get the number of categories, i.e. K
```
Here, `p` must be a real vector, of which all components are nonnegative and sum to one.

**Note:** The input vector `p` is directly used as a field of the constructed distribution, without being copied.

`Categorical` is simply a type alias describing a special case of a
`DiscreteNonParametric` distribution, so non-specialized methods defined for
`DiscreteNonParametric` apply to `Categorical` as well.

External links:

* [Categorical distribution on Wikipedia](http://en.wikipedia.org/wiki/Categorical_distribution)
"""
const Categorical{P<:Real,Ps<:AbstractVector{P}} = DiscreteNonParametric{Int,P,Base.OneTo{Int},Ps}

function Categorical{P,Ps}(p::Ps; check_args=true) where {P<:Real, Ps<:AbstractVector{P}}
    check_args && @check_args(Categorical, isprobvec(p))
    return Categorical{P,Ps}(Base.OneTo(length(p)), p, check_args=check_args)
end

Categorical(p::Ps; check_args=true) where {P<:Real, Ps<:AbstractVector{P}} =
    Categorical{P,Ps}(p, check_args=check_args)

function Categorical(k::Integer; check_args=true)
    check_args && @check_args(Categorical, k >= 1)
    return Categorical{Float64,Vector{Float64}}(Base.OneTo(k), fill(1/k, k), check_args=check_args)
end

Categorical(probabilities::Real...; check_args=true) = Categorical([probabilities...]; check_args=check_args)

### Conversions

convert(::Type{Categorical{P,Ps}}, x::AbstractVector{<:Real}) where {
    P<:Real,Ps<:AbstractVector{P}} = Categorical{P,Ps}(Ps(x))

### Parameters

ncategories(d::Categorical) = support(d).stop
params(d::Categorical{P,Ps}) where {P<:Real, Ps<:AbstractVector{P}} = (probs(d),)
partype(::Categorical{T}) where {T<:Real} = T

### Statistics

function median(d::Categorical{T}) where {T<:Real}
    k = ncategories(d)
    p = probs(d)
    cp = zero(T)
    i = 0
    while cp < 1/2 && i <= k
        i += 1
        @inbounds cp += p[i]
    end
    i
end

### Evaluation

function cdf(d::Categorical{T}, x::Int) where T<:Real
    k = ncategories(d)
    p = probs(d)
    x < 1 && return zero(T)
    x >= k && return one(T)
    c = p[1]
    for i = 2:x
        @inbounds c += p[i]
    end
    return c
end

function pdf(d::Categorical, x::Real)
    ps = probs(d)
    return insupport(d, x) ? ps[round(Int, x)] : zero(eltype(ps))
end

function _pdf!(r::AbstractArray, d::Categorical{T}, rgn::UnitRange) where {T<:Real}
    vfirst = round(Int, first(rgn))
    vlast = round(Int, last(rgn))
    vl = max(vfirst, 1)
    vr = min(vlast, ncategories(d))
    p = probs(d)
    if vl > vfirst
        for i = 1:(vl - vfirst)
            r[i] = zero(T)
        end
    end
    fm1 = vfirst - 1
    for v = vl:vr
        r[v - fm1] = p[v]
    end
    if vr < vlast
        for i = (vr - vfirst + 2):length(rgn)
            r[i] = zero(T)
        end
    end
    return r
end

# sampling

sampler(d::Categorical{P,Ps}) where {P<:Real,Ps<:AbstractVector{P}} =
   AliasTable(probs(d))

### sufficient statistics

struct CategoricalStats <: SufficientStats
    h::Vector{Float64}
end

function add_categorical_counts!(h::Vector{Float64}, x::AbstractArray{T}) where T<:Integer
    for i = 1 : length(x)
        @inbounds xi = x[i]
        h[xi] += 1.   # cannot use @inbounds, as no guarantee that x[i] is in bound
    end
    h
end

function add_categorical_counts!(h::Vector{Float64}, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Integer
    n = length(x)
    if n != length(w)
        throw(DimensionMismatch("Inconsistent array lengths."))
    end
    for i = 1 : n
        @inbounds xi = x[i]
        @inbounds wi = w[i]
        h[xi] += wi   # cannot use @inbounds, as no guarantee that x[i] is in bound
    end
    h
end

function suffstats(::Type{<:Categorical}, k::Int, x::AbstractArray{T}) where T<:Integer
    CategoricalStats(add_categorical_counts!(zeros(k), x))
end

function suffstats(::Type{<:Categorical}, k::Int, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Integer
    CategoricalStats(add_categorical_counts!(zeros(k), x, w))
end

const CategoricalData = Tuple{Int, AbstractArray}

suffstats(::Type{<:Categorical}, data::CategoricalData) = suffstats(Categorical, data...)
suffstats(::Type{<:Categorical}, data::CategoricalData, w::AbstractArray{Float64}) = suffstats(Categorical, data..., w)

# Model fitting

function fit_mle(::Type{<:Categorical}, ss::CategoricalStats)
    Categorical(pnormalize!(ss.h))
end

function fit_mle(::Type{<:Categorical}, k::Integer, x::AbstractArray{T}) where T<:Integer
    Categorical(pnormalize!(add_categorical_counts!(zeros(k), x)), check_args=false)
end

function fit_mle(::Type{<:Categorical}, k::Integer, x::AbstractArray{T}, w::AbstractArray{Float64}) where T<:Integer
    Categorical(pnormalize!(add_categorical_counts!(zeros(k), x, w)), check_args=false)
end

fit_mle(::Type{<:Categorical}, data::CategoricalData) = fit_mle(Categorical, data...)
fit_mle(::Type{<:Categorical}, data::CategoricalData, w::AbstractArray{Float64}) = fit_mle(Categorical, data..., w)

fit_mle(::Type{<:Categorical}, x::AbstractArray{T}) where {T<:Integer} = fit_mle(Categorical, maximum(x), x)
fit_mle(::Type{<:Categorical}, x::AbstractArray{T}, w::AbstractArray{Float64}) where {T<:Integer} = fit_mle(Categorical, maximum(x), x, w)

fit(::Type{<:Categorical}, data::CategoricalData) = fit_mle(Categorical, data)
fit(::Type{<:Categorical}, data::CategoricalData, w::AbstractArray{Float64}) = fit_mle(Categorical, data, w)


module DOFBasesTests

using Test
using Gridap

fun(x) = x[1]+x[2]

D = 2
T = Float64
field = AnalyticalField(fun,D)

basis = MonomialBasis(T,(1,1))

polytope = Polytope(HEX_AXIS,HEX_AXIS)
nodes, nfacenodes = Gridap.RefFEs._high_order_lagrangian_nodes_polytope(polytope,[1,1])
nodes
dofbasis = LagrangianDOFBasis(T,nodes)

field_dofs = [0.0, 1.0, 1.0, 2.0]

basis_dofs = [
 1.0 1.0 1.0 1.0;
 0.0 1.0 0.0 1.0;
 0.0 0.0 1.0 1.0;
 0.0 0.0 0.0 1.0]

test_dof_basis(dofbasis,field,basis,field_dofs,basis_dofs)

T = VectorValue{2,Float64}

fun(x) = VectorValue(x[1]+x[2],x[1])

field = AnalyticalField(fun,D)
basis = MonomialBasis(T,(1,1))

nodes, nfacenodes = Gridap.RefFEs._high_order_lagrangian_nodes_polytope(polytope,[1,1])
dofbasis = LagrangianDOFBasis(T,nodes)

# # Component major
#field_dofs = [0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 2.0, 1.0]

# Node major
field_dofs = [0.0, 1.0, 1.0, 2.0, 0.0, 1.0, 0.0, 1.0]

# Node major
basis_dofs = [
  1.0 1.0 1.0 1.0 0.0 0.0 0.0 0.0;
  0.0 0.0 0.0 0.0 1.0 1.0 1.0 1.0;
  0.0 1.0 0.0 1.0 0.0 0.0 0.0 0.0;
  0.0 0.0 0.0 0.0 0.0 1.0 0.0 1.0;
  0.0 0.0 1.0 1.0 0.0 0.0 0.0 0.0;
  0.0 0.0 0.0 0.0 0.0 0.0 1.0 1.0;
  0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0;
  0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0]

test_dof_basis(dofbasis,field,basis,field_dofs,basis_dofs)

@test dofbasis.dof_to_node == [1, 2, 3, 4, 1, 2, 3, 4]
@test dofbasis.dof_to_comp == [1, 1, 1, 1, 2, 2, 2, 2]
@test dofbasis.node_and_comp_to_dof == [1 5; 2 6; 3 7; 4 8]

end # module


### Use the lasso path for the crime train example.

include("lasso.jl")
include("utils.jl")

using DataFrames
using DelimitedFiles: readdlm
using Plots
import UnicodePlots as Uplot

function DataSummary()
    P, H = readdlm("./L1_Lasso_Module/crime-train.txt", '\t'; header=true)
    Df = DataFrame(P, reshape(H, length(H)))
    return Df
end

PrintTitle("A Summarization of the Data: The Predictors")
TheData= DataSummary()
println(describe(TheData))
println("The size of the dataframe is: $(size(TheData))")
println("The predictant is:\"$(names(TheData)[1])\"")
PrintTitle("We are gonna prepare the into a matrix and vector. ")

function PrepareTheDataMatrix(data_frame::DataFrame)
    """
        Min, max, median and mean are assumed to be important predictor
    """
    FeatureMatrix = Matrix(data_frame[!, 2:end])
    Predictants = data_frame[!, 1]
    return FeatureMatrix, Predictants
end

A, b = PrepareTheDataMatrix(TheData)

function AnalyaisWithLasso(A::Matrix, b::Matrix)
    """
        Do the lasso path thing on the data matrix,
        features are columns of the matrix, instances are rows.

        Returns:
            tuple where the first element is the data matrix and the
            second element is the column vector of all the regularization
            parameter λ.
    """
    Instance = LassoSCOP(A, b)
    Results, λs =  LassoPath(Instance)
    VisualizeLassoPath(Instance, "cirme-data-lassopath-$(rand(1:999999999)).png")

end

AnalyaisWithLasso(A, reshape(b, (length(b), 1)))


using NLsolve

"""
	StaticMatch(mtypes, ftypes, mdist, fdist, surplus)

Construct a Choo & Siow (2006) marriage market model and solve for the equilibrium.
"""
struct StaticMatch

	# m/f types
	"vector of vectors of values for each of N male traits"
	mtypes::Array{Array{T, 1}, 1} where T <: Real
	"vector of vectors of values for each of L female traits"
	ftypes::Array{Array{T, 1}, 1} where T <: Real

	# m/f masses
	"male masses: I_1 x ... x I_N"
	mdist::Array
	"female masses: J_1 x ... x J_L"
	fdist::Array

	"surplus array: (I_1 x ... x I_N) x (J_1 x ... x J_L)"
	surplus::Array

	# equilibrium
	"masses of single males: I_1 x ... x I_N"
	msingle::Array
	"masses of single females: J_1 x ... x J_L"
	fsingle::Array
	"masses of married couples: (I_1 x ... x I_N) x (J_1 x ... x J_L)"
	matches::Array
	"wife's share of marital surplus: (I_1 x ... x I_N) x (J_1 x ... x J_L)"
	wifeshare::Array

	"Inner constructor solves equilibrium and performs sanity checks"
	function StaticMatch(mtypes::Array{Array{T, 1}, 1} where T <: Real,
		                 ftypes::Array{Array{T, 1}, 1} where T <: Real,
		                 mdist::Array, fdist::Array, surplus::Array)

		# CHECK: masses of m/f must be proper probability distro
		if minimum(mdist) < 0 || minimum(fdist) < 0
			error("invalid type distribution")
		end

		# CHECK: masses of m/f match the number of types
		if ndims(mdist) != length(mtypes) || ndims(fdist) != length(ftypes)
			error("type distributions inconsistent with type vectors")
		end

		# compute equilibrium
		msingle, fsingle, matches = equilibrium(surplus, mdist, fdist)
		wifeshare = surplusdiv(matches, msingle, fsingle) ./ (2 .* surplus)

		# TEST: masses of every match outcome must be strictly positive
		negtol = 0
		if minimum(msingle) < -negtol || minimum(fsingle) < -negtol
			@warn "Non-positive mass of singles."
		end
		if minimum(matches) < -negtol
			@warn "Non-positive match mass."
		end

		# create instance
		new(mtypes, ftypes, mdist, fdist, surplus, msingle, fsingle, matches, wifeshare)
	end # constructor

end # struct

"Outer constructor that takes the production function to build the surplus array"
function StaticMatch(men::Array{Array{T, 1}, 1} where T <: Real, wom::Array{Array{T, 1}, 1} where T <: Real,
	                 mmass::Array, fmass::Array, prodfn::Function)
	# Note: prodfn(man::Vector, woman::Vector)
	surp = generate_surplus(men, wom, mmass, fmass, prodfn)
	return StaticMatch(men, wom, mmass, fmass, surp)
end

"Outer constructor for one dimensional case"
function StaticMatch(men::Vector, wom::Vector,
	                 mmass::Array, fmass::Array, prodfn::Function)
	# augment production function
	g(x::Vector, y::Vector) = prodfn(x[1], y[1])
	return StaticMatch([men,], [wom,], mmass, fmass, g)
end

"Outer constructor for one dimensional males case"
function StaticMatch(men::Vector, wom::Array{Array{T, 1}, 1} where T <: Real,
	                 mmass::Array, fmass::Array, prodfn::Function)
	return StaticMatch([men,], wom, mmass, fmass, prodfn)
end

"Outer constructor for one dimensional females case"
function StaticMatch(men::Array{Array{T, 1}, 1} where T <: Real, wom::Vector,
	                 mmass::Array, fmass::Array, prodfn::Function)
	return StaticMatch(men, [wom,], mmass, fmass, prodfn)
end

"""
Quasi-demand functions

Assuming that surplus is generated from equal gains on both sides.
This is w.l.o.g. as gains and transfers cannot be separately identified.
"""
function demand(surp::Array, transfers::Array, mmass::Array, fmass::Array)

	# un-normalized raw demands (Note: surp is defined in per-person terms)
	m_raw_dmd = exp.(surp .- transfers)
	f_raw_dmd = exp.(surp .+ transfers)

	# total raw demands per sex, by summing over opposite sex
	m_raw_dmd_tot = [sum(m_raw_dmd[i,j] for j in CartesianIndices(fmass)) for i in CartesianIndices(mmass)]
	f_raw_dmd_tot = [sum(f_raw_dmd[i,j] for i in CartesianIndices(mmass)) for j in CartesianIndices(fmass)]

	# solve for singles: normalization factors
	m_sng_dmd = mmass ./ (1 .+ m_raw_dmd_tot)
	f_sng_dmd = fmass ./ (1 .+ f_raw_dmd_tot)

	# normalized demands
	m_dmd = [m_sng_dmd[i] * m_raw_dmd[i,j] for i in CartesianIndices(mmass), j in CartesianIndices(fmass)]
	f_dmd = [f_sng_dmd[j] * f_raw_dmd[i,j] for i in CartesianIndices(mmass), j in CartesianIndices(fmass)]

	return m_sng_dmd, m_dmd, f_sng_dmd, f_dmd
end # demand

"""
Compute equilibrium shares of singles and marriages.

Solves the matching equilibrium by searching for the transfers
(analogous to prices) which clear the market.
"""
function equilibrium(surpl::Array, mmass::Array, fmass::Array)

	"Function to pass to solver to find zero."
	function eqnsys!(res::Array, tx::Array)
		msd, mdmd, fsd, fdmd = demand(surpl, tx, mmass, fmass)
		res .= mdmd .- fdmd
		return res
	end

	#initial guess of transfers/prices
	guess = zero(surpl)

	# NLsolve
	result = nlsolve(eqnsys!, guess, ftol = 1e-16, autodiff = :forward, method = :newton)

	if !converged(result)
		@warn "NLsolve failed to converge to equilibrium. Solver result:"
		@show result
	end

	eqm_transfers = reshape(result.zero, size(surpl))
	singmen, mmatches, singwom, fmatches = demand(surpl, eqm_transfers, mmass, fmass)
	matches = (mmatches .+ fmatches) ./ 2 # average out any approximation error

	return singmen, singwom, matches
end # equilibrium

"Construct production array from function."
function generate_surplus(mtypes::Array{Array{T, 1}, 1} where T <: Real,
                          ftypes::Array{Array{T, 1}, 1} where T <: Real,
                          mmass::Array, fmass::Array, prodfn::Function)
	# get dimensions
	Dm = [length(v) for v in mtypes]
	Df = [length(v) for v in ftypes]

	# initialize arrays
	surp = Array{Float64}(undef, Dm..., Df...)
	gent = Vector{Float64}(undef, length(mtypes)) # one man's vector of traits
	lady = Vector{Float64}(undef, length(ftypes))

	for coord in CartesianIndices(surp)
		for trt in 1:length(mtypes) # loop through traits in coord
			gent[trt] = mtypes[trt][coord[trt]]
		end
		for trt in 1:length(ftypes) # loop through traits in coord
			lady[trt] = ftypes[trt][coord[trt+length(mtypes)]]
		end
		surp[coord] = prodfn(gent, lady)
	end

	return surp
end # generate_surplus

"Wife's consumption out of surplus (aggregate share)."
function surplusdiv(matches::Array, sm::Array, sw::Array)
	return [log(matches[i,j]) - log(sw[j]) for i in CartesianIndices(sm), j in CartesianIndices(sw)]
end # surplusdiv

"Estimate the marital surplus from observed matches and singles."
function estimate_static_surplus(matches::Array, msingle::Array, fsingle::Array)
	return [log(matches[i,j]) - 0.5 * (log(msingle[i]) + log(fsingle[j]))
	        for i in CartesianIndices(msingle), j in CartesianIndices(fsingle)]
end # estimate_static_surplus


using Test

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

include("../utilities.jl")

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

@testset "NormInfinity" begin
    bridged_mock = MOIB.Constraint.NormInfinity{Float64}(mock)

    MOIT.basic_constraint_tests(bridged_mock, config,
        include = [(F, MOI.NormInfinityCone) for F in [
            MOI.VectorOfVariables, MOI.VectorAffineFunction{Float64}, MOI.VectorQuadraticFunction{Float64}
        ]])

    mock.optimize! = (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [1.0, 0.5, 1.0],
        (MOI.VectorAffineFunction{Float64}, MOI.Nonnegatives) => [[0.0, 1.0, 0.0, 0.0]],
        (MOI.VectorAffineFunction{Float64}, MOI.Zeros) => [[-1], [-1]])

    MOIT.norminf1vtest(bridged_mock, config)
    MOIT.norminf1ftest(bridged_mock, config)

    @testset "Test mock model" begin
        var_names = ["x", "y", "z"]
        MOI.set(mock, MOI.VariableName(), MOI.get(mock, MOI.ListOfVariableIndices()), var_names)
        nonneg = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.Nonnegatives}())
        @test length(nonneg) == 1
        MOI.set(mock, MOI.ConstraintName(), nonneg[1], "nonneg")
        zeros = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.Zeros}())
        @test length(zeros) == 2
        MOI.set(mock, MOI.ConstraintName(), zeros[1], "x_eq")
        MOI.set(mock, MOI.ConstraintName(), zeros[2], "y_eq")

        s = """
        variables: x, y, z
        nonneg: [x + -1.0y, x + -1.0z, x + y, x + z] in MathOptInterface.Nonnegatives(4)
        x_eq: [-1.0 + x] in MathOptInterface.Zeros(1)
        y_eq: [-0.5 + y] in MathOptInterface.Zeros(1)
        maxobjective: y + z
        """
        model = MOIU.Model{Float64}()
        MOIU.loadfromstring!(model, s)
        MOIU.test_models_equal(mock, model, var_names, ["nonneg", "x_eq", "y_eq"])
    end

    @testset "Test bridged model" begin
        var_names = ["x", "y", "z"]
        MOI.set(bridged_mock, MOI.VariableName(), MOI.get(bridged_mock, MOI.ListOfVariableIndices()), var_names)
        norminf = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.NormInfinityCone}())
        @test length(norminf) == 1
        MOI.set(bridged_mock, MOI.ConstraintName(), norminf[1], "norminf")
        zeros = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.Zeros}())
        @test length(zeros) == 2
        MOI.set(bridged_mock, MOI.ConstraintName(), zeros[1], "x_eq")
        MOI.set(bridged_mock, MOI.ConstraintName(), zeros[2], "y_eq")

        s = """
        variables: x, y, z
        norminf: [1.0x, y, z] in MathOptInterface.NormInfinityCone(3)
        x_eq: [-1.0 + x] in MathOptInterface.Zeros(1)
        y_eq: [-0.5 + y] in MathOptInterface.Zeros(1)
        maxobjective: y + z
        """
        model = MOIU.Model{Float64}()
        MOIU.loadfromstring!(model, s)
        MOIU.test_models_equal(bridged_mock, model, var_names, ["norminf", "x_eq", "y_eq"])
    end

    ci = first(MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.NormInfinityCone}()))

    @testset "$attr" for attr in [MOI.ConstraintPrimalStart(), MOI.ConstraintDualStart()]
        @test MOI.supports(bridged_mock, attr, typeof(ci))
        value = [4.0, 1.0, -2.0]
        MOI.set(bridged_mock, attr, ci, value)
        @test MOI.get(bridged_mock, attr, ci) ≈ value
    end

    test_delete_bridge(bridged_mock, ci, 3, ((MOI.VectorAffineFunction{Float64}, MOI.Nonnegatives, 0),))
end

@testset "NormOne" begin
    bridged_mock = MOIB.Constraint.NormOne{Float64}(mock)

    MOIT.basic_constraint_tests(bridged_mock, config,
        include = [(F, MOI.NormOneCone) for F in [
            MOI.VectorOfVariables, MOI.VectorAffineFunction{Float64}, MOI.VectorQuadraticFunction{Float64}
        ]])

    mock.optimize! = (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [1.0, 0.5, 0.5, 0.5, 0.5],
        (MOI.VectorAffineFunction{Float64}, MOI.Nonnegatives) => [[1.0, 1.0, 0.0, 0.0]],
        (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}) => [[1.0]],
        (MOI.VectorAffineFunction{Float64}, MOI.Zeros) => [[-1], [0]])

    MOIT.normone1vtest(bridged_mock, config)
    MOIT.normone1ftest(bridged_mock, config)

    var_names = ["x", "y", "z"]
    MOI.set(bridged_mock, MOI.VariableName(), MOI.get(bridged_mock, MOI.ListOfVariableIndices()), var_names)

    nonneg = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.Nonnegatives}())
    greater = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}}())
    u, v = MOI.get(mock, MOI.ListOfVariableIndices())[4:5]
    @testset "Test mock model" begin
        MOI.set(mock, MOI.VariableName(), u, "u")
        MOI.set(mock, MOI.VariableName(), v, "v")
        @test length(nonneg) == 1
        MOI.set(mock, MOI.ConstraintName(), nonneg[1], "nonneg")
        @test length(greater) == 1
        MOI.set(mock, MOI.ConstraintName(), greater[1], "greater")
        zeros = MOI.get(mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.Zeros}())
        @test length(zeros) == 2
        MOI.set(mock, MOI.ConstraintName(), zeros[1], "x_eq")
        MOI.set(mock, MOI.ConstraintName(), zeros[2], "y_eq")

        s = """
        variables: x, y, z, u, v
        nonneg: [u + -1.0y, v + -1.0z, u + y, v + z] in MathOptInterface.Nonnegatives(4)
        greater: x + -1.0u + -1.0v >= 0.0
        x_eq: [-1.0 + x] in MathOptInterface.Zeros(1)
        y_eq: [-0.5 + y] in MathOptInterface.Zeros(1)
        maxobjective: y + z
        """
        model = MOIU.Model{Float64}()
        MOIU.loadfromstring!(model, s)
        MOIU.test_models_equal(mock, model, [var_names; "u"; "v"], ["nonneg", "greater", "x_eq", "y_eq"])
    end

    @testset "Test bridged model" begin
        normone = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.NormOneCone}())
        @test length(normone) == 1
        MOI.set(bridged_mock, MOI.ConstraintName(), normone[1], "normone")
        zeros = MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.Zeros}())
        @test length(zeros) == 2
        MOI.set(bridged_mock, MOI.ConstraintName(), zeros[1], "x_eq")
        MOI.set(bridged_mock, MOI.ConstraintName(), zeros[2], "y_eq")

        s = """
        variables: x, y, z
        normone: [1.0x, y, z] in MathOptInterface.NormOneCone(3)
        x_eq: [-1.0 + x] in MathOptInterface.Zeros(1)
        y_eq: [-0.5 + y] in MathOptInterface.Zeros(1)
        maxobjective: y + z
        """
        model = MOIU.Model{Float64}()
        MOIU.loadfromstring!(model, s)
        MOIU.test_models_equal(bridged_mock, model, var_names, ["normone", "x_eq", "y_eq"])
    end

    ci = first(MOI.get(bridged_mock, MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64}, MOI.NormOneCone}()))

    @testset "$attr" for attr in [MOI.ConstraintPrimalStart(), MOI.ConstraintDualStart()]
        @test MOI.supports(bridged_mock, attr, typeof(ci))
        value = [4.0, 1.0, -2.0]
        MOI.set(bridged_mock, attr, ci, value)
        @test MOI.get(bridged_mock, attr, ci) ≈ value
        if attr isa MOI.ConstraintPrimalStart
            @test MOI.get(mock, MOI.VariablePrimalStart(), u) == 1
            @test MOI.get(mock, MOI.VariablePrimalStart(), v) == 2
            @test MOI.get(mock, attr, nonneg[1]) == [0.0, 4.0, 2.0, 0.0]
            @test MOI.get(mock, attr, greater[1]) == 1
        else
            @test MOI.get(mock, attr, nonneg[1]) == [0.0, 2.0, 1.0, 0.0]
            @test MOI.get(mock, attr, greater[1]) == 4
        end
    end

    test_delete_bridge(bridged_mock, ci, 3, (
        (MOI.ScalarAffineFunction{Float64}, MOI.GreaterThan{Float64}, 0),
        (MOI.VectorAffineFunction{Float64}, MOI.Nonnegatives, 0)))
end


"""
    degrees(G::SimpleGraph)

Compute the sequence of degrees
"""
function degrees(G::SimpleGraph)
    degrees(CombinatorialAdjacency(adjacency_matrix(G)))
end

"""
    neighbors_edges(G::SimpleGraph)

Returns deg the sequence of degrees and neig a list of list of direcetd edges numbered from 1 to 2*ne(G)
neig[i] is the list of edges j->i
"""
function neighbors_edges(G::SimpleGraph)
    deg = degrees(G)
    neig = [Vector{Int64}() for _ in vertices(G)]
    c = 1
    for (i,d) in enumerate(deg)
        neig[i] = c:(c+d-1)
        c +=d
    end
    deg, neig
end

"""
    Network

`deg` is a `2*ne(G)` vectors with the degree of the j for each edge `j->i`
`edges_arr` is a `2*ne(G)x2` array containing all directed edges `j->i` where `i` is increasing.
`neig` is a list of vectors where `neig[e]` are all indices of edges `u->j` for `u!=i` where edge `e`
 is `j->i`
"""
struct Network
    G::SimpleGraph
    deg::Vector{Int64}
    edges_arr::Array{Int64,2}
    neig::Vector{Vector{Int64}}
end

"""
    create_edges(G::SimpleGraph)

returns a 2*ne(G)x2 array containing all directed edges j->i where i is increasing.
"""
function create_edges(G::SimpleGraph)
    edges_array = Array{Int64,2}(undef,2*ne(G),2)
    c = 1
    for i in vertices(G)
        for j in G.fadjlist[i]
            edges_array[c,:] = [j i]
            c +=1
        end
    end
    edges_array
end

"""
    neig_remove(G::SimpleGraph,i::Int64,j::Int64)

returns the neighbors of i in G whitout j
"""
function neig_remove(G::SimpleGraph,i::Int64,j::Int64)
    neig = neighbors(G,i)
    if j ∉ neig
        error(j, "is not a neighbor of ", i, "in ", neig)
    end
    filter(s -> s!=j, neig)
end

"""
    create_dic(edges_arr::Array{Int64,2})

returns a dic such that dic[(i,j)] = index of edge i->j in edges_arr
"""
function create_dic(edges_arr::Array{Int64,2})
    ne, _ = size(edges_arr)
    local_dic = Dict{Tuple,Int64}()
    for i in 1:ne
        local_dic[(edges_arr[i,1],edges_arr[i,2])] = i
    end
    local_dic
end

"""
    create_neig_r(G::SimpleGraph,edges_arr::Array{Int64,2})

returns a list of vectors where neig_r[i] are all indices of edges u->v for u!=w where edge i is v->w
"""
function create_neig_r(G::SimpleGraph,edges_arr::Array{Int64,2})
    ne, _ = size(edges_arr)
    neig_r = [Vector{Int64}() for _ in 1:ne]
    local_dic = create_dic(edges_arr)
    for i in 1:ne
        neig_r[i] = map(s -> local_dic[(s,edges_arr[i,1])], neig_remove(G,edges_arr[i,1],edges_arr[i,2]))
    end
    neig_r
end

"""
    create_degrees(G::SimpleGraph,edges_arr::Array{Int64,2})

retunrs a list where degrees_G[i] is the degree of u where edge i is u->v
"""
function create_degrees(G::SimpleGraph,edges_arr::Array{Int64,2})
    ne, _ = size(edges_arr)
    degrees_G = zeros(ne)
    for i in 1:ne
        degrees_G[i] = degree(G,edges_arr[i,1])
    end
    degrees_G
end

function Network(G::SimpleGraph)
    edges_arr = create_edges(G)
    neig = create_neig_r(G,edges_arr)
    deg = create_degrees(G,edges_arr)
    return Network(G, deg, edges_arr, neig)
end

"""
    Pair_ER

for `i` in `G1`, `perm_G1G2[i]` is the corresponding vertex in `G2`. It is `0` if no corresponding vertex in `G2`.
`edge_G1G2` used for init of bias message passing (TBC).
"""
struct Pair_ER
    N1::Network
    N2::Network
    Ginter::SimpleGraph
    n::Int64
    λ::Float64
    s::Float64
    perm_G1G2::Array{Int64,1}
    perm_G2G1::Array{Int64,1}
    edge_G1G2::Array{Int64,2}
    dic_comb::Dict{Tuple{Int64, Int64}, Vector{Vector{Int64}}}
    dic_perm::Dict{Int64, Vector{Vector{Int64}}}
end

"""
    compute_ab(n,λ,s)

returns the solution of the system of equations:
a+(1-a)*b^2 = λ*s/n and (1-a)*b*(1-b) = λ*(1-s)/n
"""
function compute_ab(n,λ,s)
    λ*(s-λ/n)/(λ*(s-2)+n), λ*(1-s)/(n-λ)
end

"""
    create_cor_graphs(n,λ,s)

returns correlated ER graphs G1, G2 and their intersection. All graphs are aligned.
"""
function create_cor_graphs(n,λ,s)
    a,b = compute_ab(n,λ,s)
    G_inter = erdos_renyi(n,a)
    G_add1 = erdos_renyi(n,b)
    G_add2 = erdos_renyi(n,b)
    G_1 = union(G_inter,G_add1)
    G_2 = union(G_inter,G_add2)
    return G_1, G_2, intersect(G_1,G_2)
end

function create_perm(n, vmap1, vmap2)
    revert_vmap2 = zeros(Int64,n)
    for (i,v) in enumerate(vmap2)
        revert_vmap2[v] = i
    end
    [revert_vmap2[i] for i in vmap1]
end

function Pair_ER(n,λ,s)
    G1,G2,Ginter = create_cor_graphs(n,λ,s)
    GG1, vmap1 = largest_comp(G1)
    GG2, vmap2 = largest_comp(G2)
    GGinter, vmapinter = largest_comp(Ginter)
    perm_GinterG1 = create_perm(n, vmapinter, vmap1)
    perm_GinterG2 = create_perm(n, vmapinter, vmap2)
    perm_G1G2 = create_perm(n, vmap1, vmap2)
    perm_G2G1 = create_perm(n, vmap2, vmap1)
    N1 = Network(GG1)
    N2 = Network(GG2)
    dic_G1 = create_dic(N1.edges_arr)
    dic_G2 = create_dic(N2.edges_arr)
    Ninter= Network(GGinter)
    ne , _ = size(Ninter.edges_arr)
    edge_G1G2 = Array{Int64,2}(undef,ne,2)
    for (i,e) in enumerate(eachrow(Ninter.edges_arr))
        #println(i,e)
        edge_G1G2[i,:] = [dic_G1[(perm_GinterG1[e[1]],perm_GinterG1[e[2]])], dic_G2[(perm_GinterG2[e[1]],perm_GinterG2[e[2]])]]
    end
    deg1 = unique(sort(N1.deg))
    deg2 = unique(sort(N2.deg))
    min_deg = min(deg1[end],deg2[end])
    max_deg = max(deg1[end],deg2[end])

    dic_comb = Dict([((d,k),collect(combinations(range(1,d),k))) for d in 0:max_deg for k in 0:d])
    dic_perm = Dict([(k,collect(permutations(range(1,k)))) for k in 0:min_deg])
    return Pair_ER(N1,N2,GGinter,n,λ,s,perm_G1G2,perm_G2G1,edge_G1G2,dic_comb,dic_perm)
end

"""
    largest_comp(G::SimpleGraph)

returns the largest connected component of `G` and
for i a vertex of this component `vmap[i]` is the corresponding index in the original `G`.
"""
function largest_comp(G::SimpleGraph)
    cc = connected_components(G)
    i = argmax([length(c) for c in cc])
    sg, vmap = induced_subgraph(G,cc[i])
    return sg, vmap
end


using DateLib
using Base.Test

function yearFractionTester(StartDate,EndDate,MatlabResults,TestToll=1e-14)
for i=0:DateLib.currMaxImplemented
	@test(abs(MatlabResults[i+1]-yearfrac(StartDate,EndDate,i))<TestToll)
end
return ;
end

StartDate=DateClass(14,12,1992);
EndDate=DateClass(28,2,1996);
MatlabResults=[3.20821917808219;3.20555555555556;3.25277777777778;3.20821917808219;3.20555555555556;3.20555555555556;3.20555555555556;3.20821917808219;3.20821917808219;3.25277777777778;3.20821917808219;3.20555555555556;3.20765027322409];
yearFractionTester(StartDate,EndDate,MatlabResults);

EndDate2=DateClass(29,2,1996);
MatlabResults2=[3.21095890410959;3.20833333333333;3.25555555555556;3.21095890410959;3.20833333333333;3.20833333333333;3.20833333333333;3.21095890410959;3.21095890410959;3.25555555555556;3.21095890410959;3.20833333333333;3.21038251366122];
yearFractionTester(StartDate,EndDate2,MatlabResults2);

EndDate3=DateClass(28,2,2026);
MatlabResults3=[33.2301369863014;33.2055555555556;33.6916666666667;33.2301369863014;33.2055555555556;33.2055555555556;33.2055555555556;33.2082191780822;33.2301369863014;33.6916666666667;33.2301369863014;33.2055555555556;33.2080844374580];
TestToll=1e-13;
yearFractionTester(StartDate,EndDate3,MatlabResults3,TestToll);
println("Test YearFraction Passed")


using DelayEmbeddings, OrdinaryDiffEq

ti = time()

const diffeq = (alg = Vern9(), atol = 1e-9, rtol = 1e-9, maxiters = typemax(Int))

include("dataset_tests.jl")
include("reconstruction_tests.jl")
include("R_delay.jl")
include("R_dimension.jl")

ti = time() - ti
println("\nTest took total time of:")
println(round(ti, digits=3), " seconds or ", round(ti/60, digits=3), " minutes")