text:
@doc raw"""
    LogisticSoftMaxLikelihood(num_class::Int) -> MultiClassLikelihood

## Arguments
- `num_class::Int` : Total number of classes

    LogisticSoftMaxLikelihood(labels::AbstractVector) -> MultiClassLikelihood

## Arguments
- `labels::AbstractVector` : List of classes labels

---

The multiclass likelihood with a logistic-softmax mapping: :
```math
p(y=i|\{f_k\}_{1}^{K}) = \frac{\sigma(f_i)}{\sum_{k=1}^k \sigma(f_k)}
```
where ``\sigma`` is the logistic function.
This likelihood has the same properties as [softmax](https://en.wikipedia.org/wiki/Softmax_function).
---

For the analytical version, the likelihood is augmented multiple times.
More details can be found in the paper [Multi-Class Gaussian Process Classification Made Conjugate: Efficient Inference via Data Augmentation](https://arxiv.org/abs/1905.09670).
"""
LogisticSoftMaxLikelihood(x) = MultiClassLikelihood(LogisticSoftMaxLink(), x)

struct LogisticSoftMaxLink <: AbstractLink end

function (::LogisticSoftMaxLink)(f::AbstractVector{<:Real})
    return normalize(logistic.(f), 1)
end

function implemented(
    ::MultiClassLikelihood{<:LogisticSoftMaxLink},
    ::Union{<:AnalyticVI,<:MCIntegrationVI,<:GibbsSampling},
)
    return true
end

Base.show(io::IO, ::LogisticSoftMaxLink) = print(io, "Logistic-SoftMax Link")

## Local Updates ##
function init_local_vars(
    l::MultiClassLikelihood{<:LogisticSoftMaxLink}, batchsize::Int, T::DataType=Float64
)
    num_class = n_class(l)
    c = [ones(T, batchsize) for _ in 1:num_class] # Second moment of fₖ
    α = num_class * ones(T, batchsize) # First variational parameter of Gamma distribution
    β = num_class * ones(T, batchsize) # Second variational parameter of Gamma distribution
    θ = [rand(T, batchsize) * 2 for _ in 1:num_class] # Variational parameter of Polya-Gamma distribution
    γ = [rand(T, batchsize) for _ in 1:num_class] # Variational parameter of Poisson distribution
    return (; c, α, β, θ, γ)
end

function local_updates!(
    local_vars,
    ::MultiClassLikelihood{<:LogisticSoftMaxLink},
    y,
    μ::NTuple{N,<:AbstractVector},
    Σ::NTuple{N,<:AbstractVector},
) where {N}
    @. local_vars.c = broadcast((Σ, μ) -> sqrt.(Σ + abs2.(μ)), Σ, μ)
    for _ in 1:2
        broadcast!(
            (β, c, μ, ψα) -> 0.5 / β * exp.(ψα) .* safe_expcosh.(-0.5 * μ, 0.5 * c),
            local_vars.γ,
            Ref(local_vars.β),
            local_vars.c,
            μ,
            Ref(digamma.(local_vars.α)),
        ) # Update γ
        local_vars.α .= 1 .+ (local_vars.γ...)
    end
    broadcast!(
        (y, γ, c) -> 0.5 * (y + γ) ./ c .* tanh.(0.5 .* c),
        local_vars.θ, # target
        eachcol(y), # argument 1
        local_vars.γ, # argument 2
        local_vars.c, # argument 3
    ) # update θ
    return local_vars
end

function sample_local!(local_vars, ::MultiClassLikelihood{<:LogisticSoftMaxLink}, y, f)
    broadcast!(
        f -> rand.(Poisson.(0.5 * local_vars.α .* safe_expcosh.(-0.5 * f, 0.5 * f))),
        local_vars.γ,
        f,
    )
    local_vars.α .= rand.(Gamma.(1 .+ (local_vars.γ...), inv.(local_vars.β)))
    local_vars.θ .= broadcast(
        (y, γ, f) -> rand.(PolyaGamma.(y .+ Int.(γ), abs.(f))), eachcol(y), local_vars.γ, f
    )
    return local_vars
end

## Global Gradient Section ##

@inline function ∇E_μ(::MultiClassLikelihood{<:LogisticSoftMaxLink}, ::AOptimizer, y, state)
    return 0.5 .* (eachcol(y) .- state.γ)
end
@inline function ∇E_Σ(::MultiClassLikelihood{<:LogisticSoftMaxLink}, ::AOptimizer, y, state)
    return 0.5 .* state.θ
end

## ELBO Section ##
function expec_loglikelihood(
    ::MultiClassLikelihood{<:LogisticSoftMaxLink}, ::AnalyticVI, y, μ, Σ, state
)
    tot = -length(y) * logtwo
    tot += -sum(sum(state.γ .+ eachcol(y))) * logtwo
    tot += 0.5 * sum(zip(state.θ, state.γ, eachcol(y), μ, Σ)) do (θ, γ, y, μ, Σ)
        dot(μ, (y - γ)) - dot(θ, abs2.(μ)) - dot(θ, Σ)
    end
    return tot
end

function AugmentedKL(l::MultiClassLikelihood{<:LogisticSoftMaxLink}, state, y)
    return PolyaGammaKL(l, state, y) + PoissonKL(l, state) + GammaEntropy(l, state)
end

function PolyaGammaKL(::MultiClassLikelihood, state, y)
    return sum(broadcast(PolyaGammaKL, eachcol(y) .+ state.γ, state.c, state.θ))
end

function PoissonKL(::MultiClassLikelihood, state)
    return sum(
        broadcast(
            PoissonKL,
            state.γ,
            Ref(state.α ./ state.β),
            Ref(digamma.(state.α) .- log.(state.β)),
        ),
    )
end

##  Compute the equivalent of KL divergence between an improper prior p(λ) (``1_{[0,\\infty]}``) and a variational Gamma distribution ##
function GammaEntropy(::MultiClassLikelihood, state)
    return -sum(state.α) + sum(log, first(state.β)) - sum(first ∘ logabsgamma, state.α) -
           dot(1 .- state.α, digamma.(state.α))
end

## Numerical Gradient Section ##

function grad_samples!(
    model::AbstractGPModel{T,<:MultiClassLikelihood{<:LogisticSoftMaxLink},<:NumericalVI},
    samples::AbstractMatrix{T},
    opt_state,
    y,
    index,
) where {T}
    l = likelihood(model)
    grad_μ = zeros(T, n_latent(model))
    grad_Σ = zeros(T, n_latent(model))
    g_μ = similar(grad_μ)
    num_sample = size(samples, 1)
    @views @inbounds for i in 1:num_sample
        σ = logistic.(samples[i, :])
        samples[i, :] .= l(samples[i, :])
        s = samples[i, y][1]
        g_μ .= grad_logisticsoftmax(samples[i, :], σ, y) / s
        grad_μ += g_μ
        grad_Σ += diaghessian_logisticsoftmax(samples[i, :], σ, y) / s - abs2.(g_μ)
    end
    for k in 1:n_latent(model)
        opt_state[k].ν[index] = -grad_μ[k] / num_sample
        opt_state[k].λ[index] = grad_Σ[k] / num_sample
    end
end

function log_like_samples(
    ::AbstractGPModel{T,<:MultiClassLikelihood{<:LogisticSoftMaxLink}},
    samples::AbstractMatrix,
    y,
) where {T}
    num_sample = size(samples, 1)
    loglike = zero(T)
    for i in 1:num_sample
        σ = logistic.(samples[i, :])
        loglike += log(σ[y][1]) - log(sum(σ))
    end
    return loglike / num_sample
end

function grad_logisticsoftmax(s::AbstractVector{T}, σ::AbstractVector{T}, y) where {T<:Real}
    return s[y][1] * (y .- s) .* (one(T) .- σ)
