function freqresp(sys::DelayLtiSystem, ω::AbstractVector{T}) where {T <: Real}
    ny = noutputs(sys)
    nu = ninputs(sys)

    P_fr = freqresp(sys.P.P, ω);

    G_fr = zeros(eltype(P_fr), length(ω), ny, nu)

    for ω_idx=1:length(ω)
        P11_fr = P_fr[ω_idx, 1:ny, 1:nu]
        P12_fr = P_fr[ω_idx, 1:ny, nu+1:end]
        P21_fr = P_fr[ω_idx, ny+1:end, 1:nu]
        P22_fr = P_fr[ω_idx, ny+1:end, nu+1:end]

        delay_matrix_inv_fr = Diagonal(exp.(im*ω[ω_idx]*sys.Tau)) # Frequency response of the diagonal matrix with delays
        # Inverse of the delay matrix, so there should not be any minus signs in the exponents

        G_fr[ω_idx,:,:] .= P11_fr + P12_fr/(delay_matrix_inv_fr - P22_fr)*P21_fr # The matrix is invertible (?!)
    end

    return G_fr
end

function evalfr(sys::DelayLtiSystem, s)
    (ny, nu) = size(sys)

    P_fr = evalfr(sys.P.P, s)

    P11_fr = P_fr[1:ny, 1:nu]
    P12_fr = P_fr[1:ny, nu+1:end]
    P21_fr = P_fr[ny+1:end, 1:nu]
    P22_fr = P_fr[ny+1:end, nu+1:end]

    delay_matrix_inv_fr = Diagonal(exp.(s*sys.Tau))

    return P11_fr + P12_fr/(delay_matrix_inv_fr - P22_fr)*P21_fr
end

"""
    `y, t, x = lsim(sys::DelayLtiSystem, u, t::AbstractArray{<:Real}; x0=fill(0.0, nstates(sys)), alg=MethodOfSteps(Tsit5()), kwargs...)`

    Simulate system `sys`, over time `t`, using input signal `u`, with initial state `x0`, using method `alg` .

    Arguments:

    `t`: Has to be an `AbstractVector` with equidistant time samples (`t[i] - t[i-1]` constant)
    `u`: Function to determine control signal `ut` at a time `t`, on any of the following forms:
        Can be a constant `Number` or `Vector`, interpreted as `ut .= u` , or
        Function `ut .= u(t)`, or
        In-place function `u(ut, t)`. (Slightly more effienct)
    `kwargs...`: These are sent to `solve` from DelayDiffEq.

    Returns: times `t`, and `y` and `x` at those times.
"""
function lsim(sys::DelayLtiSystem{T,S}, u, t::AbstractArray{<:Real}; x0=fill(zero(T), nstates(sys)), alg=MethodOfSteps(Tsit5()), kwargs...) where {T,S}
    # Make u! in-place function of u
    u! = if isa(u, Number) || isa(u,AbstractVector) # Allow for u to be a constant number or vector
        (uout, t) -> uout .= u
    elseif DiffEqBase.isinplace(u, 2)               # If u is an inplace (more than 1 argument function)
        u
    else                                            # If u is a regular u(t) function
        (out, t) -> (out .= u(t))
    end
    _lsim(sys, u!, t, x0, alg; kwargs...)
end

function dde_param(dx, x, h!, p, t)
    A, B1, B2, C2, D21, Tau, u!, uout, hout, tmp = p

    u!(uout, t)     # uout = u(t)

    #dx .= A*x + B1*ut
    mul!(dx, A, x)
    mul!(tmp, B1, uout)
    dx .+= tmp

    @views for k=1:length(Tau)     # Add each of the delayed signals
        u!(uout, t-Tau[k])      # uout = u(t-tau[k])
        h!(hout, p, t-Tau[k])
        dk_delayed = dot(C2[k,:], hout) + dot(D21[k,:], uout)
        dx .+= B2[:, k] .* dk_delayed
    end
    return
end

# TODO Discontinuities in u are not handled well yet.
function _lsim(sys::DelayLtiSystem{T,S}, Base.@nospecialize(u!), t::AbstractArray{<:Real}, x0::Vector{T}, alg; kwargs...) where {T,S}
    P = sys.P

    if !iszero(P.D22)
        throw(ArgumentError("non-zero D22-matrix block is not supported")) # Due to limitations in differential equations
    end

    t0 = first(t)
    dt = t[2] - t[1]
    if !all(diff(t) .≈ dt) # QUESTION Does this work or are there precision problems?
        error("The t-vector should be uniformly spaced, t[2] - t[1] = $dt.") # Perhaps dedicated function for checking this?
    end

    # Get all matrices to save on allocations
    A, B1, B2, C1, C2, D11, D12, D21, D22 = P.A, P.B1, P.B2, P.C1, P.C2, P.D11, P.D12, P.D21, P.D22
    Tau = sys.Tau

    hout = fill(zero(T), nstates(sys))  # in place storage for h
    uout = fill(zero(T), ninputs(sys)) # in place storage for u
    tmp = similar(x0)

    h!(out, p, t) = (out .= 0)      # History function

    p = (A, B1, B2, C2, D21, Tau, u!, uout, hout, tmp)
    prob = DDEProblem{true}(dde_param, x0, h!, (T(t[1]), T(t[end])), p, constant_lags=sys.Tau)

    sol = DelayDiffEq.solve(prob, alg; saveat=t, kwargs...)

    x = sol.u::Vector{Vector{T}} # the states are labeled u in DelayDiffEq

    y = Matrix{T}(undef, noutputs(sys), length(t))
    d = Matrix{T}(undef, size(C2,1), length(t))
    # Build y signal (without d term)
    for k = 1:length(t)
        u!(uout, t[k])
        y[:,k] = C1*x[k] + D11*uout
        #d[:,k] = C2*x[k] + D21*uout
    end

    dtmp = Vector{T}(undef, size(C2,1))
    # Function to evaluate d(t)_i at an arbitrary time
    # X is continuous, so interpoate, u is not
    function dfunc!(tmp::Array{T1}, t, i) where T1
        tmp .= if t < t0
            T1(0)
        else
            sol(t) # solution object has built in interpolator of same order as solver.
        end
        u!(uout, t)
        return dot(view(C2,i,:),tmp) + dot(view(D21,i,:),uout)
    end

    # Account for the effect of the delayed d-signal on y
    for k=1:length(Tau)
        for j=1:length(t)
            di = dfunc!(tmp, t[j] - Tau[k], k)
            y[:, j] .+= view(D12,:, k) .* di
        end
    end

    return y', t, reduce(hcat ,x)'
end

# We have to default to something, look at the sys.P.P and delays
function _bounds_and_features(sys::DelayLtiSystem, plot::Val)
    ws, pz =  _bounds_and_features(sys.P.P, plot)
    logtau = log10.(abs.(sys.Tau))
    logtau = filter(x->x>4, logtau) # Ignore low frequency
    if isempty(logtau)
        return ws, pz
    end
    extreme = extrema(logtau)
    return [min(ws[1], floor(extreme[1]-0.2)), max(ws[2], ceil(extreme[2]+0.2))], pz
end

# Againm we have to do something for default vectors, more or less a copy from timeresp.jl
function _default_Ts(sys::DelayLtiSystem)
    if !isstable(sys.P.P)
        return 0.05   # Something small
    else
        ps = pole(sys.P.P)
        r = minimum([abs.(real.(ps));0]) # Find the fastest pole of sys.P.P
        r = min(r, minimum([sys.Tau;0])) # Find the fastest delay
        if r == 0.0
            r = 1.0
        end
        return 0.07/r
    end
end

iscontinuous(sys::DelayLtiSystem) = true

function Base.step(sys::DelayLtiSystem{T}, t::AbstractVector; kwargs...) where T
    nu = ninputs(sys)
    if t[1] != 0
        throw(ArgumentError("First time point must be 0 in step"))
    end
    u = (out, t) -> (t < 0 ? out .= 0 : out .= 1)
    x0=fill(zero(T), nstates(sys))
    if nu == 1
        y, tout, x = lsim(sys, u, t; x0=x0, kwargs...)
    else
        x = Array{T}(undef, length(t), nstates(sys), nu)
        y = Array{T}(undef, length(t), noutputs(sys), nu)
        for i=1:nu
            y[:,:,i], tout, x[:,:,i] = lsim(sys[:,i], u, t; x0=x0, kwargs...)
        end
    end
    return y, tout, x
end

function impulse(sys::DelayLtiSystem{T}, t::AbstractVector; alg=MethodOfSteps(BS3()), kwargs...) where T
    nu = ninputs(sys)
    iszero(sys.P.D12) || @warn("Impulse with a direct term from input to delay vector leads to poor accuracy.")
    if t[1] != 0
        throw(ArgumentError("First time point must be 0 in impulse"))
    end
    u = (out, t) -> (out .= 0)
    if nu == 1
        y, tout, x = lsim(sys, u, t; alg=alg, x0=sys.P.B[:,1], kwargs...)
    else
        x = Array{T}(undef, length(t), nstates(sys), nu)
        y = Array{T}(undef, length(t), noutputs(sys), nu)
        for i=1:nu
            y[:,:,i], tout, x[:,:,i] = lsim(sys[:,i], u, t; alg=alg, x0=sys.P.B[:,i], kwargs...)
        end
    end
    return y, tout, x
end

# Used for pade approximation
"""
`p2 = _linscale(p::Polynomial, a)`

Given a polynomial `p` and a number `a, returns the polynomials `p2` such that
`p2(s) == p(a*s)`.
"""
function _linscale(p::Polynomial, a)
    # This function should perhaps be implemented in Polynomials.jl
    coeffs_scaled = similar(p.coeffs, promote_type(eltype(p), typeof(a)))
    a_pow = 1
    coeffs_scaled[1] = p.coeffs[1]
    for k=2:length(p.coeffs)
        a_pow *= a
        coeffs_scaled[k] = p.coeffs[k]*a_pow
    end
    return Polynomial(coeffs_scaled)
