using BilevelOptimization
using BilevelOptimization.BilevelFlowProblems

using Test
using LinearAlgebra: I

import Cbc
using JuMP

test_bp() = BilevelLP(
    [1.,0.], [0.],
    [-.1 0.], -ones(1,1), [-1.],
    [-1.],
    [-.1 0.], ones(1,1), ones(1,)
)

@testset "BilevelLP type" begin
    bp = test_bp()
    @test bp isa BilevelLP{V,M} where{V<:AbstractVector{Float64},M<:AbstractMatrix{Float64}}
end

@testset "test basic problem" begin
    bp = test_bp()
    (m, x, y, λ, _) = build_blp_model(bp, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0))
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    xv = JuMP.value.(x)
    yv = JuMP.value.(y)
    @test xv[1] ≈ 0.0
    @test yv[1] ≈ 1.0
end

@testset "Integrality is registered" begin
    bp = test_bp()
    push!(bp.Jx, 2) # second useless variable is integer
    (m, x, y, λ, _) = build_blp_model(bp, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0))
    # @test JuMP.getcategory(x[2]) == :Int
    # @test JuMP.getcategory(x[1]) == :Cont
    # @test JuMP.getcategory(y[1]) == :Cont
end

@testset "Problem from Dempe-Mersha 2006" begin
    # problem taken from:
    # A.G. Mersha, S. Dempe,
    # Linear bilevel programming with upper-level constraints depending on lower-level solution.
    # Appl. Math. Computation (180), 2006
    cx = [-1.]
    cy = [-2.]
    G = zeros(2,1) .+ [-2.,1.]
    H = zeros(2,1) .+ [3.,1.]
    q = [12.,14.]
    d = [-1.]
    A = zeros(2,1) .+ [-3.,3.]
    B = ones(2,1)
    b = [-3.,30.]
    prob = BilevelLP(
        cx, cy,
        G, H, q,
        d, A, B, b
    )
    (m, x, y, λ, _) = build_blp_model(prob, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0))
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    @test JuMP.value(x[1]) ≈ 8
    @test JuMP.value(y[1]) ≈ 6

    # problem modification: passing upper-level constraints to lower level
    A = zeros(4,1) .+ [-3.,3.,-2.,1.]
    B = ones(4,1); B[3,1] = 3.
    b = [-3.,30.,12.,14.]
    prob = BilevelLP(
        cx, cy,
        G, H, q,
        d, A, B, b
    )
    (m, x, y, λ, _) = build_blp_model(prob, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0))
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    @test JuMP.value(x[1]) ≈ 6
    @test JuMP.value(y[1]) ≈ 8
end

@testset "Problem with upper bounds" begin
    # problem taken from:
    # Bilevel Programming Problems
    # S. Dempe, V. Kalashnikov, G. Pérez-Valdés, N. Kalashnykova
    # Springer 2015
    cx = [2.,1.]
    cy = [2.,-1.]
    G = zeros(0,2)
    H = zeros(0,2)
    q = Float64[]
    d = [0.,0.]
    F = [1. 0.;0. 1.]
    A = zeros(1,2)
    B = [-2. 1.]
    b = [0.]
    prob = BilevelLP(
        cx, cy,
        G, H, q,
        d, A, B, b, Int[], F
    )
    set_lower_bound(prob, BilevelOptimization.upper, 1, -1.)
    setupperbound(prob, BilevelOptimization.upper, 1, 1.)
    set_lower_bound(prob, BilevelOptimization.upper, 2, -1.)
    setupperbound(prob, BilevelOptimization.upper, 2, -0.75)
    set_lower_bound(prob, BilevelOptimization.lower, 1, -Inf64)
    setupperbound(prob, BilevelOptimization.lower, 1, 2.)
    set_lower_bound(prob, BilevelOptimization.lower, 2, 0.)
    setupperbound(prob, BilevelOptimization.lower, 2, 2.)
    @test size(prob.B) == (prob.ml,prob.nl)
    (m, x, y, λ, _) = build_blp_model(prob, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0))
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    @test all(JuMP.value.(x) .≈ (-1.,-1.))
end

@testset "Upper-level integer problem" begin
    # problem taken from:
    # Bilevel Programming Problems
    # S. Dempe, V. Kalashnikov, G. Pérez-Valdés, N. Kalashnykova
    # Springer 2015
    cx = [2.,0.]
    cy = [3.,2.,6.]
    G = [ 4.  1.
         -4. -1.]
    H = zeros(2,3)
    q = [10.,-10.]
    Jx = Int64[1, 2]
    d = [-5.,-8.,-1.]
    F = zeros(2,3)
    B = [4. 2. 0.
         2. 4. 1.
        -1. 0. 0.
         0. -1. 0.
         0. 0. -1.
         ]
    A = [-1.  0.
          0. -1.
          0.  0.
          0.  0.
          0.  0.]
    b = [0.,0.,0.,0.,0.]
    intprob = BilevelLP(
        cx, cy,
        G, H, q,
        d, A, B, b, Jx, ylowerbound = false
    )
    (m, x, y, λ, _) = build_blp_model(intprob, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0))
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    @test all(JuMP.value.(x) .≈ (2.,2.))
    @test all(JuMP.value.(y) .≈ (1//3,1//3,0))
    B = [4. 2. 0.
         2. 4. 1.]
    A = [-1.  0.
          0. -1.]
    b = [0.,0.]
    intprob = BilevelLP(
        cx, cy,
        G, H, q,
        d, A, B, b, Jx
    )
    (m, x, y, λ, _) = build_blp_model(intprob, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0))
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    @test all(JuMP.value.(x) .≈ (2.,2.))
    @test all(JuMP.value.(y) .≈ (1//3,1//3,0))
end

function test_bflow()
    init_cost = [
        0. 1. 1. 4.
        0. 0. 0. 1.
        0. 0. 0. 1.
        0. 0. 0. 0.
    ]
    taxable_edges = [
        false true true false
        false false false true
        false false false true
        false false false false
    ]
    tax_options = zeros(4,4,5)
    for i in 1:4, j in 1:4
        if taxable_edges[i,j]
            tax_options[i,j,:] .= (0.0, 0.5, 1.0, 1.5, 2.0)
        end
    end

    capacities = [
        0. 3. 2. 3.
        0. 0. 0. 2.
        0. 0. 0. 1.
        0. 0. 0. 0.
    ]
    BilevelFlowProblem(init_cost,taxable_edges,capacities,tax_options,3.)
end

@testset "Flow problem matrix construction" begin
    bfp = test_bflow()
    (B,b) = BilevelFlowProblems.flow_constraint_standard(bfp)
    @test size(B)[1] == bfp.nv-2 + 1 + bfp.nv*bfp.nv
    @test sum(B[1,:]) ≈ -3.
    @test B[1,5] ≈ -1.
    @test B[1,9] ≈ -1.
    @test B[1,13] ≈ -1.
    @test b[1] ≈ -bfp.minflow
    # flow conservation
    @test B[2,5] ≈ 1.
    @test B[2,14] ≈ -1.
    @test b[2] ≈ 0.
    @test B[3,9] ≈ 1.
    @test B[3,15] ≈ -1.
    @test b[3] ≈ 0.
    # capacity constraints
    @test all(I - B[4:end,:] .≈ 0)
    @test all(b[4:end] .≈ [bfp.capacities[i] for i in eachindex(bfp.capacities)])
end

@testset "Bilevel flow JuMP model" begin
    bfp = test_bflow()
    (m, r, y, f, λ) = build_blp_model(bfp, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0))
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    @test JuMP.objective_value(m) ≈ 6.
    for j in 1:size(r)[2]
        for i in 1:size(r)[1]
            @test JuMP.value(r[i,j]) ≈ sum(JuMP.value.(y[i,j,:]).*bfp.tax_options[i,j,:]) * JuMP.value(f[i,j])
        end
    end
end

@testset "Complementarity types" begin
    m = JuMP.Model()
    # testing non-crashing
    BilevelOptimization.add_complementarity_constraint(m, SOS1Complementarity(), [], [], [], [])
end

@testset "Basic problem with specified methods" begin
    bp = test_bp()
    (m, x, y, λ) = build_blp_model(bp, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0), comp_method = SOS1Complementarity())
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    xv = JuMP.value.(x)
    yv = JuMP.value.(y)
    @test xv[1] ≈ 0.0
    @test yv[1] ≈ 1.0
    (m, _, _, _) = build_blp_model(bp, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0), comp_method = SOS1Complementarity())
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    # arbitrary big-enough bounds
    (m, x, y, λ) = build_blp_model(bp, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0), comp_method = BoundComplementarity(100.,100.))
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    xv = JuMP.value.(x)
    yv = JuMP.value.(y)
    @test xv[1] ≈ 0.0
    @test yv[1] ≈ 1.0
end

@testset "Bilevel flow big-M bounds" begin
    bfp = test_bflow()
    (m, r, y, f, λ) = build_blp_model(bfp, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0), comp_method = BoundComplementarity(100., 100.))
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    @test JuMP.objective_value(m) ≈ 6.
    for j in 1:size(r)[2]
        for i in 1:size(r)[1]
            @test JuMP.value(r[i,j]) ≈ sum(JuMP.value.(y[i,j,:]).*bfp.tax_options[i,j,:]) * JuMP.value(f[i,j])
        end
    end
end

@testset "Bilevel flow big-M vector bounds" begin
    bfp = test_bflow()
    bounds_method = BoundComplementarity(10 .* ones(19), 30.)
    (m, r, y, f, λ) = build_blp_model(bfp, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0), comp_method = bounds_method)
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    @test JuMP.objective_value(m) ≈ 6.
    for j in 1:size(r)[2]
        for i in 1:size(r)[1]
            @test JuMP.value(r[i,j]) ≈ sum(JuMP.value.(y[i,j,:]) .* bfp.tax_options[i,j,:]) * JuMP.value(f[i,j])
        end
    end
    bounds_method = BoundComplementarity(10., 10 .* ones(19))
    (m, r, y, f, λ) = build_blp_model(bfp, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0), comp_method = bounds_method)
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.OPTIMAL
    @test JuMP.objective_value(m) ≈ 6.
    for j in 1:size(r)[2]
        for i in 1:size(r)[1]
            @test JuMP.value(r[i,j]) ≈ sum(JuMP.value.(y[i,j,:]) .* bfp.tax_options[i,j,:]) * JuMP.value(f[i,j])
        end
    end
end

@testset "Bilevel flow big-M infeasible" begin
    bfp = test_bflow()
    (m, r, y, f, λ) = build_blp_model(bfp, optimizer_with_attributes(Cbc.Optimizer, "LogLevel" => 0), comp_method = BoundComplementarity(0.1, 0.1))
    JuMP.optimize!(m)
    @test termination_status(m) === MOI.INFEASIBLE
end


export ent_spec, ent_S, entropy
"""
SchmidtMatrix{Tm <: Number, Ta <: SubArray, Tb <: SubArray, Ti <: Integer}

Struct that is used to construct the Schmidt matrix:

`|Ψ⟩ = ∑ᵢⱼ Mᵢᵢ|ψᵢ⟩ ⊗ |ψⱼ⟩`

Properties:
-----------
- `M`   : Schidt matrix.
- `A`   : View of region A.
- `B`   : View of region B.
- `base`: Base.
"""
struct SchmidtMatrix{Tm <: Number, Ta <: SubArray, Tb <: SubArray, Ti <: Integer}
    M::Matrix{Tm}
    A::Ta
    B::Tb
    base::Ti
end

Matrix(S::SchmidtMatrix) = S.M

"""
    SchmidtMatrix(T::DataType, b::AbstractBasis, Ainds::AbstractVector{Ta}) where Ta <: Integer

Construction of `SchmidtMatrix`.
"""
function SchmidtMatrix(T::DataType, b::AbstractBasis, Ainds::AbstractVector{Ta}) where Ta <: Integer
    L = length(b)
    B = base(b)
    Binds = Vector{Ta}(undef, L-length(Ainds))
    P = 1
    for i in range(one(Ta), stop=convert(Ta, L))
        if !in(i, Ainds)
            Binds[P] = i
            P += 1
        end
    end
    M = zeros(T, B^length(Ainds), B^length(Binds))
    SchmidtMatrix(M, view(digits(b), Ainds), view(digits(b), Binds), B)
end

function addto!(S::SchmidtMatrix, val::Number)
    ia = index(S.A, base=S.base)
    ib = index(S.B, base=S.base)
    S.M[ia, ib] += val
end

"""
    Vector(v::AbstractVector, b::AbstractBasis)

Convert a state in the symmetry sector to the original state in tensor-product Hilbert space.
"""
Vector(v::AbstractVector, b::AbstractBasis) = reshape(schmidt(v, 1:L, b), :)

"""
    ent_spec(v::AbstractVector, Aind::AbstractVector{<:Integer}, b::AbstractBasis)

Conpute the entanglement spectrum.
"""
ent_spec(v::AbstractVector, Aind::AbstractVector{<:Integer}, b::AbstractBasis) = svdvals(schmidt(v, Aind, b))

"""
    entropy(s::AbstractVector{<:Real}; α::Real=1, cutoff::Real=1e-20)

Compute the entropy of Schmidt values.

Inputs:
-------
- `s`     : Schmidt values.
- `α`     : Renyi index.
- `cutoff`: Cutoff of the Schmidt values.

Outputs:
--------
- `S`: Entanglement entropy.
"""
function entropy(s::AbstractVector{<:Real}; α::Real=1, cutoff::Real=1e-20)
    if isone(α)
        shannon_entropy(s, cutoff=cutoff)
    elseif iszero(α)
        renyi_zero_entropy(s, cutoff=cutoff)
    else
        renyi_entropy(s, α)
    end
end

"""
Compute Shannon entropy
"""
function shannon_entropy(s::AbstractVector{<:Real}; cutoff::Real=1e-20)
    ent = 0.0
    for si in s
        if si > cutoff
            ent -= si * log(si)
        else
            break
        end
    end
    ent
end

"""
Compute Renyi-0 entropy, which is thenumber of non-zero Schmidt values.
"""
function renyi_zero_entropy(s::AbstractVector{<:Real}; cutoff::Real=1e-20)
    N = 0
    for si in s
        if si > cutoff
            N += 1
        else
            break
        end
    end
    N
end

"""
Compute general Renyi entropy
"""
renyi_entropy(s::AbstractVector{<:Real}, α::Real) = log(sum(s.^α)) / (1-α)

"""
    ent_S(v::AbstractVector, Aind::AbstractVector{<:Integer}, b::AbstractBasis; α::Real=1, cutoff::Real=1e-20)

Conpute the entanglement entropy of a state.
"""
function ent_S(v::AbstractVector, Aind::AbstractVector{<:Integer}, b::AbstractBasis; α::Real=1, cutoff::Real=1e-20)
    s = ent_spec(v, Aind, b) .^ 2
    entropy(s, α=α, cutoff=cutoff)
end


module TestWiringLayers

using Test
using Catlab.WiringDiagrams

# Low-level graph interface
###########################

f = WiringLayer(3, 3)
@test nwires(f) == 0
@test wires(f) == []

add_wire!(f, 1 => 3)
@test has_wire(f, 1 => 3)
@test !has_wire(f, 3 => 1)
@test_throws BoundsError add_wire!(f, 4 => 1) # Source out of bounds
@test_throws BoundsError add_wire!(f, 1 => 4) # Target out of bounds