end

function diaghessian_logisticsoftmax(
    s::AbstractVector{T}, σ::AbstractVector{T}, y
) where {T<:Real}
    return s[y][1] * (one(T) .- σ) .*
           (abs2.(y - s) .* (1 .- σ) - s .* (1 .- s) .* (1 .- σ) - σ .* (y - s))
end

function hessian_logisticsoftmax(
    s::AbstractVector{T}, σ::AbstractVector{T}, y
) where {T<:Real}
    m = length(s)
    i = findfirst(y)
    hessian = zeros(T, m, m)
    @inbounds for j in 1:m
        for k in 1:m
            hessian[j, k] =
                (one(T) - σ[j]) *
                s[i] *
                (
                    (δ(T, i, k) - s[k]) * (one(T) - σ[k]) * (δ(T, i, j) - s[j]) -
                    s[j] * (δ(T, j, k) - s[k]) * (one(T) - σ[k]) -
                    δ(T, k, j) * σ[j] * (δ(T, i, j) - s[j])
                )
        end
    end
    return hessian
end


text:
"""
    @fieldequal Supertype

Define a method of `==` for all subtypes of `Supertype`
such that `==` returns true if two instances have the same field values.
"""
macro fieldequal(Supertype)
    return esc(quote
        function ==(x::T, y::T) where T <: $Supertype
            f = fieldnames(T)
            getfield.(Ref(x),f) == getfield.(Ref(y),f)
        end
    end)
end

"""
    cb(args...)

Construct a vector from `ConstantTerm`s provided as arguments.
This method is useful for working with `@formula`.

Note that `@formula` does not accept function calls with empty argument.
Hence, `cb()` cannot be used to create an empty vector `[]`
in the expression passed to `@formula`.
However, passing `cb` (without the parenthesis) to `@formula` will return `[]`
due to [`unpack`](@ref).
"""
cb(args...) = [arg isa ConstantTerm ? arg.n :
    throw(ArgumentError("only `ConstantTerm`s are accepted.")) for arg in args]

cb() = []

"""
    unpack(t::ConstantTerm)

Return the value represented by `t`.
"""
unpack(t::ConstantTerm) = t.n

"""
    unpack(t::Term)

If `t.sym` is the name of an object defined in `Main`,
either call its method with no argument if it exists
or return the object itself.
Otherwise, return `t.sym`.
"""
function unpack(t::Term)
    if isdefined(Main, t.sym)
        f = getfield(Main,  t.sym)
        hasmethod(f, Tuple{}) ? f() : f
    else
        return t.sym
    end
end

"""
    unpack(t::FunctionTerm{F})

Call the function represented by `t`.
"""
unpack(t::FunctionTerm{F}) where F = F.instance(t.args_parsed...)

"""
    kwarg(v)

Return a key-value `Pair` with the key being a keyword argument name
and the value being `v`.
The key is determined by the type of `v`.
"""
kwarg(::Any) = nothing

"""
    @unpack functionname

Define a method of `functionname` that accepts terms generated by `@formula`
as arguments.
This method can be used to translate terms into
arguments that match the other methods of `functionname`.
Each term is processed by [`unpack`](@ref) and [`kwarg`](@ref).
"""
macro unpack(functionname)
    return esc(quote
        function $functionname(ts::AbstractTerm...)
            args, kwargs = [], []
            for t in ts
                v = unpack(t)
                kv = kwarg(v)
                kv isa Nothing ? push!(args, v) : push!(kwargs, kv)
            end
            return $functionname(args...; kwargs...)
        end
    end)
end

"""
    exampledata()

Return the names of available example datasets.
"""
exampledata() = (:hrs, :nsw, :mpdta)

"""
    exampledata(name::Union{Symbol,String})

Return a `CSV.File` containing the example dataset with the specified `name`.
"""
function exampledata(name::Union{Symbol,String})
    Symbol(name) in exampledata() ||
        throw(ArgumentError("example dataset $(name) does not exist"))
    path = (@__DIR__)*"/../data/$(name).csv.gz"
    cols = open(path) |> GzipDecompressorStream |> read |> CSV.File |> VecColumnTable
    # Avoid customized index type for columns from CSV.File
    for i in 1:ncol(cols)
        col = _columns(cols)[i]
        _columns(cols)[i] = convert(Vector, col)
    end
    return cols
end

# Check whether the input data is a column table
function checktable(data)
    istable(data) ||
        throw(ArgumentError("data of type $(typeof(data)) is not `Tables.jl`-compatible"))
    Tables.columnaccess(data) ||
        throw(ArgumentError("data of type $(typeof(data)) is not a column table"))
end

ext: jl


text:
#=
gradient.jl

    Proximal gradient methods, w/ and w/o acceleration, to optimize an objective
    function that can be split into two components, one of which is
    differentiable.

@author: Quint Wiersma <q.wiersma@vu.nl>

@date: 2022/01/14
=#

"""
    update_state!(state, λ, prox!)

Update state using proximal gradient step with step size `λ`, storing the result
in `x`.

#### Arguments
  - `state::ProxGradState`  : state variables
  - `λ::Real`               : stepsize
  - `prox!::Function`       : proximal operator of ``g(x)``
"""
function update_state!(state::ProxGradState, λ::Real, prox!::Function)
    # Proximal gradient step
    state.x.= state.y .- λ .* state.∇f
    prox!(state.x, λ)

    return nothing
end

"""
    backtrack!(ls, state, f, prox!; ϵ=1e-7)

Backtracking line search to find the optimal step size `λ`.

#### Arguments
  - `ls::BackTrack`         : line search variables
  - `state::ProxGradState`  : state variables
  - `λ::Real`               : initial stepsize
  - `f::Function`           : objective function
  - `prox!::Function`       : proximal operator of ``g(x)``
  - `ϵ::Real`               : tolerance
"""
function backtrack!(ls::BackTrack, state::ProxGradState, f::Function,
                    prox!::Function; ϵ::Real=1e-7)
    # Store previous stepsize
    ls.λ_prev= ls.λ

    # Objective function value, f(y)
    f_y= f(state.y)

    # Minimal stepsize
    λ_min= ϵ * maximum(abs, state.∇f)

    # Update state
    update_state!(state, ls.λ, prox!)

    # Objective function value, f(x)
    f_x= f(state.x)

    # f̂
    state.Δ.= state.x .- state.y
    f_hat= f_y + dot(state.∇f, state.Δ) + inv(ls.λ+ls.λ)*norm(state.Δ)^2

    # tolerance
    tol= 10 * eps(f_x) * (one(f_x) + abs(f_x))

    # Backtracking line search
    while f_x > f_hat + tol && ls.λ > λ_min
        # Update stepsize
        ls.λ*= ls.β

        # Update state
        update_state!(state, ls.λ, prox!)

        # Objective function value
        f_x= f(state.x)

        # Update f̂
        state.Δ.= state.x .- state.y
        f_hat= f_y + dot(state.∇f, state.Δ) + inv(ls.λ+ls.λ)*norm(state.Δ)^2

        # Update tolerance
        tol= 10 * eps(f_x) * (one(f_x) + abs(f_x))
    end

    return nothing
end