end

# Coefficeints for Padé approximations
# PADE_Q_COEFFS = [Polynomial([binomial(N,i)*prod(N+1:2*N-i) for i=0:N]) for N=1:10]
const PADE_Q_COEFFS =  [[2, 1],
 [12, 6, 1],
 [120, 60, 12, 1],
 [1680, 840, 180, 20, 1],
 [30240, 15120, 3360, 420, 30, 1],
 [665280, 332640, 75600, 10080, 840, 42, 1],
 [17297280, 8648640, 1995840, 277200, 25200, 1512, 56, 1],
 [518918400, 259459200, 60540480, 8648640, 831600, 55440, 2520, 72, 1],
 [17643225600, 8821612800, 2075673600, 302702400, 30270240, 2162160, 110880, 3960, 90, 1],
 [670442572800, 335221286400, 79394515200, 11762150400, 1210809600, 90810720, 5045040, 205920, 5940, 110, 1]]

"""
    pade(τ::Real, N::Int)

Compute the `N`th order Padé approximation of a time-delay of length `τ`.
"""
function pade(τ::Real, N::Int)
    if !(1 <= N <= 10); error("Order of Padé approximation must be between 1 and 10. Got $N."); end

    Q = Polynomials.Polynomial(PADE_Q_COEFFS[N])

    return tf(_linscale(Q, -τ), _linscale(Q, τ)) # return Q(-τs)/Q(τs)
end

"""
    pade(G::DelayLtiSystem, N)

Approximate all time-delays in `G` by Padé approximations of degree `N`.
"""
function pade(G::DelayLtiSystem, N)
    ny, nu = size(G)
    nTau = length(G.Tau)
    X = append([ss(pade(τ,N)) for τ in G.Tau]...) # Perhaps append should be renamed blockdiag
    return lft(G.P.P, X)
end


using Random, Test, LinearAlgebra, ForwardDiff
using Turing.Variational: TruncatedADAGrad, DecayedADAGrad, apply!

θ = randn(10, 10)
@testset for opt in [TruncatedADAGrad(), DecayedADAGrad(1e-2)]
    θ_fit = randn(10, 10)
    loss(x, θ_) = mean(sum(abs2, θ*x - θ_*x; dims = 1))
    for t = 1:10^4
        x = rand(10)
        Δ = ForwardDiff.gradient(θ_ -> loss(x, θ_), θ_fit)
        Δ = apply!(opt, θ_fit, Δ)
        @. θ_fit = θ_fit - Δ
    end
    @test loss(rand(10, 100), θ_fit) < 0.01
    @test length(opt.acc) == 1
end


# QuantileLoss

using LaTeXStrings, LossFunctions, Measures, Latexify

Plots.plot(QuantileLoss(.75), c=:olive, linewidth=2)
Plots.plot!(QuantileLoss(.5), c=:orangered4, linewidth=2.5)
Plots.plot!(QuantileLoss(.25), c=:springgreen4, linewidth=2.5)


export dtw

using Spandex

struct Position
    x::Int64
    y::Int64
end

struct Solution
    pos::Position
    iter::Int64
    cost::Float64
    weight::Float64
end

function dtw(
    a::Vector{T},
    b::Vector{T},
    distance,
    window_size::Number = -1,
) where {T}
    local m = length(a)
    local n = length(b)

    if typeof(window_size) != Int
        window_size = Int64(round(min(m, n) * window_size))
    end
    if window_size < 0
        window_size = m + n + 1
    end

    local d = Matrix{T}(undef, m, n)
    fill!(d, typemax(T))
    for i = 1:m
        for j = max(1, i - window_size):min(n, i + window_size)
            d[i, j] = distance(a[i], b[j])
        end
    end

    local D = copy(d)
    for i = 1:m
        for j = max(1, i - window_size):min(n, i + window_size)
            if 1 == i && 1 == j
                continue
            end

            local d1 = typemax(T)
            if i > 1
                d1 = D[i-1, j]
            end
            local d2 = typemax(T)
            if j > 1
                d2 = D[i, j-1]
            end
            local d3 = typemax(T)
            if i > 1 && j > 1
                d3 = D[i-1, j-1]
            end

            D[i, j] += min(d1, min(d2, d3))
        end
    end

    local solution =
        find_best(d, D, Position(1, 1), Position(m, n), Int64(window_size))

    return solution.cost / Float64(solution.iter)
end

function find_best(
    d::Matrix{T},
    D::Matrix{T},
    from::Position,
    to::Position,
    window_size::Int64,
) where {T}
    local m = size(D, 1)
    local n = size(D, 2)

    local pq = PriorityQueue{Solution}((a, b) -> (a.weight < b.weight) ? a : b)
    push!(pq, Solution(from, 1, d[from.x, from.y], D[from.x, from.y]))

    local visited = Dict{Position,Float64}()
    visited[from] = D[from.x, from.y]

    local tg = Float64(n) / m
    local routes = [(1, 1), (1, 0), (0, 1)]

    local iter = 1
    local found = missing
    while pq.size > 0
        local cur = pop!(pq)
        iter = max(iter, cur.iter)

        if (iter - cur.iter) > window_size
            continue
        end

        if to == cur.pos
            found = cur
            break
        end

        for route in routes
            local next = Position(cur.pos.x + route[1], cur.pos.y + route[2])
            if next.x > m || next.y > n
                continue
            end

            local dy = next.x * tg
            local dx = next.y / tg
            if abs(next.x - dx) > window_size || abs(next.y - dy) > window_size
                continue
            end

            local new_weight = cur.weight + D[next.x, next.y]
            if haskey(visited, next)
                if visited[next] < new_weight
                    continue
                end
            end
            visited[next] = new_weight

            push!(
                pq,
                Solution(
                    next,
                    cur.iter + 1,
                    cur.cost + d[next.x, next.y],
                    new_weight,
                ),
            )
        end
    end

    return found
end


using PtFEM, Test

data = Dict(
  # Plane(ndim, nst, nxe, nye, nip, direction, finite_element(nod, nodof), axisymmetric)
  :struc_el => Plane(2, 3, 2, 2, 1, :x, Triangle(3, 2), false),
  :properties => [1.0e6 0.3;],
  :x_coords => [0.0,  0.5,  1.0],
  :y_coords => [0.0,  -0.5,  -1.0],
  :support => [
    (1, [0 1]),
    (4, [0 1]),
    (7, [0 0]),
    (8, [1 0]),
    (9, [1 0])
    ],
  :loaded_nodes => [
    (1, [0.0 -0.25]),
    (2, [0.0 -0.50]),
    (3, [0.0 -0.25])
    ]
)

@time m, dis_df, fm_df = p51(data)

arr = Vector{Float64}(m.loads[0:12])
@test  arr ≈ [0.0, -9.100000000000005e-7, 1.950000000000001e-7,
-9.100000000000006e-7, 3.900000000000002e-7, -9.1e-7, -4.5500000000000025e-7,
1.950000000000002e-7, -4.550000000000004e-7, 3.900000000000004e-7,
-4.5499999999999993e-7, 1.9500000000000038e-7, 3.9000000000000045e-7]


## Company Statement
Base.@kwdef mutable struct CompanyResults
    EBITDA::Float64 #Ebitda
    ghg_emissions::Float64 # GHG emissions in milions tonnes of CO2eq
    climate_costs::Union{Nothing,Float64} = nothing # Climate-related damages to society from the company, in Billion USD
end

company_a = CompanyResults(EBITDA = 25.2, ghg_emissions =  374)

## GHG emissions valuation
using Plots

function climate_valuation!(data::CompanyResults, scc::Float64)::CompanyResults
    data.climate_costs = - data.ghg_emissions * scc * 10^(-3) # to express in billion USD
    legend_axis = ["EBITDA", "GHG",
                "Air Pollution",
                "Water Consumption",
                "Water Pollution",
                "Land Use",
                "Waste","Incidents"]

    air_pollution = data.climate_costs * 0.2
    water_consumption = data.climate_costs * 0.1
    water_pollution = data.climate_costs * 0.08
    land_use = data.climate_costs * 0.05
    waste = data.climate_costs * 0.04
    incidents = data.climate_costs * 0.02

    list_data = [data.EBITDA,
                data.climate_costs,
                air_pollution,
                water_consumption,
                water_pollution,
                land_use,
                waste,
                incidents]
    list_data = reduce(vcat, list_data)
    println(list_data)
    display(bar(legend_axis, list_data, orientation = :horizontal, bar_width = 1, legend = nothing, xlabel = "Billion USD"))
    return data
end

climate_valuation!(company_a, 51.)

company_a.climate_costs / company_a.EBITDA * -1


export total_degree, bezout_number, total_degree_start_solutions

"""
    total_degree(
        F::System;
        parameters = nothing,
        gamma = cis(2π * rand()),
        tracker_options = TrackerOptions(),
        endgame_options = EndgameOptions(),
    )

Solve the system `F` using a total degree homotopy.
This returns a path tracker ([`EndgameTracker`](@ref) or [`OverdeterminedTracker`](@ref)) and an iterator to compute the start solutions.
If the system `F` has declared `variable_groups` then a multi-homogeneous
a start system following [^Wam93] will be constructed.

[^Wam93]: An efficient start system for multi-homogeneous polynomial continuation, Wampler, C.W. Numer. Math. (1993) 66: 517. https://doi.org/10.1007/BF01385710
"""
function total_degree(
    F::Union{System,AbstractSystem};
    compile::Union{Bool,Symbol} = COMPILE_DEFAULT[],
    kwargs...,
)
    if isa(F, AbstractSystem) || isnothing(variable_groups(F))
        return total_degree_variables(F; compile = compile, kwargs...)
    else
        return total_degree_variable_groups(F; compile = compile, kwargs...)
    end
end