add_wires!(f, [2 => 2, 3 => 1])
@test nwires(f) == 3
@test wires(f) == [1 => 3, 2 => 2, 3 => 1]

@test nwires(f, 1 => 3) == 1
add_wires!(f, [1 => 3, 1 => 3])
@test nwires(f, 1 => 3) == 3
@test wires(f) == [1 => 3, 1 => 3, 1 => 3, 2 => 2, 3 => 1]
@test out_wires(f, 1) == [1 => 3, 1 => 3, 1 => 3]
rem_wire!(f, 1 => 3)
@test nwires(f, 1 => 3) == 2
rem_wires!(f, 1 => 3)
@test !has_wire(f, 1 => 3)
@test nwires(f, 1 => 3) == 0
@test wires(f) == [2 => 2, 3 => 1]

# High-level categorical interface
##################################

A, B, C = NLayer(1), NLayer(1), NLayer(1)

# Category
@test id(A) == complete_layer(1,1)
@test compose(id(A), id(A)) == id(A)

# Symmetric monoidal category
@test otimes(id(A), id(B)) == id(otimes(A,B))
@test compose(braid(A,B), braid(B,A)) == id(otimes(A,B))

# Diagonals
@test mcopy(A) == complete_layer(1,2)
@test delete(A) == complete_layer(1,0)
@test compose(otimes(mcopy(A), mcopy(B)), otimes(id(A), braid(A,B), id(B))) ==
  mcopy(otimes(A,B))
@test compose(mcopy(A), otimes(mcopy(A), id(A))) == mcopy(A, 3)
@test compose(mcopy(A), otimes(id(A), mcopy(A))) == mcopy(A, 3)
@test compose(mcopy(A), otimes(id(A), delete(A))) == id(A)

# Codiagonals
@test mmerge(A) == complete_layer(2,1)
@test create(A) == complete_layer(0,1)
@test compose(otimes(id(A), braid(B,A), id(B)), otimes(mmerge(A), mmerge(B))) ==
  mmerge(otimes(A,B))
@test compose(otimes(mmerge(A), id(A)), mmerge(A)) == mmerge(A, 3)
@test compose(otimes(id(A), mmerge(A)), mmerge(A)) == mmerge(A, 3)
@test compose(otimes(id(A), create(A)), mmerge(A)) == id(A)

# Wiring diagrams
#################

function roundtrip(layer::WiringLayer)::WiringLayer
  ports = n::Int -> repeat([nothing], n)
  diagram = to_wiring_diagram(layer, ports(layer.ninputs), ports(layer.noutputs))
  wiring_layer_between(diagram, input_id(diagram), output_id(diagram))
end

layer = WiringLayer(3, 3)
add_wires!(layer, [1 => 1, 1 => 2, 2 => 3, 3 => 1, 3 => 3])
@test roundtrip(layer) == layer

end


"""
    AbstractApprox

Supertype of types specifying how equality or approximate equality should
be computed. See `Equal`, `Approx`, and `EachApprox`.
"""
abstract type AbstractApprox end

"""
    Equal <: AbstractApprox

Approximate equality test that actually demands strict equality.
"""
struct Equal <: AbstractApprox
end

"""
    EachApprox <: AbstractApprox

Demands that each pair of elements are approximately equal.
`kw` are keyword pairs that are forwarded to `isapprox`.
"""
struct EachApprox{T} <: AbstractApprox
    kw::T
end
EachApprox(; kws...) = EachApprox(kws)

"""
   Approx(; kws...) <: AbstractApprox

Specifies using the legacy `isapprox` interface. For example,
for `AbstractMatrix`, matrix norms are used to test closeness.
`kw` are keyword pairs that are forwarded to `isapprox`.
For example, `Approx(atol=1e-9)`.
"""
struct Approx{T} <: AbstractApprox
    kw::T
end
Approx(; kws...) = Approx(kws)

"""
    isapprox(::Equal, x, y)

Return `true` if `x == y`.
"""
Base.isapprox(::Equal, x, y) = (x == y)

"""
    isapprox(a::Union{EachApprox, Approx}, x, y)

Use the definition of approximate equality specified by `a` to determine
if `x` is approximately `y`.
"""
Base.isapprox(a::Union{EachApprox, Approx}, x, y) = isapprox(x, y; pairs(a.kw)...)

"""
    isapprox(a::EachApprox, A::AbstractArray, B::AbstractArray)

Compute element-wise approximate equality.
"""
function Base.isapprox(a::EachApprox, A::AbstractArray, B::AbstractArray)
    n1 = length(A)
    n2 = length(B)
    if n1 != n2
        return false
    end
    for (x, y) in zip(A, B)
        if ! isapprox(a, x, y)
            return false
        end
    end
    return true
end

"""
    UpToPhase <: AbstractApprox

Demands that each pair of elements are approximately equal up to a phase,
that is a number whose absolute value is one.
`kw` are keyword pairs that are forwarded to `isapprox`.
"""
struct UpToPhase{T} <: AbstractApprox
    kw::T
end
UpToPhase(; kws...) = UpToPhase(kws)

function Base.isapprox(a::UpToPhase, x::Number, y::Number)
    aa = Approx(;a.kw...)
    if isapprox(aa, x, zero(x))
        return isapprox(aa, y, zero(y))
    elseif isapprox(aa, y, zero(y))
        return isapprox(aa, x, zero(x))
    end
    return isunitary(x / y, aa)
end

function Base.isapprox(_app::UpToPhase,  A::AbstractArray, B::AbstractArray)
    n1 = length(A)
    n2 = length(B)
    if n1 != n2
        return false
    end
    app = Approx(;_app.kw...)
    seen_non_zero_flag = false
    z = zero(eltype(A)) # TODO use promotion
    for (a, b) in zip(A, B)
        if iszero(a, app)
            !iszero(b, app) && return false
        elseif iszero(b, app)
            !iszero(a, app) && return false
        else
            if ! seen_non_zero_flag
                z = a/b
                isunitary(z, app) || return false
                seen_non_zero_flag = true
            else
                isapprox(app, a/b, z) || return false
            end
        end
    end
    return true
end


include("basic_omnivory_module.jl")
include("press.jl")
using PyPlot
pygui(true)

# PULSE SIMULATIONS for ranges of K, aCP, eCP and mP.

# range of K
println("Simulations for a range of K values")
res_K = []
rg_K = 2:0.05:4
for k in rg_K
    par = ModelPar()
    par.K_base = par.K = k
    res_K = push!(res_K, press(par, 0.1, 0.5))
end

# range of a_CP
println("Simulations for a range of aCP values")
res_aCP = []
rg_aCP = 0.3:.0025:.55
for a in rg_aCP
    par = ModelPar()
    par.a_CP = a
    res_aCP = push!(res_aCP, press(par, 0.1, 0.5))
end

# range of e_CP
println("Simulations for a range of eCP values")
res_eCP = []
rg_eCP = 0.35:.0025:0.7
for e in rg_eCP
    par = ModelPar()
    par.e_CP = e
    res_eCP = push!(res_eCP, press(par, 0.1, 0.5))
end

# range of m_P
println("Simulations for a range of mP values")
res_mP = []
rg_mP = 0.15:.005:.3
for m in rg_mP
    par = ModelPar()
    par.m_P = m
    res_mP = push!(res_mP, press(par, 0.1, 0.5))
end

# As this takes some time to run you may want to save this (it takes ~500Mo),
# if so uncomment the lines below.
# using FileIO, JLD2
# # mkdir("res")
# save(
#     "res/res_S2.jld2",
#     "rg_K", rg_K, "res_K", res_K,
#     "rg_aCP", rg_aCP, "res_aCP", res_aCP,
#     "rg_eCP", rg_eCP, "res_eCP", res_eCP,
#     "rg_mP", rg_mP, "res_mP", res_mP
#     )

# FIGURE
println("Drawing Figure S2")

fig = figure(figsize = (12, 8))
ind = [2 5 8]

##----------- K
tly = ["-1/Re(λ1)" "Degree of Overshoot" "Max-Min"]
tlx = ["" "" "K"]
for i in 1:3
    subplot(3, 4, (i - 1) * 4 + 1)
    plot_sa_unit(rg_K, res_K, ind[i], false, tlx[i], tly[i])
end

##----------- aCP
tlx = ["" "" L"a_{CP}"]
for i in 1:3
    subplot(3, 4, (i-1) * 4 + 2)
    plot_sa_unit(rg_aCP, res_aCP, ind[i], false, tlx[i], "")
end

##----------- eCP
tlx = ["" "" L"e_{CP}"]
for i in 1:3
    subplot(3, 4, (i-1) * 4 + 3)
    plot_sa_unit(rg_eCP, res_eCP, ind[i], false, tlx[i], "")
end

##----------- mP
tlx = ["" "" L"m_P"]
for i in 1:3
    subplot(3, 4, (i-1) * 4 + 4)
    plot_sa_unit_rev(rg_mP, res_mP, ind[i], i == 1, tlx[i], "")
    val = [0.15, 0.20, 0.25, 0.30]
    xticks(val, reverse(val))
end

tight_layout()
savefig("fig/fig_S2.svg")


function compute_one_form(t, x, p, Θ)
    @assert size(x) == size(p)
    @assert length(t) == size(x,2) == size(p,2)

    tp = zeros(eltype(x), size(x,1))

    for k in 1:size(x,3)
        for j in 1:size(x,2)
            Θ(t[j], x[:,j,k], tp)
            p[:,j,k] .= tp
        end
    end
end

function compute_velocity(γ, γ̇)
    @assert size(γ) == size(γ̇)

    n = size(γ,2)
    k = collect(0:div(n-1,2)+1)
    κ = 2π .* 1im .* k ./ n

    for l in 1:size(γ,1)
        γ̂ = rfft(γ[l,:])
        γ̇[l,:] .= irfft(γ̂ .* κ, n)
    end
end

function compute_correction(t, x, λ, σ, g)
    @assert size(x) == size(λ) == size(σ)
    @assert length(t) == size(x,2) == size(σ,2)

    tσ = zeros(eltype(x), size(x,1))

    for k in 1:size(x,3)
        for j in 1:size(x,2)
            g(t[j], x[:,j,k], λ[:,j,k], tσ)
            σ[:,j,k] .= tσ
        end
    end
end

function compute_loop_integral(p::Array{T,2}, v::Array{T,2}) where {T}
    local result = zero(T)
    local error  = zero(T)

    for i in 1:size(v,2)
        for l in 1:size(v,1)
            result, error = compensated_summation(p[l,i] * v[l,i], result, error)
        end
    end

    return result
end


using NearestNeighbors, Statistics, Distances

export estimate_dimension, stochastic_indicator
export Euclidean, Chebyshev, Cityblock

#####################################################################################
#                                Estimate Dimension                                 #
#####################################################################################
"""
    estimate_dimension(s::AbstractVector, τ::Int, γs = 1:5, method = "afnn"; kwargs...)

Compute a quantity that can estimate an optimal amount of
temporal neighbors `γ` to be used in [`reconstruct`](@ref) or [`embed`](@ref).

## Description
Given the scalar timeseries `s` and the embedding delay `τ` compute a quantity
for each `γ ∈ γs` based on the "nearest neighbors" in the embedded time series.

The quantity that is calculated depends on the algorithm defined by the string `method`:

* `"afnn"` (default) is Cao's "Averaged False Nearest Neighbors" method [1], which
    gives a ratio of distances between nearest neighbors. This ratio saturates
    around `1.0` near the optimal value of `γ` (see [`afnn`](@ref)).
* `"fnn"` is Kennel's "False Nearest Neighbors" method [2], which gives the
    number of points that cease to be "nearest neighbors" when the dimension
    increases. This number drops down to zero near the optimal value of `γ`.
    This method accepts the keyword arguments `rtol` and `atol`, which stand
    for the "tolerances" required by Kennel's algorithm (see [`fnn`](@ref)).
* `"f1nn"` is Krakovská's "False First Nearest Neighbors" method [3], which
    gives the ratio of pairs of points that cease to be "nearest neighbors"
    when the dimension increases. This number drops down to zero near the
    optimal value of `γ` (see [`f1nn`](@ref)).

`"afnn"` and `"f1nn"` also support the `metric` keyword, which can be any of
`Cityblock(), Euclidean(), Chebyshev()`. This metric is used both
for computing the nearest neighbors (`KDTree`s) as well as the distances necessary for
Cao's method (eqs. (2, 3) of [1]). Defaults to `Euclidean()` (note that [1] used
`Chebyshev`).

Please be aware that in **DynamicalSystems.jl** `γ` stands for the amount of temporal
neighbors and not the embedding dimension (`D = γ + 1`, see also [`embed`](@ref)).

## References

[1] : Liangyue Cao, [Physica D, pp. 43-50 (1997)](https://www.sciencedirect.com/science/article/pii/S0167278997001188?via%3Dihub)

[2] : M. Kennel *et al.*, [Phys. Review A **45**(6), 3403-3411](https://journals.aps.org/pra/abstract/10.1103/PhysRevA.45.3403)
(1992).

[3] : Anna Krakovská *et al.*, [J. Complex Sys. 932750 (2015)](https://doi.org/10.1155/2015/932750)
"""
function estimate_dimension(s::AbstractVector, τ::Int, γs = 1:5, method = "afnn";
    metric = Euclidean(), kwargs...)

    if method == "afnn"
        return afnn(s, τ, γs, metric)
    elseif method == "fnn"
        return fnn(s, τ, γs; kwargs...)
    elseif method == "f1nn"
        return f1nn(s, τ, γs, metric)
    end
end

"""
    afnn(s::AbstractVector, τ:Int, γs = 1:5, metric=Euclidean())

Compute the parameter E₁ of Cao's "averaged false nearest neighbors" method for
determining the minimum embedding dimension of the time series `s`, with
a sequence of `τ`-delayed temporal neighbors [1].

## Description
Given the scalar timeseries `s` and the embedding delay `τ` compute the
values of `E₁` for each `γ ∈ γs`, according to Cao's Method (eq. 3 of [1]).

This quantity is a ratio of the averaged distances between the nearest neighbors
of the reconstructed time series, which quantifies the increment of those
distances when the number of temporal neighbors changes from `γ` to `γ+1`.

Return the vector of all computed `E₁`s. To estimate a good value for `γ` from this,
find `γ` for which the value `E₁` saturates at some value around 1.

*Note: This method does not work for datasets with perfectly periodic signals.*

See also: [`estimate_dimension`](@ref), [`fnn`](@ref), [`f1nn`](@ref).

## References

[1] : Liangyue Cao, [Physica D, pp. 43-50 (1997)](https://www.sciencedirect.com/science/article/pii/S0167278997001188?via%3Dihub)
"""
function afnn(s::AbstractVector{T}, τ::Int, γs = 1:5, metric=Euclidean()) where {T}
    E1s = zeros(length(γs))
    aafter = 0.0
    aprev = _average_a(s, γs[1], τ, metric)
    for (i, γ) ∈ enumerate(γs)
        aafter = _average_a(s, γ+1, τ, metric)
        E1s[i] = aafter/aprev
        aprev = aafter
    end
    return E1s
end
# then use function `saturation_point(γs, E1s)` from ChaosTools