"""
    prox_grad(x0, f, ∇f!, prox!; style="none", β=.5, ϵ=1e-7, max_iter=1000)

Minimize an objective function ``f(x) + g(x)``, where ``f(x)`` is differentibale
while ``g(x)`` is not, using the proximal gradient method.

#### Arguments
  - `x0::AbstractVector`    : initial parameter values (n x 1)
  - `f::Function`           : ``f(x)``
  - `∇f!::Function`         : gradient of `f`
  - `prox!::Function`       : proximal operator of ``g(x)``
  - `style::AbstractString` : acceleration style
  - `β::Real`               : line search parameter
  - `ϵ::Real`               : tolerance
  - `max_iter::Integer`     : max number of iterations

#### Returns
  - `x::AbstractVector` : minimizer (optimal parameter values) (n x 1)
"""
function prox_grad(x0::AbstractVector, f::Function, ∇f!::Function, prox!::Function;
                    style::AbstractString="none", β::Real=.5, ϵ::Real=1e-7,
                    max_iter::Integer=1000)
    # Initialize state and line search
    state= ProxGradState(copy(x0), copy(x0), similar(x0), similar(x0), similar(x0))
    ls= BackTrack(one(Float64), one(Float64), β)

    # Initialize acceleration
    if style == "none"
        acc= NoAccel(zero(Float64))
    elseif style == "simple"
        acc= Simple(zero(Float64), one(Int64))
    elseif style == "nesterov"
        acc= Nesterov(ω= zero(Float64), θ= one(Float64), ls= ls)
    end

    # Initialize gradient change
    change= Inf
    # Initialize iteration counter
    iter= 1
    # Proximal gradient method
    while change > ϵ && iter < max_iter
        # Store change in state
        state.Δ.= state.x .- state.x_prev

        # Store current parameters
        copyto!(state.x_prev, state.x)

        # Extrapolation
        update_acc!(acc, state)
        state.y.= state.x .+ acc.ω .* state.Δ

        # Current gradient
        ∇f!(state.∇f, state.y)

        # Backtracking linesearch
        backtrack!(ls, state, f, prox!, ϵ=ϵ)

        # gradient change
        change= norm(state.∇f)

        # Update iteration counter
        iter+=1
    end

    return state.x
end

"""
    prox_grad!(x, f, ∇f!, prox!; style="none", β=.5, ϵ=1e-7, max_iter=1000)

Minimize an objective function ``f(x) + g(x)``, where ``f(x)`` is differentibale
while ``g(x)`` is not, using the proximal gradient method. Storing the result in
`x`. See also `prox_grad`.
"""
function prox_grad!(x::AbstractVector, f::Function, ∇f!::Function, prox!::Function;
                    style::AbstractString="none", β::Real=.5, ϵ::Real=1e-7,
                    max_iter::Integer=1000)
    # Initialize state and line search
    state= ProxGradState(x, copy(x), similar(x), similar(x), similar(x))
    ls= BackTrack(one(Float64), one(Float64), β)

    # Initialize acceleration
    if style == "none"
        acc= NoAccel(zero(Float64))
    elseif style == "simple"
        acc= Simple(zero(Float64), one(Int64))
    elseif style == "nesterov"
        acc= Nesterov(ω= zero(Float64), θ= one(Float64), ls= ls)
    end

    # Initialize gradient change
    change= Inf
    # Initialize iteration counter
    iter= 1
    # Proximal gradient method
    while change > ϵ && iter < max_iter
        # Store change in state
        state.Δ.= state.x .- state.x_prev

        # Store current parameters
        copyto!(state.x_prev, state.x)

        # Extrapolation
        update_acc!(acc, state)
        state.y.= state.x .+ acc.ω .* state.Δ

        # Current gradient
        ∇f!(state.∇f, state.y)

        # Backtracking linesearch
        backtrack!(ls, state, f, prox!, ϵ=ϵ)

        # gradient change
        change= norm(state.∇f, Inf)

        # Update iteration counter
        iter+=1
    end

    return nothing
end


text:
import Snopt

#####
##### Model
#####

println("===== Building Model =====")
crc3, parameters, environment, alphas, stepsymbol, cl_stall, nx, nu = MT.build_crc3(; nwings)
crc3_model = MT.AircraftDynamics(crc3, parameters, environment, stepsymbol, cl_stall, nx, nu)

#####
##### trim
#####

# xu_xdot = [
#     1. px,
#     2. py,
#     3. vx,
#     4. vy,
#     5. theta,
#     6. theta_dot,
#     7. energy,
#     8. px_dot,
#     9. py_dot,
#     10. vx_dot,
#     11. vy_dot,
#     12. theta_dot,
#     13. theta_dot_dot,
#     14. power
# ]

x_trim_guess = [
    0.0, # x
    0.0, # y
    0.0, # vx
    0.0, # vy
    pi/2, # theta
    0.0, # theta_dot
    0 # energy
]

u_trim_guess = ones(nwings) * 10000 .* pi/30

xu_trim_guess = vcat(x_trim_guess, u_trim_guess)

x_xdot_trimmed = [
    # 0.0, # px
    # 0.0, # py
    # 0.0, # vx
    # 0.0, # vy
    pi/2, # theta
    # 0.0, # theta_dot
    # 0.0, # energy
    # 0.0, # px_dot
    # 0.0, # py_dot
    0.0, # vx_dot
    0.0, # vy_dot
    # 0.0, # theta_dot
    0.0, # theta_dot_dot
    # 0.0, # power
]

i_x_xdot_trimmed = [5,10,11,13]
# i_xu_xdot_trimmed = [3, 4, 6, 7, 8, 9, 10, 11, 12, 13]

lxu = [
    0.0,
    0.0,
    -10.0,
    -10.0,
    -pi,
    -10,
    0.0,
    0.0,
    0.0
]

uxu = [
    0.0,
    0.0,
    10.0,
    10.0,
    pi,
    10.0,
    10.0,
    10000.0 * pi/30,
    10000.0 * pi/30
]

println("===== Begin Optimization =====")

xopt, fopt, info = MT.trim(crc3_model, xu_trim_guess, x_xdot_trimmed, i_x_xdot_trimmed, lxu, uxu;
    derivatives = MT.SNOW.CentralFD(),
    # derivatives = MT.SNOW.ComplexStep(),
    # solver = MT.SNOW.SNOPT(),
    solver = MT.SNOW.IPOPT(),
    objective = (x, xdot, u) -> MT.LA.norm(x[3:4]) / 1000
)

println("xstar = ", xopt)
println("fstar = ", fopt)
println("info = ", info)

max_stars_repo_name: byuflowlab/MultifidelityTrajectories

max_issues_repo_name: byuflowlab/MultifidelityTrajectories

max_forks_repo_name: byuflowlab/MultifidelityTrajectories


text:
## Label propagation algorithm for detecting communities ##

function _count_neighbor_labels(g::Graph, label::Vector{Int}, n::Integer)
    count = Dict{Int,Int}()
    for i = neighbors(g, n)
        l = label[i]
        count[l] = get(count, l, 0) + 1
    end
    return count
end

"""
    label_propagation(g::Graph)

A simple implementation of the label propagation algorithm.

# References

Raghavan, Usha Nandini, Réka Albert, and Soundar Kumara.
“Near Linear Time Algorithm to Detect Community Structures in Large-Scale Networks.”
Physical Review E 76, no. 3 (2007).
"""
function label_propagation(g::Graph)
    label = collect(nodes(g))
    finished = false
    while !finished
        order = shuffle(nodes(g))
        # Assign new labels
        for i = order
            count = _count_neighbor_labels(g, label, i)
            maxlabel, maxcount = 0, 0
            for l = keys(count)
                if count[l] > maxcount || count[l] == maxcount && order[l] < order[maxlabel]
                    maxlabel = l
                    maxcount = count[l]
                end
            end
            label[i] = maxlabel
        end
        # Check if labels don't need to be changed
        finished = true
        for i = nodes(g)
            count = _count_neighbor_labels(g, label, i)
            lcount = get(count, label[i], 0)
            if any(c > lcount for c = values(count))
                finished = false
                break
            end
        end
    end
    return relabel(label)
end