function total_degree_variables(
    F::Union{System,AbstractSystem};
    target_parameters = nothing,
    γ = cis(2π * rand()),
    gamma = γ,
    tracker_options = TrackerOptions(),
    endgame_options = EndgameOptions(),
    compile::Union{Bool,Symbol} = COMPILE_DEFAULT[],
    kwargs...,
)
    unsupported_kwargs(kwargs)
    m, n = size(F)

    @unique_var x[1:n] t s[1:n]
    support, coeffs = support_coefficients(System(F(x, target_parameters), x))

    homogeneous = true
    D = zeros(Int, length(support))
    for (k, A) in enumerate(support)
        d = 0
        for j = 1:size(A, 2)
            dⱼ = 0
            for i = 1:size(A, 1)
                dⱼ += A[i, j]
            end
            homogeneous = homogeneous && (j == 1 || dⱼ == d)
            d = max(d, dⱼ)
        end
        D[k] = d
    end
    scaling = maximum.(abs ∘ float, coeffs)

    m ≥ (n - homogeneous) || throw(FiniteException(n - homogeneous - m))

    overdetermined = m > n - homogeneous
    # sort by descending degree for overdetermined systems
    if overdetermined && F isa System
        F = deepcopy(F)
        perm = sortperm(D; rev = true)
        permute!(F.expressions, perm)
        permute!(D, perm)
        permute!(scaling, perm)
    end

    if F isa System
        F = fixed(F; compile = compile)
    end
    if target_parameters !== nothing
        F = fix_parameters(F, target_parameters; compile = compile)
    end
    # if homogeneous put on on an affine chart
    if homogeneous && overdetermined
        F = on_affine_chart(F)
        push!(D, 1)
        push!(scaling, 1.0)
        homogeneous = false
        F = square_up(F)
        scaling = [LA.I F.A] * scaling
        D = D[1:n]
    elseif overdetermined
        F = square_up(F)
        scaling = [LA.I F.A] * scaling
        D = D[1:n]
    end
    if homogeneous
        G = fixed(
            System(s[1:n-1] .* (x[1:n-1] .^ D[1:n-1] .- x[end] .^ D[1:n-1]), x, s[1:n-1]);
            compile = compile,
        )
    else
        G = fixed(
            System(s .* (x .^ D .- 1), x, s);
            compile = compile,
            optimizations = false,
        )
    end
    G = fix_parameters(G, scaling)

    H = StraightLineHomotopy(G, F; γ = gamma)
    if homogeneous
        H = on_affine_chart(H)
    end
    T = EndgameTracker(H, tracker_options = tracker_options, options = endgame_options)
    if overdetermined
        T = OverdeterminedTracker(T, F)
    end
    starts = total_degree_start_solutions(D; homogeneous = homogeneous)

    T, starts
end

function total_degree_variable_groups(
    F::System;
    γ = cis(2π * rand()),
    gamma = γ,
    target_parameters = nothing,
    tracker_options = TrackerOptions(),
    endgame_options = EndgameOptions(),
    compile::Union{Bool,Symbol} = COMPILE_DEFAULT[],
    kwargs...,
)
    unsupported_kwargs(kwargs)
    m, n = size(F)
    homogeneous = is_homogeneous(F)
    D = multi_degrees(F)
    vargroups = variable_groups(F)
    projective_dims = length.(vargroups) .- homogeneous
    M = length(projective_dims)

    m ≥ (n - M * homogeneous) || throw(FiniteException(n - M * homogeneous - m))

    overdetermined = m > n - M * homogeneous

    F = fixed(F; compile = compile)
    if target_parameters !== nothing
        F = fix_parameters(F, target_parameters; compile = compile)
    end

    if homogeneous && overdetermined
        F = on_affine_chart(F, projective_dims)
        D = [D LA.I]
        homogeneous = false
        F = square_up(F)
        D = max.(D[:, 1:n], maximum(D[:, n+1:end], dims = 2))
        projective_dims .+= 1
    elseif overdetermined
        F = square_up(F)
        D = max.(D[:, 1:n], maximum(D[:, n+1:end], dims = 2))
    end

    g, C = multi_homogeneous_system(D, vargroups; homogeneous = homogeneous)
    P = parameters(g)
    p = randn(length(P))
    G = fix_parameters(fixed(g; compile = compile), p)
    starts = MultiBezoutSolutionsIterator(
        D,
        C(P => p),
        projective_dims;
        homogeneous = homogeneous,
    )
    H = StraightLineHomotopy(G, F; gamma = gamma)
    if homogeneous
        H = on_affine_chart(H, projective_dims)
    end
    T = EndgameTracker(H, tracker_options = tracker_options, options = endgame_options)
    if overdetermined
        T = OverdeterminedTracker(T, F)
    end
    T, starts
end

function multi_homogeneous_system(D, vargroups; homogeneous::Bool)
    M, n = size(D)
    projective_dims = length.(vargroups) .- homogeneous
    @unique_var c[1:n+M, 1:n]
    C = zeros(Expression, n + M, n)
    for i = 1:n
        j = 1
        for kⱼ in projective_dims
            for l = 1:kⱼ
                C[j, i] = begin
                    if (i == l && i ≤ kⱼ)
                        1
                    elseif i ≠ l && i ≤ kⱼ
                        0
                    else
                        c[j, i]
                    end
                end
                j += 1
            end
            C[j, i] = 1
            j += 1
        end
    end

    P = variables(C)

    if homogeneous
        Z = vargroups
    else
        Z = vcat.(vargroups, 1)
    end

    m, n = size(D)
    g = map(1:n) do i
        s = Ref(1)
        prod(1:m) do j
            dᵢⱼ = D[j, i]
            kⱼ = length(Z[j]) - 1
            if dᵢⱼ == 0
                s[] += kⱼ + 1
                return 1
            end
            bᵢⱼ = sum(1:kⱼ) do l
                cz = C[s[], i] * Z[j][l]
                s[] += 1
                cz
            end
            s[] += 1
            bᵢⱼ^dᵢⱼ - Z[j][end]^dᵢⱼ
        end
    end
    System(g; variable_groups = vargroups, parameters = P), C
end

struct TotalDegreeStartSolutionsIterator{Iter}
    degrees::Vector{Int}
    homogeneous::Bool
    iterator::Iter
end
function TotalDegreeStartSolutionsIterator(degrees, homogeneous::Bool)
    iterator = Iterators.product(map(d -> 0:(d-1), degrees)...)
    TotalDegreeStartSolutionsIterator(degrees, homogeneous, iterator)
end
function Base.show(io::IO, iter::TotalDegreeStartSolutionsIterator)
    print(io, "$(length(iter)) total degree start solutions for degrees $(iter.degrees)")
end

function Base.iterate(iter::TotalDegreeStartSolutionsIterator)
    indices, state = iterate(iter.iterator)
    _value(iter, indices), state
end
function Base.iterate(iter::TotalDegreeStartSolutionsIterator, state)
    it = iterate(iter.iterator, state)
    it === nothing && return nothing
    _value(iter, first(it)), last(it)
end

function _value(iter::TotalDegreeStartSolutionsIterator, indices)
    x = map((k, d) -> cis(2π * k / d), indices, iter.degrees)
    iter.homogeneous && push!(x, 1)
    x
end

Base.length(iter::TotalDegreeStartSolutionsIterator) = length(iter.iterator)
Base.eltype(iter::Type{<:TotalDegreeStartSolutionsIterator}) = Vector{Complex{Float64}}

"""
    total_degree_start_solutions(degrees)

Returns an iterator of the start solutions of the system
```math
\\begin{array}{rl}
    z_1^{d_1} - 1 &= 0 \\\\
    z_2^{d_2} - 1 &= 0 \\\\
    &\\vdots \\\\
    z_n^{d_n} - 1 &= 0 \\\\
\\end{array}
```
where ``d_i`` is `degrees[i]`.

## Example
```julia-repl
julia> iter = total_degree_start_solutions([2, 2])
4 total degree start solutions for degrees [2, 2]

julia> collect(iter)
4-element Array{Array{Complex{Float64},1},1}:
 [1.0 + 0.0im, 1.0 + 0.0im]
 [-1.0 + 1.2246467991473532e-16im, 1.0 + 0.0im]
 [1.0 + 0.0im, -1.0 + 1.2246467991473532e-16im]
 [-1.0 + 1.2246467991473532e-16im, -1.0 + 1.2246467991473532e-16im]
```
"""
total_degree_start_solutions(
    degrees::AbstractVector{<:Integer};
    homogeneous::Bool = false,
) = TotalDegreeStartSolutionsIterator(degrees, homogeneous)

###
### Multihomogeneous
###

struct MultiBezoutIndicesIterator
    D::Matrix{Int} # Degree matrix
    k::Vector{Int} # (projective) dimensions

    function MultiBezoutIndicesIterator(D::Matrix{Int}, k::Vector{Int})
        @assert size(D, 1) == length(k)
        new(D, k)
    end
end
Base.IteratorSize(::Type{MultiBezoutIndicesIterator}) = Base.SizeUnknown()
Base.eltype(::Type{MultiBezoutIndicesIterator}) = Tuple{Int,Vector{Int}}
function Base.iterate(
    iter::MultiBezoutIndicesIterator,
    state = [i for i = 1:length(iter.k) for j = 1:iter.k[i]],
)
    d = 0
    m, n = size(iter.D)
    p = state
    while d == 0
        state[1] > n && return nothing
        p, state = nextpermutation(1:m, state)
        let p = p, D = iter.D
            d = prod(i -> D[p[i], i], 1:n)::Int
        end
    end

    (d, p), state
end