function _average_a(s::AbstractVector{T},γ,τ,metric) where {T}
    #Sum over all a(i,d) of the Ddim Reconstructed space, equation (2)
    Rγ = reconstruct(s[1:end-τ],γ,τ)
    tree2 = KDTree(Rγ, metric)
    nind = (x = knn(tree2, Rγ.data, 2)[1]; [ind[1] for ind in x])
    e = 0.0
    @inbounds for (i,j) ∈ enumerate(nind)
        δ = evaluate(metric, Rγ[i], Rγ[j])
        #If Rγ[i] and Rγ[j] are still identical, choose the next nearest neighbor
        if δ == 0.0
            j = knn(tree2, Rγ[i], 3, true)[1][end]
            δ = evaluate(metric, Rγ[i], Rγ[j])
        end
        e += _increase_distance(δ,s,i,j,γ,τ,metric)/δ
    end
    return e / (length(Rγ)-1)
end

# Function to increase the distance (p-norm) between two points `(i,j)` of
# the embedded time `s`series, by adding one temporal neighbor
_increase_distance(δ, s, i::Int, j::Int, γ::Int, τ::Int, ::Chebyshev) =
    max(δ, abs(s[i+γ*τ+τ] - s[j+γ*τ+τ]))
_increase_distance(δ, s, i::Int, j::Int, γ::Int, τ::Int, ::Euclidean) =
    sqrt(δ*δ + abs2(s[i+γ*τ+τ] - s[j+γ*τ+τ]) )
_increase_distance(δ, s, i::Int, j::Int, γ::Int, τ::Int, ::Cityblock) =
    δ + abs(s[i+γ*τ+τ] - s[j+γ*τ+τ])

"""
    fnn(s::AbstractVector, τ:Int, γs = 1:5; rtol=10.0, atol=2.0)

Calculate the number of "false nearest neighbors" (FNN) of the datasets created
from `s` with a sequence of `τ`-delayed temporal neighbors.

## Description
Given a dataset made by embedding `s` with `γ` temporal neighbors and delay `τ`,
the "false nearest neighbors" (FNN) are the pairs of points that are nearest to
each other at dimension `γ`, but are separated at dimension `γ+1`. Kennel's
criteria for detecting FNN are based on a threshold for the relative increment
of the distance between the nearest neighbors (`rtol`, eq. 4 in [1]), and
another threshold for the ratio between the increased distance and the
"size of the attractor" (`atol`, eq. 5 in [1]). These thresholds are given
as keyword arguments.

The returned value is a vector with the number of FNN for each `γ ∈ γs`. The
optimal value for `γ` is found at the point where the number of FNN approaches
zero.

See also: [`estimate_dimension`](@ref), [`afnn`](@ref), [`f1nn`](@ref).

## References

[1] : M. Kennel *et al.*, "Determining embedding dimension for phase-space
reconstruction using a geometrical construction", *Phys. Review A 45*(6), 3403-3411
(1992).
"""
function fnn(s::AbstractVector, τ::Int, γs = 1:5; rtol=10.0, atol=2.0)
    rtol2 = rtol^2
    Ra = std(s, corrected=false)
    nfnn = zeros(length(γs))
    @inbounds for (k, γ) ∈ enumerate(γs)
        y = reconstruct(s[1:end-τ],γ,τ)
        tree = KDTree(y)
        nind = (x = knn(tree, y.data, 2)[1]; [ind[1] for ind in x])
        for (i,j) ∈ enumerate(nind)
            δ = norm(y[i]-y[j], 2)
            # If y[i] and y[j] are still identical, choose the next nearest neighbor
            # as in Cao's algorithm (not suggested by Kennel, but still advisable)
            if δ == 0.0
                j = knn(tree, y[i], 3, true)[1][end]
                δ = norm(y[i]-y[j])
            end
            δ1 = _increase_distance(δ,s,i,j,γ,τ,Euclidean())
            cond_1 = ((δ1/δ)^2 - 1 > rtol2) # equation (4) of Kennel
            cond_2 = (δ1/Ra > atol)         # equation (5) of Kennel
            if cond_1 | cond_2
                nfnn[k] += 1
            end
        end
    end
    return nfnn
end

"""
    f1nn(s::AbstractVector, τ:Int, γs = 1:5, metric = Euclidean())

Calculate the ratio of "false first nearest neighbors" (FFNN) of the datasets created
from `s` with a sequence of `τ`-delayed temporal neighbors.

## Description
Given a dataset made by embedding `s` with `γ` temporal neighbors and delay `τ`,
the "first nearest neighbors" (FFNN) are the pairs of points that are nearest to
each other at dimension `γ` that cease to be nearest neighbors at dimension
`γ+1` [1].

The returned value is a vector with the ratio between the number of FFNN and
the number of points in the dataset for each `γ ∈ γs`. The optimal value for `γ`
is found at the point where this ratio approaches zero.

See also: [`estimate_dimension`](@ref), [`afnn`](@ref), [`fnn`](@ref).

## References

[1] : Anna Krakovská *et al.*, "Use of false nearest neighbours for selecting
variables and embedding parameters for state space reconstruction", *J Complex
Sys* 932750 (2015), DOI: 10.1155/2015/932750
"""
function f1nn(s::AbstractVector, τ::Int, γs = 1:5, metric = Euclidean())
    f1nn_ratio = zeros(length(γs))
    γ_prev = 0 # to recall what γ has been analyzed before
    Rγ = reconstruct(s[1:end-τ],γs[1],τ) # this is for the first iteration
    for (i, γ) ∈ enumerate(γs)
        if i>1 && γ!=γ_prev+1
            # Re-calculate the series with γ delayed dims if γ does not follow
            # the dimension of the previous iteration
            Rγ = reconstruct(s[1:end-τ],γ,τ)
        end
        (nf1nn, Rγ) = _compare_first_nn(s,γ,τ,Rγ,metric)
        f1nn_ratio[i] = nf1nn/length(Rγ)
        # Trim Rγ for the next iteration
        Rγ = Rγ[1:end-τ,:]
        γ_prev = γ
    end
    return f1nn_ratio
end

function _compare_first_nn(s, γ::Int, τ::Int, Rγ::Dataset{D,T}, metric) where {D,T}
    # This function compares the first nearest neighbors of `s`
    # embedded with Dimensions `γ` and `γ+1` (the former given as input)
    tree = KDTree(Rγ,metric)
    Rγ1 = reconstruct(s,γ+1,τ)
    tree1 = KDTree(Rγ1,metric)
    nf1nn = 0
    # For each point `i`, the fnn of `Rγ` is `j`, and the fnn of `Rγ1` is `k`
    nind = (x = knn(tree, Rγ.data, 2)[1]; [ind[1] for ind in x])
    @inbounds for  (i,j) ∈ enumerate(nind)
        k = knn(tree1, Rγ1.data[i], 2, true)[1][end]
        if j != k
            nf1nn += 1
        end
    end
    # `Rγ1` is returned to re-use it if necessary
    return (nf1nn, Rγ1)
end

"""
    stochastic_indicator(s::AbstractVector, τ:Int, γs = 1:4) -> E₂s

Compute an estimator for apparent randomness in a reconstruction with `γs` temporal
neighbors.

## Description
Given the scalar timeseries `s` and the embedding delay `τ` compute the
values of `E₂` for each `γ ∈ γs`, according to Cao's Method (eq. 5 of [1]).

Use this function to confirm that the
input signal is not random and validate the results of [`estimate_dimension`](@ref).
In the case of random signals, it should be `E₂ ≈ 1 ∀ γ`.

## References

[1] : Liangyue Cao, [Physica D, pp. 43-50 (1997)](https://www.sciencedirect.com/science/article/pii/S0167278997001188?via%3Dihub)
"""
function stochastic_indicator(s::AbstractVector{T},τ, γs=1:4) where T # E2, equation (5)
    #This function tries to tell the difference between deterministic
    #and stochastic signals
    #Calculate E* for Dimension γ+1
    E2s = Float64[]
    for γ ∈ γs
        Rγ1 = reconstruct(s,γ+1,τ)
        tree1 = KDTree(Rγ1[1:end-1-τ])
        method = FixedMassNeighborhood(2)

        Es1 = 0.
        nind = (x = neighborhood(Rγ1[1:end-τ], tree1, method); [ind[1] for ind in x])
        for  (i,j) ∈ enumerate(nind)
            Es1 += abs(Rγ1[i+τ][end] - Rγ1[j+τ][end]) / length(Rγ1)
        end

        #Calculate E* for Dimension γ
        Rγ = reconstruct(s,γ,τ)
        tree2 = KDTree(Rγ[1:end-1-τ])
        Es2 = 0.
        nind = (x = neighborhood(Rγ[1:end-τ], tree2, method); [ind[1] for ind in x])
        for  (i,j) ∈ enumerate(nind)
            Es2 += abs(Rγ[i+τ][end] - Rγ[j+τ][end]) / length(Rγ)
        end
        push!(E2s, Es1/Es2)
    end
    return E2s
end


using Gadfly

plot([sin, cos], 0, 25)


@testset "Config" begin
    Impl.@polyvar x y z

    f = Polynomial{Float64}(x^2+3y+z*x+y^2+12*x^7*z^5)
    g = Polynomial{Float64}(4*y*x^2-3z+z*x^4+y+y^2-y*x)

    ∇f = [differentiate(f, i) for i=1:3]
    ∇g = [differentiate(g, i) for i=1:3]

    w = rand(3)
    u = zeros(3)

    cfg = GradientConfig(f, w)
    @test f(w) ≈ evaluate(f, w, cfg)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient(f, w, cfg)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient!(u, f, w, cfg)
    @test [p(w) for p in ∇f] ≈ u

    cfg = GradientConfig(f, Complex128)
    wc = rand(Complex128, 3)
    uc = zeros(Complex128, 3)
    @test f(wc) ≈ evaluate(f, wc, cfg)
    @test [p(wc) for p in ∇f] ≈ FixedPolynomials.gradient(f, wc, cfg)
    @test [p(wc) for p in ∇f] ≈ FixedPolynomials.gradient!(uc, f, wc, cfg)
    @test [p(wc) for p in ∇f] ≈ uc

    cfg = GradientConfig(f)
    @test f(w) ≈ evaluate(f, w, cfg)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient(f, w, cfg)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient!(u, f, w, cfg)
    @test [p(w) for p in ∇f] ≈ u

    cfg2 = deepcopy(cfg)
    @test cfg2 !== cfg

    r = GradientDiffResult(cfg)

    FixedPolynomials.gradient!(r, f, w, cfg)
    @test f(w) ≈ value(r)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient(r)

    v = zeros(4)
    r = GradientDiffResult((@view v[1:3]))
    FixedPolynomials.gradient!(r, f, w, cfg)
    @test f(w) ≈ value(r)
    @test [p(w) for p in ∇f] ≈ FixedPolynomials.gradient(r)

    @test_throws MethodError GradientConfig([f, f], w)
    @test_throws MethodError JacobianConfig(f, w)

    u = zeros(2)
    cfg = JacobianConfig([f, g], w)
    @test [f(w), g(w)] ≈ evaluate([f, g], w, cfg)
    @test [f(w), g(w)] ≈ evaluate!(u, [f, g], w, cfg)
    @test [f(w), g(w)] ≈ u

    cfg2 = deepcopy(cfg)
    @test cfg2 !== cfg

    U = zeros(2, 3)
    DF = vcat(RowVector([p(w) for p in ∇f]), [p(w) for p in ∇g] |> RowVector)
    @test DF ≈ jacobian([f, g], w, cfg)
    @test DF ≈ jacobian!(U, [f, g], w, cfg)
    @test DF ≈ U

    r = JacobianDiffResult(cfg)
    jacobian!(r, [f, g], w, cfg)
    @test DF ≈ jacobian(r)
    @test [f(w), g(w)] ≈ value(r)

    v = zeros(3)
    V = zeros(3, 3)

    r = JacobianDiffResult((@view v[1:2]), (@view V[1:2, :]))
    jacobian!(r, [f, g], w, cfg)
    @test DF ≈ jacobian(r)
    @test [f(w), g(w)] ≈ value(r)
end


#======================================
Methods using semidefinite programming
======================================#
using SumOfSquares, DynamicPolynomials, SemialgebraicSets, SDPA

"""
    new_sos(backend; kwargs...)

Return a new (empty) sum-of-squares optimization problem for the given backend.

### Input

- `backend` -- the backend (also called JuMP `solver`)
- `kwargs`  -- additional keyword arguments

### Output

An instance of `SOSModel` for the given backend and options.
"""
@inline function new_sos(backend, kwargs...)
    𝑂 = Dict(kwargs)

    if VERSION < v"1.6"
        if :QUIET ∈ keys(𝑂)
            # for mosek solver
            SOSModel(with_optimizer(backend, QUIET=𝑂[:QUIET]))
        else
            SOSModel(with_optimizer(backend))
        end
    else
        if :QUIET ∈ keys(𝑂)
            # for mosek solver
            SOSModel(backend, QUIET=𝑂[:QUIET])
        else
            SOSModel(backend)
        end
    end
end

"""
    enclose_SumOfSquares(f::Function, dom::Interval_or_IntervalBox,
                         order::Int=5, backend=SDPA.Optimizer; kwargs...)

Compute a range enclosure using sum-of-squares optimization.

### Input

- `f`       -- function
- `dom`     -- hyperrectangular domain, either a unidimensional `Interval` or
               a multidimensional `IntervalBox`
- `order`   -- (optional, default: `5`) maximum degree of the SDP relaxation
- `backend` -- (optional, default: `SDPA.Optimizer`) the optimization backend
               (aka JuMP solver)
- `kwargs`  -- additional keyword arguments

### Output

An interval representing the range enclosure (minimum and maximum) of `f` over
its domain `dom`.

### Algorithm

The range enclosure is computed using polynomial optimization methods.
We refer to the documentation and examples of `SumOfSquares.jl` for details.

### Notes

Use the `backend` keyword argument to pass an SDP solver backend of your choice;
additional arguments to the solver can be passed in `kwargs`.

For instance, to use `SDPA` (default choice), use it as:

```julia
julia> backend = SDPA.Optimizer

julia> enclose_SumOfSquares(f, dom, order, backend=backend)
...
```

To use `MOSEK` in non-verbose mode, let

```julia
julia> using MosekTools

julia> backend = MosekTools.Mosek.Optimizer;

julia> enclose_SumOfSquares(f, dom, order, backend=backend, QUIET=true)
...
```

To get the runtime, use `MOI.get(model, MOI.SolveTime())`.
"""
function enclose_SumOfSquares(f::Function, dom::Interval_or_IntervalBox;
                              order::Int=5, backend=SDPA.Optimizer, kwargs...)
    _enclose_SumOfSquares(f, dom, order, backend; kwargs...)
end

function _enclose_SumOfSquares(f::Function, dom::Interval, order::Int,
                               backend; kwargs...)

    # polynomial variables
    @polyvar x
    p = f(x)

    # box constraints
    B = @set inf(dom) <= x && x <= sup(dom)

    # ============
    # Upper bound
    # ============
    model = new_sos(backend, kwargs...)
    @variable(model, γ) # JuMP decision variable
    @constraint(model, p <= γ, domain=B, maxdegree=order)
    @objective(model, Min, γ)
    optimize!(model)
    upper_bound = objective_value(model)

    # ============
    # Lower bound
    # ============
    model = new_sos(backend, kwargs...)
    @variable(model, γ) # JuMP decision variable
    @constraint(model, p >= γ, domain=B, maxdegree=order)
    @objective(model, Max, γ)
    optimize!(model)
    lower_bound = objective_value(model)

    return Interval(lower_bound, upper_bound)