text:
@polly function kernel_mvt(x1, x2, y1, y2, A)
    n = size(x1,1)
    for i = 1:n, j = 1:n
        x1[i] = x1[i] + A[i,j] * y1[j]
    end
    for i = 1:n, j = 1:n
        x2[i] = x2[i] + A[j,i] * y2[j]
    end
end

let
    n = 2000

    A = zeros(Float32, n, n)
    x1 = zeros(Float32, n)
    x2 = zeros(Float32, n)
    y1 = zeros(Float32, n)
    y2 = zeros(Float32, n)

    for i = 1:n
        x1[i] = (i % n) / n
        x2[i] = ((i + 1) % n) / n
        y1[i] = ((i + 3) % n) / n
        y2[i] = ((i + 4) % n) / n
        for j = 1:n
            A[i,j] = (i*j % n) / n
        end
    end

    SUITE["mvt"] = @benchmarkable kernel_mvt(x1, x2, y1, y2, A) setup = (x1 = copy($x1); x2 = copy($x2); y1 = copy($y1); y2 = copy($y2); A = copy($A))
end


text:
function gen_lattice_hexagonal(a; coa=8.0/3.0)
    LL = zeros(3,3)
    LL[:,1] = [1.0, 0.0, 0.0]
    LL[:,2] = [-cos(pi/3.0), sin(pi/3.0), 0.0]
    LL[:,3] = [0.0, 0.0, coa]
    return a*LL
end

function gen_lattice_fcc(a)
    LL = zeros(3,3)
    LL[:,1] = [ 0.0, 1.0, 1.0]
    LL[:,2] = [ 1.0, 0.0, 1.0]
    LL[:,3] = [ 1.0, 1.0, 0.0]
    return 0.5*a*LL
end


text:
module EGraphs

include("../docstrings.jl")

using RuntimeGeneratedFunctions
RuntimeGeneratedFunctions.init(@__MODULE__)

import ..Rule
import ..iscall
import ..istree
import ..getfunsym
import ..getfunargs
import ..setfunsym!
import ..setfunargs!

using ..Util

include("enode.jl")
export isenode

include("egg.jl")
export find
export EClass
export EGraph
export AbstractAnalysis
export merge!
export addexpr!
export addanalysis!
export rebuild!

include("analysis.jl")

include("ematch.jl")
include("Schedulers/Schedulers.jl")

include("saturation_report.jl")
include("saturation.jl")
export saturate!
include("equality.jl")
export areequal
export @areequal
export @areequalg

include("extraction.jl")
export extract!
export ExtractionAnalysis
export astsize
export astsize_inv
export @extract

end


text:
#empty
#code was put into __init__()


text:
using ITensors
using ITensorsGateEvolution

function main()

N = 10

osRand = ProductOps()
for n in 1:N
  osRand *= ("rand", n)
end

osX = ProductOps()
for n in 1:N
  osX *= ("X", n)
end

osZ = ProductOps()
for n in 1:N
  osZ *= ("Z", n)
end

osSw = ProductOps()
for n in 1:N-1
  osSw *= ("Sw", n, n+1)
end

osCx = ProductOps()
for n in 1:N-1
  osCx *= ("Cx", n, n+1)
end

osRand = ProductOps()
for n in 1:N-1
  osRand *= ("rand", n, n+1)
end

osT = ProductOps()
for n in 1:N-4
  osT *= ("T", n, n+2, n+4)
end

os = osRand * osX * osSw * osZ * osCx * osT

@show os

s = siteinds("qubit", N)
gates = ops(s, os)

ψ0 = productMPS(s, "0")

# Apply the gates
ψ = apply(gates, ψ0; cutoff = 1e-15, maxdim = 100)
@show dim(s[1])^(N ÷ 2)
@show maxlinkdim(ψ)

prodψ = apply(gates, prod(ψ0))
@show prod(ψ) ≈ prodψ
@show norm(prod(ψ) - prodψ) / norm(prodψ)

return ψ

end

ψ = main();


text:
using Images, ImageFeatures, CoordinateTransformations

img1 = Gray.(load("cat-3417184_640.jpg"))
img2 = Gray.(load("cat-3417184_640_watermarked.jpg"))

# rotation img2 around the center
rot = recenter(RotMatrix(5pi/6), [size(img2)...] .÷ 2)  # a rotation around the center
tform = rot ∘ Translation(-50, -40)
img2 = warp(img2, tform, size(img2))

orb_params = ORB(num_keypoints = 1000)

desc_1, ret_keypoints_1 = create_descriptor(img1, orb_params)
desc_2, ret_keypoints_2 = create_descriptor(img2, orb_params)

matches = match_keypoints(ret_keypoints_1, ret_keypoints_2, desc_1, desc_2, 0.2)

grid = Gray.(hcat(img1, img2))
offset = CartesianIndex(0, size(img1, 2))
map(m -> draw!(grid, LineSegment(m[1], m[2] + offset)), matches)
grid

imshow(grid)


text:
immutable Pareto <: ContinuousUnivariateDistribution
    scale::Float64
    shape::Float64
    function Pareto(sc::Real, sh::Real)
        sc > zero(sc) && sh > zero(sh) || error("shape and scale must be positive")
        new(float64(sc), float64(sh))
    end
    Pareto() = new(1.0, 1.0)
end

Pareto(scale::Real) = Pareto(scale, 1.0)

cdf(d::Pareto, q::Real) = q >= d.scale ? 1.0 - (d.scale / q)^d.shape : 0.0

entropy(d::Pareto) = log(d.shape / d.scale) + 1. / d.scale + 1.

insupport(d::Pareto, x::Number) = isreal(x) && isfinite(x) && x > d.scale

function kurtosis(d::Pareto)
    a = d.shape
    a > 4.0 || error("Kurtosis undefined for Pareto w/ shape <= 4")
    (6.0 * (a^3 + a^2 - 6.0 * a - 2.0)) / (a * (a - 3.0) * (a - 4.0))
end

mean(d::Pareto) = d.shape <= 1.0 ? Inf : (d.scale * d.shape) / (d.scale - 1.0)

median(d::Pareto) = d.scale * 2.0^d.shape

mode(d::Pareto) = d.scale
modes(d::Pareto) = [d.scale]

function pdf(d::Pareto, q::Real)
    q >= d.scale ? (d.shape * d.scale^d.shape) / (q^(d.shape + 1.0)) : 0.0
end

quantile(d::Pareto, p::Real) = d.scale / (1.0 - p)^(1.0 / d.shape)

rand(d::Pareto) = d.shape / (rand()^(1.0 / d.scale))

function skewness(d::Pareto)
    a = d.shape
    a > 3.0 || error("Skewness undefined for Pareto w/ shape <= 3")
    ((2.0 * (1.0 + a)) / (a - 3.0)) * sqrt((a - 2.0) / a)
end

function var(d::Pareto)
    d.scale < 2.0 ? Inf : (d.shape^2 * d.scale) / ((d.scale - 1.0)^2 * (d.scale - 2.0))
end


text:
using DiffEqJump, DiffEqBase

rate = (u,p,t) -> p[1]*u[1]*u[2]
affect! = function (integrator)
  integrator.u[1] -= 1
  integrator.u[2] += 1
end
jump = ConstantRateJump(rate,affect!)

rate = (u,p,t) -> p[2]*u[2]
affect! = function (integrator)
  integrator.u[2] -= 1
  integrator.u[3] += 1
end
jump2 = ConstantRateJump(rate,affect!)

u0 = [999,1,0]
p = (.1/1000,.01)
tspan = (0.0,2500.0)

dprob  = DiscreteProblem(u0,tspan,p)
jprob = JumpProblem(dprob,Direct(),jump,jump2,save_positions=(false,false))
sol = solve(jprob, SSAStepper())
@test sol[1,end] == 0