# Adopted from Combinatorics.jl:
# https://github.com/JuliaMath/Combinatorics.jl/blob/8d9571402319799b29da2005a65b627e8771c1e4/src/permutations.jl#L47
function nextpermutation(m, state)
    n = length(state)
    perm = [m[state[i]] for i = 1:n]
    s = copy(state)
    i = n - 1
    while i >= 1 && s[i] >= s[i+1]
        i -= 1
    end
    if i > 0
        j = n
        while j > i && s[i] >= s[j]
            j -= 1
        end
        s[i], s[j] = s[j], s[i]
        reverse!(s, i + 1)
    else
        s[1] = n + 1
    end
    return (perm, s)
end

struct MultiBezoutSolutionsIterator
    indices::MultiBezoutIndicesIterator
    coeffs::Matrix{Float64}
    homogeneous::Bool
    # precomputed stuff and cache
    roots_of_unity::Matrix{ComplexF64} # Lower Triangular matrix storing the precomputed roots
    A::Vector{Matrix{Float64}}
    b::Vector{Vector{ComplexF64}}
    ranges::Vector{UnitRange{Int}}
end

function MultiBezoutSolutionsIterator(
    indices::MultiBezoutIndicesIterator,
    coeffs::Matrix;
    homogeneous::Bool,
)
    d_max = maximum(indices.D)
    roots_of_unity = zeros(ComplexF64, d_max, d_max)
    for i = 1:d_max, j = 0:i-1
        roots_of_unity[i, j+1] = cis(2π * j / i)
    end
    A = map(kⱼ -> zeros(kⱼ, kⱼ), indices.k)
    b = map(kⱼ -> zeros(ComplexF64, kⱼ), indices.k)
    r_stop = Ref(1)
    ranges = map(indices.k) do kⱼ
        range = r_stop[]:(kⱼ+r_stop[]-1)
        r_stop[] += kⱼ + 1
        range
    end
    MultiBezoutSolutionsIterator(indices, coeffs, homogeneous, roots_of_unity, A, b, ranges)
end
function MultiBezoutSolutionsIterator(
    D::Matrix,
    C::Matrix,
    proj_dims::Vector{Int};
    kwargs...,
)
    MultiBezoutSolutionsIterator(MultiBezoutIndicesIterator(D, proj_dims), C; kwargs...)
end

function Base.show(io::IO, iter::MultiBezoutSolutionsIterator)
    print(io, "Solutions iterator for a multi-homogeneous start system")
end
Base.show(io::IO, ::MIME"application/prs.juno.inline", x::MultiBezoutSolutionsIterator) = x
Base.IteratorSize(::Type{<:MultiBezoutSolutionsIterator}) = Base.SizeUnknown()
Base.eltype(::Type{MultiBezoutSolutionsIterator}) = Vector{ComplexF64}

function Base.iterate(iter::MultiBezoutSolutionsIterator)
    (_, perm), indices_state = iterate(iter.indices)
    dᵢ = ntuple(i -> iter.indices.D[perm[i], i], length(perm))
    d_iter = Iterators.product(map(dᵢ -> 1:dᵢ, dᵢ)...)
    q, d_state = iterate(d_iter)

    x = compute_solution(iter, perm, q, dᵢ)
    x, (perm, indices_state, dᵢ, d_iter, d_state)
end

function Base.iterate(iter::MultiBezoutSolutionsIterator, state)
    perm, indices_state, dᵢ, d_iter, d_state = state

    next_d_state = iterate(d_iter, d_state)
    if next_d_state !== nothing
        q, d_state = next_d_state
        x = compute_solution(iter, perm, q, dᵢ)
        return x, (perm, indices_state, dᵢ, d_iter, d_state)
    end

    # We exhausted our current batch. move indices forward
    indices_newstate = iterate(iter.indices, indices_state)
    # check whether we are completely done
    indices_newstate === nothing && return nothing

    (_, perm), indices_state = indices_newstate
    # Create new d_iter
    new_dᵢ::typeof(dᵢ) = let D = iter.indices.D, perm = perm
        ntuple(i -> D[perm[i], i], length(perm))
    end
    new_d_iter::typeof(d_iter) = Iterators.product(map(dᵢⱼ -> 1:dᵢⱼ, new_dᵢ)...)
    q, d_state = iterate(new_d_iter)
    x = compute_solution(iter, perm, q, new_dᵢ)
    return x, (perm, indices_state, new_dᵢ, new_d_iter, d_state)
end

function compute_solution(iter::MultiBezoutSolutionsIterator, perm, q, dᵢ)
    m, n = size(iter.indices.D)
    # for each variable group we have to solve a linear system
    t = 1
    if iter.homogeneous
        data = zeros(ComplexF64, n + m)
    else
        data = zeros(ComplexF64, n)
    end
    for j = 1:m
        kⱼ = iter.indices.k[j]
        Aⱼ = iter.A[j]
        bⱼ = iter.b[j]
        s = 1
        for i = 1:n
            perm[i] == j || continue
            range = iter.ranges[j]
            for (t, l) in enumerate(range)
                Aⱼ[s, t] = iter.coeffs[l, i]
            end
            bⱼ[s] = iter.roots_of_unity[dᵢ[i], q[i]]
            s += 1
        end

        LinearAlgebra.ldiv!(LinearAlgebra.lu!(Aⱼ), bⱼ)

        data[t:(t+kⱼ-1)] .= bⱼ
        if iter.homogeneous
            data[t+kⱼ] = 1
            t += kⱼ + 1
        else
            t += kⱼ
        end
    end

    data
end

function paths_to_track(f, ::Val{:total_degree})
    target_parameters = nparameters(f) > 0 ? zeros(nparameters(f)) : nothing
    _, starts = total_degree(f; target_parameters = target_parameters)
    if Base.IteratorSize(typeof(starts)) == Base.SizeUnknown()
        k = 0
        for s in starts
            k += 1
        end
        k
    else
        length(starts)
    end
end
@deprecate bezout_number(f::Union{System,AbstractSystem}) paths_to_track(
    f;
    start_system = :total_degree,
)


using HDF5, JLD, KDTrees, DataFrames
using Dates

function testingTimes (start::Int64, finish::Int64, radius::Float64)
	startingHour = start
	endingHour = finish
	radius = 140.0
	nodePositions = load("../Cities/Manhattan/manhattan.jld","positions")

	trainingNumRides = zeros(Int64, length(nodePositions), length(nodePositions))
	trainingAvgTime = zeros(Float32, length(nodePositions), length(nodePositions))
	testingNumRides = zeros(Int64, length(nodePositions), length(nodePositions))
	testingAvgTime = zeros(Float32, length(nodePositions), length(nodePositions))

	nodes = Array(Float64, (2,length(nodePositions)))
	for (i,p) in enumerate(nodePositions)
	    nodes[1,i] = p[1]
	    nodes[2,i] = p[2]
	end

	tree = KDTree(nodes)

	function coordinatesToNode(x::Float64, y::Float64)
		return inball(tree, [x, y], radius, false)
	end

	function computeDistance(px::Float64, py::Float64, dx::Float64, dy::Float64)
		return sqrt((px - dx) ^ 2 + (py - dy) ^ 2)
	end

	for i = 1:12
		println(i)
		df = 0
		@time df = readtable("../../Data/reduced_trip_data_$(i).csv");
		for j = 1:nrow(df)
			if j % 10000 == 0
				print("\r $(j * 100 / nrow(df))%            ")
			end
			pickup = df[j, :pickup_datetime]
			dropoff = df[j, :dropoff_datetime]
			pickupTime = DateTime(pickup, "y-m-d H:M:S")
			dropoffTime = DateTime(dropoff, "y-m-d H:M:S")
			pickupFlag = 1 <= dayofweek(pickupTime) <= 5 && startingHour <= hour(pickupTime) <= endingHour
			dropoffFlag = 1 <= dayofweek(dropoffTime) <= 5 && startingHour <= hour(dropoffTime) <= endingHour
			if pickupFlag && dropoffFlag
				if computeDistance(df[j, :pickup_x], df[j, :pickup_y], df[j, :dropoff_x], df[j, :dropoff_y]) > 140.0
					nodeSrcList = coordinatesToNode(df[j, :pickup_x], df[j, :pickup_y])
					nodeDstList = coordinatesToNode(df[j, :dropoff_x], df[j, :dropoff_y])
					for nodeSrc in nodeSrcList, nodeDst in nodeDstList
						if rand() > 0.2
							trainingNumRides[nodeSrc, nodeDst] += 1
							prevTime = trainingAvgTime[nodeSrc, nodeDst]
							newTime = (dropoffTime - pickupTime).value / 1000.0
							trainingAvgTime[nodeSrc, nodeDst] = (newTime + prevTime * (trainingNumRides[nodeSrc, nodeDst] - 1)) / trainingNumRides[nodeSrc, nodeDst]
						else
							testingNumRides[nodeSrc, nodeDst] += 1
							prevTime = testingAvgTime[nodeSrc, nodeDst]
							newTime = (dropoffTime - pickupTime).value / 1000.0
							testingAvgTime[nodeSrc, nodeDst] = (newTime + prevTime * (testingNumRides[nodeSrc, nodeDst] - 1)) / testingNumRides[nodeSrc, nodeDst]
						end
					end
				end
			end
		end
	end

	trainingDF = DataFrame(node1 = Int64[], node2 = Int64[], numberOfRides = Int64[], averageTime = Float32[])
	testingDF = DataFrame(node1 = Int64[], node2 = Int64[], numberOfRides = Int64[], averageTime = Float32[])
	for src = 1:length(nodePositions), dst = 1:length(nodePositions)
		if trainingNumRides[src, dst] > 0
			push!(trainingDF, [src dst trainingNumRides[src, dst] trainingAvgTime[src, dst]])
		end
		if testingNumRides[src, dst] > 0
			push!(testingDF, [src dst testingNumRides[src, dst] testingAvgTime[src, dst]])
		end
	end

	startingString = string(startingHour)
	endingString = string(endingHour)
	if startingHour < 10
		startingString = "0" * startingString
	end
	if endingHour < 10
		endingString = "0" * endingString
	end
	writetable("../../Travel_times/training_r$(int(radius))_wd_$(startingString)$(endingString).csv", trainingDF)
	writetable("../../Travel_times/testing_r$(int(radius))_wd_$(startingString)$(endingString).csv", testingDF)