end

function _enclose_SumOfSquares(f::Function, dom::IntervalBox{N}, order::Int,
                               backend; kwargs...) where {N}

    # polynomial variables
    @polyvar x[1:N]
    p = f(x...)

    # box constraints
    Bi =[@set inf(dom[i]) <= x[i] && x[i] <= sup(dom[i]) for i in 1:N]
    B = reduce(intersect, Bi)

    # ============
    # Upper bound
    # ============
    model = new_sos(backend, kwargs...)
    @variable(model, γ) # JuMP decision variable
    @constraint(model, p <= γ, domain=B, maxdegree=order)
    @objective(model, Min, γ)
    optimize!(model)
    upper_bound = objective_value(model)

    # ============
    # Lower bound
    # ============
    model = new_sos(backend, kwargs...)
    @variable(model, γ) # JuMP decision variable
    @constraint(model, p >= γ, domain=B, maxdegree=order)
    @objective(model, Max, γ)
    optimize!(model)
    lower_bound = objective_value(model)

    return Interval(lower_bound, upper_bound)
end


@doc raw"""
```
l = segment_image(Algorithm::GrowCut, feature_vector::AbstractArray, seeds::AbstractArray{Int, 2}, t1::Int=9, t2::Int=9; max_iter::Int=1000, converge_at::Int=1, rule::Tuple{Function, Function}=(colour_diff,find_maxC))
```

Using a provided `feature_vector` (default rule uses RGB image) and an array of
seeds, `GrowCut` will use Cellular Automata to segment an image into different
regions. The algorithm takes two thresholds which can be used to control the
smoothness of the segmentation.

# Output

Returns a Array of Integers `l` that specifies the label assigned to each pixel.

# Details

The algorithm uses Cellular Automata to iteratively apply a rule to each pixel
which then competes with its neighbours to label the pixel. By default the
algorithm takes an RGB image for the `feature_vector` however a custom rule can
be defined for other `feature_vector` types. To improve performance after the
first iteration only the neighbourhood surrounding a pixel that changed in the
previous iteration is checked.

# Arguments

The function arguments are described in more detail below.

##  `feature_vector`

An `AbstractArray` containing the feature vector for each pixel. The default rule
assumes this is an RGB value for the pixel. Note: RGB images should be converted
to RGB{Float64} to prevent overflow.

##  `seeds`

An `Array{Int, 2}` specifying initial labels. Any pixel that is a seed should be
set to the desired integer label and all other positions should be 0. This can
be generated using `set_seed_pixels`

## `t1` and `t2`

`Int`s that specifiy the thresholds to be used for smoothing. `t1` is the maximum
number of "enemies" that a pixel can be surrounded by and still attack other
pixels. `t2` is the maximum number of "enemies" before the pixel is forced to
become the weakest enemy pixel. Note: these values can produce scenarios where
the algorithm is unable to converge and will instead reach its maximum number of
iterations.

## `maxiter`

An `Int` that specifies the maximum number of iterations before execution is
aborted. If left unspecified a default value of 1000 is used. This should never
be reached under normal conditions.

## `converge_at`

An `Int` that specifies the minimum number of changes across the image for
convergence. If left unspecified a default value of 1 is used.

## `rule`

A `Tuple{Function, Function}` that specifies the rule to be used for the
function. The first element is the rule that determines the strength of the
attacking pixel and the second element is used to calculate any extra
information needed by the rule. By default `colour_diff` and `find_maxC` are used. A
template for custom rules can be found using `help_rule_template` or generated into
a new file using `create_rule`.

# Example

Segment the "toucan" image using `GrowCut`

```julia
using Images, PictureSegmentation
img = testimage("toucan")

#convert to RGB{Float64}
img=RGB{Float64}.(img)

#create array of seeds in initialized to 0
seeds = zeros(Int,axes(img))

#use set_seed_pixels to set the inital seeds
set_seed_pixels(clicks, img, suppress_feedback=false)

#segment the image
l=segment_image(GrowCut(),img,clicks,9,9)

#create and display segmented image from labels
img2 = convert_labels(l)
img2
```

# Advanced

By default GrowCut uses the rule:

```math
\begin{aligned}
\textbf{for } \forall p &  \in  P                               &\\
                        & l_p^{t+1} =  l_p^{t}                  &\\
                        & \theta_p^{t+1} =  \theta_p^{t}        &\\
                        & \textbf{for } \forall q  \in  N(P)    & \\
                        &    \qquad    \textbf{if } g(\| RGB_p - RGB_q \|_2) \cdot \theta_q^t > \theta_p^t \\
                        &       \qquad \qquad l_p^{t+1} =  l_q^{t}  & \\
                        & \qquad  \qquad  \theta_p^{t+1} = g(\| RGB_p - RGB_q \|_2) \cdot \theta_q^t  \\
                         & \qquad     \textbf{end }  & \\
                        & \textbf{end }                         &\\
\textbf{end }\\
\end{aligned}
```
Where q is a pixel in the neighbourhood of p, l is the label of each pixel, θ is
the current strength of each pixel and RGB is the RGB vector of the pixel. g is
defined as ``g(x)=1-\dfrac{x}{\text{max}\|RGB\|}``.

An example of a custom rule for Lab colour is defined below:

```math
\begin{aligned}
\textbf{for } \forall p &  \in  P                               &\\
                        & l_p^{t+1} =  l_p^{t}                  &\\
                        & \theta_p^{t+1} =  \theta_p^{t}        &\\
                        & \textbf{for } \forall q  \in  N(P)    & \\
                        &    \qquad    \textbf{if } g(\| Lab_p - Lab_q \|_2) \cdot \theta_q^t > \theta_p^t \\
                        &       \qquad \qquad l_p^{t+1} =  l_q^{t}  & \\
                        & \qquad  \qquad  \theta_p^{t+1} = g(\| Lab_p - Lab_q \|_2) \cdot \theta_q^t  \\
                         & \qquad     \textbf{end }  & \\
                        & \textbf{end }                         &\\
\textbf{end }\\
\end{aligned}
```

where ``g(x)=1-\dfrac{x}{\text{max}\|Lab\|}``.
This would be implemented as per follows:

```julia

#calculate maxLab as needed by rule
function calculate_maxLab(LAB_color_vector::AbstractArray)
    maxLab = 0
    for i in CartesianIndices(feature_vector)
        current_Lab = sqrt(LAB_color_vector[i].l^2 + LAB_color_vector[i].a^2 + LAB_color_vector[i].b^2)
        if current_Lab > maxLab
            maxC = current_Lab
        end
    end
    return maxLab
end

#calculates the adjusted strength as per rule
function example_LAB_rule(θ::Array{Float64,1}, LAB_color_vector::AbstractArray, target_pixel::CartesianIndex, attacking_pixel::CartesianIndex, maxLab)

    LabTarget = LAB_color_vector[target_pixel]
    LabAttacker = LAB_color_vector[attacking_pixel]

    #calculate difference between LAB vectors
    LABdiff=Lab(LabTarget.l-LabAttacker.l,LabTarget.a-LabAttacker.a,LabTarget.b-LabAttacker.b)

    #find the absolute of the difference
    absLABdiff=sqrt(LABdiff.l^2+LABdiff.a^2+LABdiff.b^2)

    #calculate strength
    adjusted_strength = (1-(absLABdiff/maxLab))*θ[attacking_pixel]
    return adjusted_strength
end
```

The rule can then be passed to `Growcut` as per following:

```
l=segment_image(GrowCut(), img, clicks, rule=(example_LAB_rule,calculate_maxLab))
```

# Reference
1. V. Vezhnevets and V. Konouchine, "“GrowCut” - Interactive Multi-Label N-D Image Segmentation By Cellular Automata", in Graphicon, Novosibirsk Akademgorodok, Russia, 2005.
"""
function segment_image(Algorithm::GrowCut, feature_vector::AbstractArray, seeds::AbstractArray{Int, 2}, t1::Int=9, t2::Int=9; max_iter::Int=1000, converge_at::Int=1, rule::Tuple{Function, Function}=(colour_differance, find_maximum))
    # initialize values
    calculate_adjusted_strength = rule[1]
    extra_function = rule[2]
    extras = extra_function(feature_vector)
    labels = copy(seeds)
    E = zeros(Int, axes(labels))
    θ = set_strength(labels)
    labelsₜ₊₁ = copy(seeds)
    Eₜ₊₁ = zeros(Int, axes(labels))
    θₜ₊₁ = set_strength(labels)
    regions = maximum(seeds)
    iter = 0
    count = length(labels)
    changed = Array{Tuple{UnitRange{Int},UnitRange{Int}},1}(undef,1)
    changedₜ₊₁ = Array{Tuple{UnitRange{Int},UnitRange{Int}},1}(undef,0)
    rₘ, cₘ = size(labels)
    changed[1] = 1:rₘ, 1:cₘ

    # overflow warning
    if typeof(feature_vector) == Array{RGB,2} && typeof(feature_vector) != Array{RGB{Float64},2}
        @warn "feature_vector is RGB and not of type RGB{Float64}. This may cause overflow errors."
    end

    # iterate till convergence or maximum iterations reached
    while count > converge_at && iter < max_iter
        iter += 1
        count = 0
        copyto!(θₜ₊₁, θ)
        copyto!(labelsₜ₊₁, labels)
        copyto!(Eₜ₊₁, E)
        # loop through active pixel neighbourhoods (all pixels on first iteration)
        for j in CartesianIndices(changed)
            @inbounds for i in CartesianIndices(changed[j])
                r, c=i.I
                r₀ = max(r-1, 1)
                r₁ = min(r+1, rₘ)
                c₀ = max(c-1, 1)
                c₁ = min(c+1, cₘ)
                N₈ = r₀:r₁, c₀:c₁
                Eₜ₊₁[i] = count_enemy(labels,i,N₈)
                if E[i] < t2
                    θₜ₊₁, labelsₜ₊₁, count, changedₜ₊₁ = update_pixel(θ, θₜ₊₁, labels, labelsₜ₊₁, feature_vector, E, i, extras, t1, count, N₈, changedₜ₊₁, calculate_adjusted_strength)
                else
                    θₜ₊₁, labelsₜ₊₁, count, changedₜ₊₁ = occupy_pixel(θ, θₜ₊₁, labels, labelsₜ₊₁, feature_vector, i, extras,count, N₈, changedₜ₊₁, calculate_adjusted_strength)
                end
            end
        end

        # update values
        changed = changedₜ₊₁
        changedₜ₊₁ = empty(changedₜ₊₁)
        E = Eₜ₊₁
        θ = θₜ₊₁
        labels = labelsₜ₊₁
    end
    return labels
end


using Turing
using Turing.RandomMeasures
using Random
using AutoGibbs
using BenchmarkTools
using Plots
using StatsPlots

include("models.jl")
include("models_comparison.jl")

function benchmark_models(combinations...; N=100, bparams...)
    old_prog = Turing.PROGRESS[]
    Turing.turnprogress(false)
    suite = BenchmarkGroup()

    benchmarks = map(combinations) do (name, (m, alg))
        suite[name] = @benchmarkable sample($m, $alg, $N)
    end

    tune!(suite)
    results = run(suite;bparams...)

    Turing.turnprogress(old_prog)
    return results
end

function compare_gmm()
    data = [randn(10) .- 0.5; randn(5) .+ 1]
    m1 = gmm_example(data)
    m2 = gmm_tarray_example(data)

    cond_gmm_z = StaticConditional(m1, :z)
    alg1 = Gibbs(cond_gmm_z, MH(:w, :μ))
    alg2 = Gibbs(cond_gmm_z, HMC(0.01, 10, :w, :μ))
    alg3 = Gibbs(PG(20, :z), MH(:w, :μ))
    alg4 = Gibbs(PG(20, :z), HMC(0.01, 10, :w, :μ))

    b = benchmark_models("AG+MH" => (m1, alg1),
                         "AG+HMC" => (m1, alg2),
                         "PG+MH" => (m2, alg3),
                         "PG+HMC" => (m2, alg4);
                         N=1000, verbose=true, samples=10)
    display(b)
end

compare_gmm()


Pkg.add("MetadataTools")
# Pkg.test("MetadataTools")
Pkg.add("MetaGraphs")
Pkg.add("GraphPlot")
using MetadataTools
using LightGraphs
using MetaGraphs
using GraphPlot

pkgmeta = get_all_pkg()
pg = make_dep_graph(pkgmeta)
pkg_graph = get_pkg_dep_graph("LightGraphs", pg)
pg
pg.pkgnames
pg.adjlist

g = DiGraph(length(pg.adjlist))
mg = MetaDiGraph(g)
for i in 1:length(pg.adjlist)
    for j in pg.adjlist[i]
        add_edge!(mg, j, i)
    end
end
for v in 1:nv(mg)
    val = pg.pkgnames[v]
    set_prop!(mg, v, :name, val)
end

lgv = find(x->x=="LightGraphs", pg.pkgnames)[1]
t = bfs_tree(mg, lgv)
tvertices = [v for v in vertices(t) if degree(t, v) > 0]
subg = induced_subgraph(mg, tvertices)[1]
tnames = pg.pkgnames[tvertices]
tnamesmeta = [prop[:name] for (v, prop) in subg.vprops]
sort(tnamesmeta) == sort(tnames)

fname = "graph.plot.svg"
gplot_kwargs = Dict(:nodelabel=>tnames,
                    # :nodesize=> Float64[sqrt(degree(subg, v))/10 for v in vertices(subg)])
                    :nodesize=>100)
draw(SVG(fname, 8inch, 8inch), gplot(subg; gplot_kwargs...))


using Gnuplot
x = y = -15:0.33:15
fz(x,y) = sin.(sqrt.(x.^2 + y.^2))./sqrt.(x.^2+y.^2)
fxy = [fz(x,y) for x in x, y in y]
@gsp x y fxy "w l lc 'orange'"
save(term="pngcairo size 800,800", output="./code/plt3d_ex3.png") # hide


funs = (lmul!, rmul!, ldiv!, rdiv!)
altfuns = ( (D,F)->F*D, (D,F)->D*F, (D,F)->F\D, (D,F)->D/F)
#triangles = subtypes(LinearAlgebra.AbstractTriangular)
triangles = (LowerTriangular,UpperTriangular,UnitLowerTriangular,UnitUpperTriangular)
flips = (identity,Transpose,Adjoint)

let
    count = 0
    n = 4
    for (fun,altfun) in zip(funs,altfuns)
        for (Dtype,Ftype) in Iterators.product(triangles,triangles)
            for (Dflip,Fflip) in Iterators.product(flips,flips)
                if Dflip != identity || rand() > 1/2    # test real and complex, but complex only for non-transposed desitinations
                    D = Dflip(Dtype(randn(n,n)))
                    F = Fflip(Ftype(randn(n,n)))
                else
                    D = Dflip(Dtype(complex.(randn(n,n),randn(n,n))))
                    F = Fflip(Ftype(complex.(randn(n,n),randn(n,n))))
                end

                println("Testing: triangle_update!(",typeof(D),", ",typeof(F),", ",fun," )")

                D0 = copyto!(similar(D),D)
                F0 = copyto!(similar(F),F)
                α = rand() > 1/2 ? 1 : π
                LHS = α == 1 ? triangle_update!(D,F,fun) : triangle_update!(D,F,fun,α)
                @test F == F0  # check F not clobbered
                @test LHS ≈ α*altfun(D0,F)
                count += 1
            end
        end
    end
    println("Perfomed ",count," tests on triangle_update!")