u02 = [1000,1,0]
p2 = (.1/1000,0.0)
dprob2 = remake(dprob, u0=u02, p=p2)
jprob2 = remake(jprob, prob=dprob2)
sol2 = solve(jprob2, SSAStepper())
@test sol2[2,end] == 1001

tspan2 = (0.0, 25000.0)
jprob3 = remake(jprob, p=p2, tspan=tspan2)
sol3 = solve(jprob3, SSAStepper())
@test sol3[2,end] == 1000
@test sol3.t[end] == 25000.0

# test error handling
@test_throws ErrorException jprob4 = remake(jprob, prob=dprob2, p=p2)
@test_throws ErrorException jprob5 = remake(jprob, aggregator=RSSA())


text:
#### JuMPSubproblem Implementation ####

"""
Implementation of the [`AbstractSubproblem`](@ref) interface using JuMP.
"""
struct JuMPSubproblem <: AbstractSubproblem
    model::JuMP.Model
    scenario::ScenarioID
    stage_map::Dict{StageID, Vector{JuMP.VariableRef}}
    vars::Dict{VariableID, JuMP.VariableRef}
    w_vars::Dict{VariableID, JuMP.VariableRef}
    xhat_vars::Dict{VariableID, JuMP.VariableRef}
end

function JuMPSubproblem(m::JuMP.Model,
                        scid::ScenarioID,
                        stage_map::Dict{StageID, Vector{JuMP.VariableRef}}
                        )::JuMPSubproblem
    return JuMPSubproblem(m, scid, stage_map,
                          Dict{VariableID, JuMP.VariableRef}(),
                          Dict{VariableID, JuMP.VariableRef}(),
                          Dict{VariableID, JuMP.VariableRef}()
                          )
end

## JuMPSubproblem Internal Structs ##

struct JSVariable
    ref::JuMP.VariableRef
    name::String
    node_id::NodeID
end

## JuMPSubproblem Penalty Parameter Functions ##

function get_penalty_value(r::Float64,
                           vid::VariableID,
                           )::Float64
    return r
end

function get_penalty_value(r::Dict{VariableID,Float64},
                           vid::VariableID
                           )::Float64
    return r[vid]
end

## AbstractSubproblem Interface Functions ##

function add_lagrange_terms(js::JuMPSubproblem,
                            vids::Vector{VariableID}
                            )::Nothing

    obj = JuMP.objective_function(js.model,
                                  JuMP.GenericQuadExpr{Float64, JuMP.VariableRef}
                                  )

    jvi = JuMP.VariableInfo(false, NaN,   # lower_bound
                            false, NaN,   # upper_bound
                            true, 0.0,    # fixed
                            false, NaN,   # start value
                            false, false) # binary, integer

    for vid in vids
        var = js.vars[vid]
        w_ref = JuMP.add_variable(js.model, JuMP.build_variable(error, jvi))
        JuMP.add_to_expression!(obj, w_ref, var)
        js.w_vars[vid] = w_ref
    end

    JuMP.set_objective_function(js.model, obj)

    return
end

function add_ph_objective_terms(js::JuMPSubproblem,
                                vids::Vector{VariableID},
                                r::Union{Float64,Dict{VariableID,Float64}},
                                )::Nothing
    add_lagrange_terms(js, vids)
    add_proximal_terms(js, vids, r)
    return
end

function objective_value(js::JuMPSubproblem)::Float64
    return JuMP.objective_value(js.model)
end

function report_penalty_info(js::JuMPSubproblem,
                             vids::Vector{VariableID},
                             ::Type{ProportionalPenaltyParameter},
                             )::Dict{VariableID,Float64}
    return objective_coefficients(js, vids)
end

function report_penalty_info(js::JuMPSubproblem,
                             vids::Vector{VariableID},
                             ::Type{SEPPenaltyParameter},
                             )::Dict{VariableID,Float64}
    return objective_coefficients(js, vids)
end

function report_values(js::JuMPSubproblem,
                       vars::Vector{VariableID}
                       )::Dict{VariableID,Float64}
    val_dict = Dict{VariableID, Float64}()
    for vid in vars
        val_dict[vid] = JuMP.value(js.vars[vid])
    end
    return val_dict
end

function report_variable_info(js::JuMPSubproblem,
                              st::ScenarioTree
                              )::Dict{VariableID, VariableInfo}

    var_info = Dict{VariableID, VariableInfo}()

    for node in scenario_nodes(st, js.scenario)

        stid = stage(node)

        for (k, var) in enumerate(js.stage_map[stid])
            vid = VariableID(js.scenario, stid, Index(k))
            js.vars[vid] = var
            var_info[vid] = VariableInfo(JuMP.name(var),
                                         JuMP.is_integer(var) || JuMP.is_binary(var),
                                         )
        end

    end

    return var_info
end

function solve_subproblem(js::JuMPSubproblem)::MOI.TerminationStatusCode
    JuMP.optimize!(js.model)
    return JuMP.termination_status(js.model)
end

function update_lagrange_terms(js::JuMPSubproblem,
                               w_vals::Dict{VariableID,Float64}
                               )::Nothing

    for (wid, wval) in pairs(w_vals)
        JuMP.fix(js.w_vars[wid], wval, force=true)
    end

    return

end

function update_ph_terms(js::JuMPSubproblem,
                         w_vals::Dict{VariableID,Float64},
                         xhat_vals::Dict{VariableID,Float64}
                         )::Nothing

    update_lagrange_terms(js, w_vals)
    update_proximal_terms(js, xhat_vals)

    return

end

function warm_start(js::JuMPSubproblem)::Nothing
    for var in JuMP.all_variables(js.model)
        if !JuMP.is_fixed(var)
            JuMP.set_start_value(var, JuMP.value(var))
        end
    end
    return
end

function ef_copy_model(efm::JuMP.Model,
                       js::JuMPSubproblem,
                       scid::ScenarioID,
                       tree::ScenarioTree,
                       node_var_map::Dict{NodeID, Set{JSVariable}}
                       )::Dict{NodeID, Set{JSVariable}}

    (snode_var_map, s_var_map) = _ef_copy_variables(efm, js, scid, tree, node_var_map)
    processed = Set(keys(node_var_map))
    _ef_copy_constraints(efm, js, s_var_map, processed)
    _ef_copy_objective(efm, js, s_var_map, tree.prob_map[scid])

    return snode_var_map
end

function ef_node_dict_constructor(::Type{JuMPSubproblem})
    return Dict{NodeID, Set{JSVariable}}()
end

## JuMPSubproblem Specific Functions ##

function add_proximal_terms(js::JuMPSubproblem,
                            vids::Vector{VariableID},
                            r::Union{Float64,Dict{VariableID,Float64}}
                            )::Nothing

    obj = JuMP.objective_function(js.model,
                                  JuMP.GenericQuadExpr{Float64, JuMP.VariableRef}
                                  )

    jvi = JuMP.VariableInfo(false, NaN,   # lower_bound
                            false, NaN,   # upper_bound
                            true, 0.0,    # fixed
                            false, NaN,   # start value
                            false, false) # binary, integer

    for vid in vids
        var = js.vars[vid]
        xhat_ref = JuMP.add_variable(js.model, JuMP.build_variable(error, jvi))
        rho = get_penalty_value(r, vid)
        JuMP.add_to_expression!(obj, 0.5 * rho * (var - xhat_ref)^2)
        js.xhat_vars[vid] = xhat_ref
    end

    JuMP.set_objective_function(js.model, obj)

    return

end

function fix_variables(js::JuMPSubproblem,
                       vids::Dict{VariableID,Float64}
                       )::Nothing

    for (vid, val) in pairs(vids)
        JuMP.fix(js.vars[vid], val, force=true)
    end

    return

end