end

# testingTimes(12, 14, 140.0)
# testingTimes(2, 4, 140.0)
# testingTimes(7, 9, 140.0)


# push!(LOAD_PATH, joinpath(@__DIR__, "..", ".."))
using JLD2

# jlfile = jldopen("data/local.jld2", "a+")
# jlfile = jldopen("data/nonlocal.jld2", "a+")
jlfile = jldopen("data/nonlocal_symmetric.jld2", "a+")
# jlfile = jldopen("data/nonlocal_more_velocity.jld2")
EF¹¹ = jlfile["diffusivity"]["K11"]
EF¹² = jlfile["diffusivity"]["K12"]
EF²¹ = jlfile["diffusivity"]["K21"]
EF²² = jlfile["diffusivity"]["K22"]
x = jlfile["grid"]["x"]
z = jlfile["grid"]["z"]
N = length(x)
M = length(z)
tmpE = copy(EF¹¹)
EF¹¹ = reshape(permutedims(reshape(EF¹¹, (N, M, N, M)), (2, 1, 4, 3)), (N * M, N * M))
EF¹² = reshape(permutedims(reshape(EF¹², (N, M, N, M)), (2, 1, 4, 3)), (N * M, N * M))
EF²¹ = reshape(permutedims(reshape(EF²¹, (N, M, N, M)), (2, 1, 4, 3)), (N * M, N * M))
EF²² = reshape(permutedims(reshape(EF²², (N, M, N, M)), (2, 1, 4, 3)), (N * M, N * M))
# 2 goes to 1, 1 goes to 2, 4 goes to 3, 3 goes to 4
##

using GLMakie
fig = Figure(resolution = (1800, 1300), title = "Nonlocal Operators")
titlestring = "Kˣˣ"
ax1 = Axis(fig[1, 1], title = titlestring, titlesize = 30)
titlestring = "Kˣᶻ"
ax2 = Axis(fig[1, 3], title = titlestring, titlesize = 30)
titlestring = "Kᶻˣ"
ax3 = Axis(fig[2, 1], title = titlestring, titlesize = 30)
titlestring = "Kᶻᶻ"
ax4 = Axis(fig[2, 3], title = titlestring, titlesize = 30)

colormap = :thermal
colormap2 = :balance
hm1 = heatmap!(ax1, EF¹¹, colormap = colormap)
ax1.yreversed = true

hm2 = heatmap!(ax2, EF¹², colormap = colormap2)
ax2.yreversed = true

hm3 = heatmap!(ax3, EF²¹, colormap = colormap2)
ax3.yreversed = true

hm4 = heatmap!(ax4, EF²², colormap = colormap)
ax4.yreversed = true

Colorbar(fig[1, 2], hm1, height = Relative(3 / 4), width = 25, ticklabelsize = 30,
    labelsize = 30, ticksize = 25, tickalign = 1,)
Colorbar(fig[2, 2], hm2, height = Relative(3 / 4), width = 25, ticklabelsize = 30,
    labelsize = 30, ticksize = 25, tickalign = 1,)
Colorbar(fig[1, 4], hm3, height = Relative(3 / 4), width = 25, ticklabelsize = 30,
    labelsize = 30, ticksize = 25, tickalign = 1,)
Colorbar(fig[2, 4], hm4, height = Relative(3 / 4), width = 25, ticklabelsize = 30,
    labelsize = 30, ticksize = 25, tickalign = 1,)
display(fig)


function turing_inference(
    prob::DiffEqBase.DEProblem,
    alg,
    t,
    data,
    priors;
    likelihood_dist_priors = [InverseGamma(2, 3)],
    likelihood = (u,p,t,σ) -> MvNormal(u, σ[1]*ones(length(u))),
    num_samples=1000, sampler = Turing.NUTS(0.65),
    syms = [Turing.@varname(theta[i]) for i in 1:length(priors)],
    kwargs...,
)
    N = length(priors)
    Turing.@model mf(x, ::Type{T} = Float64) where {T <: Real} = begin
        theta = Vector{T}(undef, length(priors))
        for i in 1:length(priors)
            theta[i] ~ NamedDist(priors[i], syms[i])
        end
        σ = Vector{T}(undef, length(likelihood_dist_priors))
        for i in 1:length(likelihood_dist_priors)
            σ[i] ~ likelihood_dist_priors[i]
        end
        p_tmp = remake(prob, u0 = convert.(T, (prob.u0)), p = theta)
        _saveat = t === nothing ? Float64[] : t
        sol_tmp = solve(p_tmp, alg; saveat = _saveat, kwargs...)

        if sol_tmp isa DiffEqBase.AbstractEnsembleSolution
            failure = any((s.retcode != :Success for s in sol_tmp)) && any((s.retcode != :Terminated for s in sol_tmp))
        else
            failure = sol_tmp.retcode != :Success && sol_tmp != :Terminated
        end

        if failure
            @logpdf() = -Inf
            return
        end
        if sol_tmp isa DiffEqBase.AbstractNoTimeSolution
            res = sol_tmp.u
            x ~ likelihood(res, theta, Inf, σ)
        else
            u = sol_tmp.u
            for i = 1:length(t)
                res = sol_tmp.u[i]
                _t = sol_tmp.t[i]
                x[:,i] ~ likelihood(res, theta, _t, σ)
            end
        end
        return
    end

    # Instantiate a Model object.
    model = mf(data)
    chn = sample(model, sampler, num_samples)
    return chn
end


"""
    AbstractSystem

Abstract supertype for all system types.
"""
abstract type AbstractSystem end

"""
    statedim(s::AbstractSystem)

Returns the dimension of the state space of system `s`.
"""
function statedim end

"""
    stateset(s::AbstractSystem)

Returns the set of allowed states of system `s`.
"""
function stateset end

"""
    inputdim(s::AbstractSystem)

Returns the dimension of the input space of system `s`.
"""
function inputdim end

"""
    inputset(s::AbstractSystem)

Returns the set of allowed inputs of system `s`.
"""
function inputset end

"""
    AbstractDiscreteSystem

Abstract supertype for all discrete system types.
"""
abstract type AbstractDiscreteSystem <: AbstractSystem end

"""
    AbstractContinuousSystem

Abstract supertype for all continuous system types.
"""
abstract type AbstractContinuousSystem <: AbstractSystem end

"""
    islinear(s::AbstractSystem)

Specifies if the dynamics of system `s` is specified by linear equations.

### Notes

We adopt the notion from [Section 2.7, 1]. For example, the system with inputs
``x' = f(t, x, u) = A x + B u`` is linear, since the function ``f(t, ⋅, ⋅)`` is
linear in ``(x, u)`` for each ``t ∈ \\mathbb{R}``. On the other hand,
``x' = f(t, x, u) = A x + B u + c`` is affine but not linear, since it is not
linear in ``(x, u)``.

This function uses the information of the type, not the value. So, if a system
type allows an instance that is not linear, it returns `false` by default.
For example, polynomial systems can be nonlinear; hence `islinear`
returns `false`.

[1] Sontag, Eduardo D. *Mathematical control theory: deterministic finite dimensional
systems.* Vol. 6. Springer Science & Business Media, 2013.
"""
function islinear(::AbstractSystem) end

"""
    isaffine(s::AbstractSystem)

Specifies if the dynamics of system `s` is specified by affine equations.

### Notes

An affine system is the composition of a linear system and a translation.
See [`islinear(::AbstractSystem)`](@ref) for the notion of linear system adopted
in this library.

This function uses the information of the type, not the value. So, if a system
type allows an instance that is not affine, it returns `false` by default.
For example, polynomial systems can be nonlinear; hence `isaffine` is `false`.
"""
function isaffine(::AbstractSystem) end

"""
    AbstractMap

Abstract supertype for all map types.
"""
abstract type AbstractMap end

"""
    outputdim(m::AbstractMap)

Returns the dimension of the output space of the map `m`.
"""
function outputdim end

"""
    outputmap(s::SystemWithOutput)

Returns the output map of a system with output.
"""
function outputmap end

"""
    islinear(m::AbstractMap)

Specifies if the map `m` is linear or not.

### Notes

A map is *linear* if it preserves the operations of scalar multiplication and
vector addition.
"""
function islinear(::AbstractMap) end

"""
    isaffine(m::AbstractMap)

Specifies if the map `m` is affine or not.

### Notes

An affine map is the composition of a linear map and a translation.
See also [`islinear(::AbstractMap)`](@ref).
"""
function isaffine(::AbstractMap) end

"""
    apply(m::AbstractMap, args...)

Apply the rule specified by the map to the given arguments.
"""
function apply end


using Flux
using Flux: onecold, onehotbatch
using Flux.Losses: logitcrossentropy
using GraphNeuralNetworks
using MLDatasets: Cora
using Statistics, Random
using CUDA
CUDA.allowscalar(false)

function eval_loss_accuracy(X, y, ids, model, g)
    ŷ = model(g, X)
    l = logitcrossentropy(ŷ[:,ids], y[:,ids])
    acc = mean(onecold(ŷ[:,ids]) .== onecold(y[:,ids]))
    return (loss = round(l, digits=4), acc = round(acc*100, digits=2))
end

# arguments for the `train` function
Base.@kwdef mutable struct Args
    η = 5f-3             # learning rate
    epochs = 10         # number of epochs
    seed = 17             # set seed > 0 for reproducibility
    usecuda = false      # if true use cuda (if available)
    nhidden = 64        # dimension of hidden features
end