end


using jFEMTools
import Tensors

mesh = unitSquareMesh2(RectangleCell, (3,3));

# forcing function
rhs(x::Tensors.Vec{2}) = 2*π^2*sin(π*x[1])*sin(π*x[2])

degree = 2;
dim = 2;
element = PoissonVirtualElement(dim,degree);
u = TrialFunction(VEMFunctionSpace(mesh,element))
dofs = DofHandler(mesh, [u]);
operators = VEMOperators(dofs, u;load = rhs);

K = assemble_stiffnessMat(operators);
b = assemble_load(operators);

# Boundary condition
g(x::Tensors.Vec{2}) = sin(π * x[2])*sin(π * x[1]);
dbc = Dirichlet(dofs,u,"boundary",g);
apply!(K,b,dbc);

#Solve
x = K\b

#Plot solution
using Makie
import AbstractPlotting
include("src/plot_recipes.jl")
#popdisplay(AbstractPlotting.PlotDisplay())
#AbstractPlotting.inline!(true)
scene = Scene(resolution = (400, 200))
# Plot approximation
vi = jFEMTools.vertexdofs(dofs, u)
plot!(scene, mesh, color = x[vi])

#Plot exact solution
vv = get_vertices_matrix(mesh)
xx = [g(vv) for vv in mesh.vertices];
plot!(scene, mesh, color = xx)
#display(scene)


using Unfolding
using CSV
using DataFrames
using Test

@testset "Unfolding.jl" begin
    # Reading data
    df_samps = CSV.read("samples.csv", DataFrame)
    df_block = CSV.read("block_model.csv", DataFrame)
    println("- Samples loaded")

    # Get coordinate points as matrix
    input_block = coords(df_block, columns=["XC","YC","ZC"])
    input_samps = coords(df_samps, columns=["X","Y","Z"])

    # Get reference surface points for unfolding
    ref_surface = getreference(input_block, axis=["X","Y"])
    println("- Reference surface extracted")

    # Get transformed coordinates of blocks and samples after unfolding
    optim = (neigh=70,)
    unf_block = unfold(input_block, ref_surface, optim=optim)
    unf_samps = unfold(input_samps, input_block, unf_block, optim=optim)
    println("- Unfolding finished")

    # Write new XT, YT and ZT columns with the transformed coordinates
    for (i,c) in enumerate([:XT,:YT,:ZT])
        df_samps[:,c] = unf_samps[i,:]
        df_block[:,c] = unf_block[i,:]
    end

    to_vtk("test_out_vtk", unf_samps, (test=rand(size(unf_samps,2)),))

    all_input = hcat(input_samps,input_block)
    all_unf   = hcat(unf_samps,unf_block)
    error     = errors(:dists, all_input, all_unf)
    good, bad = errors(:ids, all_input, all_unf)

    @test length(good) > (30 * length(bad))

    println("- Quick test of unfold with multiple inputs")
    s1, s2 = unfold([input_samps[:,1:80],input_samps[:,80:end]], ref_surface)
    s1, s2 = unfold([input_samps[:,1:80],input_samps[:,80:end]], input_block, unf_block)
end


"""
    nonlinear objective
"""
struct NonlinearObjective <: Objective
    f # objectivce
    g # objective gradient
    gradient_type
end

function nonlinear_objective(f)
    NonlinearObjective(f, x -> zero(x), :auto)
end

function nonlinear_objective(f, g)
    NonlinearObjective(f, g, :gradient)
end

function objective(Z, obj::NonlinearObjective, idx, T)
    return obj.f(Z)
end

function objective_gradient!(∇J, Z, obj::NonlinearObjective, idx, T)
    if obj.gradient_type == :auto
        ∇J .+= ForwardDiff.gradient(obj.f, Z)
    else
        ∇J .+= obj.g(Z)
    end
    return nothing
end


using ReservoirComputing
using GaussianProcesses

#model parameters
const approx_res_size = 30
const radius = 1.2
const activation = tanh
const degree = 6
const sigma = 0.1
const beta = 0.1
const alpha = 1.0
const nla_type = NLADefault()
const in_size = 1
const out_size = 1
extended_states = true
const h_steps = 2

const train_len = 50
const predict_len = 12
data = ones(Float64, in_size, 100)
train = data[:, 1:1+train_len-1]
test = data[:, train_len:train_len+predict_len-1]

esn = ESN(approx_res_size,
    train,
    degree,
    radius,
    activation = activation,
    sigma = sigma,
    alpha = alpha,
    nla_type = nla_type,
    extended_states = extended_states)

W_out = ESNtrain(esn, beta)
@test size(W_out) == (out_size, esn.res_size+out_size)
#test predict
output = ESNpredict(esn, predict_len, W_out)
@test size(output) == (out_size, predict_len)
#test h steps predict
output = ESNpredict_h_steps(esn, predict_len, h_steps, test, W_out)
@test size(output) == (out_size, predict_len)

#test esgp
mean = MeanZero()
kernel = Lin(1.0)

gp = ESGPtrain(esn, mean, kernel, lognoise = -2.0, optimize = false)

esgp_output, sigma_esgp = ESGPpredict(esn, predict_len, gp)
@test size(output) == (out_size, predict_len)
esgp_output, sigma_esgp = ESGPpredict_h_steps(esn, predict_len, h_steps, test, gp)
@test size(output) == (out_size, predict_len)

#model parameters
const first_activation = tanh
const second_activation = identity
const first_lambda = 0.45
const second_lambda = 0.3

esn = dafESN(approx_res_size,
            train,
            degree,
            radius,
            first_lambda,
            second_lambda,
            first_activation = first_activation,
            second_activation = second_activation,
            sigma = sigma,
            alpha = alpha,
            nla_type = nla_type,
            extended_states = extended_states)

W_out = ESNtrain(esn, beta)
@test size(W_out) == (out_size, esn.res_size+out_size)
#test predict
output = dafESNpredict(esn, predict_len, W_out)
@test size(output) == (out_size, predict_len)
#test h steps predict
output = dafESNpredict_h_steps(esn, predict_len, h_steps, test, W_out)
@test size(output) == (out_size, predict_len)


# Find Patterns Forming Clumps in a String

const NUCLEOTIDES = ('A', 'C', 'G', 'T')
const NUCLEOTIDE_ORDER = Dict{Char, Int64}('A' => 0, 'C' => 1, 'G' => 2, 'T' => 3)

convert_to_intlist(line) = map(x -> parse(Int, x), split(line))

pattern_to_number(pattern) = foldl((acc, nucleotide) -> 4 * acc + NUCLEOTIDE_ORDER[nucleotide], pattern; init=0)

function number_to_pattern(encoding, k)
    pattern = Char[]
    for _ in 1:k
        encoding, remainder = divrem(encoding, 4)
        push!(pattern, NUCLEOTIDES[remainder+1])
    end
    join(reverse(pattern))
end

function compute_frequencies(text, k)
    frequencies = zeros(Int32, 4 ^ k)
    for ix in 1:(length(text)-k+1)
        pattern = text[ix:ix+k-1]
        code = pattern_to_number(pattern)
        frequencies[code+1] += 1
    end
    frequencies
end

function collect_clump_forming_k_mers(clump_array, k)
    k_mers_in_clump = String[]
    for (ix, isclump) in enumerate(clump_array)
        if isclump
            k_mer = number_to_pattern(ix - 1, k)
            push!(k_mers_in_clump, k_mer)
        end
    end
    k_mers_in_clump
end

function find_clumps_in_text(genome, k, l, t)
    clump_array = falses(4 ^ k)
    frequencies = compute_frequencies(genome[1:l], k)
    for (ix, freq) in enumerate(frequencies)
        if freq >= t
            clump_array[ix] = true
        end
    end

    for ix in 2:(length(genome)-l)
        start_k_mer = genome[ix-1:ix-2+k]
        start_code = pattern_to_number(start_k_mer)
        frequencies[start_code+1] -= 1
        end_k_mer = genome[ix+l-k:ix+l-1]
        end_code = pattern_to_number(end_k_mer)
        frequencies[end_code+1] += 1
        if frequencies[end_code+1] >= t
            clump_array[end_code+1] = true
        end
    end

    collect_clump_forming_k_mers(clump_array, k)
end

function main()
    genome = readline()
    k, l, t = convert_to_intlist(readline())
    result = find_clumps_in_text(genome, k, l, t)
    println(join(result, " "))
end

main()


using Test

using CompScienceMeshes
using BEAST

const e0 = point(0.0,0.0,0.0)
const e1 = point(1.0,0.0,0.0)
const e2 = point(0.0,1.0,0.0)
const e3 = point(0.0,0.0,1.0)

p = simplex(
    [
        point(0.0,0.0,0.0),
        point(1.0,0.0,0.0)
    ],
    Val{1}
)

q = simplex(
    [
        point(0.0,0.0,0.0),
        point(0.5,0.0,0.0)
    ],
    Val{1}
)

f = BEAST.LagrangeRefSpace{Float64,1,2,2}()
x = neighborhood(p, [0.0])
v = f(x)

@test v[1].value == 0
@test v[2].value == 1

@test v[1].derivative == -1
@test v[2].derivative == +1

Q = restrict(f, p, q)
@test Q == [
    1.0 0.5
    0.0 0.5]

# Test restriction of RT elements
ni, no = 6, 7;
ui = transpose([triangleGaussA[ni] triangleGaussB[ni] ]);
uo = transpose([triangleGaussA[no] triangleGaussB[no] ]);
wi = triangleGaussW[ni];
wo = triangleGaussW[no];
# universe = Universe(1.0, ui, wi, uo, wo);

p = simplex([2*e0,2*e1,2*e2], Val{2})
x = neighborhood(p,[0.5, 0.5])
ϕ = BEAST.RTRefSpace{Float64}()
v = ϕ(x)

Q = restrict(ϕ, p, p)
@test Q == Matrix(LinearAlgebra.I, 3, 3)

q = simplex([e0+e1, 2*e1, e1+e2], Val{2})
Q = restrict(ϕ, p, q)
@test Q == [
    2 -1 0
    0  1 0
    0 -1 2] // 4


function continuous!(f, r, diam)
    dr = r[2] - r[1]
    m = 0

    for (i, j) in enumerate(r)
        if j < diam
            m = i
        end
    end

    fpd = -274.0 * f[m + 1] + 600.0 * f[m + 2] - 600.0 * f[m + 3]
    fpd += 400.0 * f[m + 4] - 150.0 * f[m + 5] + 24.0 * f[m + 6]
    fpd /= 120.0 * dr

    fppd = 225.0 * f[m + 1] - 770.0 * f[m + 2] + 1070.0 * f[m + 3]
    fppd += -780.0 * f[m + 4] + 305.0 * f[m + 5] - 50.0 * f[m + 6]
    fppd /= 60.0 * dr^2

    fd = fpd * (diam - r[m + 1]) + f[m + 1]

    fpi = 274.0 * f[m] - 600.0 * f[m - 1] + 600.0 * f[m - 2]
    fpi += -400.0 * f[m - 3] + 150.0 * f[m - 4] - 24.0 * f[m - 5]
    fpi /= 120.0 * dr

    fppi = 225.0 * f[m] - 770.0 * f[m - 1] + 1070.0 * f[m - 2]
    fppi += -780.0 * f[m - 3] + 305.0 * f[m - 4] - 50.0 * f[m - 5]
    fppi /= 60.0 * dr^2

    fi = fpi * (diam - r[m]) + f[m]
    delf = fd - fi
    delfp = fpd - fpi
    delfpp = fppd - fppi
    a0 = delf - diam * delfp + diam^2 * delfpp / 2.0
    a1 = delfp - diam * delfpp
    a2 = delfpp / 2.0

    for i in 1:m
        f[i] += a0 + a1 * r[i] + a2 * r[i]^2
    end

    return a0, a1, a2
end

function remove_continuous!(f, q, diam, params)
    (a0, a1, a2) = params

    for i in eachindex(f)
        if q[i] > 0.0
            part1 = a0 + 2.0 * diam * a1 + 3.0 * diam^2 * a2
            part1 -= 6.0 * a2 / q[i]^2
            part1 *= sin(q[i] * diam) / q[i]^3

            part2 = -diam * a0 + 2.0 * a1 / q[i]^2 - diam^2 * a1
            part2 += 6.0 * diam * a2 / q[i]^2
            part2 = (part2 - diam^3 * a2) * cos(q[i] * diam) / q[i]^2

            part3 = -2.0 * a1 / q[i]^4

            f[i] -= 4.0 * π * (part1 + part2 + part3)
        else
            f[i] = 0.0
        end
    end

    return nothing
end

@inline function extrapolation(g, f, rho, drho)
    a = (f .- g) ./ drho
    b = @. f - a * rho
    result = @. a * (rho + drho) + b

    return result
end

function inner_product(f, df)
    (integral, _) = romberg(df, f)

    return integral
end

function pairwise_inner_product(x::AbstractMatrix, dx::AbstractRange)
    n = size(x, 1)
    result_matrix = zeros(eltype(x), (n, n))

    # First, deal with the diagonal
    for i in 1:n
        temp_prod = x[i, :] .* x[i, :]
        result_matrix[i, i] = inner_product(temp_prod, dx)
    end

    for i in 1:n
        for j in (i + 1):n
            temp_prod = x[i, :] .* x[j, :]
            result_matrix[i, j] = inner_product(temp_prod, dx)
            result_matrix[j, i] = result_matrix[i, j] # Fill in the symmetric part
        end
    end

    return result_matrix
end


"""
A Stochastic Gradient optimisation algorithm based on relativistic Hamiltonian dynamics

X. Lu, V. Perrone, L. Hasenclever, Y. W. Teh, Sebastian J. Vollmer relativistic Monte Carlo https://arxiv.org/abs/1609.04388

Fields:
    - x: position
    - p: momentum
    - ntiers: number of iteration
    - mass: diagonal mass matrix
    - c: speed of light
    - D: diagonal preconditioner
    - independent_momenta: true means kinetic energy is applied to each component of momenta
                           false mean kinetic energy is applied to euclidean length of momenta
Constructors:
    - construct using `RelFrozenState(x::Array{Float64};stepsize = 0.001, p=:none, mass=[1.0],c=[1.0], D=[1.0],  independent_momenta=true)`
"""

type RelFrozenState <: SamplerState
    x::Array{Float64}
    p::Array{Float64}

    stepsize::Float64
    mass::Array{Float64,1}
    c::Array{Float64,1}
    D::Array{Float64,1}

    independent_momenta::Bool
    function RelFrozenState(x::Array{Float64};stepsize = 0.001, p=:none, mass=[1.0],c=[1.0], D=[1.0],  independent_momenta=true)
        if isa(mass,Number)
          mass = mass * ones(length(x))
        end
        if p == :none
            p = sample_rel_p(mass, c, length(x))
        end
    new(x,p,stepsize,mass,c,D,independent_momenta)
    end
end