function objective_coefficients(js::JuMPSubproblem,
                                vids::Vector{VariableID},
                                )::Dict{VariableID,Float64}
    pen_dict = Dict{VariableID,Float64}()
    obj = JuMP.objective_function(js.model)
    for vid in vids
        pen_dict[vid] = JuMP.coefficient(obj, js.vars[vid])
    end
    return pen_dict
end

function update_proximal_terms(js::JuMPSubproblem,
                               xhat_vals::Dict{VariableID,Float64}
                               )::Nothing

    for (xhid, xhval) in pairs(xhat_vals)
        JuMP.fix(js.xhat_vars[xhid], xhval, force=true)
    end

    return

end

## JuMPSubproblem Internal Functions ##

function _build_var_info(vref::JuMP.VariableRef)
    hlb = JuMP.has_lower_bound(vref)
    hub = JuMP.has_upper_bound(vref)
    hf = JuMP.is_fixed(vref)
    ib = JuMP.is_binary(vref)
    ii = JuMP.is_integer(vref)

    return JuMP.VariableInfo(hlb,
                             hlb ? JuMP.lower_bound(vref) : 0,
                             hub,
                             hub ? JuMP.upper_bound(vref) : 0,
                             hf,
                             hf ? JuMP.fix_value(vref) : 0,
                             false, # Some solvers don't accept starting values
                             0,
                             ib,
                             ii)
end

function _ef_add_variables(model::JuMP.Model,
                           js::JuMPSubproblem,
                           s::ScenarioID,
                           node::ScenarioNode,
                           )

    smod = js.model

    var_map = Dict{JuMP.VariableRef, JSVariable}()
    new_vars = Set{JSVariable}()

    for vref in js.stage_map[node.stage]
        info = _build_var_info(vref)
        var = JuMP.name(vref)
        vname = var * "_{" * stringify(value.(scenario_bundle(node))) * "}"
        new_vref = JuMP.add_variable(model,
                                     JuMP.build_variable(error, info),
                                     vname)
        jsv = JSVariable(new_vref, vname, node.id)
        var_map[vref] = jsv
        push!(new_vars, jsv)
    end

    return (var_map, new_vars)
end

function _ef_map_variables(js::JuMPSubproblem,
                           node::ScenarioNode,
                           new_vars::Set{JSVariable},
                           )
    var_map = Dict{JuMP.VariableRef, JSVariable}()

    for vref in js.stage_map[stage(node)]

        var = JuMP.name(vref)

        for jsv in new_vars

            @assert jsv.node_id == node.id

            if occursin(var, jsv.name)
                var_map[vref] = jsv
            end

        end

    end

    return var_map
end

function _ef_copy_variables(model::JuMP.Model,
                            js::JuMPSubproblem,
                            s::ScenarioID,
                            tree::ScenarioTree,
                            node_var_map::Dict{NodeID, Set{JSVariable}},
                            )

    # Below for mapping variables in the subproblem model `smod` into variables for
    # the extensive form model `model`
    s_var_map = Dict{JuMP.VariableRef, JSVariable}()

    # For saving updates to node_var_map and passing back up
    snode_var_map = Dict{NodeID, Set{JSVariable}}()

    smod = js.model

    for node in scenario_nodes(tree, s)

        # For the given model `smod`, either create extensive variables corresponding
        # to this node or map them onto existing extensive variables.
        if !haskey(node_var_map, id(node))
            (var_map, new_vars) = _ef_add_variables(model, js, s, node)
            snode_var_map[node.id] = new_vars
        else
            var_map = _ef_map_variables(js,
                                        node,
                                        node_var_map[node.id])
        end

        @assert(isempty(intersect(keys(s_var_map), keys(var_map))))
        merge!(s_var_map, var_map)
    end

    return (snode_var_map, s_var_map)
end

function _ef_convert_and_add_expr(add_to::JuMP.QuadExpr,
                                  convert::JuMP.AffExpr,
                                  s_var_map::Dict{JuMP.VariableRef,JSVariable},
                                  scalar::Real,
                                  )::Set{NodeID}

    nodes = Set{NodeID}()

    JuMP.add_to_expression!(add_to, scalar * JuMP.constant(convert))

    for (coef, var) in JuMP.linear_terms(convert)
        vi = s_var_map[var]
        nvar = vi.ref
        JuMP.add_to_expression!(add_to, scalar*coef, nvar)

        push!(nodes, vi.node_id)
    end

    return nodes
end

function _ef_convert_and_add_expr(add_to::JuMP.QuadExpr,
                                  convert::JuMP.QuadExpr,
                                  s_var_map::Dict{JuMP.VariableRef,JSVariable},
                                  scalar::Real,
                                  )::Set{NodeID}

    nodes = _ef_convert_and_add_expr(add_to, convert.aff, s_var_map, scalar)

    for (coef, var1, var2) in JuMP.quad_terms(convert)
        vi1 = s_var_map[var1]
        vi2 = s_var_map[var2]

        nvar1 = vi1.ref
        nvar2 = vi2.ref
        JuMP.add_to_expression!(add_to, scalar*coef, nvar1, nvar2)

        push!(nodes, vi1.node_id)
        push!(nodes, vi2.node_id)
    end

    return nodes
end

function _ef_copy_constraints(model::JuMP.Model,
                              js::JuMPSubproblem,
                              s_var_map::Dict{JuMP.VariableRef,JSVariable},
                              processed::Set{NodeID},
                              )::Nothing

    smod = js.model
    constraint_list = JuMP.list_of_constraint_types(smod)

    for (func,set) in constraint_list

        if func == JuMP.VariableRef
            # These constraints are handled by the variable bounds
            # which are copied during copy variable creation so
            # we skip them
            continue
        end

        for cref in JuMP.all_constraints(smod, func, set)

            cobj = JuMP.constraint_object(cref)
            expr = zero(JuMP.QuadExpr)
            nodes = _ef_convert_and_add_expr(expr,
                                             JuMP.jump_function(cobj),
                                             s_var_map,
                                             1)

            # If all variables in the expression are from processed nodes,
            # then this constraint has already been added to the model
            # and can be skipped.
            if !issubset(nodes, processed)
                JuMP.drop_zeros!(expr)
                JuMP.@constraint(model, expr in JuMP.moi_set(cobj))
            end
        end
    end

    return
end

function _ef_copy_objective(model::JuMP.Model,
                            js::JuMPSubproblem,
                            s_var_map::Dict{JuMP.VariableRef,JSVariable},
                            prob::Real
                            )::Nothing

    add_obj = JuMP.objective_function(js.model)
    obj = JuMP.objective_function(model)
    _ef_convert_and_add_expr(obj, add_obj, s_var_map, prob)
    JuMP.drop_zeros!(obj)
    JuMP.set_objective_function(model, obj)

    return
end


text:
print_bool=false

include("../functions/MTiXsens_FillMissing.jl")

AccFree_MTiXsens=Dict()
Xsensor_reduce=Dict()

for experiment=1:size(filenameMTiXsens,1)
    #Extract FreeAcc from Xsens
    AccFree_MTiXsens_wthPacketCounter=Dict()
    for n in sensorXsensNames
        #extract useful column of MTiXsens datas and fill missing row datas
        AccFree_MTiXsens_wthPacketCounter[n] = MTiXsens_FillMissing(MTiXsens[filenameMTiXsens[experiment,1]][n][:,[1,5,6,7]])
    end
    AccFree_MTiXsens[filenameMTiXsens[experiment,1]]=AccFree_MTiXsens_wthPacketCounter
    #[filenameMTiXsens[experiment,1]]

    if print_bool
        println(string("Extract FreeAcc finished, experiment : "),experiment)
    end
end