function train(Layer; verbose=false, kws...)
    args = Args(; kws...)
    args.seed > 0 && Random.seed!(args.seed)

    if args.usecuda && CUDA.functional()
        device = Flux.gpu
        args.seed > 0 && CUDA.seed!(args.seed)
    else
        device = Flux.cpu
    end

    # LOAD DATA
    data = Cora.dataset()
    g = GNNGraph(data.adjacency_list, graph_type=GRAPH_T) |> device
    X = data.node_features |> device
    y = onehotbatch(data.node_labels, 1:data.num_classes) |> device
    train_ids = data.train_indices |> device
    val_ids = data.val_indices |> device
    test_ids = data.test_indices |> device
    ytrain = y[:,train_ids]

    nin, nhidden, nout = size(X,1), args.nhidden, data.num_classes

    ## DEFINE MODEL
    model = GNNChain(Layer(nin, nhidden),
                    #  Dropout(0.5),
                     Layer(nhidden, nhidden),
                     Dense(nhidden, nout))  |> device

    ps = Flux.params(model)
    opt = ADAM(args.η)

    ## TRAINING
    function report(epoch)
        train = eval_loss_accuracy(X, y, train_ids, model, g)
        test = eval_loss_accuracy(X, y, test_ids, model, g)
        println("Epoch: $epoch   Train: $(train)   Test: $(test)")
    end

    verbose && report(0)
    @time for epoch in 1:args.epochs
        gs = Flux.gradient(ps) do
            ŷ = model(g, X)
            logitcrossentropy(ŷ[:,train_ids], ytrain)
        end
        Flux.Optimise.update!(opt, ps, gs)
        verbose && report(epoch)
    end

    train_res = eval_loss_accuracy(X, y, train_ids, model, g)
    test_res = eval_loss_accuracy(X, y, test_ids, model, g)
    return train_res, test_res
end

function train_many(; usecuda=false)
    for (layer, Layer) in [
                ("GCNConv", (nin, nout) -> GCNConv(nin => nout, relu)),
                ("ResGatedGraphConv", (nin, nout) -> ResGatedGraphConv(nin => nout, relu)),
                ("GraphConv", (nin, nout) -> GraphConv(nin => nout, relu, aggr=mean)),
                ("SAGEConv", (nin, nout) -> SAGEConv(nin => nout, relu)),
                ("GATConv", (nin, nout) -> GATConv(nin => nout, relu)),
                ("GINConv", (nin, nout) -> GINConv(Dense(nin, nout, relu), 0.01, aggr=mean)),
                ## ("ChebConv", (nin, nout) -> ChebConv(nin => nout, 2)), # not working on gpu
                ## ("NNConv", (nin, nout) -> NNConv(nin => nout)),  # needs edge features
                ## ("GatedGraphConv", (nin, nout) -> GatedGraphConv(nout, 2)),  # needs nin = nout
                ## ("EdgeConv",(nin, nout) -> EdgeConv(Dense(2nin, nout, relu))), # Fits the traning set but does not generalize well
                ]

        @show layer
        @time train_res, test_res = train(Layer; usecuda, verbose=false)
        @test train_res.acc > 94
        @test test_res.acc > 70
    end
end

## if GRAPH_T != :dense # some erratic errors with :dense
train_many(usecuda=false)
if TEST_GPU
    train_many(usecuda=true)
end


println("Advent of Code 2017, day 6")

f = open("input6")
input = read(f,String)
close(f)

banks = [parse(Int,n) for n in split(input)]
n = length(banks)

# Part 1

function next(j)
  if j<n return j+1 else return 1 end
end

function redistribute(bk)
  b = copy(bk)
  (m,j) = findmax(b)
  b[j] = 0
  while m > 0
    j = next(j)
    b[j] += 1
    m -= 1
  end
  return b
end

cycles = [ banks ]
b = redistribute(banks)
i = 1
while !(b in cycles)
  append!(cycles,[b])
  global b = redistribute(b)
  global i += 1
end

println("Part 1: ",i)

# Part 2

index = findfirst(x -> x==b,cycles)

println("Part 2: ",i-index+1)


module ReinterpretTests

using Test
using Gridap.TensorValues

v = VectorValue(1.3,2.1)
V = fill(v,5)
A = reinterpret(V)
R = [1.3 1.3 1.3 1.3 1.3; 2.1 2.1 2.1 2.1 2.1]
@test A == R

v = TensorValue(2.0,3.5,1.0,4.0)
V = fill(v,5)
A = reinterpret(V)
R = [2.0 2.0 2.0 2.0 2.0; 3.5 3.5 3.5 3.5 3.5; 1.0 1.0 1.0 1.0 1.0; 4.0 4.0 4.0 4.0 4.0]
@test A == R

v = SymFourthOrderTensorValue(1,2,3, 4,5,6, 7,8,9)
V = fill(v,5)
A = reinterpret(V)
R = [1 1 1 1 1; 2 2 2 2 2; 3 3 3 3 3; 4 4 4 4 4; 5 5 5 5 5; 6 6 6 6 6; 7 7 7 7 7; 8 8 8 8 8; 9 9 9 9 9]
@test A == R

v = SymTensorValue(1,2,3)
V = fill(v,5)
A = reinterpret(V)
R = [1 1 1 1 1; 2 2 2 2 2; 3 3 3 3 3]
@test A == R

end # module ReinterpretTests


################################################################################

# Setup

# Note: This version is against the true solution, so adaptive timestepping is
# used. This means that the option which is set to get the correct timepoints
# is saveat, not tstops!

using DifferentialEquations, RecursiveArrayTools
using NLopt
using BlackBoxOptim

Xiang2015Bounds = Tuple{Float64, Float64}[(9, 11), (20, 30), (2, 3)] # for local optimizations
xlow_bounds = [9.0,20.0,2.0]
xhigh_bounds = [11.0,30.0,3.0]
LooserBounds = Tuple{Float64, Float64}[(0, 22), (0, 60), (0, 6)] # for global optimization
GloIniPar = [0.0, 0.5, 0.1] # for global optimizations
LocIniPar = [9.0, 20.0, 2.0] # for local optimization

g1 = @ode_def_nohes LorenzExample begin
  dx = σ*(y-x)
  dy = x*(ρ-z) - y
  dz = x*y - β*z
end σ=>10.0 ρ=>28.0 β=>2.6666

r0 = [0.1;0.0;0.0]
tspan = (0.0,30.0)
prob = ODEProblem(g1,r0,tspan)
tspan2 = (0.0,3.0)
prob_short = ODEProblem(g1,r0,tspan2)

dt = 30.0/3000
tf = 30.0
tinterval = 0:dt:tf
t  = collect(tinterval)

h = 0.01
M = 300
tstart = 0.0
tstop = tstart + M * h
tinterval_short = 0:h:tstop
t_short = collect(tinterval_short)

# Generate Data
data_sol_short = solve(prob_short,Vern7(),saveat=t_short,reltol=1e-12,abstol=1e-12)
data_short = convert(Array,data_sol_short)
data_sol = solve(prob,Vern7(),saveat=t,reltol=1e-12,abstol=1e-12)
data = convert(Array,data_sol)

# Plot the data

# using Plots
# plot(data_sol_short) # the short solution
# plot(data_sol) # the longer solution

################################################################################

# Using BlackBoxOptim

# Note: Euler uses tstops to hit the estimation timepoints exactly since it's not adaptive
obj_short = build_loss_objective(prob_short,Euler(),L2Loss(t_short,data_short),tstops=t_short)
res1 = bboptimize(obj_short;SearchRange = Xiang2015Bounds, MaxSteps = 11e3)
# Euler could not recover the correct results since its error is too high!

obj_short = build_loss_objective(prob_short,Tsit5(),L2Loss(t_short,data_short))
res1 = bboptimize(obj_short;SearchRange = Xiang2015Bounds, MaxSteps = 11e3)
# Tolernace is still too high to get close enough

obj_short = build_loss_objective(prob_short,Tsit5(),L2Loss(t_short,data_short),reltol=1e-9)
res1 = bboptimize(obj_short;SearchRange = Xiang2015Bounds, MaxSteps = 11e3)
# With the tolerance lower, it achieves the correct solution in 4.5 seconds.

obj_short = build_loss_objective(prob_short,Vern7(),L2Loss(t_short,data_short),reltol=1e-12,abstol=1e-12)
res1 = bboptimize(obj_short;SearchRange = Xiang2015Bounds, MaxSteps = 11e3)
# But too much accuracy in the numerical solution of the ODE actually leads to
# slower convergence, since each step takes longer!

# Fastest BlackBoxOptim: 4.5 seconds

################################################################################

# using NLopt
obj_short = build_loss_objective(prob_short,Tsit5(),L2Loss(t_short,data_short),reltol=1e-14)

opt = Opt(:GN_ORIG_DIRECT_L, 3)
lower_bounds!(opt,LocIniPar)
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,LocIniPar)

opt = Opt(:GN_CRS2_LM, 3)
lower_bounds!(opt,LocIniPar)
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,LocIniPar)

opt = Opt(:GN_ISRES, 3)
lower_bounds!(opt,LocIniPar)
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,LocIniPar)

opt = Opt(:LN_BOBYQA, 3)
lower_bounds!(opt,LocIniPar)
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,LocIniPar)
# This one took 0.04 seconds! Wow!

opt = Opt(:LN_NELDERMEAD, 3)
lower_bounds!(opt,LocIniPar)
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,LocIniPar)

opt = Opt(:LD_SLSQP, 3)
lower_bounds!(opt,LocIniPar)
upper_bounds!(opt,[11.0,30.0,3.0])
min_objective!(opt, obj_short)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,LocIniPar)

################################################################################

# Longer version

obj = build_loss_objective(prob,Euler(),L2Loss(t,data),tstops = t)
# res1 = bboptimize(obj;SearchRange = Xiang2015Bounds, MaxSteps = 8e3)
# Once again, Euler fails to convergence its error is too high

obj = build_loss_objective(prob,Vern7(),L2Loss(t,data),reltol=1e-14)
# res1 = bboptimize(obj;SearchRange = Xiang2015Bounds, MaxSteps = 8e3)
# BB with Tsit5 converges just fine in 14.5 seconds