function sample!(s::RelFrozenState, sgrad)
  D = s.D
  stepsize = s.stepsize

  m = s.mass
  c = s.c
  independent_momenta=s.independent_momenta
  independent_momenta ? tmp = m .* sqrt(s.p.^2 ./ (m.^2 .* c.^2) + 1) : tmp = m .* sqrt(s.p'*s.p ./ (m.^2 .* c.^2) + 1)
  p_grad = s.p./tmp
  s.p[:] += stepsize.*(sgrad(s.x)-D[1]*p_grad)
  independent_momenta ? tmp = m .* sqrt(s.p.^2 ./ (m.^2 .* c.^2) + 1)  :  tmp = m .* sqrt(s.p'*s.p ./ (m.^2 .* c.^2) + 1)
  s.x[:] += stepsize.*s.p./tmp
end


export getSensTMatVec

function getSensTMatVec(v::Vector,m::Vector,pFor::FWIparam)

    # extract pointers
    M  				= pFor.Mesh
    omega 			= pFor.omega
    gamma 			= pFor.gamma
    Q     			= pFor.Sources
    P     			= pFor.Receivers
    Ainv			= pFor.ForwardSolver

	batchSize 		= pFor.forwardSolveBatchSize;
	select  		= pFor.sourceSelection;

	nsrc = size(Q,2); nrec = size(P,2);

	if length(select) > 0
		nsrc = length(select);
	end

	if batchSize > nsrc
		batchSize = nsrc;
	end

	numBatches 	= ceil(Int64,nsrc/batchSize);
	An2cc = getNodalAverageMatrix(M);
	m_nodal = An2cc'*m;
	gamma_nodal = An2cc'*gamma;
	n_nodal = prod(M.n.+1);
	H = spzeros(ComplexF64,n_nodal,n_nodal);
	if isa(Ainv,ShiftedLaplacianMultigridSolver)
		H = GetHelmholtzOperator(M,m_nodal,omega, gamma_nodal, true,useSommerfeldBC);
		Ainv.helmParam = HelmholtzParam(M,gamma_nodal,m_nodal,omega,true,useSommerfeldBC);
		H = H + GetHelmholtzShiftOP(m_nodal, omega,Ainv.shift[1]);
		H = sparse(H');
		# H is actually shifted laplacian now...
	elseif isa(Ainv,JuliaSolver)
		H = GetHelmholtzOperator(M,m_nodal,omega, gamma_nodal, true,useSommerfeldBC);
	end

	JTv = zeros(Float64,prod(M.n))
	Vdatashape = reshape(v,nrec,nsrc);

	for k_batch = 1:numBatches
		batchIdxs = (k_batch-1)*batchSize + 1 : min(k_batch*batchSize,nsrc);
		V = P*Vdatashape[:,batchIdxs];
		V,Ainv = solveLinearSystem(H,V,Ainv,1);
		if pFor.useFilesForFields
			filename = getFieldsFileName(omega);
			file     = matopen(filename);
			U        = read(file,string("Ubatch_",k_batch));
			close(file);
		else
			U = pFor.Fields[:,batchIdxs];
		end
		# Original code:  V = conj((1.0.-1im*vec(gamma./omega)).*U).*V;
		V = An2cc*(conj(U).*(1.0.+1im*vec(gamma_nodal./omega)).*V);
		JTv .+= omega^2*vec(real(sum(V,dims=2)));
	end

	if isa(Ainv,ShiftedLaplacianMultigridSolver)
		clear!(Ainv.MG);
	end

    return JTv
end


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 9e3f6282-3575-11eb-35e8-69e4cf57b6f4
using LinearAlgebra, Plots, PlutoUI

# ╔═╡ 16df5992-3574-11eb-3682-0981f8882895
md"""
# Generalized optimal transport

Traditional optimal transport find the optimal transportation map given fixed marginals. For many applications, this requirement is too stringent, for example, because the marginals are not balanced. So, one can define **generalized optimal transport** as:

$$\max_Q\, \langle M, Q \rangle + \epsilon H(Q) + \alpha D_\text{KL}(\mathbf{a}\mid Q \mathbb{1}) + \beta D_\text{KL}(\mathbf{b}\mid Q^\intercal \mathbb{1})\,.$$

One notes that the **Kullback-Leibler** is now used as a soft-constraints to drive the marginals of $Q$ to match $\mathbf{a}$ and $\mathbf{b}$, these themselfs don't have to be normalized.

One can solve this problem using the **generalized Sinkhorn algorithm**.

"""

# ╔═╡ 244cb164-3575-11eb-3b9d-837acaa0e9e9
function gen_sinkhorn(M, a, b; ϵ=1.0, α=1.0, β=1.0)
	K = exp.(M ./ ϵ)
	u, v = ones(length(a)), ones(length(b))
	for i in 1:10
		u .= (a ./ K * v) .^ (α / (α + ϵ))
		v .= (b ./ K' * u) .^ (β / (β + ϵ))
	end
	K .*= u
	K .*= v'
	K ./= sum(K)
	return K
end

# ╔═╡ 62aa2bf2-3576-11eb-0cc6-e31f7ad771d0
n, m = 20, 10

# ╔═╡ 3486c424-3576-11eb-35a3-0bcc30b65445
M = randn(n, m)

# ╔═╡ ac95ae3a-357b-11eb-1e16-b56d76851314
normalize(x) = x ./ sum(x)

# ╔═╡ 5ecc17ca-3576-11eb-00ce-2fe7c5d8f55b
a = rand(n) #|> normalize

# ╔═╡ 72fc2c12-3576-11eb-23d1-c1d14da6de89
b = rand(m) #|> normalize

# ╔═╡ 3d0844ac-3577-11eb-3086-112eb6425f79
plot(bar(a), bar(b), layout=(2,1))

# ╔═╡ d9ffdbac-3576-11eb-24f1-55d257632108
@bind α Slider(0:0.01:10, default=1)

# ╔═╡ f22bc95c-3576-11eb-2a14-9d92a4497f68
@bind β Slider(0:0.01:10, default=1)

# ╔═╡ cd231944-357b-11eb-24b6-dbbb08cb40d3
@bind ϵ Slider(0:0.01:10, default=1)

# ╔═╡ 7702f944-3576-11eb-08e4-0f1f61b1d3bd
Q = gen_sinkhorn(M, a, b; ϵ, α, β)

# ╔═╡ c80cf7c2-3576-11eb-0e81-05121d909271
sum(Q)

# ╔═╡ d0a700bc-3576-11eb-2d6c-2354903c6a7d
heatmap(Q)

# ╔═╡ 84a3ccea-357b-11eb-0924-b9f6fe8c481f
ã = sum(Q, dims=2)[:]

# ╔═╡ 8fd9f8dc-357b-11eb-0fa8-812f0ea106e5
b̃ = sum(Q, dims=1)[:]

# ╔═╡ 7c05eabe-357b-11eb-0f67-4fa981d4d725
plot(bar(ã), bar(b̃), layout=(2,1))

# ╔═╡ Cell order:
# ╠═16df5992-3574-11eb-3682-0981f8882895
# ╠═9e3f6282-3575-11eb-35e8-69e4cf57b6f4
# ╠═244cb164-3575-11eb-3b9d-837acaa0e9e9
# ╠═62aa2bf2-3576-11eb-0cc6-e31f7ad771d0
# ╠═3486c424-3576-11eb-35a3-0bcc30b65445
# ╠═ac95ae3a-357b-11eb-1e16-b56d76851314
# ╠═5ecc17ca-3576-11eb-00ce-2fe7c5d8f55b
# ╠═72fc2c12-3576-11eb-23d1-c1d14da6de89
# ╠═3d0844ac-3577-11eb-3086-112eb6425f79
# ╠═d9ffdbac-3576-11eb-24f1-55d257632108
# ╠═f22bc95c-3576-11eb-2a14-9d92a4497f68
# ╠═cd231944-357b-11eb-24b6-dbbb08cb40d3
# ╠═7702f944-3576-11eb-08e4-0f1f61b1d3bd
# ╠═c80cf7c2-3576-11eb-0e81-05121d909271
# ╠═d0a700bc-3576-11eb-2d6c-2354903c6a7d
# ╠═84a3ccea-357b-11eb-0924-b9f6fe8c481f
# ╠═8fd9f8dc-357b-11eb-0fa8-812f0ea106e5
# ╠═7c05eabe-357b-11eb-0f67-4fa981d4d725


include("sec_6_2.jl")

using PyPlot

r = 3
M = 2r
ppI = 40
store = Store(α, r, M, ppI)

N = 3
t, U = FODEdG!(λ, tmax, f, u0, N, r, M, store)

figure(3)
tt = range(0, tmax, length=201)
exact_u = [ u(t, λ, u0, f, M, store) for t in tt ]
pcwise_t, pcwise_U = evaluate_pcwise_poly!(U, t, ppI, store)
line1 = plot(tt, exact_u, "k:")
line2 = [ plot(pcwise_t[:,n], pcwise_U[:,n], "c") for n = 1:N ]
legend((line1[1], line2[1][1]), (L"$u$", L"$U$"))
grid(true)
xlabel(L"$t$", fontsize=12)
xticks((0.0, 2/3, 4/3, 2), ("0", "2/3", "4/3", "2"))
savefig("fig3.eps")


#=
BasicPrim:
- Julia version: 1.6.2
- Author: samuel
- Date: 2021-09-10
=#

using SuiteSparseGraphBLAS
using SparseArrays

function stoppingCondition(m, n)

infinito = true

for i = 1:n
        if(m[i] == 0.0)
            return false
    end
end
    return infinito
end

#A is the input matrix, n is the number of nodes, m is the source matrix Output: minimum spanning tree cost
function b_prim(A, n, m)

    d = GBVector{Float64}(n)

    weight = 0.0

    d = A[1,:]
    print("INITIAL WEIGHT: ")
    print(weight)
    print("\n\n")
    while(stoppingCondition(m, n) == false)

        u = argmin(m'+d)
        m[u[2]] = Inf
        weight = weight + d[u[2]]
        print("WEIGHT: ")
        print(weight)
        d = emul(d, A[u[1],:],  BinaryOps.MIN)
        print("  ITERATION FINISHED")
        print("\n\n\n")
    end

    print("COMPLETED")
    print(" ALL THE NODES HAVE BEEN PROCESSED (INF)")
    print("\n\n\n")
    print(m')
    return weight

end


@testset "Quadratic" begin
    p, q = synthetic_gradient(SynthSphere(50))
    z = zeros(Float64, size(p))
    Z = Quadratic(z=z)(p, q)
    @test Z[38,38] ≈ -10.398166437049104
    @test Z[1,1] ≈ -10.704233100241964
    @test Z[76,76] ≈ 36.889627123013
end


module Light

export Polar, Stokes, PolEllipse, rotate, mean, convert, measurements2stokes
import Base: mean, convert, +, *

immutable Stokes{T <: AbstractFloat}
	s0::T
	s1::T
	s2::T
	s3::T
	function Stokes(s0,s1,s2,s3)
		if (s0 < 0) || (s1*s1 + s2*s2 + s3*s3 > s0*s0)
			s0 = s1 = s2 = s3 = T(NaN)
		end
		new(s0,s1,s2,s3)
	end
end
Stokes{T <: AbstractFloat}(s0::T,s1::T,s2::T,s3::T) = Stokes{T}(s0,s1,s2,s3)

immutable Polar{T <: AbstractFloat}
	I::T
	dolp::T
	aop::T
	docp::T
	function Polar(I,dolp,aop,docp)
		if (I < 0) || (dolp*dolp + docp*docp > 1)
			I = dolp = aop = docp = T(NaN)
		end
		new(I,dolp,aop,docp)
	end
end
Polar{T <: AbstractFloat}(I::T,dolp::T,aop::T,docp::T) = Polar{T}(I,dolp,aop,docp)

#stokes2polar(s0,s1,s2,s3) = (s0, sqrt(s1*s1 + s2*s2)/s0, atan2(s2,s1)/2, s3/s0)

immutable PolEllipse{T <: AbstractFloat}
	polarization::T
	a::T
	b::T
	angle::T
	lefthand::Bool
end
#PolEllipse{T <: AbstractFloat}(polarization::T,a::T,angle::T,lefthand::Bool) = PolEllipse{T}(polarization,ellipticity,angle,lefthand)

function PolEllipse(s::Stokes)
	tmp = s.s1*s.s1 + s.s2*s.s2
	s0P = sqrt(tmp + s.s3*s.s3)
	dolp = sqrt(tmp)
	R = (s0P - dolp)/(s0P + dolp)
	aop = atan2(s.s2,s.s1)/2
	h = s.s3 < 0
	P = s0P/s.s0
	a = P*R
	b = P
	return PolEllipse(P,a,b,aop,s.s3 < 0)
end

#=function convert(::Type{PolEllipse}, a::Polar)
	P = sqrt(a.dolp*a.dolp + a.docp*a.docp)
	A = sqrt(.5*(1 + a.dolp))
	B = sqrt(.5*(1 - a.dolp))
	return PolEllipse(P,A/B,a.aop,a.docp < 0)
end

convert(::Type{PolEllipse}, a::Stokes) = PolEllipse(Polar(a))=#

convert(::Type{Stokes}, a::Polar) = Stokes(a.I, a.dolp*cos(2a.aop)*a.I, a.dolp*sin(2a.aop)*a.I, a.docp*a.I)
convert(::Type{Polar}, a::Stokes) = Polar(a.s0, sqrt(a.s1*a.s1 + a.s2*a.s2)/a.s0, atan2(a.s2,a.s1)/2, a.s3/a.s0)

*(a::Stokes,b::AbstractFloat) = Stokes(a.s0*b, a.s1*b, a.s2*b, a.s3*b)

+(a::Stokes,b::Stokes) = Stokes(a.s0 + b.s0, a.s1 + b.s1, a.s2 + b.s2, a.s3 + b.s3)

rotate(a::Polar,b::AbstractFloat) = Polar(a.I,a.dolp,a.aop + b,a.docp)
rotate(a::Stokes,b::AbstractFloat) = Stokes(rotate(Polar(a),b))

function mean(a::Vector{Polar})
	I = mean(map(x -> x.I,a))
	dolp = mean(map(x -> x.dolp,a))
	α = map(x -> x.aop,a)
	aop = atan2(mean(sin(α)),mean(cos(α)))
	docp = mean(map(x -> x.docp,a))
	return Polar(I,dolp,aop,docp)
end

function measurements2stokes(ad, v, d, h, r, l)
	s0 = (v + h + d + ad + r + l)/3
	s1 = v - h
	s2 = d - ad
	s3 = r - l
	return (s0,s1,s2,s3)
end

convert(::Type{Stokes}, ad, v, d, h, r, l) = Stokes(measurements2stokes(ad, v, d, h, r, l)...)
convert(::Type{Polar}, ad, v, d, h, r, l) = Polar(Stokes(ad, v, d, h, r, l))

end


# NOTE distance in potentials is ångströms when fitting.
# But bohrs when CP2K calculates them. So that is why there is
# 0.52917721090 -terms in few places.

"""
GeneralAngle <: AbstractClusterPotential

General potential with customizable powers for both radial polynomials and cosθ.

# Fields
- `constants::Array{Float64,2}` : constants
- `ppowers::Vector{Int}`        : powers for radial polynomials
- `cpowers::Vector{Int}`        : powers for cosθ

To create potential use:
```julia
GeneralAngle(ppowers::AbstractVector{Int}, cpowers::AbstractVector{Int})
```
"""
mutable struct GeneralAngle <: AbstractClusterPotential
    constants::Array{Float64,2}
    "Powers for radial polynomials"
    ppowers::Vector{Int}
    "Powers for cosθ"
    cpowers::Vector{Int}
    function GeneralAngle(ppowers::AbstractVector{Int}, cpowers::AbstractVector{Int})
        new(zeros(length(cpowers), length(ppowers)), ppowers, cpowers)
    end
end

function Base.show(io::IO, potential::GeneralAngle; energy_unit="cm^-1")
    print(io, "GeneralAngle constants: ",potential.constants)
end

function calculate_potential(potential::GeneralAngle,
                                     cluster::AbstractCluster, indices)
    r = distances(cluster, indices[1], indices[2])
    θ = cluster_angle(cluster, indices[1], indices[2], indices[3])
    return potential(r,θ)
end

function potential_variables(potential::GeneralAngle,
                            cluster::AbstractCluster, indices)
    # Convert from Å to bohr
    r = distances(cluster, indices[1], indices[2]) ./ 0.52917721090
    θ = cluster_angle(cluster, indices[1], indices[2], indices[3])
    rr = r.^potential.ppowers
    cθ = cos(θ).^potential.cpowers
    out = [  x .* cθ for x in rr ]
    return vcat(out...)'
end

function get_potential!(potential::GeneralAngle, constants...)
    @assert length(constants) == length(potential.constants)
    for i in eachindex(constants)
        potential.constants[i] = constants[i]
    end
end

function (p::GeneralAngle)(r::Number,α::Number)
    # input is Å calculation in bohr
    c = cos.(α)
    cp = c.^p.cpowers
    rp = (r./0.52917721090).^p.ppowers
    sum(p.constants .* (cp * rp'))
end


module Examples

using CombinatorialSpaces
using CombinatorialSpaces.ExteriorCalculus
using Catlab
using Catlab.Present
using Catlab.Programs
using Catlab.WiringDiagrams
using Catlab.Graphics
using Catlab.CategoricalAlgebra
using Decapodes.Simulations
using Decapodes.Diagrams
using Decapodes.Schedules
using CombinatorialSpaces: ∧

using LinearAlgebra

import Catlab.Graphics: to_graphviz

export dual, sym2func, to_graphviz, gen_dec_rules

function dual(s::EmbeddedDeltaSet2D{O, P}) where {O, P}
  sd = EmbeddedDeltaDualComplex2D{O, eltype(P), P}(s)
  subdivide_duals!(sd, Barycenter())
  sd
end

sym2func(sd) = begin
  Dict(:d₀=>Dict(:operator => d(Val{0}, sd), :type => MatrixFunc()),
       :d₁=>Dict(:operator => d(Val{1}, sd), :type => MatrixFunc()),
       :dual_d₀=>Dict(:operator => dual_derivative(Val{0}, sd),
                      :type => MatrixFunc()),
       :dual_d₁=>Dict(:operator => dual_derivative(Val{1}, sd),
                      :type => MatrixFunc()),
       :⋆₀=>Dict(:operator => hodge_star(Val{0}, sd), :type => MatrixFunc()),
       :⋆₁=>Dict(:operator => hodge_star(Val{1}, sd), :type => MatrixFunc()),
       :⋆₂=>Dict(:operator => hodge_star(Val{2}, sd), :type => MatrixFunc()),
       :⋆₀⁻¹=>Dict(:operator => inv_hodge_star(Val{0}, sd), :type => MatrixFunc()),
       :⋆₁⁻¹=>Dict(:operator => inv_hodge_star(Val{1}, sd; hodge=DiagonalHodge()),
                   :type => MatrixFunc()),
       :⋆₂⁻¹=>Dict(:operator => inv_hodge_star(Val{2}, sd), :type => MatrixFunc()),
       :∧₁₀=>Dict(:operator => (γ, α,β)->(γ .= ∧(Tuple{1,0}, sd, α, β)),
                  :type=>InPlaceFunc()),
       :∧₁₁=>Dict(:operator => (γ, α,β)->(γ .= ∧(Tuple{1,1}, sd, α, β)),
                  :type=>InPlaceFunc()),
       :plus => Dict(:operator => (x,y)->(x+y), :type => ElementwiseFunc()),
       :L₀ => Dict(:operator => (v,u)->(lie_derivative_flat(Val{2}, sd, v, u)),
                   :type => ArbitraryFunc()))
#       :Δ₀=>Dict(:operator => Δ(Val{0}, sd), :type => MatrixFunc()),
#       :Δ₁=>Dict(:operator => Δ(Val{1}, sd), :type => MatrixFunc()),
end

function expand_dwd(dwd_orig, patterns)
  dwd = deepcopy(dwd_orig)
  has_updated = true
  while(has_updated)
    updates = map(1:nboxes(dwd)) do b
      v = box(dwd, b).value
      if v ∈ keys(patterns)
        dwd.diagram[b, :value] = patterns[v]
        true
      else
        false
      end
    end
    dwd = substitute(dwd)
    has_updated = any(updates)
  end
  dwd
end

function contract_matrices!(dwd, s2f)
  has_updated = true
  is_matrix(b) = begin
    b >= 1 || return false
    v = Symbol(first(split("$(box(dwd, b).value)", "⋅")))
    v != :remove && (s2f[v][:type] isa MatrixFunc) &&
      length(input_ports(dwd, b)) == 1 &&
      length(output_ports(dwd, b)) == 1
  end
  while(has_updated)
    has_updated = false
    # Join parallel matrix boxes and mark for removal
    for b in 1:nboxes(dwd)
      if is_matrix(b)
        ows = out_wires(dwd, b)
        if length(ows) == 1
          ow = ows[1]
          if is_matrix(ow.target.box)
            has_updated = true
            tb = ow.target.box
            dwd.diagram[b, :value] = Symbol(dwd.diagram[tb, :value], :⋅, dwd.diagram[b, :value])
            dwd.diagram[tb, :value] = :remove
            add_wires!(dwd, map(w -> Wire(w.value, Port(b, OutputPort, 1), w.target),
                                out_wires(dwd, tb)))
            dwd.diagram[incident(dwd.diagram, b, :out_port_box)[1], :out_port_type] = output_ports(dwd, tb)[1]
          end
        end
      end
    end
    rem_boxes!(dwd, findall(b -> dwd.diagram[b, :value] == :remove, 1:nboxes(dwd)))
  end
end

function boundary_inds(::Type{Val{1}}, s)
  collect(findall(x -> x != 0, boundary(Val{2},s) * fill(1,ntriangles(s))))
end

function boundary_inds(::Type{Val{0}}, s)
  ∂1_inds = boundary_inds(Val{1}, s)
  unique(vcat(s[∂1_inds,:src],s[∂1_inds,:tgt]))
end

function boundary_inds(::Type{Val{2}}, s)
  ∂1_inds = boundary_inds(Val{1}, s)
  inds = map([:∂e0, :∂e1, :∂e2]) do esym
    vcat(incident(s, ∂1_inds, esym)...)
  end
  unique(vcat(inds...))
end

function bound_edges(s, ∂₀)
  te = vcat(incident(s, ∂₀, :tgt)...)
  se = vcat(incident(s, ∂₀, :src)...)
  intersect(te, se)
end

function adj_edges(s, ∂₀)
  te = vcat(incident(s, ∂₀, :tgt)...)
  se = vcat(incident(s, ∂₀, :src)...)
  unique(vcat(te, se))
end

function gen_dec_rules()
  @present ExtendedOperators(FreeExtCalc2D) begin
    X::Space
    neg₁̃::Hom(DualForm1(X), DualForm1(X)) # negative
    half₁̃::Hom(DualForm1(X), DualForm1(X)) # half
    proj₁_¹⁰₁::Hom(Form1(X) ⊗ Form0(X), Form1(X))
    proj₂_¹⁰₀::Hom(Form1(X) ⊗ Form0(X), Form0(X))
    proj₁_¹²₁::Hom(Form1(X) ⊗ Form2(X), Form1(X))
    proj₂_¹²₂::Hom(Form1(X) ⊗ Form2(X), Form2(X))
    proj₁_¹²̃₁::Hom(Form1(X) ⊗ DualForm2(X), Form1(X))
    proj₂_¹²̃₂̃::Hom(Form1(X) ⊗ DualForm2(X), DualForm2(X))
    proj₁_¹¹̃₁::Hom(Form1(X) ⊗ DualForm1(X), Form1(X))
    proj₂_¹¹̃₁̃::Hom(Form1(X) ⊗ DualForm1(X), DualForm1(X))
    proj₁_¹̃¹̃₁̃::Hom(DualForm1(X) ⊗ DualForm1(X), DualForm1(X))
    proj₂_¹̃¹̃₁̃::Hom(DualForm1(X) ⊗ DualForm1(X), DualForm1(X))
    proj₁_¹¹₁::Hom(Form1(X)⊗Form1(X), Form1(X))
    proj₂_¹¹₁::Hom(Form1(X)⊗Form1(X), Form1(X))
    sum₁̃::Hom(DualForm1(X)⊗DualForm1(X), DualForm1(X))
    sum₁::Hom(Form1(X)⊗Form1(X), Form1(X))
    L₀::Hom(Form1(X)⊗DualForm2(X), DualForm2(X))
    L₁::Hom(Form1(X)⊗DualForm1(X), DualForm1(X))
    i₀::Hom(Form1(X)⊗DualForm2(X), DualForm1(X))
    i₁::Hom(Form1(X)⊗DualForm1(X), DualForm0(X))
  end

  Lie0Imp = @decapode ExtendedOperators begin
    (dḞ2, dF2)::DualForm2{X}
    F1::Form1{X}

    dḞ2 ==  dual_d₁{X}(i₀(F1, dF2))
  end
  lie0_imp = diag2dwd(Lie0Imp, in_vars = [:F1, :dF2], out_vars = [:dḞ2])

  Lie1Imp = @decapode ExtendedOperators begin
    (dḞ1, dF1)::DualForm1{X}
    F1::Form1{X}
    dḞ1 == i₀(F1, dual_d₁{X}(dF1)) + dual_d₀{X}(i₁(F1, dF1))
  end
  lie1_imp = diag2dwd(Lie1Imp, in_vars = [:F1, :dF1], out_vars = [:dḞ1])

  I0Imp = @decapode ExtendedOperators begin
    F1::Form1{X}
    dF1::DualForm1{X}
    dF2::DualForm2{X}
    dF1 == ⋆₁{X}(∧₁₀{X}(F1, ⋆₀⁻¹{X}(dF2)))
  end
  i0_imp = diag2dwd(I0Imp, in_vars = [:F1, :dF2], out_vars = [:dF1])

  I1Imp = @decapode ExtendedOperators begin
    F1::Form1{X}
    dF1::DualForm1{X}
    dF0::DualForm0{X}
    dF0 == ⋆₂{X}(∧₁₁{X}(F1, ⋆₁⁻¹{X}(dF1)))
  end
  i1_imp = diag2dwd(I1Imp, in_vars = [:F1, :dF1], out_vars = [:dF0])

  δ₁Imp = @decapode ExtendedOperators begin
    F1::Form1{X}
    F0::Form0{X}
    F0 == ⋆₀⁻¹{X}(dual_d₁{X}(⋆₁{X}(F1)))
  end
  δ₁_imp = diag2dwd(δ₁Imp, in_vars = [:F1], out_vars = [:F0])

  δ₂Imp = @decapode ExtendedOperators begin
    F1::Form1{X}
    F2::Form2{X}
    F1 == ⋆₁⁻¹{X}(dual_d₀{X}(⋆₂{X}(F2)))
  end
  δ₂_imp = diag2dwd(δ₂Imp, in_vars = [:F2], out_vars = [:F1])

  Δ0Imp = @decapode ExtendedOperators begin
    (F0, Ḟ0)::Form0{X}
    Ḟ0 == δ₁{X}(d₀{X}(F0))
  end
  Δ0_imp = diag2dwd(Δ0Imp, in_vars = [:F0], out_vars = [:Ḟ0])

  Δ1Imp = @decapode ExtendedOperators begin
    (F1, Ḟ1)::Form1{X}
    Ḟ1 == δ₂{X}(d₁{X}(F1)) + d₀{X}(δ₁{X}(F1))
  end
  Δ1_imp = diag2dwd(Δ1Imp, in_vars = [:F1], out_vars = [:Ḟ1])

  Dict(:L₀ => lie0_imp, :i₀ => i0_imp, :L₁ => lie1_imp, :i₁ => i1_imp,
       :δ₁ => δ₁_imp, :δ₂ => δ₂_imp, :Δ₀ => Δ0_imp, :Δ₁ => Δ1_imp)
end
end


# # Eady instability
#
# This script simulates the instability of a sheared, free-surface
# flow using `ClimateMachine.Ocean.HydrostaticBoussinesqSuperModel`.

using Printf

ENV["GKSwstype"] = "nul"
using Plots
using JLD2
using Plots
using ClimateMachine
using CUDA

ClimateMachine.init()

using ClimateMachine.Ocean
using ClimateMachine.Ocean.Domains
using ClimateMachine.Ocean.Fields

using ClimateMachine.GenericCallbacks: EveryXSimulationTime
using ClimateMachine.GenericCallbacks: EveryXSimulationSteps
using ClimateMachine.Ocean: current_step, Δt, current_time
using CLIMAParameters: AbstractEarthParameterSet, Planet

# # The Eady instability problem
#
# Our objective is to model the instability
# of a sinusoidal geostrophic jet with the streamfunction
#
# ```math
# ψ(y, z) = α λ cos(y / λ) (z + H)
# ```
#
# where ``α`` is geostrophic shear, ``λ`` is the width
# of the sinusoid, and ``H`` is depth. We set ``λ = L / π``,
# where ``L`` is the width of the domain, which fits a
# half wavelength of the sinusoid.
#
# The jet is vertically sheared with the total velocity
# field ``u = - ∂_y ψ``, such that
#
# ```math
# u = α sin(y / λ) (z + H) \, .
# ```
#
# The total buoyancy field is ``b = f ∂_z ψ + N² z``, where
# ``f ∂_z ψ`` is the geostrophic component of buoyancy and
# ``N² z`` reflects a stable stratification with buoyancy
# frequency ``N``. We thus have
#
# ```math
# b = α λ cos(y / λ) + N² z
# ```
#
# The baroclinic component of ``u`` is the "thermal wind"
# associated with ``b``. The barotropic component of ``u``,
# ``⟨u⟩ = \frac{1}{H} ∫ u \rm{d} z``, however, is
# geostrophically-balanced with the sea surface displacement
# via the north-south (``y``) momentum balance,
#
# ```math
# f ⟨u⟩ = - g ∂_y η \, ,
# ```
#
# where ``g`` is gravitational acceleration. Using
# ``⟨u⟩ = α H sin(y / λ) / 2``, we obtain
#
# ```math
# η = α f λ H / (2 g) cos(y / λ) \, .
# ```
#
# We add a periodic barotropic perturbation to
# this unstable state to simulate the growth of
# instability. The periodic perturbation has the
# streamfunction
#
# ```math
# ψ̃(x, y) = \cos(k x) \cos(ℓ y) \exp((y - y_0)^2 / 2L̃^2) \, ,
# ```
#
# where ``k = 12π / L``, ``ℓ = 6π / L``, ``y_0 = L/2``, and
# ``L̃ = 5 \times 10^4`` m. The velocity field and height
# perturbation associated with this streamfunction is
#
# ```math
# \begin{gather}
# ũ = - ∂_y ψ̃ = (ℓ \tan(ℓ y) + (y - y₀) / L̃^2 ) ψ̃
# ṽ = + ∂_x ψ̃ = k \tan(k x) ψ̃
# η = + f ψ̃  / g
# \end{gather}
#```
#
# We set the amplitude of the perturbation to 1% of the
# basic state.
#
# # Parameters
#
# We now choose parameters for the Eady instability problem.

Nh = 32
Nz = 4
Np = 4

L = 1e6 # Domain width (m)
H = 1e3 # Domain height (m)
f = 1e-4 # Coriolis parameter (s⁻¹)
α = 5f # Geostrophic shear (s⁻¹)
N² = 1e-5 # Initial buoyancy gradient (s⁻²)
νh = κh = 1e3 # Horizontal viscosity and diffusivity (m² s⁻¹)
νz = κh = 1e-2 # Vertical viscosity and diffusivity (m² s⁻¹)

minutes = 60.0
hours = 60minutes
days = 24hours
years = 365days

stop_time = 60days # Simulation stop time

struct EarthParameters <: AbstractEarthParameterSet end
Planet.grav(::EarthParameters) = 0.1
g = Planet.grav(EarthParameters())

# # The domain

domain = RectangularDomain(
    Ne = (Int(Nh / Np), Int(Nh / Np), Int(Nz / Np)),
    Np = Np,
    x = (0, L),
    y = (0, L),
    z = (-H, 0),
    periodicity = (true, false, false),
)

# # Boundary conditions

no_slip = OceanBC(Impenetrable(NoSlip()), Insulating())
free_surface = OceanBC(Penetrable(FreeSlip()), Insulating())

# # Initial conditions
#
# We idealize θ as buoyancy. We initialize the flow with a
# vertically-sheared, horizontally-sinusoidal jet in geostrophic
# balance. We add a small amount of noise to the buoyancy field.

## Geostrophic jet:
λ = L / π # Sinusoidal jet width (m)

U(y, z) = +α * sin(y / λ) * (z + H)
B(y, z) = +α * f * λ * cos(y / λ) + N² * z
h(y, z) = α * f * λ * H / 2g * cos(y / λ)

# Geostrophic mode-3 perturbation:
width = 5e4
kx = 12π / L
ky = 6π / L
y₀ = L / 2

ψ̃(x, y) = cos(kx * x) * cos(ky * y) * exp(-(y - y₀)^2 / 2width^2)

# ũ = - ∂y ψ̃
ũ(x, y) = (ky * tan(ky * y) + (y - y₀)^2 / width^2) * ψ̃(x, y)

# ṽ = ∂x ψ̃
ṽ(x, y) = -kx * tan(kx * x) * ψ̃(x, y)

a = 0.01
uᵢ(x, y, z) = U(y, z) + a * ũ(x, y)
vᵢ(x, y, z) = a * ṽ(x, y)
ηᵢ(x, y, z) = h(y, z) + a * f / g * ψ̃(x, y)
θᵢ(x, y, z) = B(y, z)

initial_conditions = InitialConditions(u = uᵢ, v = vᵢ, θ = θᵢ, η = ηᵢ)

minutes = 60.0

model = Ocean.HydrostaticBoussinesqSuperModel(
    domain = domain,
    time_step = 5minutes,
    initial_conditions = initial_conditions,
    parameters = EarthParameters(),
    # array_type = CuArray,
    buoyancy = (αᵀ = 1 / g,),
    coriolis = (f₀ = f, β = 0),
    turbulence_closure = (νʰ = νh, κʰ = κh, νᶻ = νz, κᶻ = νz),
    rusanov_wave_speeds = (cʰ = sqrt(g * H), cᶻ = 1e-2),
    boundary_tags = ((0, 0), (1, 1), (1, 2)),
    boundary_conditions = (no_slip, free_surface),
)

# We prepare a callback that periodically fetches the horizontal velocity and
# tracer concentration for later animation,

fetched_states = []

realdata = Array(model.state.realdata)
u = SpectralElementField(domain, model.grid, view(realdata, :, 1, :))

volume = assemble(u)
x = volume.x[:, 1, 1]
y = volume.y[1, :, 1]

start_time = time_ns()

data_fetcher = EveryXSimulationTime(days) do

    realdata = Array(model.state.realdata)
    u = SpectralElementField(domain, model.grid, view(realdata, :, 1, :))
    v = SpectralElementField(domain, model.grid, view(realdata, :, 2, :))
    η = SpectralElementField(domain, model.grid, view(realdata, :, 3, :))
    θ = SpectralElementField(domain, model.grid, view(realdata, :, 4, :))

    umax = maximum(abs, u)
    elapsed = (time_ns() - start_time) * 1e-9

    step = @sprintf("Step: %d", current_step(model))
    sim_time = @sprintf("time: %.2f days", current_time(model) / days)
    wall_time = @sprintf("wall time: %.2f min", elapsed / 60)

    @info "$step, $sim_time, $wall_time"

    isnan(umax) && error("NaN'd out.")

    u_assembly = assemble(data.(u.elements))
    v_assembly = assemble(data.(v.elements))
    θ_assembly = assemble(data.(θ.elements))

    push!(
        fetched_states,
        (
            u = u_assembly,
            v = v_assembly,
            θ = θ_assembly,
            time = current_time(model),
        ),
    )
end

# and then run the simulation.

model.solver_configuration.timeend = stop_time

try
    result = ClimateMachine.invoke!(
        model.solver_configuration;
        user_callbacks = [data_fetcher],
    )
catch err
    @warn "Simulation stopped early because " print(showerror, err)
end

# Finally, we make an animation of the evolving shear instability.

animation = @animate for (i, state) in enumerate(fetched_states)
    local u
    local v
    local θ

    @info "Plotting frame $i of $(length(fetched_states))..."

    kwargs = (
        xlim = domain.x,
        ylim = domain.y,
        linewidth = 0,
        aspectratio = 1,
        xlabel = "x (m)",
        ylabel = "y (m)",
    )

    u = state.u[:, :, end]
    v = state.v[:, :, end]
    θ = state.θ[:, :, end]

    ulim = 1.2 * α * H
    θmin = minimum(θ)
    θmax = maximum(θ)

    ulevels = range(-ulim, ulim, length = 31)
    θlevels = range(θmin, θmax, length = 31)

    u_plot = heatmap(
        x,
        y,
        clamp.(u, -ulim, ulim)';
        #levels = ulevels,
        color = :balance,
        clim = (-ulim, ulim),
        kwargs...,
    )

    v_plot = contourf(
        x,
        y,
        clamp.(v, -ulim, ulim)';
        #levels = ulevels,
        color = :balance,
        clim = (-ulim, ulim),
        kwargs...,
    )

    θ_plot = heatmap(
        x,
        y,
        clamp.(θ, θmin, θmax)';
        #levels = θlevels,
        color = :thermal,
        clim = (θmin, θmax),
        kwargs...,
    )

    u_title =
        @sprintf("surface u (m s⁻¹) at t = %.2f days", state.time / days)
    v_title =
        @sprintf("surface v (m s⁻¹) at t = %.2f days", state.time / days)
    θ_title =
        @sprintf("surface b (m s⁻²) at t = %.2f days", state.time / days)

    plot(
        u_plot,
        v_plot,
        θ_plot,
        title = [u_title v_title θ_title],
        layout = (1, 3),
        size = (1500, 400),
    )
end

name =
    @sprintf("eady_dg_Np%d_Nh%d_Nz%d_αf%.2e_ν%.2e.gif", Np, Nh, Nz, α / f, νh)

gif(animation, name, fps = 8)


export initAddEcho
export addEcho

"""
    arcrc(T₆₀::Number, Lᵣ::NTuple{3,Number})

                  β * V
`T₆₀ = ———————————————————————`\n
          - S * ln(1-α) + m₄ * V
where β = 0.160465, m₄ = 0.0104.
V is volume unit m³, S is area unit m²
"""
function arcrc(T₆₀::Number, Lᵣ::NTuple{3,Number})
    # acoustical reflaction coefficient and counts
    (T₆₀>3.00) ? (T₆₀=3.00) : (T₆₀=T₆₀) # max time
    (T₆₀<0.05) ? (T₆₀=0.05) : (T₆₀=T₆₀) # min time
    x,y,z = Lᵣ
    V = x * y * z                 # volume
    S = (x*y + x*z + y*z)*2       # area
    β  = 0.1611                   # 24log(10)/344
    m₄ = 0.0104                   # 4m value
    L  = 4 * V / S                # mean free propagatation length
    α = 1 - exp((m₄ - β/T₆₀)*V/S) # acoustical absorption coefficient
    C = floor(Int, 344 * T₆₀ / L) # reflaction counts
    γ = sqrt(1-α)                 # reflaction coefficient
    return (-γ, C÷6)
end

"""
    y = directconv(x, h)
convolves vectors x and h. The resulting vector is length length(x)+length(h)-1 which
has a O(Lx*Lh) complexity.
"""
function directconv(x, h)
    Lx = veclen(x)
    Lh = veclen(h)
    L  = Lx + Lh - 1
    y  = zeros(eltype(x), L)
    c  = h[end:-1:1]
    # boundary left y[1:Lh-1]
    for i=1:Lh-1
        y[i] = sum(x[1:i] .* h[i:-1:1])
    end
    # middle part y[Lh:Lx]
    for i=Lh:Lx
        y[i] = sum(x[i-Lh+1:i] .* c)
    end
    # right part y[Lx+1:Ly]
    for i=0:Lh-2
        y[L-i] = sum(x[Lx-i:Lx] .* h[Lh:-1:Lh-i])
    end
    return y
end

"""
    y = conv(x, h)
convolves signals x and h. The resulting signal has length (Lx + Lh - 1) which
has nearly O(Lx*Log(Lx)) complexity when Lx>>Lh. This function uses fft to reduce complexity.
"""
function conv(x, h)
    Lx = veclen(x)
    Lh = veclen(h)
    L  = Lx + Lh - 1
    xₑ = zeros(eltype(x), L)
    hₑ = zeros(eltype(h), L)
    xₑ[1:Lx] = x
    hₑ[1:Lh] = h
    return real(ifft( fft(xₑ) .* fft(hₑ) ))
end

"""
    y = rir(fs::Number, T60::Number, room::NTuple{3,Number}, src::NTuple{3,Number}, mic::NTuple{3,Number})
Generate room impulse response.
"""
function rir(fs::Number, T60::Number, room::NTuple{3,Number}, src::NTuple{3,Number}, mic::NTuple{3,Number})
    γ, C = arcrc(T60, room)
    N = 2*C + 1
    L = floor(Int, fs*T60)
    T = eltype(T60)
    x = zeros(T, N)
    y = zeros(T, N)
    z = zeros(T, N)
    h = zeros(T, L)
    xr,yr,zr = room  # room's size in meters
    xs,ys,zs = src   # source's coordinates in meters
    xm,ym,zm = mic   # micphone's coordinates in meters
    for i=-C:1:C
        if i&1==1 # odd number
            x[i+C+1] = xs + i*xr - xm;
            y[i+C+1] = ys + i*yr - ym;
            z[i+C+1] = zs + i*zr - zm;
        else # even number
            x[i+C+1] = -xs + (i+1)*xr - xm;
            y[i+C+1] = -ys + (i+1)*yr - ym;
            z[i+C+1] = -zs + (i+1)*zr - zm;
        end
    end
    maxi = 0
    for i=-C:1:C
        for j=-C:1:C
            for k=-C:1:C
                d = sqrt( x[i+C+1]^2 + y[j+C+1]^2 + z[k+C+1]^2 )
                n = floor(Int, d / 344.0 * fs)
                if n>L; continue; end
                if maxi<n; maxi=n;end
                m = abs(i) + abs(j) + abs(k)
                h[n] += γ^m / d
            end
        end
    end
    return h[1:maxi]
end

"""
    y = addEcho(wav::Array, fs::Number, T60::Number,
                room::NTuple{3,Number},
                src::NTuple{3,Number},
                mic::NTuple{3,Number})
Generate impulse response function online, and add reverberation effect to wav.
## Arguments
- `wav`: sound samples
- `fs`: sound's sampling rate
- `T60`: effective reverberation time
- `room`: room's size in meters, e.g., (3, 4, 2.5)
- `src`: source's coordinates in meters, e.g., (1, 2, 1.2)
- `mic`: micphone's coordinates in meters, e.g., (1.6, 2.6, 1.0)
"""
function addEcho(wav::Array, fs::Number, T60::Number,
                 room::NTuple{3,Number},
                 src::NTuple{3,Number},
                 mic::NTuple{3,Number})
    return conv(wav, rir(fs, T60, room, src, mic))
end

function initAddEcho(path::String, period::Int)
    counter = 1
    h = nothing
    FILES = readdir(path)
    function addecho(speech::Array)
        if counter == 1
            file = rand(FILES,1)[1]
            @assert endswith(file, "wav") "$path should only keep *.wav files"
            h, fs = wavread(joinpath(path, file))
        end # read another Room Impulse Response every period
        (counter == period) ? (counter=1) : (counter+=1)
        return conv(speech, h)
    end
    return addecho
end

"""
    initAddEcho(fs::Number, T₆₀Span::NTuple{2,Number}, roomSpan::NTuple{6,Number}) -> addecho(wav::Array)
init reverberation effect function.
+ `fs` sampling rate
+ `T₆₀Span` T₆₀ range e.g. (0.05, 0.5)
+ `roomSpan`room size e.g. (MinL, MaxL, MinW, MaxW, MinH, MaxH)
"""
function initAddEcho(fs::Number, T₆₀Span::NTuple{2,Number}, roomSpan::NTuple{6,Number})
    MinT₆₀, MaxT₆₀ = T₆₀Span
    MinL, MaxL, MinW, MaxW, MinH, MaxH = roomSpan
    @assert MinT₆₀ <= MaxT₆₀
    @assert MinL <= MaxL
    @assert MinW <= MaxW
    @assert MinH <= MaxH
    function addecho(wav::Array)
        T₆₀ = rand()*(MaxT₆₀ - MinT₆₀) + MinT₆₀
        Lx  = rand()*(MaxL - MinL) + MinL
        Ly  = rand()*(MaxW - MinW) + MinW
        Lz  = rand()*(MaxH - MinH) + MinH
        room = (Lx, Ly, Lz)
        src  = (rand()*Lx, rand()*Ly, rand()*Lz)
        mic  = (rand()*Lx, rand()*Ly, rand()*Lz)
        return conv(wav, rir(fs, T₆₀, room, src, mic))
    end
    return addecho
end


🦀 = parse.(Int, split(read("input.txt", String), ","))
fuel1, fuel2 = typemax(Int), typemax(Int)
for loc in minimum(🦀):maximum(🦀)
    fuel1 = min(fuel1, sum(x -> abs(x - loc), 🦀))
    fuel2 = min(fuel2, sum(x -> sum(1:abs(x - loc)), 🦀))
end
@info "Part 1: minimum fuel usage = $fuel1"
@info "Part 2: minimum fuel usage = $fuel2"