for experiment=1:size(filenameXsensor,1)
    #Extract from Xsensor    Average Pressure	Peak Pressure	Contact Area (in²)	Est. Load (N)	COP Row	COP Column
    #siege first and dossier second
    ###WARNING###
    #Pay attention to column in xsensor experiment files
    #in this experiment there is a std dev column between Est. Load and COP
    #############
    #Xsensor_reduce[filenameXsensor[experiment,1]]=tryparse.(Float64,replace.(Xsensor[filenameXsensor[experiment,1]][:,[9,10,11,12,13,14,18,19,20,21,22,23]],[','=>'.']))
    if size(Xsensor[filenameXsensor[1,1]],2)>16
        if tryparse.(Float64,replace.(Xsensor[filenameXsensor[experiment,1]][1,12],[','=>'.']))[1] < 800
            Xsensor_reduce[filenameXsensor[experiment,1]]=tryparse.(Float64,replace.(Xsensor[filenameXsensor[experiment,1]][:,[19,20,21,22,24,25,9,10,11,12,14,15]],[','=>'.']))
        else
            Xsensor_reduce[filenameXsensor[experiment,1]]=tryparse.(Float64,replace.(Xsensor[filenameXsensor[experiment,1]][:,[9,10,11,12,14,15,19,20,21,22,24,25]],[','=>'.']))
        end
    else
        Xsensor_reduce[filenameXsensor[experiment,1]]=tryparse.(Float64,replace.(Xsensor[filenameXsensor[experiment,1]][:,[9,10,11,12,14,15]],[','=>'.']))

    end

    if print_bool
        println(string("Extract Xsensor finished, experiment : "),experiment)
    end

end

println("Extract data finished")