opt = Opt(:GN_ORIG_DIRECT_L, 3)
lower_bounds!(opt,zeros(3))
upper_bounds!(opt,[22.0,60.0,6.0])
min_objective!(opt, obj)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,GloIniPar)

opt = Opt(:GN_CRS2_LM, 3)
lower_bounds!(opt,zeros(3))
upper_bounds!(opt,[22.0,60.0,6.0])
min_objective!(opt, obj)
xtol_rel!(opt,1e-12)
maxeval!(opt, Int(1e5))
@time (minf,minx,ret) = NLopt.optimize(opt,GloIniPar)

opt = Opt(:GN_ISRES, 3)
lower_bounds!(opt,zeros(3))
upper_bounds!(opt,[22.0,60.0,6.0])
min_objective!(opt, obj)
xtol_rel!(opt,1e-12)
maxeval!(opt, Int(1e6))
@time (minf,minx,ret) = NLopt.optimize(opt,GloIniPar)

opt = Opt(:GN_ESCH, 3)
lower_bounds!(opt,zeros(3))
upper_bounds!(opt,[22.0,60.0,6.0])
min_objective!(opt, obj)
xtol_rel!(opt,1e-12)
maxeval!(opt, Int(1e5))
@time (minf,minx,ret) = NLopt.optimize(opt,GloIniPar)

################################################################

# Local Optimizers fail on the hard problem, as expected
LocIniPar = [11.5312, 26.0192, 2.79983]
opt = Opt(:LN_BOBYQA, 3)
lower_bounds!(opt,xlow_bounds)
upper_bounds!(opt,[13.0,30.0,3.0])
min_objective!(opt, obj)
xtol_rel!(opt,1e-12)
maxeval!(opt, Int(1e5))
@time (minf,minx,ret) = NLopt.optimize(opt,LocIniPar)
# Converges in 0.05 seconds!

opt = Opt(:LN_NELDERMEAD, 3)
lower_bounds!(opt,xlow_bounds)
upper_bounds!(opt,[13.0,30.0,3.0])
min_objective!(opt, obj)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,LocIniPar)

opt = Opt(:LD_SLSQP, 3)
lower_bounds!(opt,xlow_bounds)
upper_bounds!(opt,[13.0,30.0,3.0])
min_objective!(opt, obj)
xtol_rel!(opt,1e-12)
maxeval!(opt, 10000)
@time (minf,minx,ret) = NLopt.optimize(opt,LocIniPar)

################################################################################

# Conclusion:

# 1) When doing these experiments, making sure the "data" is correct is key.
#    The lorenz_test.jl results are all skewed because the actual Euler solution
#    that made the data is skewed. But, since it's a fixed timestep, it was able
#    to find the parameters that made the data, given that it's hitting all the
#    same points. That's not the case here. Because there's adaptive timestepping,
#    we're getting rid of the ability to "overfit", i.e. simply match the data
#    we generated. Instead, we have to hit the same solution, albeit at different
#    stepping points, making it less prone to the bias.
# 2) When trying to hit the real data, having a low enough tolerance on the
#    numerical solution is key. If the numerical solution is too rough, then
#    we can never actually hone in on the true parameters since even with the
#    true parameters we will erroneously induce numerical error. Maybe this
#    could be adaptive?
# 3) Too low of a tolerance in the numerical solution is unnecessary and will
#    slow down the convergence of the estimation.
# 4) The estimation method makes a huge difference in the timings. Here, BBO
#    doesn't do very well. Some of the methods in NLopt, like :LN_BOBYQA, solve
#    the problem in <0.05 seconds. This is an algorithm that can scale.
# 5) Fitting shorter timespans is easier... maybe this could be some kind of
#    adaptivity trick as well?


# ---
# title: 847. Shortest Path Visiting All Nodes
# id: problem847
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Dynamic Programming, Breadth-first Search
# link: <https://leetcode.com/problems/shortest-path-visiting-all-nodes/description/>
# hidden: true
# ---
#
# An undirected, connected graph of N nodes (labeled `0, 1, 2, ..., N-1`) is
# given as `graph`.
#
# `graph.length = N`, and `j != i` is in the list `graph[i]` exactly once, if
# and only if nodes `i` and `j` are connected.
#
# Return the length of the shortest path that visits every node. You may start
# and stop at any node, you may revisit nodes multiple times, and you may reuse
# edges.
#
#
#
# **Example 1:**
#
#
#
#     Input: [[1,2,3],[0],[0],[0]]
#     Output: 4
#     **Explanation** : One possible path is [1,0,2,0,3]
#
# **Example 2:**
#
#
#
#     Input: [[1],[0,2,4],[1,3,4],[2],[1,2]]
#     Output: 4
#     **Explanation** : One possible path is [0,1,4,2,3]
#
#
#
#
# **Note:**
#
#   1. `1 <= graph.length <= 12`
#   2. `0 <= graph[i].length < graph.length`
#
#
## @lc code=start
using LeetCode

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


"""
Take data in `input` and replace variable symbols in `modelex`
expression with the arrays β and ζ that hold the unknown parameters and
estimates of measured parameters, respectively.
"""
function _makemodel(input::LSQInput, modelex::Expr; debug=false)

    modelname = gensym("model")

    # substitute symbols in model expression with entries in beta and zeta,
    # according to mapping in input
    subssyms = let
        substitutions = [k => :(ζ[$v]) for (k, v) in input.zindex]
        append!(substitutions, [k => :(β[$v]) for (k, v) in input.βindex])
        Dict(substitutions)
    end
    modelex = subs(modelex, subssyms)

    modelcode = quote
        function $modelname(β::Vector{Float64}, ζ::Vector{Float64})
            $modelex
        end
    end

    # print generated model to identify problems
    if debug
        println(modelcode)
    end

    eval(modelcode)
end

"""
    yourmodel = @withinput input begin
        # your model expression
    end

Declare a model using the data in `input`.

All variables that are not assigned to must be present in
the input data, i.e.

```
y = x + 1
```

will throw an error when the model is evaluated if `x` is not included
in the input data.

**The model has to return the values of the constraints**, e.g.
if two constraints are present:

```
y - a*x = 0
a - 2b  = 0
```

The model has to return:

```
return [
    y - a*x
    a - 2b
]
```

See README and tests for further examples.
"""
macro withinput(input::Symbol, modelex::Expr)

    modelex = Expr(:quote, modelex)
    esc(quote
        ConstrainedLeastSquares._makemodel($input, $modelex)
    end)
end

"""
Like `@withinput`, but prints the generated model.
"""
macro withinput_debug(input::Symbol, modelex::Expr)

    modelex = Expr(:quote, modelex)
    quote
        makemodel($input, $modelex; debug=true)
    end
end


# script to create all plots

for file in ("plot_intro_exact_figure_1a.jl",
             "plot_intro_approximate_figure_1b.jl",
             "plot_intro_recurrence_figure_1c.jl",
             "plot_homogenize_figure_2a_and_figure_8.jl",
             "plot_shrink_delta_figure_2b.jl",
             "plot_delta_figure_3.jl",
             "plot_X0_figure_4.jl",
             "plot_U_figure_5.jl",
             # figure 6 is generated from the experiment runtime data
             "plot_underapproximation_figure_7.jl")
    println("Running $file ...")
    include(file)
end

nothing


"""
    shift_constant(set::MOI.AbstractScalarSet,
                   offset)

Returns a new scalar set `new_set` such that `func`-in-`set` is equivalent to
`func + offset`-in-`new_set`.

## Examples

The call `shift_constant(MOI.Interval(-2, 3), 1)` is equal to
`MOI.Interval(-1, 4)`.
"""
function shift_constant(set::Union{MOI.LessThan{T},
                                   MOI.GreaterThan{T},
                                   MOI.EqualTo{T}},
                        offset::T) where T
    return typeof(set)(MOI.constant(set) + offset)
end
function shift_constant(set::MOI.Interval, offset)
    return MOI.Interval(set.lower + offset, set.upper + offset)
end

const ScalarLinearSet{T} = Union{MOI.EqualTo{T}, MOI.LessThan{T}, MOI.GreaterThan{T}}
const VectorLinearSet = Union{MOI.Zeros, MOI.Nonnegatives, MOI.Nonpositives}

vector_set_type(::Type{<:MOI.EqualTo}) = MOI.Zeros
vector_set_type(::Type{<:MOI.LessThan}) = MOI.Nonpositives
vector_set_type(::Type{<:MOI.GreaterThan}) = MOI.Nonnegatives

scalar_set_type(::Type{<:MOI.Zeros}, T::Type) = MOI.EqualTo{T}
scalar_set_type(::Type{<:MOI.Nonpositives}, T::Type) = MOI.LessThan{T}
scalar_set_type(::Type{<:MOI.Nonnegatives}, T::Type) = MOI.GreaterThan{T}

"""
    is_diagonal_vectorized_index(index::Base.Integer)

Return whether `index` is the index of a diagonal element in a
[`MOI.AbstractSymmetricMatrixSetTriangle`](@ref) set.
"""
function is_diagonal_vectorized_index(index::Base.Integer)
    perfect_square = 1 + 8index
    return isqrt(perfect_square)^2 == perfect_square
end

"""
    side_dimension_for_vectorized_dimension(n::Integer)

Return the dimension `d` such that
`MOI.dimension(MOI.PositiveSemidefiniteConeTriangle(d))` is `n`.
"""
function side_dimension_for_vectorized_dimension(n::Base.Integer)
    return div(isqrt(1 + 8n), 2)
end


#=
Julia has a few packages for visualization, depending on your tastes.

\toc

## Plots

Plots is probably the most well-known, and also one of the oldest.
It provides an abstraction layer and you can select the backend that actually does the plotting.
See the backend list [here](https://docs.juliaplots.org/latest/backends/).

When you install Plots it also install GR, which is quite fast and has some decent graphics.
GR is the default choice, but if you alternate from it, you can activate it again with `gr()`.

The PyPlot package provides another common backend.
It uses Python's matplotlib, which provides very nice graphics, although it is a bit slower than GR.
Change to PyPlot using `pyplot()`.

You can also use UnicodePlots for plotting directly into the REPL.
Change using `unicodeplots()`.

PlotlyJS is also quite well-liked, because it usually allows some interactivity in the plots.
Use `plotlyjs()` to change.

Some packages are more fickle than others to install.
GR, for instance, decided to not work with GitHub Actions for this site.
For that reason, I use mostly PyPlot.

Here are some examples with PyPlot.
Try them with other backends.
=#
using LinearAlgebra, Plots, Random
pyplot()

Random.seed!(0)
n = 100
x = range(-1, 1, length=n)
y = 5 * exp.(-10x.^2) + x + 0.2 * randn(n)
Etr, Ete = Float64[], Float64[]
Ite = shuffle(1:n)[1:2 * div(n, 5)]
Itr = setdiff(1:n, Ite)
pmax = 6
plt = plot(title="ML stuff", layout=grid(1, 2), size=(800,400))
scatter!(plt[1], x, y, leg=false)
for p = 1:pmax
  X = [x[i]^j for i = 1:n, j = 0:p]
  β = X[Itr,:] \ y[Itr]
  plot!(plt[1], t -> β[1] + sum(β[j+1] * t^j for j = 1:p), -1, 1)
  r = y - X * β
  push!(Etr, norm(r[Itr])^2 / length(Itr))
  push!(Ete, norm(r[Ite])^2 / length(Ite))
end
plot!(title="MSE")
plot!(plt[2], Etr, m=(3,:blue,:square), lab="Train")
plot!(plt[2], Ete, m=(3,:red,:circle), lab="Test")
png(joinpath(@OUTPUT, "vis-plots-1"))

# \fig{vis-plots-1.png}

X = [
  0.4 * randn(div(n, 2), 2) .+ 3;
  1.2 * randn(div(n, 2), 2) .- 1;
]
h = 1.0
K(x) = exp(-norm(x)^2 / h) / h
f(x) = sum(K(x - X[i,:]) for i = 1:n) / n

plt = plot(layout=grid(2,2), leg=false)
scatter!(plt[1], X[:,1], X[:,2], title="scatter")
contour!(plt[2], -5:0.1:5, -5:0.1:5, (x,y) -> f([x;y]), title="countour", levels=50)
surface!(plt[3], -5:0.1:5, -5:0.1:5, (x,y) -> f([x;y]), title="surface")
surface!(plt[4], -5:0.1:5, -5:0.1:5, (x,y) -> f([x;y]), title="surface", camera=(70,10))
plt
png(joinpath(@OUTPUT, "vis-plots-2"))

# \fig{vis-plots-2.png}

θg = range(0, 2π, length=120)
anim = Animation()
n = 3
plot(leg=false, axis=false, grid=false, size=(600,600))
for i = 1:6
  θr = range(0, 2π, length=n + 1)
  plot!(cos.(θr), sin.(θr), c=:blue, opacity=0.5, lw=3)
  radius = 1 / cos(π / n)
  plot!(radius * cos.(θr .+ π / n), radius * sin.(θr .+ π / n), c=:red, opacity=0.5, lw=3)
  xlims!(-2, 2)
  ylims!(-2, 2)
  frame(anim)
  global n = 2n
end

gif(anim, joinpath(@OUTPUT, "vis-plots-anim-1.gif"), fps=1)

# \fig{vis-plots-anim-1.gif}

#=
## PyPlot

PyPlot uses Python's matplotlib, as said above.
We can use it as a backend to Plots, but it can also be used directly.
The main advantage are that the syntax is closer to matplotlib's syntax, and Plots doesn't support everything that PyPlot does.
=#

using PyPlot
x = range(0, 2π, length=1000)
y = sin.(3 * x + 4 * cos.(2 * x))
PyPlot.figure()
PyPlot.plot(x, y, color="red", linewidth=2.0, linestyle="--")
PyPlot.title("A sinusoidally modulated sinusoid")
PyPlot.savefig(joinpath(@OUTPUT, "vis-pyplot-1"))

# \fig{vis-pyplot-1.png}

#=
## Makie

Makie is a complete alternative to Plots.
It tries to achieve high performance using heavier backends, such as OpenGL.

## StatsPlots

This package provides "recipes for Plots.jl" - i.e., new plotting functions - for statistical plots.
=#
using RDatasets, StatsPlots

df = dataset("datasets", "iris")
@df df cornerplot(cols(1:4), compact=true)
png(joinpath(@OUTPUT, "vis-statsplots-1"))

# \fig{vis-statsplots-1.png}

plt = plot(layout=grid(2, 2))
for i = 1:4
  StatsPlots.violin!(plt[i], df.Species, df[!,i], lab="", ylabel=names(df)[i])
  StatsPlots.boxplot!(plt[i], df.Species, df[!,i], lab="", fillalpha=0.6, m=(0))
  StatsPlots.dotplot!(plt[i], df.Species, df[!,i], lab="", m=(:white, 2), opacity=0.6)
end
png(plt, joinpath(@OUTPUT, "vis-statsplots-2"))

# \fig{vis-statsplots-2.png}

#=
## GraphRecipes

Plot recipes for Graphs.
=#
using GraphRecipes
g = rand(0:1, 5, 5)
graphplot(g)
png(joinpath(@OUTPUT, "vis-graphrecipes-1"))

# \fig{vis-graphrecipes-1.png}


"""
    Prechelt

> Find a sequence of words such that the sequence of characters in these words exactly corresponds to the sequence of
> digits in the phone number. All possible solutions must be found and printed. The solutions are created word by word
> and if no word from the dictionary can be inserted at some point during that process, a single digit from the phone
> number can appear in the result at that position. Many phone numbers have no solution at all. Here is an example of
> the program output for the phone number “3586-75,” when the dictionary contained the words "Dali," "um," "Sao," "da,"
> "Pik," and 73,108 others:
>
>     3586-75: Dali um
>     3586-75: Sao 6 um
>     3586-75: da Pik 5
>
> A list of partial solutions needs to be maintained by the program while processing each number, and the dictionary
> must be embedded in a supporting data structure (such as a 10-ary digit tree) for efficient access.
>
>  -- <cite>Prechelt, Lutz. (1999). Comparing Java vs. C/C++ Efficiency Differences to Interpersonal Differences. Communications of the ACM. 42. 10.1145/317665.317683.</quote>

This implementation uses the [Aho-Corasick algorithm](https://en.wikipedia.org/wiki/Aho%E2%80%93Corasick_algorithm)
with two small specializations

    * words are mapped to digits as they are inserted in the Aho-Corasick data structure
    * the data structure uses fixed length arrays to minimize memory allocations

```jldoctest
julia> replaceall("358", ["Sao", "dal"])
1-element Vector{String}:
 "dal"

```
"""

const DIGITS = Vector{Int}(undef, 128)
DIGITS[codeunits("0123456789")] = 1:10
DIGITS[codeunits("ejnqrwxdsyftamcivbkulopghz")] = DIGITS[codeunits("01112223334455666777888999")]
DIGITS[codeunits("EJNQRWXDSYFTAMCIVBKULOPGHZ")] = DIGITS[codeunits("01112223334455666777888999")]

mutable struct Node
    depth::Int
    keys::Vector{String}
    next::Vector{Union{Missing,Node}}
    fail::Node

    function Node(depth::Int=0)
        node = new()
        node.depth = depth
        node.keys = String[]
        node.next = Vector{Union{Missing,Node}}(missing, 10)
        node
    end
end

digits(s::String) = DIGITS[codeunits(s)]

function insert!(node::Node, word::String)
    for digit ∈ digits(word)
        next = node.next[digit]
        if ismissing(node.next[digit])
            node.next[digit] = next = Node(node.depth + 1)
        end
        node = next
    end
    push!(node.keys, word)
end

function add_fail_transition!(node::Node)
    fail = node.fail
    for (index, next) ∈ enumerate(node.next)
        ismissing(next) && continue
        next.fail = coalesce(fail.next[index], fail)
    end
    for next ∈ skipmissing(node.next)
        add_fail_transition!(next)
    end
end

function search(node::Node, index::Int)
    coalesce(node.next[index], node.depth == 0 ? node : search(node.fail, index))
end

function aho_corasick(words::Union{Base.EachLine,Vector{String}})
    node = Node()
    node.fail = node
    for word ∈ words
        insert!(node, word)
    end
    for next ∈ skipmissing(node.next)
        next.fail = node
    end
    for next ∈ skipmissing(node.next)
        add_fail_transition!(next)
    end
    node
end

# varargs suppress StaticLint's "Possible method call error"
⊗(x...) = kron(x...)  # https://julialang.org/blog/2017/01/moredots/

function replaceall(number::String, node::Node)
    """
    ```jldoctest
    julia> replaceall("358675", ["Dali", "um", "Sao", "da", "Pik"])
    3-element Vector{String}:
     "Dalium"
     "Sao6um"
     "daPik5"

    ```
    """
    heads = Vector{Vector{String}}(undef, length(number))
    for (k, digit) ∈ enumerate(digits(number))
        node = next = search(node, digit)
        while length(next.keys) > 0
            rest = next.keys
            j = k - node.depth
            i = j - 1

            heads[k] = (
                isassigned(heads, i)
                ? [heads[j] ⊗ rest; heads[i] ⊗ (number[i:i] .* rest)]
                : isassigned(heads, j)
                ? heads[j] ⊗ rest
                : rest
            )
            next = next.fail
        end
    end
    filter(s -> count(r"\d", s) <= 1, [heads[end]; heads[end - 1] .* number[end:end]])
end

replaceall(number::String, words::Vector{String}) = replaceall(number, aho_corasick(words))