text:
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Usage:
julia --project=[...] hard_example.jl n delta alpha
"""

include("src/AGD_distance_restart.jl")

gr()
fntsm = Plots.font("serif", pointsize=12)
fntlg = Plots.font("serif", pointsize=18)
default(titlefont=fntlg, guidefont=fntlg, tickfont=fntsm, legendfont=fntsm)

n = parse(Int64, ARGS[1])
delta = parse(Float64, ARGS[2])
alpha = parse(Float64, ARGS[3])

function h_val(x::Float64)
  if x >= -delta
    return 0.5 * x^2
  else
    return -delta * x - 0.5 * delta^2
  end
end
function h_grad(x::Float64)
  if x >= -delta
    return x
  else
    return -delta
  end
end

vec = collect(1:n) .+ 0.0
function f_val(x::Vector{Float64})
  return sum(h_val.(x) .* vec) + 0.5 * alpha * sum(x.^2)
end
function f_grad(x::Vector{Float64})
  return h_grad.(x) .* vec + alpha * x
end

hard_function = differentiable_function(
  f_val,
  f_grad,
  n,
)

v_0 = -ones(n)
f_min_threshold = 1e-10
parameters = default_basic_agd_parameters(
  verbose=false,
  max_iter=40000,
  function_value_optimality_threshold=f_min_threshold,
);
basic_history, fval_restart_history, our_adaptive_history = run_algs(
  hard_function, v_0, parameters
);

optimality_threshold = 1e-8

BEST_FIXED_FREQUENCY = true
if BEST_FIXED_FREQUENCY
  best_fixed_frequency_result =
    pick_best_fixed_restart_frequency(hard_function, v_0, parameters;
      set_of_frequencies = [128, 256, 512, 1024, 2048, 4096, 8192],
      approximate_optimality_threshold = optimality_threshold)
  best_fixed_frequency_history = best_fixed_frequency_result.history
  print("best_fixed_frequency = ", best_fixed_frequency_result.frequency)
end

plot_algs(
  basic_history,
  fval_restart_history,
  our_adaptive_history,
  y_min=optimality_threshold,
)

if BEST_FIXED_FREQUENCY
  plot_history!(best_fixed_frequency_history;
      label="Best fixed ($(best_fixed_frequency_result.frequency))",
      color=3,
    )
end

println("")
println("Plot name: hard-example")
savefig(joinpath(@__DIR__, "plots/hard-example.pdf"))

# plot h_vals function for a larger delta value
delta = 1e-1
function h_val(x::Float64)
  if x >= -delta
    return 0.5 * x^2
  else
    return -delta * x - 0.5 * delta^2
  end
end

x_vals = collect(range(-1,stop=1,length=200))
y_vals = h_val.(x_vals)
plot(x_vals,y_vals, xlabel="x", ylabel="h(x)", legend=false)

savefig(joinpath(@__DIR__, "plots/h_func.pdf"))


text:
#==
Note this is a straight translation of the Lark
equivalent file and has not been tested or
simplified for Julia
==#

abstract type ParserPuppet end

struct NormalParserPuppet <: ParserPuppet
    parser
    parser_state
    lexer_state
end

# In lark, the immutable puppet keeps its
# result in a class variable so it can be
# hashed. I don't think we have this issue
# in Julia.
struct ImmutableParserPuppet <: ParserPuppet
    parser
    parser_state
    lexer_state
end

feed_token(p::ParserPuppet,token) = begin
    feed_token(p.parser_state,token, token.type_ == "\$END")
end

# TODO: proper copies

isequal(p1::ParserPuppet,p2) = false
isequal(p1::ParserPuppet,p2::ParserPuppet) = begin
    p1.parser_state == p2.parser_state && p1.lexer_state == p2.lexer_state
end

copy(p::ParserPuppet) = typeof(p)(p.parser,copy(p.parser_state),copy(p.lexer_state))

as_immutable(p::ParserPuppet) = ImmutableParserPuppet(p.parser,p.parser_state,p.lexer_state)

pretty(p::ParserPuppet) = begin
    out = ["Puppet choices:"]
    for (k,v) in p.choices()
        push!(out,"\t- $k -> $v")
    end
    push!(out,"stack size: $(length(p.parser_state.state_stack))")
    return join(out,"\n")
end

choices(p::ParserPuppet) = p.parser_state.parse_conf.parse_table.states[p.parser_state.position]

accepts(p::ParserPuppet) = begin
    accepts = Set()
    for t in keys(p.choices())
        if isupper(t)
            new_puppet = copy(p)
            try
                feed_token(new_puppet,Token(t,""))
            catch e
                if e isa UnexpectedToken
                    continue
                else
                    push!(accepts,t)
                end
            end
        end
    end
    return accepts
end

resume_parse(p::ParserPuppet) = begin
    p.parser.parse_from_state(p.parser_state)
end

# TODO: make the immutableparserpuppet work
feed_token(i::ImmutableParserPuppet) = begin
    c = copy(i)
end


text:
using Test
using ThreeBodyDecay

let m1sq = 1.0, m2sq = 2.0, iϵ = 1e-8im
  sv = (√m1sq+√m2sq)^2 .+ rand(10) # rand -> [0, 1]
  v1 = iRho.(sv .+ iϵ, 1.0, 2.0)
  v2 = 1im .* imag.(ChewMandestam.(sv .+ iϵ,1.0,2.0))
  @test max(abs.(v1 - v2)...) < 10*imag(iϵ)
end

max_stars_repo_name: mmikhasenko/ThreeBodyDecay

max_issues_repo_name: mmikhasenko/ThreeBodyDecay

max_forks_repo_name: mmikhasenko/ThreeBodyDecay


text:
import Plots
using PrettyTables
using JLD2
using IterativeSolvers
# export residuals,niterations
# using Formatting
# using Printf
Plots.pyplot()

residuals(hist) = hist.data[:resnorm]
residuals(x::Nothing) = nothing
niterations(x::Nothing) = nothing
residuals(x::Array) = x
niterations(x) = length(residuals(x))==0 ? nothing : length(residuals(x))

function plotbar(table,fn)
    plt=Plots.bar(Vector{String}(table[:,1]),Vector{Float64}(table[:,4]),title="Poisson Solver Duration (s)",
        xrotation=90,label="",size=(600,400),tickfont=Plots.font(12),guidefont=Plots.font(12))
    Plots.display(plt)
    Plots.savefig(plt,fn)
end

function plotallconvergence(itable,fn)
    plt=nothing

    for i=1:size(itable,1)
        resi=residuals(itable[i,7])
        ls=:solid
        if itable[i,5]>1.e-6
            # @show "nc",itable[i,1],itable[i,5]
            ls=:dot
        else
            # @show "co",itable[i,1],itable[i,5]
            ls=:solid
        end
        if i==1
            plt=Plots.plot(resi,yaxis=:log,label=itable[i,1],
                legend=:outertopright,  linewidth = 3,linestyle=ls,
                ylabel="Residual norm",xlabel="Iteration #")
        else
            plt=Plots.plot!(plt,resi,yaxis=:log,label=itable[i,1],
                legend=:outertopright,  linewidth = 3,linestyle=ls,
                size=(600,400),tickfont=Plots.font(12),guidefont=Plots.font(12),legendfont=Plots.font(11))
        end
    end
    Plots.display(plt)
    Plots.savefig(plt,fn)
end

function plotfast(itable,fn)
    markers = filter((m->begin
            m in Plots.supported_markers()
        end), Plots._shape_keys)
    markers = reshape(markers, 1, length(markers))
    # @show markers
    plt2=nothing
    for i=1:size(itable,1)
        resi=residuals(itable[i,7])
        ls=:solid
        if itable[i,6]<100
            # @show itable[i,1],itable[i,6]
            if plt2==nothing
                plt2=Plots.plot(resi,yaxis=:log,label=itable[i,1],
                    linewidth = 3,linestyle=ls,
                    ylabel="Residual norm",xlabel="Iteration #", reuse=false,
                    marker=markers[i],markersize=6)
            else
                plt2=Plots.plot!(plt2,resi,yaxis=:log,label=itable[i,1],
                    linewidth = 3,linestyle=ls,
                    size=(600,400),tickfont=Plots.font(12),guidefont=Plots.font(12),legendfont=Plots.font(11),
                    marker=markers[i],markersize=6)
            end
        end
    end
    display(plt2)
    Plots.savefig(plt2,fn)
end

function plotconvergence(itable,fn)
    markers = filter((m->begin
            m in Plots.supported_markers()
        end), Plots._shape_keys)
    markers = reshape(markers, 1, length(markers))
    # @show markers
    plt2=nothing
    for i=1:size(itable,1)
        resi=residuals(itable[i,7])
        # @show label=itable[i,1]
        ls=:solid
        # @show itable[i,1],itable[i,6],resi
        if plt2==nothing
            # @show itable[i,1],itable[i,6]
            plt2=Plots.plot(resi,yaxis=:log,label=itable[i,1],
                linewidth = 3,linestyle=ls,
                ylabel="Residual norm",xlabel="Iteration #",
                size=(600,400),tickfont=Plots.font(12),guidefont=Plots.font(12),legendfont=Plots.font(11),
                marker=markers[i],markersize=6)
        else
            # @show itable[i,1],itable[i,6]
            plt2=Plots.plot!(plt2,resi,yaxis=:log,label=itable[i,1],
                linewidth = 3,linestyle=ls,
                ylabel="Residual norm",xlabel="Iteration #",
                size=(600,400),tickfont=Plots.font(12),guidefont=Plots.font(12),legendfont=Plots.font(11),
                marker=markers[i],markersize=6)
        end
    end
    display(plt2)
    Plots.savefig(plt2,fn)
end

function measureplot(msm,sname)
    header=keys(msm[1])[1:end]
    nrows=length(msm)
    ncols=length(header)
    table=Matrix{Any}(undef,nrows,ncols)
    iterations=Vector{Union{Nothing,Vector{Float64}}}(undef,nrows)
    for i=1:nrows
        for j=1:ncols
            table[i,j]=msm[i][header[j]]
        end
        table[i]=replace(table[i],"Poisson"=>"")
        table[i]=replace(table[i],"Preconditioner"=>"")
        table[i]=replace(table[i],"SparseCG"=>"PCG")
    end
    table[7,1]="PCGAMG{SmoothA}"
    table[8,1]="PCGAMG{RugeStuben}"

    #Remove some measurements :
    removed=["GMG_CGSmoother"]
    table=table[filter(i->!(table[i,1] in removed),1:size(table,1)),:]

    #add Total Time and niter
    header=[header[1:3]...,"Total Time",header[4],"niters",header[5]]
    table=hcat(table[:,1:3],table[:,2].+table[:,3],table[:,4],niterations.(residuals(table[:,5])),table[:,5])
    #Sort
    sp=sortperm(table[:, 4])
    table=table[sp,:]
    # iterations=iterations[sp,:]
    nr,nc=size(table)
    println("all methods")
    pretty_table(table[:,1:6],collect(header[1:6]);formatters=ft_printf("%5.3E",2:5),alignment=:l)
    println("converged methods")
    ctable=table[filter(i->(!isnan(table[i,5]) && table[i,5]<1.e-6),1:size(table,1)),:]
    pretty_table(ctable[:,1:6],collect(header[1:6]);formatters=ft_printf("%5.3E",2:5),alignment=:l)

    println("non converged methods")
    nctable=table[filter(i->(isnan(table[i,5]) ||  table[i,5]>1.e-6),1:size(table,1)),:]
    pretty_table(nctable[:,1:6],collect(header[1:6]);formatters=ft_printf("%5.3E",2:5),alignment=:l)
    println("converged iterative methods")
    ictable=ctable[filter(i->ctable[i,6]!==nothing,1:size(ctable,1)),:]
    pretty_table(ictable[:,1:6],collect(header[1:6]);formatters=ft_printf("%5.3E",2:5),alignment=:l)
    println("efficient converged iterative methods")
    iectable=ictable[filter(i->ictable[i,4]<10.0,1:size(ictable,1)),:]
    pretty_table(ictable[:,1:6],collect(header[1:6]);formatters=ft_printf("%5.3E",2:5),alignment=:l)
    println("iterative methods")
    itable=table[filter(i->table[i,6]!==nothing,1:size(table,1)),:]
    pretty_table(itable[:,1:6],collect(header[1:6]);formatters=ft_printf("%5.3E",2:5),alignment=:l)
    println("fast methods")
    ftable=ictable[filter(i->(ictable[i,6]<25 && ictable[i,4]<10.0),1:size(ictable,1)),:]
    pretty_table(ftable[:,1:6],collect(header[1:6]);formatters=ft_printf("%5.3E",2:5),alignment=:l)

    sp=sortperm(itable[:, 6]+0.0001*itable[:,5])
    itable=itable[reverse(sp),:]
    pretty_table(itable[:,1:6],collect(header[1:6]);formatters=ft_printf("%5.3E",2:5),alignment=:l)

    cftable=table[filter(i->(!isnan(table[i,5]) && table[i,5]<1.e-6 && table[i,4]<10.0),1:size(table,1)),:]

    plotbar(cftable,"time_converge_"*sname*"_all.svg")
    plotbar(iectable,"time_converge_"*sname*"_mixed.svg")
    plotbar(ftable,"time_converge_"*sname*"_fast.svg")

    plotallconvergence(itable,"allconvergence_"*sname*".svg")
    plotfast(itable,"fastconvergence_"*sname*".svg")
    plotconvergence(ftable,"convergence_"*sname*".svg")

    table,header
end

function test(sname)
    fnjld2="msm"*sname*".jld2"
    @load fnjld2 msm
    t,h=measureplot(msm,sname)
    nothing
end