export PseudoPeriodicTwin

import TimeseriesSurrogates: Surrogate, SurrogateGenerator, surrogenerator
import StatsBase: sample, pweights
import Distances: Euclidean, PreMetric, pairwise
import DelayEmbeddings: Dataset, embed

"""
    PseudoPeriodicTwin(d::Int, τ::Int, δ = 0.2, ρ = 0.1, metric = Euclidean())
    PseudoPeriodicTwin(δ = 0.2, ρ = 0.1, metric = Euclidean())

A pseudoperiodic twin surrogate[^Miralles2015], which is a fusion of the
twin surrogate[^Thiel2006] and the pseudo-periodic surrogate[^Small2001].

## Input parameters

A delay reconstruction of the input timeseries is constructed using embedding
dimension `d` and embedding delay `τ`. The threshold `δ ∈ (0, 1]` determines which
points are "close" (neighbors) or not, and is expressed as a fraction of the
attractor diameter, as determined by the input data. The authors of the original
twin surrogate paper recommend `0.05 ≤ δ ≤ 0.2`[^Thiel2006].

If you have pre-embedded your timeseries, and timeseries is already a `::Dataset`, use the
three-argument constructor (so that no delay reconstruction is performed).
If you want a surrogate for a scalar-valued timeseries, use the five-argument constructor
to also provide the embedding delay `τ` and embedding dimension `d`.

## Null hypothesis

Pseudo-periodic twin surrogates generate signals similar to the original data if the
original signal is (quasi-)periodic. If the original signal is not
(quasi-)periodic, then these surrogates will have different
recurrence plots than the original signal, but preserve the overall
shape of the attractor. Thus, `PseudoPeriodicTwin` surrogates can be used to test
null hypothesis that the observed timeseries (or orbit) is consistent
with a quasi-periodic orbit[^Miralles2015].

## Returns

A `d`-dimensional surrogate orbit (a `Dataset`) is returned. Sample the first
column of this dataset if a scalar-valued surrogate is desired.

[^Small2001]: Small et al., Surrogate test for pseudoperiodic timeseries data, [Physical Review Letters, 87(18)](https://doi.org/10.1103/PhysRevLett.87.188101)
[^Thiel2006]: Thiel, Marco, et al. "Twin surrogates to test for complex synchronisation." EPL (Europhysics Letters) 75.4 (2006): 535.
[^Miralles2015]: Miralles, R., et al. "Characterization of the complexity in short oscillating timeseries: An application to seismic airgun detonations." The Journal of the Acoustical Society of America 138.3 (2015): 1595-1603.
"""
struct PseudoPeriodicTwin{T<:Real, P<:Real, D<:PreMetric} <: Surrogate
    d::Union{Nothing, Int}
    τ::Union{Nothing, Int}
    δ::T
    ρ::P
    metric::D

    function PseudoPeriodicTwin(d::Int, τ::Int, δ::T = 0.2, ρ::P = 0.1, metric::D = Euclidean()) where {T, P, D}
        new{T, P, D}(d, τ, δ, ρ, metric)
    end

    function PseudoPeriodicTwin(δ::T = 0.2, ρ::P = 0.1, metric::D = Euclidean()) where {T, P, D}
        new{T, P, D}(nothing, nothing, δ, ρ, metric)
    end
end

"""
    _prepare_embed(x::AbstractVector, d, τ) → Dataset
    _prepare_embed(x::Dataset, d, τ) → Dataset

Prepate input data for surrogate generation. If input is a vector, embed it using
the provided parameters. If input as a dataset, we assume it already represents an
orbit.
"""
function _prepare_embed end
_prepare_embed(x::AbstractVector, d, τ) = embed(x, d, τ)
_prepare_embed(x::Dataset, d, τ) = x

function surrogenerator(x::Union{AbstractVector, Dataset}, pp::PseudoPeriodicTwin, rng = Random.default_rng())
    d, τ, δ, metric = getfield.(Ref(pp), (:d, :τ, :δ, :metric))
    ρ = getfield.(Ref(pp), (:ρ))

    pts = _prepare_embed(x, d, τ)

    Nx = length(x)
    Npts = length(pts)

    dists = pairwise(metric, Matrix(pts), dims = 1)
    normalisedδ = δ*maximum(dists)

    T = eltype(pts)
    R = zeros(T, Npts, Npts)

    # Recurrence matrix
    for j = 1:Npts
        for i = 1:Npts
            R[i, j] = normalisedδ - dists[i, j] >= 0 ? 1.0 : 0.0
        end
    end

    # Identify twins
    #println("R contains $(count(R .== 1)/length(R)*100)% black dots")
    twins_i = Vector{Int}(undef, 0)
    twins_j = Vector{Int}(undef, 0)
    for j = 1:Npts
        for i = 1:Npts
            if i !== j && all(R[:, i] .≈ R[:, j])
                push!(twins_i, i)
                push!(twins_j, j)
            end
        end
    end
    #println("Found $(length(twins)) twins")

    twins = Dict{Int,Vector{Int}}()

    # For every point that has a twin, store the indices of all of its twins
    for twi in unique(twins_i)
        twins[twi] = twins_j[findall(twins_i .== twi)]
    end
    for twj in unique(twins_j)
        twins[twj] = twins_i[findall(twins_j .== twj)]
    end

    # Sampling weights (exclude the point itself)
    W = [pweights(exp.(-dists[setdiff(1:Npts, i), i] / ρ)) for i = 1:Npts]

    # The surrogate will be a vector of vectors (if pts is a Dataset, then
    # the eltype is SVector).
    PT = eltype(pts.data)
    s = Vector{PT}(undef, Nx)

    init = (pts = pts, Nx = Nx, Npts = Npts, dists = dists, R = R, twins = twins, W = W)
    return SurrogateGenerator(pp, x, s, init, rng)
end

function (sg::SurrogateGenerator{<:PseudoPeriodicTwin})()
    pts, Nx, Npts, twins, W = getfield.(Ref(sg.init), (:pts, :Nx, :Npts, :twins, :W))
    ρ = getfield.(Ref(sg.method), (:ρ))
    s = sg.s

    # Randomly pick a point from the state space as the starting point for the surrogate.
    n = 1
    i = rand(sg.rng, 1:Npts)
    s[n] = pts[i]

    while n < Nx
        # Look for possible twins of the point xᵢ. If any twin exists, jump to one of them
        # with probability 1/nⱼ, where nⱼ are the number of twins for the point xᵢ.
        if haskey(twins, i)
            # sample uniformly (with probability 1/ntargettwins) over possible target twins
            j = twins[i][rand(sg.rng, 1:length(twins[i]))]
            s[n] = pts[j]
            i = j
            n += 1
        end

        # The orbit moves on from the current point xᵢ to a randomly selected point
        # on the attractor. Closer points are more likely to be selected, and points
        # further away are less likely to be selected. The sampling probability
        # of the next point xⱼ decreases exponentially with increasing distance
        # from the current point xᵢ. Probabilities for jumping from xᵢ to any other
        # point have been pre-computed, and are stored in W[i].
        j = sample(sg.rng, 1:Npts, W[i])
        s[n] = pts[j]
        i = j
        n += 1
    end

    s[Nx] = pts[sample(1:Npts, W[i])[1]]

    return Dataset(s)
end


using TensorCore
using LinearAlgebra
using Test

@testset "Ambiguities" begin
    @test isempty(detect_ambiguities(TensorCore, Base, Core, LinearAlgebra))
end

@testset "TensorCore.jl" begin
    for T in (Int, Float32, Float64, BigFloat)
        a = [T[1, 2], T[-3, 7]]
        b = [T[5, 11], T[-13, 17]]
        @test map(⋅, a, b) == map(dot, a, b) == [27, 158]
        @test map(⊙, a, b) == map(hadamard, a, b) == [a[1].*b[1], a[2].*b[2]]
        @test map(⊗, a, b) == map(tensor, a, b) == [a[1]*transpose(b[1]), a[2]*transpose(b[2])]
        @test hadamard!(fill(typemax(Int), 2), T[1, 2], T[-3, 7]) == [-3, 14]
        @test tensor!(fill(typemax(Int), 2, 2), T[1, 2], T[-3, 7]) == [-3 7; -6 14]
    end

    @test_throws DimensionMismatch [1,2] ⊙ [3]
    @test_throws DimensionMismatch hadamard!([0, 0, 0], [1,2], [-3,7])
    @test_throws DimensionMismatch hadamard!([0, 0], [1,2], [-3])
    @test_throws DimensionMismatch hadamard!([0, 0], [1], [-3,7])
    @test_throws DimensionMismatch tensor!(Matrix{Int}(undef, 2, 2), [1], [-3,7])
    @test_throws DimensionMismatch tensor!(Matrix{Int}(undef, 2, 2), [1,2], [-3])

    u, v = [2+2im, 3+5im], [1-3im, 7+3im]
    @test u ⋅ v == conj(u[1])*v[1] + conj(u[2])*v[2]
    @test u ⊙ v == [u[1]*v[1], u[2]*v[2]]
    @test u ⊗ v == [u[1]*v[1] u[1]*v[2]; u[2]*v[1] u[2]*v[2]]
    @test hadamard(u, v) == u ⊙ v
    @test tensor(u, v)   == u ⊗ v
    dest = similar(u)
    @test hadamard!(dest, u, v) == u ⊙ v
    dest = Matrix{Complex{Int}}(undef, 2, 2)
    @test tensor!(dest, u, v) == u ⊗ v

    for (A, B, b) in (([1 2; 3 4], [5 6; 7 8], [5,6]),
                      ([1+0.8im 2+0.7im; 3+0.6im 4+0.5im],
                       [5+0.4im 6+0.3im; 7+0.2im 8+0.1im],
                       [5+0.6im,6+0.3im]))
        @test A ⊗ b == cat(A*b[1], A*b[2]; dims=3)
        @test A ⊗ B == cat(cat(A*B[1,1], A*B[2,1]; dims=3),
                           cat(A*B[1,2], A*B[2,2]; dims=3); dims=4)
    end

    A, B = reshape(1:27, 3, 3, 3), reshape(1:4, 2, 2)
    @test A ⊗ B == [a*b for a in A, b in B]

    # Adjoint/transpose is a dual vector, not an AbstractMatrix
    v = [1,2]
    @test_throws ErrorException v ⊗ v'
    @test_throws ErrorException v ⊗ transpose(v)
    @test_throws ErrorException v' ⊗ v
    @test_throws ErrorException transpose(v) ⊗ v
    @test_throws ErrorException v' ⊗ v'
    @test_throws ErrorException transpose(v) ⊗ transpose(v)
    @test_throws ErrorException v' ⊗ transpose(v)
    @test_throws ErrorException transpose(v) ⊗ v'
    @test_throws ErrorException A ⊗ v'
    @test_throws ErrorException A ⊗ transpose(v)

    # Docs comparison to `kron`
    v, w = [1,2,3], [5,7]
    @test kron(v,w) == vec(w ⊗ v)
    @test w ⊗ v == reshape(kron(v,w), (length(w), length(v)))
end


### A Pluto.jl notebook ###
# v0.12.11

using Markdown
using InteractiveUtils

# ╔═╡ 61dc86f0-238f-11eb-054f-41a45b78e94d
begin
	using  MLJ
	using  MLJLinearModels
	using  CSV
	import RDatasets: dataset
    import DataFrames: describe, select, Not, rename!,head,DataFrame, nename!
	import Plots: plot,plot!,scatter,scatter!,histogram,histogram!
end

# ╔═╡ ff5bd480-29c0-11eb-3112-cb35c17acc60
md" ### Libs"

# ╔═╡ a2efe1c0-29c8-11eb-287d-49b90e597876
md"### Sobre o Problema

url: https://archive.ics.uci.edu/ml/datasets/banknote+authentication

Data Set Information:

Data were extracted from images that were taken from genuine and forged banknote-like specimens. For digitization, an industrial camera usually used for print inspection was used. The final images have 400x 400 pixels. Due to the object lens and distance to the investigated object gray-scale pictures with a resolution of about 660 dpi were gained. Wavelet Transform tool were used to extract features from images.

Attribute Information:

1. variance of Wavelet Transformed image (continuous)
2. skewness of Wavelet Transformed image (continuous)
3. curtosis of Wavelet Transformed image (continuous)
4. entropy of image (continuous)
5. class (integer)
"

# ╔═╡ 8e226310-29c0-11eb-1093-f95e53c029b6
md"### Carga de Dados"

# ╔═╡ 114c91a0-2a90-11eb-365b-4d360bb51ebe
md" É interessante sempre mostrarmos os primeiros exemplos dos dados para ter uma ideia dos valores. Entretanto, como visto no curso, há várias análises possíveis para se fazer como identificar nulos, outliers, ... Tudo isto, é bom ser tratado."

# ╔═╡ 1c623cf0-2d17-11eb-2f20-ad2b49b221de
begin
	url = "https://archive.ics.uci.edu/ml/machine-learning-databases/00267/data_banknote_authentication.txt"
	tmp = download(url)
end

# ╔═╡ 049db9d0-2d19-11eb-24f1-0ba111942cb8
begin
	banknote = DataFrame(CSV.File(tmp; header=false))
	rename!(banknote, Dict(:Column1 => "Variance", :Column2 => "Skewness",:Column3 => "Curtosis",:Column4 => "Entropy",:Column5 => "Target"))
end

# ╔═╡ 9a810b20-29c0-11eb-1ca3-d155108c6e52
md"### Estatísticas dos Dados"

# ╔═╡ 91187080-2a8f-11eb-3001-1973d885e3f0
md"É sempre bom tiramos umas estatísticas dos dados para entender melhor o comportamento das variáveis do nosso problema. Qual a variância de cada variável? Elas estão na mesma ordem de grandeza? Máixmo e minímo ?"

# ╔═╡ 2033e842-294f-11eb-2b38-1969d7440b88
describe(banknote, :mean, :std, :min, :max, :eltype)

# ╔═╡ 2f06c4a0-294f-11eb-2cb1-2da452f98b0f
begin
	data = coerce(banknote, autotype(banknote, :discrete_to_continuous));
	data = coerce(data,:Target => OrderedFactor)#Multiclass)
end

# ╔═╡ d2a9dfe0-29c0-11eb-36a9-e7e3c1d17717
md"### Selecionando Dados"

# ╔═╡ 49b99b50-2a90-11eb-2b54-2704b914334d
md"É nesta fase que você vai escolher seu target e vai selecionar os atributos para trabalhar."

# ╔═╡ 47a80140-294f-11eb-2880-ef7cfb0a8997
begin
	y = data.Target;
    X = select(data, Not(:Target));
end

# ╔═╡ 429fc960-2d25-11eb-3c06-614daa2fc78f
schema(X)

# ╔═╡ e3d83ea0-29a3-11eb-10f5-b5f94102e1f1
md"### Verificando Algoritmos Disponíveis"

# ╔═╡ 60083560-2a90-11eb-373f-3b8a1ba538a0
md"Essa é uma característica bem legal do MLJ. Como nosso dados possuem atributos numéricos e o target numérico, ele me mostra todos os algoritmos que conseguem rodar para este tipo de dado."

# ╔═╡ 868b73b0-29a9-11eb-0ef8-971446a3c919
#models(matching(X, y))
begin
	task(model) = matching(model, X, y) && model.prediction_type == :probabilistic
	models(task)
end

# ╔═╡ 87cedb30-2a90-11eb-26dd-6162d49edd88
md"Abaixo realizamos uma carga no nosso regressor linear."

# ╔═╡ 9abff260-2a90-11eb-026f-59a85ec6b64c
md"O MLJ nos permite montar um pipeline. Pipeline é uma forma de montar uma sequência de operações poupando-lhe código. No nosso caso, queremos padronizar X e y antes de realizarmos o treino do regressor. Podemos fazer isso em uma única linha."

# ╔═╡ fcde2920-29a6-11eb-37f7-4befad3ee3d2
predictor = @pipeline  Standardizer LogisticClassifier(lambda = 0.5)

# ╔═╡ 323d2620-29c0-11eb-258a-4703b4f5a0d2
md"### Partição dos Dados"

# ╔═╡ e57688a2-2a90-11eb-2272-3d9898bf4375
md"Realizamos aqui o que conhecemos como hold-out. Particionamos em treino e teste, que é a forma mais simples de se fazer uma validação. Sabemos que existem outras no curso. No caso, deixei 70% para treino e 30% para teste. Como podem notar, a função gera índices qu correspondem às linhas dos exemplos. Ele não tira cópia dos dados."

# ╔═╡ fca39122-29a6-11eb-3b25-c7f0df7fc68f
train, test  = partition(eachindex(y), 0.7, shuffle=true);

# ╔═╡ 287ae7e0-2a91-11eb-0d5b-b138443bdc29
md"Um conceito interessante do MLJ é o machine. Esta função faz uma mapeamento do preditor ao dado. Quando aplicamos machine, estamos gerando um nó num grafo de computação e um estado associado. Existem algmas vantagens neste paradigma como: associar um mesmo preditor a vários dados distintos que podem ser processados em diferentes recursos; evitar re execução de um mesmo preditor e um mesmo dado."

# ╔═╡ fc474050-29a6-11eb-1fbd-87a88d38525a
pred_machine = machine(predictor, X, y);

# ╔═╡ 4449f8d0-29b5-11eb-3f83-c352d034ecb2
md"### Treino do Classificador"

# ╔═╡ 91918630-2a91-11eb-1113-43b18086f208
md" Aqui temos enfim, o treino do nosso regressor para o data X, apenas no conjunto de treino. Isto porque estamos passando os índices do exemplos."

# ╔═╡ 03caa2a0-2d2a-11eb-0742-413f150b9056
fit!(pred_machine, rows=train);

# ╔═╡ 50f36580-29b5-11eb-06f2-cf4b9e142979
md"### Análise da Predição"

# ╔═╡ ef445730-2a91-11eb-1552-c3e7d6e5ff34
md"Aqui analisamos cada predição e cada target. Este tipo de análise é mais apropriada quando nosso conjunto de dados possui uma relação de ordem temporal entre os exemplos (que não é o caso). Entretanto, é uma análise boa a se fazer. Podemos notar que há um atendência de subestimar os valores mais altos do que valores mais baixos. Você pode analisar o que pode estar gerando dificldade nisto analisando alguns exemplos individualmente."

# ╔═╡ 36bd3310-2dbd-11eb-0e46-bf6c0b1e80a9
md"Aqui pegamos as classes e probabilidades"

# ╔═╡ 152f7100-29a7-11eb-2300-735ca72b15c3
preds = predict(pred_machine, rows=test)

# ╔═╡ 275bab42-2dbd-11eb-1bd7-a1c0cb2965a7
md"Pegando probabilidades da predição por classe"

# ╔═╡ 1336f910-2d9c-11eb-222b-ade11f5b7b51
begin
	probs_fraude  = broadcast(pdf,preds,1.0);
	probs_nfraude = broadcast(pdf,preds,0.0);
end

# ╔═╡ d4536cb0-2d9b-11eb-16f7-993d0065c683
#mode.(ŷ)
ŷ = predict_mode(pred_machine,rows = test)

# ╔═╡ 35b079a0-2d9f-11eb-0fa6-1bb075367008
measures(matching(y))

# ╔═╡ dbb70eb2-2dc0-11eb-3da4-754b2e0de392
md"#### Métricas para o limiar padrão (0.5)"

# ╔═╡ 32c66e80-2d9e-11eb-220e-093e86fdc72f
cm = confusion_matrix(ŷ,y[test])

# ╔═╡ 3593e090-2e05-11eb-0557-13c977ff89b2
false_positive_rate(cm),false_negative_rate(cm)

# ╔═╡ 463a4470-2e05-11eb-0cac-2f56d62bfa20
true_positive_rate(cm),true_negative_rate(cm) # recall

# ╔═╡ 1fedd0a0-2e07-11eb-1f26-8575ef999663
positive_predictive_value(cm),negative_predictive_value(cm) # precision

# ╔═╡ 6a6e0550-2d9e-11eb-1136-119a2997dcae
accuracy(cm)

# ╔═╡ 0af080a0-2e05-11eb-0069-8d908435f94b
balanced_accuracy(ŷ,y[test])

# ╔═╡ c8288560-2dc8-11eb-14f6-19ae67fb2357
begin
	f1 = FScore{1.0}()
	f1(cm)
end

# ╔═╡ 00989c50-2dc9-11eb-31fc-0d8ab39b5a92
matthews_correlation(cm)

# ╔═╡ 77646ed0-2dc5-11eb-35c5-4b1582878184
md"#### Métricas para limiares distintos"

# ╔═╡ f3292f10-2d9d-11eb-1e76-8de80adce2e6
tprs, fprs, ts = roc_curve(preds,y[test]);

# ╔═╡ 96c67790-2e07-11eb-0000-19536e55469e
ts[180]

# ╔═╡ 903e2090-2dc5-11eb-2855-bfc16c6145b6
md" Área embaixo da curva"

# ╔═╡ 58f2c540-2d9e-11eb-3a14-0931bcd2cc5d
auc(preds,y[test])

# ╔═╡ c29c1380-29b5-11eb-254c-2d026e15f903
plot(tprs,fprs,label="LG Predictions",ylabel="Verdadeiros Positivos",xlabel="Falsos Positivos")

# ╔═╡ 89d06380-29bf-11eb-0d5b-4533c94745d5
md"#### Distribuição das probabilidades das predições"

# ╔═╡ 8b3ac8f0-29bf-11eb-274e-f7b07eb35e7c
histogram(probs_fraude,bins=10,label="Fraude")

# ╔═╡ 1e35c150-29de-11eb-0069-cb75699b6d61
histogram(probs_nfraude,bins=10,label="Não Fraude")

# ╔═╡ 76b4c570-29b5-11eb-32d5-75da0001a894
md"### Parâmetros encontrados"

# ╔═╡ 46a9fc40-2a8e-11eb-301f-077f9a6d50fe
md" Muitas vezes é interessante analisarmos a relevância dos coeficientes por váriosmotivos. Isto porque possibilita descobrir quais são as variáveis que ajudam mais num problema, ajudando inclusive na tomada de decisão de um ser humano. Adicionalmente, pode ser usado para compactar o modelo (ambientes com poucos recursos), eliminando atributos pouco relevantes. Finalmente, pode-se utilizar para verificar se seu modelo usa alguma variável que seja discrimnatória em termos de sociedade (ex. idade para tomada de crédito). "

# ╔═╡ 8d50cc90-294f-11eb-012b-3b973cbd1f1e
begin
	fp = fitted_params(pred_machine)
	coefs = fp.fitted_params_given_machine[fp.machines[2]]
	#fp.machines
end

# ╔═╡ e102cfc0-295c-11eb-2c5b-834000ca2077
#begin
#	res = ŷ .- y[test]
#	plot(res)
#	scatter!(res)
#end

# ╔═╡ 08b31090-2e01-11eb-27a7-c3758492386e
begin
	names = [string(e[1]) for e in coefs[2] ]
	vals  = [e[2] for e in coefs[2] ]
	plot(names,vals,linetype = :bar,label="Coefs")
end

# ╔═╡ Cell order:
# ╟─ff5bd480-29c0-11eb-3112-cb35c17acc60
# ╠═61dc86f0-238f-11eb-054f-41a45b78e94d
# ╟─a2efe1c0-29c8-11eb-287d-49b90e597876
# ╟─8e226310-29c0-11eb-1093-f95e53c029b6
# ╟─114c91a0-2a90-11eb-365b-4d360bb51ebe
# ╠═1c623cf0-2d17-11eb-2f20-ad2b49b221de
# ╠═049db9d0-2d19-11eb-24f1-0ba111942cb8
# ╟─9a810b20-29c0-11eb-1ca3-d155108c6e52
# ╟─91187080-2a8f-11eb-3001-1973d885e3f0
# ╠═2033e842-294f-11eb-2b38-1969d7440b88
# ╠═2f06c4a0-294f-11eb-2cb1-2da452f98b0f
# ╟─d2a9dfe0-29c0-11eb-36a9-e7e3c1d17717
# ╟─49b99b50-2a90-11eb-2b54-2704b914334d
# ╠═47a80140-294f-11eb-2880-ef7cfb0a8997
# ╠═429fc960-2d25-11eb-3c06-614daa2fc78f
# ╟─e3d83ea0-29a3-11eb-10f5-b5f94102e1f1
# ╟─60083560-2a90-11eb-373f-3b8a1ba538a0
# ╠═868b73b0-29a9-11eb-0ef8-971446a3c919
# ╟─87cedb30-2a90-11eb-26dd-6162d49edd88
# ╟─9abff260-2a90-11eb-026f-59a85ec6b64c
# ╠═fcde2920-29a6-11eb-37f7-4befad3ee3d2
# ╟─323d2620-29c0-11eb-258a-4703b4f5a0d2
# ╟─e57688a2-2a90-11eb-2272-3d9898bf4375
# ╠═fca39122-29a6-11eb-3b25-c7f0df7fc68f
# ╠═287ae7e0-2a91-11eb-0d5b-b138443bdc29
# ╠═fc474050-29a6-11eb-1fbd-87a88d38525a
# ╟─4449f8d0-29b5-11eb-3f83-c352d034ecb2
# ╟─91918630-2a91-11eb-1113-43b18086f208
# ╠═03caa2a0-2d2a-11eb-0742-413f150b9056
# ╟─50f36580-29b5-11eb-06f2-cf4b9e142979
# ╟─ef445730-2a91-11eb-1552-c3e7d6e5ff34
# ╟─36bd3310-2dbd-11eb-0e46-bf6c0b1e80a9
# ╠═152f7100-29a7-11eb-2300-735ca72b15c3
# ╟─275bab42-2dbd-11eb-1bd7-a1c0cb2965a7
# ╠═1336f910-2d9c-11eb-222b-ade11f5b7b51
# ╠═d4536cb0-2d9b-11eb-16f7-993d0065c683
# ╠═35b079a0-2d9f-11eb-0fa6-1bb075367008
# ╟─dbb70eb2-2dc0-11eb-3da4-754b2e0de392
# ╠═32c66e80-2d9e-11eb-220e-093e86fdc72f
# ╠═3593e090-2e05-11eb-0557-13c977ff89b2
# ╠═463a4470-2e05-11eb-0cac-2f56d62bfa20
# ╠═1fedd0a0-2e07-11eb-1f26-8575ef999663
# ╠═6a6e0550-2d9e-11eb-1136-119a2997dcae
# ╠═0af080a0-2e05-11eb-0069-8d908435f94b
# ╠═c8288560-2dc8-11eb-14f6-19ae67fb2357
# ╠═00989c50-2dc9-11eb-31fc-0d8ab39b5a92
# ╠═77646ed0-2dc5-11eb-35c5-4b1582878184
# ╠═f3292f10-2d9d-11eb-1e76-8de80adce2e6
# ╠═96c67790-2e07-11eb-0000-19536e55469e
# ╟─903e2090-2dc5-11eb-2855-bfc16c6145b6
# ╠═58f2c540-2d9e-11eb-3a14-0931bcd2cc5d
# ╟─c29c1380-29b5-11eb-254c-2d026e15f903
# ╟─89d06380-29bf-11eb-0d5b-4533c94745d5
# ╟─8b3ac8f0-29bf-11eb-274e-f7b07eb35e7c
# ╟─1e35c150-29de-11eb-0069-cb75699b6d61
# ╟─76b4c570-29b5-11eb-32d5-75da0001a894
# ╟─46a9fc40-2a8e-11eb-301f-077f9a6d50fe
# ╠═8d50cc90-294f-11eb-012b-3b973cbd1f1e
# ╟─e102cfc0-295c-11eb-2c5b-834000ca2077
# ╟─08b31090-2e01-11eb-27a7-c3758492386e


"""
    IndependentMOKernel(k::Kernel)

Kernel for multiple independent outputs with kernel `k` each.

# Definition

For inputs ``x, x'`` and output dimensions ``p_x, p_{x'}'``, the kernel ``\\widetilde{k}``
for independent outputs with kernel ``k`` each is defined as
```math
\\widetilde{k}\\big((x, p_x), (x', p_{x'})\\big) = \\begin{cases}
    k(x, x') & \\text{if } p_x = p_{x'}, \\\\
    0 & \\text{otherwise}.
\\end{cases}
```
Mathematically, it is equivalent to a matrix-valued kernel defined as
```math
\\widetilde{K}(x, x') = \\mathrm{diag}\\big(k(x, x'), \\ldots, k(x, x')\\big) \\in \\mathbb{R}^{m \\times m},
```
where ``m`` is the number of outputs.
"""
struct IndependentMOKernel{Tkernel<:Kernel} <: MOKernel
    kernel::Tkernel
end

function (κ::IndependentMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,Int})
    return κ.kernel(x, y) * (px == py)
end

function _kernelmatrix_kron_helper(::MOInputIsotopicByFeatures, Kfeatures, B)
    return kron(Kfeatures, B)
end

function _kernelmatrix_kron_helper(::MOInputIsotopicByOutputs, Kfeatures, B)
    return kron(B, Kfeatures)
end

function kernelmatrix(
    k::IndependentMOKernel, x::MOI, y::MOI
) where {MOI<:IsotopicMOInputsUnion}
    @assert x.out_dim == y.out_dim
    Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
    mtype = eltype(Kfeatures)
    return _kernelmatrix_kron_helper(x, Kfeatures, Eye{mtype}(x.out_dim))
end

if VERSION >= v"1.6"
    function _kernelmatrix_kron_helper!(K, ::MOInputIsotopicByFeatures, Kfeatures, B)
        return kron!(K, Kfeatures, B)
    end

    function _kernelmatrix_kron_helper!(K, ::MOInputIsotopicByOutputs, Kfeatures, B)
        return kron!(K, B, Kfeatures)
    end

    function kernelmatrix!(
        K::AbstractMatrix, k::IndependentMOKernel, x::MOI, y::MOI
    ) where {MOI<:IsotopicMOInputsUnion}
        @assert x.out_dim == y.out_dim
        Ktmp = kernelmatrix(k.kernel, x.x, y.x)
        mtype = eltype(Ktmp)
        return _kernelmatrix_kron_helper!(
            K, x, Ktmp, Matrix{mtype}(I, x.out_dim, x.out_dim)
        )
    end
end

function Base.show(io::IO, k::IndependentMOKernel)
    return print(io, string("Independent Multi-Output Kernel\n\t", string(k.kernel)))
end


#References: Sandoval, Steven, and Phillip L. De Leon. "The Instantaneous Spectrum: A General Framework for Time-Frequency Analysis." IEEE Transactions on Signal Processing 66.21 (2018): 5679-5693.

"""
    𝐒 = numSet([𝐂₀,𝐂₁,𝐂₂])

Create a numerical component set 'numSet' by providing a vector of
numerical canonical triplets.

# Examples
```@example
using ISA
ψ₀ = AMFMcomp(t->exp(-t^2),t->2.0,0.0)
ψ₁ = AMFMcomp(t->1.0,t->10*t,0.1)
ψ₂ = AMFMcomp(t->0.8*cos(2t),t->10 + 7.5*sin(t),π)
fs = 16_000
t = 0:1/fs:1
𝚿₀ = numComp( ψ₀(t), fs )
𝚿₁ = numComp( ψ₁(t), fs )
𝚿₂ = numComp( ψ₂(t), fs )
𝐂₀ = AMFMdemod(𝚿₀)
𝐂₁ = AMFMdemod(𝚿₁)
𝐂₂ = AMFMdemod(𝚿₂)
𝐒 = numSet([𝐂₀,𝐂₁,𝐂₂])
```
"""
struct numSet
  S::Vector{numTriplet}
end

# DISPLAY
function displayS(x)
  text = "\n"
  for i in 1:length(x.S)
    T = typeof(x.S[i])
    text *= "numerical canonical triplet #$(i)\n"
    for (name, typ) in zip(fieldnames(T), T.types)
        text *= "$name is $typ\n"
    end
    text *= "\n"
  end
  return text
end
Base.show(io::IO, x::numSet) = print(io, "numerical component set with $(length(x.S)) components$(displayS(x))")

"""
    𝐳 = numModel([𝚿₀,𝚿₁,𝚿₂])

Create a 'numModel' by providing a vector of numerical components
'numComp'.

# Examples
```@example
using ISA
ψ₀ = AMFMcomp(t->exp(-t^2),t->2.0,0.0)
ψ₁ = AMFMcomp(t->1.0,t->10*t,0.1)
ψ₂ = AMFMcomp(t->0.8*cos(2t),t->10 + 7.5*sin(t),π)
fs = 16_000
t = 0:1/fs:1
𝚿₀ = numComp( ψ₀(t), fs )
𝚿₁ = numComp( ψ₁(t), fs )
𝚿₂ = numComp( ψ₂(t), fs )
𝐳 = numModel([𝚿₀,𝚿₁,𝚿₂])
```
"""
struct numModel
  𝚿ₖ::Vector{numComp}
end

# METHODS
function (𝐳::numModel)(t::Real)
  out = 0.0
  for 𝚿 ∈ 𝐳.𝚿ₖ
    if 𝚿.t[1]<=t<=𝚿.t[end]
      out += 𝚿(t)
    end
  end
  return out
end
function (𝐳::numModel)(t::Vector{<:Real})
  return 𝐳.(t)
end
function (𝐳::numModel)(t::StepRangeLen)
  return 𝐳.(collect(t))
end
function (𝐳::numModel)(t::UnitRange)
  return 𝐳.(collect(t))
end

# DISPLAY
Base.show(io::IO, x::numModel) = print(io, "numerical AM-FM model with $(length(x.𝚿ₖ)) components")


using Distributions
using JSON

n = 3
p = 0.2
d = Binomial( n,p )

x = [ -5, -2.5, 0, 2.5, 5 ]

dmgf(t) = mgf(d, t )
y = map( dmgf, x )
println( y )

data = Dict([
	("n", n),
	("p", p),
	("data", x),
	("expected", y)
])

outfile = open("./test/fixtures/number.json", "w")
JSON.json(data)

write( outfile, JSON.json(data) )


#=
# MNIST

We begin by importing the required packages.
We load MNIST via the MLDatasets.jl package.
=#

#=
Some references

<https://www.jmlr.org/beta/papers/v17/14-237.html>,

<http://www.cs.toronto.edu/~tang/papers/RbmZM.pdf>,

<https://doi.org/10.1007/978-3-642-35289-8_3>

=#

import CairoMakie
import Makie
import MLDatasets
import Flux
import RestrictedBoltzmannMachines as RBMs
import WhitenedRBMs as WhiteRBMs
using Statistics: mean, var, std
using Random: bitrand
using ValueHistories: MVHistory
using RestrictedBoltzmannMachines: BinaryRBM
using WhitenedRBMs: whiten, blacken, whiten_visible_from_data
nothing #hide

#=
Useful function to plot MNIST digits.
=#

"""
    imggrid(A)

Given a four dimensional tensor `A` of size `(width, height, ncols, nrows)`
containing `width x height` images in a grid of `nrows x ncols`, this returns
a matrix of size `(width * ncols, height * nrows)`, that can be plotted in a heatmap
to display all images.
"""
function imggrid(A::AbstractArray{<:Any,4})
    reshape(permutedims(A, (1,3,2,4)), size(A,1)*size(A,3), size(A,2)*size(A,4))
end

#=
Now load the MNIST dataset.
=#

Float = Float32
train_x, train_y = MLDatasets.MNIST.traindata()
digit = 0 # the digit we work with
train_x = Array{Float}(train_x[:, :, train_y .== digit] .> 0.5)
train_y = train_y[train_y .== digit]
println(length(train_y), " train samples")

#=
Initialize and train a normal RBM
=#

rbm = BinaryRBM(Float, (28,28), 200)
RBMs.initialize!(rbm, train_x)
rbm = whiten_visible_from_data(whiten(rbm), train_x, WhiteRBMs.Identity(); ϵ=1f-3)
rbm.hidden.θ .= 0
nothing #hide

# Train init

batchsize = 256
optim = Flux.ADAM()
vm = bitrand(28, 28, batchsize) # fantasy chains
history = MVHistory()
push!(history, :lpl, mean(RBMs.log_pseudolikelihood(rbm, train_x)))
nothing #hide

# Pseudolikelihood before training

mean(@time RBMs.log_pseudolikelihood(rbm, train_x))

# Train

@time for epoch in 1:100 # track pseudolikelihood every 5 epochs
    WhiteRBMs.pcd!(
        rbm, train_x; epochs=5, vm, history, batchsize, optim, ϵv=1f-3,
        transform_v=WhiteRBMs.Identity(), transform_h=WhiteRBMs.Identity()
    )
    push!(history, :lpl, mean(RBMs.log_pseudolikelihood(rbm, train_x)))
end
nothing #hide

# Convert to equivalent RBM (without affine transforms)

rbm = WhiteRBMs.blacken(rbm)
nothing #hide

# Plot log-pseudolikelihood of train data during learning.

fig = Makie.Figure(resolution=(600, 300))
ax = Makie.Axis(fig[1,1], xlabel="epochs", ylabel="pseudolikelihood")
Makie.lines!(ax, get(history, :lpl)...)
fig

# Seconds per epoch.

fig = Makie.Figure(resolution=(600, 300))
ax = Makie.Axis(fig[1,1], xlabel="epoch", ylabel="seconds")
Makie.lines!(ax, get(history, :Δt)...)
fig

# Log-pseudolikelihood vs. computation time instead of epoch count.

fig = Makie.Figure(resolution=(600, 300))
ax = Makie.Axis(fig[1,1], xlabel="seconds", ylabel="pseudolikelihood")
Makie.lines!(ax, cumsum([0; get(history, :Δt)[2]])[1:5:end], get(history, :lpl)[2])
fig

# Now we do the Gibbs sampling to generate RBM digits.

nrows, ncols = 10, 15
nsteps = 5000
fantasy_F = zeros(nrows*ncols, nsteps)
fantasy_x = bitrand(28,28,nrows*ncols)
fantasy_F[:,1] .= RBMs.free_energy(rbm, fantasy_x)
@time for t in 2:nsteps
    fantasy_x .= RBMs.sample_v_from_v(rbm, fantasy_x)
    fantasy_F[:,t] .= RBMs.free_energy(rbm, fantasy_x)
end
nothing #hide

# Check equilibration of sampling

fig = Makie.Figure(resolution=(400,300))
ax = Makie.Axis(fig[1,1], xlabel="sampling time", ylabel="free energy")
fantasy_F_μ = vec(mean(fantasy_F; dims=1))
fantasy_F_σ = vec(std(fantasy_F; dims=1))
Makie.band!(ax, 1:nsteps, fantasy_F_μ - fantasy_F_σ/2, fantasy_F_μ + fantasy_F_σ/2)
Makie.lines!(ax, 1:nsteps, fantasy_F_μ)
fig

# Plot the resulting samples.

fig = Makie.Figure(resolution=(40ncols, 40nrows))
ax = Makie.Axis(fig[1,1], yreversed=true)
Makie.image!(ax, imggrid(reshape(fantasy_x, 28, 28, ncols, nrows)), colorrange=(1,0))
Makie.hidedecorations!(ax)
Makie.hidespines!(ax)
fig


# Records the statistics of two state variables X and Y:
# - covariance cov(X,Y)
# - correlation coefficient
# - marginal variance var(X) and var(Y)
# - joint average E(XY)
# - marginal average E(X) and E(Y)

# Uses the shifted-data method to avoid loss of precision.

mutable struct JointStatsRecorder{T,X,XF,Y,YF}
    x::X
    y::Y
    Kx::Vector{T}
    Ky::Vector{T}
    Σx::Vector{T}
    Σy::Vector{T}
    Σxx::Vector{T}
    Σyy::Vector{T}
    Σxy::Matrix{T}
    n::Int
end

function JointStatsRecorder(x, field_x::Symbol, y, field_y::Symbol)
    data_x = getfield(x, field_x)
    data_y = getfield(y, field_y)
    T = promote_type(eltype(data_x), eltype(data_y))
    X = typeof(x)
    Y = typeof(y)
    XF = field_x
    YF = field_y
    JointStatsRecorder{T,X,XF,Y,YF}(x, y,
        zeros(T, length(data_x)),
        zeros(T, length(data_y)),
        zeros(T, length(data_x)),
        zeros(T, length(data_y)),
        zeros(T, length(data_x)),
        zeros(T, length(data_y)),
        zeros(T, (length(data_x), length(data_y))),
        0)
end

function covariance(rc::JointStatsRecorder)
    return (rc.Σxy .- (rc.Σx.*rc.Σy')/rc.n) / rc.n
end

function variances(rc::JointStatsRecorder)
    varx = (rc.Σxx .- (rc.Σx.^2)/rc.n) / rc.n
    vary = (rc.Σyy .- (rc.Σy.^2)/rc.n) / rc.n
    return (varx, vary)
end

function corrcoef(rc::JointStatsRecorder)
    cov = covariance(rc)
    varx, vary = variances(rc)
    return cov ./ (sqrt(varx).*sqrt(vary'))
end

function marginal_averages(rc::JointStatsRecorder)
    Ex = rc.Σx / rc.n .+ rc.Kx
    Ey = rc.Σy / rc.n .+ rc.Ky
    return (Ex, Ey)
end

function sta(rc::JointStatsRecorder)
    rc.Σxy / rc.n .+ rc.Kx .+ rc.Ky
end

function joint_average(rc::JointStatsRecorder)
    Ex, Ey = marginal_averages(rc)
    return covariance(rc) .+ (Ex .* Ey')
end

function reset!(rc::JointStatsRecorder)
    rc.Kx .= 0
    rc.Σx .= 0
    rc.Σxx .= 0
    rc.Ky .= 0
    rc.Σy .= 0
    rc.Σyy .= 0
    rc.Σxy .= 0
    rc.n = 0
end

@generated function record!{T,X,XF,Y,YF}(rc::JointStatsRecorder{T,X,XF,Y,YF}, step)
    quote
        data_x = rc.x.$XF
        data_y = rc.y.$YF

        if rc.n == 0
            rc.Kx .= data_x
            rc.Ky .= data_y
        end

        rc.n += 1

        @inbounds @simd for i in eachindex(data_x)
            dx = data_x[i] - rc.Kx[i]
            rc.Σx[i] += dx
            rc.Σxx[i] += dx^2
        end
        @inbounds @simd for j in eachindex(data_y)
            dy = data_y[j] - rc.Ky[j]
            rc.Σy[j] += dy
            rc.Σyy[j] += dy^2
        end
        @inbounds for i in eachindex(data_x)
            @simd for j in eachindex(data_y)
                dx = data_x[i] - rc.Kx[i]
                dy = data_y[j] - rc.Ky[j]
                rc.Σxy[i,j] += dx*dy
            end
        end
    end
end


using FactCheck, Base.Test, Images, Colors, FixedPointNumbers

facts("Exposure") do

	context("Histogram") do

        # Many of these test integer values, but of course most images
        # should have pixel values between 0 and 1.  Still, it doesn't
        # hurt to get the integer case right too.
        img = 1:10
        bins, hist = imhist(img, 10)
        @fact length(hist) --> length(bins)+1
        @fact bins --> 1.0:1.0:11.0
        @fact hist --> [0,1, 1, 1, 1, 1, 1, 1, 1, 1, 1,0]
        bins, hist = imhist(img, 5, 2, 6)
        @fact length(hist) --> length(bins)+1
        @fact bins --> 2.0:1.0:7.0
        @fact hist --> [1, 1, 1, 1, 1, 1, 4]

        img = reshape(0:99, 10, 10)
        bins, hist = imhist(img, 10)
        @fact length(hist) --> length(bins)+1
        @fact bins --> 0.0:10.0:100.0
        @fact hist --> [0, 10, 10, 10, 10, 10, 10, 10, 10, 10, 10, 0]
        bins, hist = imhist(img, 7, 25, 59)
        @fact length(hist) --> length(bins)+1
        @fact bins --> 25.0:5.0:60.0
        @fact hist --> [25, 5, 5, 5, 5, 5, 5, 5, 40]

        # Test the more typical case
        img = reinterpret(Gray{U8}, [0x20,0x40,0x80,0xd0])
        @fact imhist(img, 5) --> (0.0:0.2:1.0,[0,1,1,1,0,1,0])
        img = reinterpret(Gray{U8}, [0x00,0x40,0x80,0xd0])
        @fact imhist(img, 5) --> (0.0:0.2:1.0,[0,1,1,1,0,1,0])
        img = reinterpret(Gray{U8}, [0x00,0x40,0x80,0xff])
        @fact imhist(img, 6) --> (0.0:0.2:1.2,[0,1,1,1,0,0,1,0])

    end

	context("Histogram Equalisation") do

        #DataTypes
        img = ones(Images.Gray{Float64}, 10, 10)
        ret = histeq(img, 100)
        @fact img == ret --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.Gray{Images.U8}, 10, 10)
        ret = histeq(img, 100)
        @fact img == ret --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.Gray{Images.U16}, 10, 10)
        ret = histeq(img, 100)
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.AGray{Images.U8}, 10, 10)
        ret = histeq(img, 100)
        @fact img == ret --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.RGB{Images.U8}, 10, 10)
        ret = histeq(img, 100)
        @fact img == ret --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.RGB{Images.U16}, 10, 10)
        ret = histeq(img, 100)
        @fact img == ret --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.RGB{Float64}, 10, 10)
        ret = histeq(img, 100)
        @fact all(map((i, r) -> isapprox(i, r), img, ret)) --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.ARGB{Images.U8}, 10, 10)
        ret = histeq(img, 100)
        @fact img == ret --> true
        @fact eltype(ret) == eltype(img) --> true

        #Working

        img = zeros(10, 10)
        for i in 1:10
            img[i, :] = 10 * (i - 1)
        end
        @fact img == Images.histeq(img, 10, 0, 90) --> true

        ret = Images.histeq(img, 2, 0, 90)
        @fact all(ret[1:5, :] .== 0) --> true
        @fact all(ret[6:10, :] .== 90) --> true

        ret = Images.histeq(img, 5, 0, 90)
        for i in 1:2:10
            @fact all(ret[i:i+1, :] .== 22.5 * floor(i / 2)) --> true
        end

        img = [0.0,  21.0,  29.0,  38.0,  38.0,  51.0,  66.0,  79.0,  79.0,  91.0,
                  1.0,  21.0,  28.0,  39.0,  39.0,  52.0,  63.0,  77.0,  77.0,  91.0,
                  2.0,  22.0,  21.0,  35.0,  35.0,  56.0,  64.0,  73.0,  73.0,  94.0,
                  3.0,  21.0,  22.0,  33.0,  33.0,  53.0,  63.0,  72.0,  72.0,  93.0,
                  4.0,  22.0,  27.0,  32.0,  32.0,  52.0,  64.0,  78.0,  78.0,  92.0,
                 11.0,  21.0,  31.0,  31.0,  42.0,  66.0,  71.0,  71.0,  88.0,  91.0,
                 12.0,  22.0,  32.0,  32.0,  44.0,  62.0,  75.0,  75.0,  88.0,  91.0,
                 13.0,  23.0,  34.0,  34.0,  49.0,  67.0,  74.0,  74.0,  82.0,  94.0,
                 14.0,  26.0,  35.0,  35.0,  43.0,  68.0,  74.0,  74.0,  83.0,  93.0,
                 15.0,  27.0,  36.0,  36.0,  44.0,  69.0,  74.0,  74.0,  86.0,  92.0]
        img = reshape(img, 10, 10)'

        ret = Images.histeq(img, 10, 0, 99)
        cdf = cumsum(imhist(img, 10)[2][2:end-1])
        @fact all(ret[1:cdf[1]] .== 0.0) --> true
        for i in 1:(size(cdf)[1]-1)
            @fact all(ret[cdf[i] + 1 : cdf[i + 1]] .== (cdf[i + 1] - cdf[1]) * 99.0 / (cdf[end] - cdf[1])) --> true
        end

    end

	context("Gamma Correction") do

        img = ones(Images.Gray{Float64}, 10, 10)
        ret = adjust_gamma(img, 1)
        @fact img == ret --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.Gray{Images.U8}, 10, 10)
        ret = adjust_gamma(img, 1)
        @fact img == ret --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.Gray{Images.U16}, 10, 10)
        ret = adjust_gamma(img, 1)
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.AGray{Images.U8}, 10, 10)
        ret = adjust_gamma(img, 1)
        @fact img == ret --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.RGB{Images.U8}, 10, 10)
        ret = adjust_gamma(img, 1)
        @fact img == ret --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.RGB{Images.U16}, 10, 10)
        ret = adjust_gamma(img, 1)
        @fact img == ret --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.RGB{Float64}, 10, 10)
        ret = adjust_gamma(img, 1)
        @fact all(map((i, r) -> isapprox(i, r), img, ret)) --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.ARGB{Images.U8}, 10, 10)
        ret = adjust_gamma(img, 1)
        @fact img == ret --> true

        #Working

        img = reshape(1:1:100, 10, 10)
        ret = Images.adjust_gamma(img, 2)
        @fact ret == img .^ 2 --> true
        ret = Images.adjust_gamma(img, 0.3)
        @fact ret == img .^ 0.3 --> true
        ret = Images.adjust_gamma(img, 1, 1, 100)
        @fact ret == img --> true

        img = zeros(10, 10)
        ret = Images.adjust_gamma(img, 2)
        @fact all(ret .== 0) --> true

        a = ARGB(0.2, 0.3, 0.4, 0.5)
        r = Images._gamma_pixel_rescale(a, 1)
        @fact isapprox(a, r, rtol = 0.0001) --> true
        r = Images._gamma_pixel_rescale(a, 2)
        @fact alpha(r) --> alpha(a)

        b = AGray(0.2, 0.6)
        r = Images._gamma_pixel_rescale(b, 1)
        @fact b --> r
        r = Images._gamma_pixel_rescale(b, 2)
        @fact alpha(r) --> alpha(b)
        @fact isapprox(r.val, b.val ^ 2) --> true

    end

	context("Histogram Matching") do

        #DataTypes
        img = ones(Images.Gray{Float64}, 10, 10)
        ret = histmatch(img, img)
        @fact all(ret .== zero(eltype(img))) --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.Gray{Images.U8}, 10, 10)
        ret = histmatch(img, img)
        @fact all(ret .== zero(eltype(img))) --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.Gray{Images.U16}, 10, 10)
        ret = histmatch(img, img)
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.AGray{Images.U8}, 10, 10)
        ret = histmatch(img, img)
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.RGB{Images.U8}, 10, 10)
        ret = histmatch(img, img)
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.RGB{Images.U16}, 10, 10)
        ret = histmatch(img, img)
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.RGB{Float64}, 10, 10)
        ret = histmatch(img, img)
        @fact all(map((i, r) -> isapprox(zero(RGB), r, atol = 0.001), img, ret)) --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.ARGB{Images.U8}, 10, 10)
        ret = histmatch(img, img)
        @fact eltype(ret) == eltype(img) --> true

        img = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
        edges, hist = imhist(img, 2)
        himg = Images._histmatch(img, edges, hist)
        @fact himg == [0, 0, 0, 0, 0, 5, 5, 5, 5, 5] --> true
        edges, hist = imhist(img, 5)
        himg = Images._histmatch(img, edges, hist)
        @fact himg == [0, 0, 2, 2, 4, 4, 6, 6, 8, 8] --> true
    end

    context("CLAHE") do

        #DataTypes
        img = ones(Images.Gray{Float64}, 10, 10)
        ret = clahe(img, 100)
        @fact size(ret) == size(img) --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.Gray{Images.U8}, 10, 10)
        ret = clahe(img, 100)
        @fact size(ret) == size(img) --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.Gray{Images.U16}, 10, 10)
        ret = clahe(img, 100)
        @fact size(ret) == size(img) --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.AGray{Images.U8}, 10, 10)
        ret = clahe(img, 100)
        @fact size(ret) == size(img) --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.RGB{Images.U8}, 10, 10)
        ret = clahe(img, 100)
        @fact size(ret) == size(img) --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.RGB{Images.U16}, 10, 10)
        ret = clahe(img, 100)
        @fact size(ret) == size(img) --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.RGB{Float64}, 10, 10)
        ret = clahe(img, 100)
        @fact size(ret) == size(img) --> true
        @fact eltype(ret) == eltype(img) --> true

        img = ones(Images.ARGB{Images.U8}, 10, 10)
        ret = clahe(img, 100)
        @fact size(ret) == size(img) --> true
        @fact eltype(ret) == eltype(img) --> true

        #Working

        cdf = Array(0.1:0.1:1.0)
        res_pix = Images._clahe_pixel_rescale(0.5654, cdf, cdf, cdf, cdf, 0.1:0.1:1.0, 1, 4, 10, 10)
        @fact isapprox(res_pix, 0.5) --> true
        res_pix = Images._clahe_pixel_rescale(0.2344, cdf, cdf, cdf, cdf, 0.1:0.1:1.0, 1, 4, 10, 10)
        @fact isapprox(res_pix, 0.2) --> true
        res_pix = Images._clahe_pixel_rescale(0.8123, cdf, cdf, cdf, cdf, 0.1:0.1:1.0, 1, 4, 10, 10)
        @fact isapprox(res_pix, 0.8) --> true

        cdf2 = [0.0, 0.0, 0.1, 0.1, 0.2, 0.3, 0.5, 0.6, 0.8, 1.0]

        res_pix = Images._clahe_pixel_rescale(0.3, cdf2, cdf, 0.1:0.1:1, 3, 10)
        @fact isapprox(res_pix, 0.144444444) --> true

        res_pix = Images._clahe_pixel_rescale(0.3, cdf2, cdf, 0.1:0.1:1, 6, 10)
        @fact isapprox(res_pix, 0.211111111) --> true

        res_pix = Images._clahe_pixel_rescale(0.5654, cdf2, 0.1:0.1:1.0)
        @fact isapprox(res_pix, 0.2) --> true

        cdf3 = [0.0, 0.0, 0.1, 0.3, 0.5, 0.5, 0.6, 0.9, 1.0, 1.0]
        cdf4 = [0.0, 0.1, 0.3, 0.4, 0.5, 0.6, 0.6, 0.6, 0.8, 1.0]

        res_pix = Images._clahe_pixel_rescale(0.8123, cdf, cdf2, cdf3, cdf4, 0.1:0.1:1.0, 3, 4, 10, 10)
        @fact isapprox(res_pix, 0.781481481) --> true
        res_pix = Images._clahe_pixel_rescale(0.3, cdf, cdf2, cdf3, cdf4, 0.1:0.1:1.0, 5, 4, 10, 10)
        @fact isapprox(res_pix, 0.2037037037) --> true

        img = [ 0.907976  0.0601377   0.747873   0.938683   0.935189  0.49492    0.508545  0.394573   0.991136   0.210218
                0.219737  0.727806    0.606995   0.25272    0.279777  0.584529   0.708244  0.958645   0.979217   0.418678
                0.793384  0.314803    0.324222   0.840858   0.49438   0.102003   0.724399  0.133663   0.312085   0.662211
                0.905336  0.423511    0.564616   0.692544   0.367656  0.630386   0.133439  0.529039   0.0175596  0.644065
                0.255092  0.00242822  0.723724   0.673323   0.244508  0.518068   0.204353  0.34222    0.106092   0.0331161
                0.432963  0.383491    0.903465   0.579605   0.719874  0.571533   0.728544  0.1864     0.950187   0.470226
                0.877475  0.554114    0.987133   0.947148   0.710115  0.131948   0.711611  0.0221843  0.470008   0.952806
                0.231911  0.177463    0.742054   0.0333307  0.481319  0.716638   0.332057  0.978177   0.0610481  0.439462
                0.935457  0.159602    0.178357   0.163585   0.275052  0.0557963  0.066368  0.199349   0.9238     0.85161
                0.692148  0.503563    0.0918827  0.0206237  0.702344  0.546088   0.04163   0.174103   0.45499    0.90019
                ]
        hist_eq_img = Images.clahe(img, xblocks = 2, yblocks = 2, clip = 5)
        ret_expected = [ 0.972222  0.154261  0.672734  0.962918  0.960518  0.603715  0.559332  0.684868  0.919632  0.11213
                        0.26227   0.577172  0.615628  0.420623  0.356921  0.581208  0.845865  0.888908  0.908469  0.518953
                        0.79315   0.398521  0.454306  0.692093  0.404984  0.287649  0.619472  0.47354   0.442781  0.826681
                        0.875951  0.325109  0.561585  0.746554  0.355356  0.542615  0.308523  0.303799  0.12818   0.662404
                        0.218908  0.145377  0.627466  0.703423  0.255746  0.543485  0.270196  0.240487  0.119202  0.0762326
                        0.443491  0.419296  0.857736  0.778583  0.78256   0.660936  0.66722   0.420796  0.793904  0.50695
                        0.765105  0.583873  0.870947  0.825773  0.755839  0.346564  0.598017  0.329962  0.564487  0.897893
                        0.512372  0.328346  0.50314   0.339795  0.493125  0.476082  0.45704   0.597256  0.385295  0.814294
                        0.866771  0.256846  0.140199  0.118129  0.287474  0.175723  0.115673  0.343809  0.735832  0.928661
                        0.833503  0.481259  0.135795  0.116606  0.737954  0.57691   0.103692  0.132731  0.539984  0.999525
                        ]
        @fact all(map((i, j) -> isapprox(i, j, rtol = 0.00001), hist_eq_img, ret_expected)) --> true

    end

    context("Other") do

        # Issue #282
        img = convert(Images.Image{Gray{UFixed8}}, eye(2,2))
        imgs = Images.imstretch(img, 0.3, 0.4)
        @fact data(imgs) --> roughly(1./(1 + (0.3./(eye(2,2) + eps())).^0.4))

        img = convert(Images.Image{Gray{UFixed16}}, [0.01164 0.01118; 0.01036 0.01187])
        @fact data(imadjustintensity(img,[0.0103761, 0.0252166]))[2,1] --> 0.0
        @fact eltype(imadjustintensity(img)) --> Gray{UFixed16}

        img = convert(Images.Image{Gray{UFixed16}}, [0.01164 0.01118; 0.01036 0.01187])
        @fact Images.complement(Gray(0.5)) == Gray(0.5) --> true
        @fact Images.complement(Gray(0.2)) == Gray(0.8) --> true
        @fact all(imcomplement(img) .== 1 - img) --> true

        hist = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

        clipped_hist = cliphist(hist, 2)
        @fact all(clipped_hist .== 2.0) --> true
        clipped_hist = cliphist(hist, 3)
        @fact all(clipped_hist .== 3.0) --> true
        clipped_hist = cliphist(hist, 4)
        @fact all(clipped_hist .== 4.0) --> true
        clipped_hist = cliphist(hist, 5)
        @fact all(clipped_hist .== 5.0) --> true
        clipped_hist = cliphist(hist, 6)
        @fact clipped_hist == vec([3, 4, 6, 6, 6, 6, 6, 6, 6, 6]) --> true
        clipped_hist = cliphist(hist, 7)
        @fact clipped_hist == vec([2.6, 3.6, 3.6, 4.6, 5.6, 7, 7, 7, 7, 7]) --> true
        clipped_hist = cliphist(hist, 8)
        @fact clipped_hist == vec([2.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3, 8, 8, 8]) --> true
        clipped_hist = cliphist(hist, 9)
        @fact clipped_hist == vec([2.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1, 9, 9]) --> true
        clipped_hist = cliphist(hist, 10)
        @fact clipped_hist == hist --> true
    end

end


immutable IdentityMap{T}<:AbstractLinearMap{T} # T will be determined from maps to which I is added
    M::Int
end
IdentityMap(T::Type,M::Int)=IdentityMap{T}(M)
IdentityMap(T::Type,M::Int,N::Int)=(M==N ? IdentityMap{T}(M) : error("IdenityMap needs to be square"))
IdentityMap(T::Type,sz::Tuple{Int,Int})=(sz[1]==sz[2] ? IdentityMap{T}(sz[1]) : error("IdenityMap needs to be square"))
IdentityMap(M::Int)=IdentityMap(Bool,M)
IdentityMap(M::Int,N::Int)=IdentityMap(Bool,M,N)
IdentityMap(sz::Tuple{Int,Int})=IdentityMap(Bool,sz)

# properties
Base.size(A::IdentityMap,n)=(n==1 || n==2 ? A.M : error("AbstractLinearMap objects have only 2 dimensions"))
Base.size(A::IdentityMap)=(A.M,A.M)
Base.isreal(::IdentityMap)=true
Base.issym(::IdentityMap)=true
Base.ishermitian(::IdentityMap)=true
Base.isposdef(::IdentityMap)=true

# multiplication with vector
Base.A_mul_B!(y::AbstractVector,A::IdentityMap,x::AbstractVector)=(length(x)==length(y)==A.M ? copy!(y,x) : throw(DimensionMismatch("A_mul_B!")))
*(A::IdentityMap,x::AbstractVector)=x

Base.At_mul_B!(y::AbstractVector,A::IdentityMap,x::AbstractVector)=(length(x)==length(y)==A.M ? copy!(y,x) : throw(DimensionMismatch("At_mul_B!")))
Base.At_mul_B(A::IdentityMap,x::AbstractVector)=(length(x)==A.M ? x : throw(DimensionMismatch("At_mul_B")))

Base.Ac_mul_B!(y::AbstractVector,A::IdentityMap,x::AbstractVector)=(length(x)==length(y)==A.M ? copy!(y,x) : throw(tMismatch("Ac_mul_B!")))
Base.Ac_mul_B(A::IdentityMap,x::AbstractVector)=(length(x)==A.M ? x : throw(DimensionMismatch("Ac_mul_B")))

# combine AbstractLinearMap and UniformScaling objects in linear combinations
+{T}(A1::AbstractLinearMap,A2::UniformScaling{T})=A1+A2[1,1]*IdentityMap{T}(size(A1,1))
+{T}(A1::UniformScaling{T},A2::AbstractLinearMap)=A1[1,1]*IdentityMap{T}(size(A2,1))+A2


# rank adaptation
function normal_component(U, Z, C, dY; tol = 1e-8)
    # U = basis
    # Z = coefficients
    # C = covariance matrix, Z'*Z for example (assumed to be of very small dimension)
    # dY = dynamics, f(UZ') for example
    return (I - U*U')*dY*(I-Z*pinv(C, atol = tol)*Z')
end

function normal_component(LRA::TwoFactorRepresentation, C, dY; tol = 1e-8)
    # C = covariance matrix, Z'*Z for example (assumed to be of very small dimension)
    # dY = dynamics, f(UZ') for example
    return normal_component(LRA.U, LRA.Z, C, dY; tol = tol)
end

function normal_component(LRA::TwoFactorRepresentation, dY; tol = 1e-8)
    # C = covariance matrix, Z'*Z for example (assumed to be of very small dimension)
    # dY = dynamics, f(UZ') for example
    return normal_component(LRA.U, LRA.Z, LRA.Z'*LRA.Z, dY; tol = tol)
end

function gd_truncated_svd!(Φ::TwoFactorRepresentation, Ψ::TwoFactorRepresentation, alg::GradientDescent, orth_alg)
    #UZ => lower rank approximation
    #Ψ => larger rank matrix to be approximated
    @unpack μ, rtol, atol, maxiter = alg
    iter = 0
    C = Symmetric(Φ.Z', Φ.Z)
    UtU = Ψ.U'*Φ.U
    ZtZ = Ψ.Z'*Φ.Z
    k1 = ZtZ / C
    k2 = Ψ.Z*UtU
    k3 = Φ.U*UtU
    k4 = k3*k1
    k5 = Ψ.U*k1

    dZ = Φ.Z - k2
    dU = k4 - k5
    while (iter < maxiter &&
          norm(dU)^2 + norm(dZ)^2 > atol^2 &&
          (norm(dU)^2 + norm(dZ)^2)/(norm(U)^2 + norm(Z)^2) > rtol^2)

        Φ.U .-= μ*dU
        Ψ.Z .-= μ*dZ
        orthonormalize!(Φ, orth_alg)
        iter += 1

        C = Symmetric(Φ.Z', Φ.Z)
        mul!(UtU, Ψ.U',Φ.U)
        mul!(ZtZ, Ψ.Z',Φ.Z)
        k1 = ZtZ / C
        mul!(k2,Ψ.Z,UtU)
        mul!(k3,Φ.U,UtU)
        mul!(k4,k3,k1)
        mul!(k5,Ψ.U,k1)
        dZ = Φ.Z - k2
        dU = k4 - k5
    end
    if iter == maxiter
        @warn "Orthonormalization did not converge.
               Iterations exceeded maxiters = $maxiter.
               Residual: $(norm(dU)^2+norm(dZ)^2)"
    end
end


using ITensors
using Symbolics

using ITensors: contraction_cost, optimal_contraction_sequence

function tensor_network(; m, k, d)
  l = Index(m, "l")
  r = Index(m, "r")
  h₁ = Index(k, "h₁")
  h₂ = Index(k, "h₂")
  h₃ = Index(k, "h₃")
  s₁ = Index(d, "s₁")
  s₂ = Index(d, "s₂")

  ψ = emptyITensor(l, s₁, s₂, r)
  L = emptyITensor(dag(l), l', h₁)
  H₁ = emptyITensor(dag(s₁), s₁', dag(h₁), h₂)
  H₂ = emptyITensor(dag(s₂), s₂', dag(h₂), h₃)
  R = emptyITensor(dag(r), r', h₃)
  return [ψ, L, H₁, H₂, R]
end

function main()
  mrange = 50:10:80
  krange = 50:10:80
  sequence_costs = Matrix{Any}(undef, length(mrange), length(krange))
  for iₘ in eachindex(mrange), iₖ in eachindex(krange)
    m_val = mrange[iₘ]
    k_val = krange[iₖ]
    d_val = 4

    TN = tensor_network(; m = m_val, k = k_val, d = d_val)
    sequence = optimal_contraction_sequence(TN)
    cost = contraction_cost(TN; sequence = sequence)

    @variables m, k, d
    TN_symbolic = tensor_network(; m = m, k = k, d = d)
    cost_symbolic = contraction_cost(TN_symbolic; sequence = sequence)
    sequence_cost = (dims = (m = m_val, k = k_val, d = d_val), sequence = sequence, cost = cost, symbolic_cost = cost_symbolic)
    sequence_costs[iₘ, iₖ] = sequence_cost
  end
  return sequence_costs
end

sequence_costs = main()

# Analyze the results.
println("Index dimensions")
display(getindex.(sequence_costs, :dims))

println("\nContraction sequences")
display(getindex.(sequence_costs, :sequence))

println("\nSymbolic contraction cost with d = 4")
# Fix d to a certain value (such as 4 for a Hubbard site)
@variables d
var_sub = Dict(d => 4)
display(substitute.(sum.(getindex.(sequence_costs, :symbolic_cost)), (var_sub,)))


using Gen
using FunctionalCollections: PersistentVector
import Distributions

####################
# generative model #
####################

function x_mean(x_prev::Real, t::Int)
    (x_prev / 2.) + 25 * (x_prev / (1 + x_prev * x_prev)) + 8 * cos(1.2 * t)
end

y_mean(x::Real) = (x * x / 20.)

@gen function hmm(var_x, var_y, T::Int)
    xs = Vector{Float64}(undef, T)
    ys = Vector{Float64}(undef, T)
    xs[1] = @trace(normal(0, 5), :x => 1)
    ys[1] = @trace(normal(y_mean(xs[1]), sqrt(var_y)), :y => 1)
    for t=2:T
        xs[t] = @trace(normal(x_mean(xs[t-1], t), sqrt(var_x)), :x => t)
        ys[t] = @trace(normal(y_mean(xs[t]), sqrt(var_y)), :y => t)
    end
    return (xs, ys)
 end

struct State
    x::Float64
    y::Float64
end

struct Params
    var_x::Float64
    var_y::Float64
end

# the kernel accepts two special arguments (the time step t and the return
# value of the previous tieration, followed by the rest of the arguments,
# specified as an argument to unfold)
@gen (static) function kernel(t::Int, state::State, params::Params)
    x::Float64 = @trace(normal(t > 1 ? x_mean(state.x, t) : 0.,
                              t > 1 ? sqrt(params.var_x) : 5.), :x)
    y::Float64 = @trace(normal(y_mean(x), sqrt(params.var_y)), :y)
    ret::State = State(x, y)
    return ret
end

# arguments: (T::Int, init_state, common)
# i.e. the kernel needs to accept three arguments..
hmm2 = Unfold(kernel)

#################################
# sequential Monte Carlo in hmm #
#################################

#const num_particles = 200
#const ess = 100

# implement of SMC using forward simulation as the proposal
# it's okay to be specialized to the model 'hmm'

@gen function single_step_observer(t::Int, y::Float64)
    @trace(dirac(y), :y => t)
end

@gen (static) function obs_sub(y::Float64)
    @trace(dirac(y), :y)
end

single_step_observer2 = at_dynamic(obs_sub, Int)

#@gen (static) function single_step_observer2(t::Int, y::Float64)
    #@trace(at_dynamic(obs_sub,Int)(t, (y,)))
    #@trace(dirac(y), t => :y)
#end

function logsumexp(arr)
    min_arr = maximum(arr)
    min_arr + log(sum(exp.(arr .- min_arr)))
end

function effective_sample_size(log_weights::Vector{Float64})
    # assumes weights are normalized
    log_ess = -logsumexp(2. * log_weights)
    exp(log_ess)
end

# TODO need to implement generate and extend for the unfold generator type
# it should have a custom argchange argument that is indicates whether only the
# time length argument has changed, which will allow extend to only visit new
# time steps

function smc(var_x, var_y, T::Int, N, ess_threshold, ys::AbstractArray{Float64,1})
    log_unnormalized_weights = Vector{Float64}(undef, N)
    log_ml_estimate = 0.
    obs = get_choices(simulate(single_step_observer2, (1, (ys[1],))))
    traces = Vector{Gen.get_trace_type(hmm2)}(undef, N)
    next_traces = Vector{Gen.get_trace_type(hmm2)}(undef, N)
    args = (1, State(NaN, NaN), Params(var_x, var_y))
    for i=1:N
        (traces[i], log_unnormalized_weights[i]) = generate(hmm2, args, obs)
    end
    num_resamples = 0
    args_change = UnfoldCustomArgDiff(true, false, false)
    for t=2:T
        #println(t)
        log_total_weight = logsumexp(log_unnormalized_weights)
        log_normalized_weights = log_unnormalized_weights .- log_total_weight
        if effective_sample_size(log_normalized_weights) < ess_threshold
            weights = exp.(log_normalized_weights)
            parents = rand(Distributions.Categorical(weights / sum(weights)), N)
            #parents = map((i) -> categorical(weights / sum(weights)), 1:N)
            log_ml_estimate += log_total_weight - log(N)
            log_unnormalized_weights = zeros(N)
            num_resamples += 1
        else
            parents = 1:N
        end
        obs = get_choices(simulate(single_step_observer2, (t, (ys[t],))))
        args = (t, State(NaN, NaN), Params(var_x, var_y))
        for i=1:N
            parent = parents[i]
            (next_traces[i], weight) = extend(trace[i], hmm2, args, args_change, obs)
            log_unnormalized_weights[i] += weight
        end
        tmp = traces
        traces = next_traces
        next_traces = tmp
    end
    log_total_weight = logsumexp(log_unnormalized_weights)
    log_normalized_weights = log_unnormalized_weights .- log_total_weight
    log_ml_estimate += log_total_weight - log(N)
    log_ml_estimate # just return the log ML estimate
end

##########################################
# collapsed generative model (handcoded) #
##########################################

using Gen: VectorDistTrace, VectorDistTraceChoiceMap
import Gen: get_call_record, has_choices, get_choices, simulate, assess

struct CollapsedHMMTrace
    vector::VectorDistTrace{Float64}
end

get_call_record(trace::CollapsedHMMTrace) = trace.vector.call
has_choices(trace::CollapsedHMMTrace) = length(trace.vector.call.retval) > 0
get_choices(trace::CollapsedHMMTrace) = CollapsedHMMChoiceMap(get_choices(trace.vector))

struct CollapsedHMMChoiceMap <: ChoiceMap
    y_assignment::VectorDistTraceChoiceMap
end

# addrs are: :y => i

Base.isempty(assignment::CollapsedHMMChoiceMap) = (length(assignment.trace.call.retval) > 0)
Gen.get_address_schema(::Type{CollapsedHMMChoiceMap}) = DynamicAddressSchema()
Gen.has_internal_node(assignment::CollapsedHMMChoiceMap, addr) = (addr == :y)
function Gen.get_internal_node(assignment::CollapsedHMMChoiceMap, addr)
    addr == :y ? assignment.y_assignment : throw(KeyError(addr))
end
Gen.get_internal_nodes(assignment::CollapsedHMMChoiceMap) = ((:y, assignment.y_assignment),)

struct CollapsedHMM <: GenerativeFunction{PersistentVector{Float64},CollapsedHMMTrace} end
collapsed_hmm = CollapsedHMM()

function unbiased_logpdf_est(args, ys::PersistentVector{Float64})
    (var_x, var_y, T, num_particles, ess) = args
	local lml_estimate::Float64
	lml_estimate = smc(var_x, var_y, T, num_particles, ess, ys)
    retval = PersistentVector{Float64}(ys)
    vector = VectorDistTrace(retval, args, lml_estimate, length(ys))
    CollapsedHMMTrace(vector)
end

function Gen.generate(generator::CollapsedHMM, args, constraints)
    (var_x, var_y, T, num_particles, ess) = args
    if isempty(constraints)
        error("Unsupported constraints")
    end
    ys = PersistentVector{Float64}(get_leaf_node(constraints, :y => t) for t=1:T)
    if var_x < 0 || var_y < 0
        retval = PersistentVector{Float64}(ys)
        vector = VectorDistTrace(retval, args, -Inf, length(ys))
        trace = CollapsedHMMTrace(vector)
        weight = -Inf
    else
        trace = unbiased_logpdf_est(args, ys)
        weight = get_call_record(trace).score
    end
    (trace, weight)
end

function Gen.update(generator::CollapsedHMM, new_args, args_change, trace, constraints)
    (var_x, var_y, T, num_particles, ess) = new_args
    if !isempty(constraints)
        error("Unsupported constraints")
    end
    ys = trace.vector.values
    if var_x < 0 || var_y < 0
        retval = get_call_record(trace).retval
        T != length(retval) && error("cannot change length")
        vector = VectorDistTrace(retval, new_args, -Inf, length(retval))
        new_trace = CollapsedHMMTrace(vector)
        weight = -Inf
    else
        new_trace = unbiased_logpdf_est(new_args, ys)
        weight = get_call_record(new_trace).score - get_call_record(trace).score
    end
    (new_trace, weight, EmptyChoiceMap(), nothing)
end

@gen (static) function model_collapsed(T::Int)
    var_x::Float64 = @trace(gamma(1, 1), :var_x)
    var_y::Float64 = @trace(gamma(1, 1), :var_y)
    @trace(collapsed_hmm(var_x, var_y, T, 4096, 2048), :hmm)
end


using AlgebraicPetri
using OrdinaryDiffEq
using DiffEqFlux
#using Catalyst
using LinearAlgebra: dot

"""
We will eventually be more sophisticated and use the actual projection
maps of a pullback, but for now we check whether "I" is in the name

Returns vector int
"""
get_infected_states(g::AbstractLabelledPetriNet) =
   [i for (i,s) in enumerate(g[:sname]) if occursin("I", string(s))]

function get_bounds(g::AbstractLabelledPetriNet)
  lb, ub = Float64[],Float64[]
  for (sym,_) in g[:tname]
    sym_= string(sym)
    if occursin("inf", sym_)
      push!(lb, log(1e-8)); push!(ub, log(1e-5))
    elseif occursin("rec", sym_)
      push!(lb, log(1e-4)); push!(ub, log(1e-1))
    else
      push!(lb, log(1e-10)); push!(ub, log(1))
    end
  end
  for (sym,_) in g[:sname]
    sym_ = string(sym)
    if occursin("I", sym_)
      push!(lb, log(1e2)); push!(ub, log(1e6))
    elseif occursin("R", sym_)
      push!(lb, log(1e-10)); push!(ub, log(1))
    else
      push!(lb, log(1e2)); push!(ub, log(1e8))
    end
  end
  return lb => ub
end
"""
state - Vector{Float}
"""
sum_infected_states(g::AbstractLabelledPetriNet, states) =
  sum([s for (i, s) in enumerate(states) if i ∈ get_infected_states(g)])

"""Generate a sequence of new infections"""
function generate_data(rxn::AbstractLabelledPetriNet, rxn_rates, initial_conc)
  o = ODEProblem(ReactionSystem(rxn),  initial_conc,
                 (0.0,100.0), vcat(rxn_rates, initial_conc))
  trajectory = solve(o, Tsit5())
  t_steps = [trajectory(i) for i in 0:100]
  return [sum_infected_states(rxn, t) for t in t_steps]
end

"""
Produces function from Vector{Float} -> Float.
- The input is the logged params (exponentiate to get the real params)

Outer function inputs:
prob - current problem (ODE problem)
times - each timestep (indices of timestep correspond to indices of data array)
data - Vector{Float} representing Infected at each step - target to hit
states/rates - dictionary whatever you want to be fixed
"""
function make_loss(g::AbstractLabelledPetriNet, prob, data; states=Dict(), rates=Dict())
  function loss(p)
    cur_p = exp.(p)
    u0 = exp.(p[(1:ns(g)) .+ nt(g)])
    for (k,v) in rates
      cur_p[k] = v
    end
    for (k,v) in states
        u0[k] = v
    end

    prob′ = remake(prob, p=cur_p, u0=u0, tspan=(0.0,100.0))
    sol = solve(prob′, Tsit5())
    sol_steps = [sol(i) for i in 0:100]
    prediction = [sum_infected_states(g, s) for s in sol_steps]
    log(sum(abs2, data .- prediction)), sol
  end
end

"""
Train a Petri net against data. Return the optimized parameters and minimum loss.
"""
function train(g::AbstractLabelledPetriNet,data::Vector{Float64})::Pair{Float64, Vector{Float64}}
  param_guess = vcat(repeat([1e-10], nt(g)), repeat([1e-3], ns(g)))
  #param_guess = [1e-10,1e-10,1e6,1e6,1e6] # correct [1e-8,1e-4,1e7,1e5,1e-10]
  prob = ODEProblem(ReactionSystem(g),
                    zeros(ns(g)), # initial state
                    (0.0,100.0),
                    zeros(ns(g)+nt(g))) # current variables in model
  lb, ub = get_bounds(g)
  l_func = make_loss(g, prob, data)
  p = DiffEqFlux.sciml_train(l_func,lb,ADAM(0.1), lower_bounds=lb, upper_bounds=ub, maxiters = 10000)
  return l_func(p.u)[1] => exp.(p.u)
end

function eval_petri_fn(true_net::AbstractLabelledPetriNet, true_rates::Vector{Float64}, true_initial_conc::Vector{Float64})::Function
  data = generate_data(true_net, true_rates, true_initial_conc)
  function model_search_loss_fn(g::AbstractLabelledPetriNet)#::Float64
    return first(train(g, data))
    #return train(g, data)
  end
  return model_search_loss_fn
end

#l_func = make_loss(model, prob, week_times, week_avg; calc_inf=calc_inf, rates=Dict(5=>1/14, 6=>1/14), states=Dict(2=>1e-9, 4=>1e-9));

#p = DiffEqFlux.sciml_train(l_func,log.(param_guess),ADAM(0.1),maxiters = 1000)


# Linear Regression with L1 penalty
PlotResults=false

useproc("GPU") # GPU is significantly faster for larger systems
#useproc("CPU")

StartCode()
X=ADnode()
w=ADvariable()
Y=ADnode()
loss=meanSquareLoss(X*w,Y) # mixed regularisation term
net=EndCode() # defines the graph

# instantiate parameter nodes and inputs:

D=100 # dimension of weight vector
N=100 # number of datapoints
truew=randn(D,1).*(rand(D,1).>0.8)
net.value[w]=truew
net.value[X]=randn(N,D)
net.value[Y]=net.value[X]*net.value[w] # make a realisable problem
net.value[w]=randn(D,1) # start with a random w

net=compile(net) # compile and preallocate memory

@gpu CUDArt.init([0]) # let the user do device management

@gpu net=convert(net,PROC,Float32)
#@gpu net=convert(net,PROC,Float64)

if 1==1

gradcheck(net)

# Training:
    tic()
println("Training: using $(net.gpu==true? "GPU" : "CPU") ")
ParsToUpdate=Parameters(net)
nupdates=500
LearningRate=0.01
error=Array(Float64,0)
velo=NesterovInit(net) # Nesterov velocity
for i=1:nupdates
    ADforward!(net)
    ADbackward!(net)
    push!(error,extract(net.value[net.FunctionNode])[1])
    printover("iteration $i: training loss = $(error[i])")
    for par in ParsToUpdate
        #GradientDescentUpdate!(net.value[par],net.gradient[par],LearningRate)
        NesterovGradientDescentUpdate!(net.value[par],net.gradient[par],velo[par],LearningRate,i)
    end
end
@gpu CUDArt.device_synchronize()
println()
    toc()

#println("\ntrue weights   ",round(100*truew')/100)
#println("learned weights",round(100*extract(net.value[w])')/100)

if PlotResults
    figure(1);plot(truew);title("true weights")
    figure(2);plot(extract(net.value[w])); title("learned weights")
end

end


# # SIMP example: Point Load Cantilever
#-
#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`simp.ipynb`](@__NBVIEWER_ROOT_URL__/examples/simp.ipynb)
#-
# ## Commented Program
#
# What follows is a program spliced with comments.
#md # The full program, without comments, can be found in the next [section](@ref simp-plain-program).

using TopOpt

# ### Define the problem
E = 1.0 # Young’s modulus
v = 0.3 # Poisson’s ratio
f = 1.0; # downward force

nels = (30, 10, 10)
problem = PointLoadCantilever(Val{:Linear}, nels, (1.0, 1.0, 1.0), E, v, f);

# See also the detailed API of `PointLoadCantilever`:
#md # ```@docs
#md # TopOpt.TopOptProblems.PointLoadCantilever
#md # ```

# ### Parameter settings
V = 0.3 # volume fraction
xmin = 1e-6 # minimum density
rmin = 2.0; # density filter radius

# ### Define a finite element solver
penalty = TopOpt.PowerPenalty(3.0)
solver = FEASolver(Direct, problem; xmin=xmin, penalty=penalty)

# ### Define compliance objective
comp = TopOpt.Compliance(problem, solver)
filter = DensityFilter(solver; rmin=rmin)
obj = x -> comp(filter(x))

# ### Define volume constraint
volfrac = TopOpt.Volume(problem, solver)
constr = x -> volfrac(filter(x)) - V

# You can enable the iteration printouts with `Nonconvex.NonconvexCore.show_residuals[] = true`

# ### Define subproblem optimizer
mma_options =
    options = MMAOptions(;
        maxiter=3000, tol=Nonconvex.Tolerance(; x=1e-3, f=1e-3, kkt=0.001)
    )
convcriteria = Nonconvex.KKTCriteria()
x0 = fill(V, length(solver.vars))
optimizer = Optimizer(
    obj, constr, x0, MMA87(); options=mma_options, convcriteria=convcriteria
)

# ### Define SIMP optimizer
simp = SIMP(optimizer, solver, penalty.p);

# ### Solve
result = simp(x0);

@show result.convstate
@show optimizer.workspace.iter
@show result.objval

# ### (Optional) Visualize the result using Makie.jl
# Need to run `using Pkg; Pkg.add(Makie)` first
# ```julia
# using TopOpt.TopOptProblems.Visualization: visualize
# fig = visualize(problem; topology = result.topology,
#     problem; topology = result.topology, default_exagg_scale = 0.07,
#     scale_range = 10.0, vector_linewidth = 3, vector_arrowsize = 0.5,
# )
# Makie.display(fig)
# ```

# or convert it to a Mesh
# Need to run `using Pkg; Pkg.add(GeometryBasics)` first
# ```julia
# import Makie, GeometryBasics
# result_mesh = GeometryBasics.Mesh(problem, result.topology);
# Makie.mesh(result_mesh)
# ```

#md # ## [Plain Program](@id simp-plain-program)
#md #
#md # Below follows a version of the program without any comments.
#md # The file is also available here: [simp.jl](simp.jl)
#md #
#md # ```julia
#md # @__CODE__
#md # ```


export bissect_Cub
function bissect_Cub(h :: AbstractNLPModel;
                    t1 :: Float64 = h.meta.x0[1],
                    t2 :: Float64 = 100.0,
                    tol :: Float64 = 1e-7,
                    maxiter :: Int = 50,
                    verbose :: Bool = false)

    (length(h.meta.x0) > 1) && warn("Not a 1-D problem ")
  (tₐ, tᵦ) = trouve_intervalle(h, t1, 2.5)

  γ = 0.8; t = tᵦ; tₚ = tₐ; tqnp = tₐ
  hₖ = 0; hkm1 = 0; gkm1 = 0; hₚ = 0
  iter = 0

  hₖ = obj(h, [t])[1]; gₖ = grad(h, [t])[1]
  hkm1 = obj(h, [tqnp])[1]; gkm1 = grad(h, [tqnp])[1]

  verbose &&
    @printf(" iter        tₚ        t         dN         ")
  verbose &&
    @printf("gₖ          gplus        \n")
  verbose &&
    @printf(" %7.2e %7.2e %7.2e  %7.2e  %7.2e  %7.2e\n",
             iter, tₚ, t, 0.0, gₖ,  0.0)

  while ((abs(gₖ) > tol) && (iter < maxiter)) || (iter == 0)
    s = t - tqnp; y = gₖ - gkm1; α = -s
    z = gₖ + gkm1 + 3 * (hₖ - hkm1) / α
    discr = z^2 - gₖ * gkm1; denom = gₖ + gkm1 + 2 * z
    if (discr > 0.0) && (abs(denom) > eps(Float64))
      #si on peut on utilise l'interpolation cubique
      w = sqrt(discr)
      dN = -s * (gₖ + z + sign(α) * w) / (denom)
    else #on se rabat sur une étape de sécante
      dN = -gₖ * s / y
    end

    if ((tₚ - t) * dN > 0.0) & (dN / (tₚ - t) < γ)
      tplus = t + dN
      hplus = obj(h, [tplus])[1]
      gplus = grad(h, [tplus])[1]
      verbose && println("N")
    else
      tplus = (t + tₚ) / 2
      hplus = obj(h, [tplus])[1]
      gplus = grad(h, [tplus])[1]
      verbose && println("B")
    end

    if t > tₚ
      if gplus < 0.0
        tₚ = t; tqnp = t; t = tplus
      else
        tqnp = t; t = tplus
      end
    else
      if gplus > 0.0
        tₚ = t; tqnp = t; t = tplus
      else
        tqnp = t; t = tplus
      end
    end

    #mise à jour des valeurs
    hkm1 = hₖ; gkm1 = gₖ; hₖ = hplus; gₖ = gplus
    iter = iter + 1
    verbose &&
        @printf(" %7.2e %7.2e %7.2e  %7.2e  %7.2e  %7.2e\n",
                iter, tₚ, t, dN, gₖ, gplus)
  end

  topt = t
  gopt = grad(h, [topt])[1]
  status = :NotSolved
  (abs(gopt) < tol) && (status = :Optimal)
  (iter >= maxiter) && (status = :Tired)
  tired = iter > maxiter
  optimal = abs(gopt) < tol
  return (topt, obj(h, [topt])[1], norm(gopt, Inf), iter, optimal, tired, status, h.counters.neval_obj, h.counters.neval_grad, h.counters.neval_hess)

end # function


for i in 1:2:9
    println(i)
end #> 1 3 5 7 9

typeof(1:1000) #> UnitRange{Int64}

a = split("A,B,C,D",",")
typeof(a) #> Array{SubString{String},1}
show(a) #> SubString{String}["A","B","C","D"]

arr = [100, 25, 37]
arra = Any[100, 25, "ABC"] #> element type Any
arr[1] #> 100
arr[end] #> 37
# arr[6]
#> ERROR: BoundsError: attempt to access 3-element Array{Int64,1} at index [6]
println(eltype(arr)) #> Int64
println(length(arr)) #> 3
println(ndims(arr)) #> 1
println(size(arr,1)) #> 3
println(size(arr)) #> (3,)

arr2 = Array{Int64}(undef, 5)
show(arr2) #> [1, 2, 3, 4, 5]
sizehint!(arr2, 10^5)

arr3 = Float64[] #> 0-element Array{Float64,1}
push!(arr3, 1.0) #> 1-element Array{Float64,1}

arr4 = collect(1:7) #> 7-element Array{Int64,1}
show(arr4) #> [1, 2, 3, 4, 5, 6, 7]

# for in loop and changing the array:
da = [1,2,3,4,5]
for n in da
    n *= 2
end
da #> 5-element Array{Int64,1}: 1 2 3 4 5
for i in 1:length(da)
    da[i] *= 2
end
da #> 5-element Array{Int64,1}: 2 4 6 8 10

arr4 = [1,2,3,4,5,6,7] #> 7-element Array{Int64,1}: [1,2,3,4,5,6,7]
join(arr4, ", ") #> "1, 2, 3, 4, 5, 6, 7"
arr4[1:3] # => [1, 2, 3]
arr4[4:end] # => [4, 5, 6, 7]

arr = [1,2,3,4,5]
arr[2:4] = [8,9,10]
println(arr) #> 1 8 9 10 5

zeros(5)
ones(4)
ones(3, 2)
eqa = range(0, step=10, length=5) #> 0:10:40
println()
show(eqa) #> 0:10:40
fill!(arr, 42) #> [42, 42, 42, 42, 42]
println()
println(Array{Any}(undef, 4)) #> Any[#undef,#undef,#undef,#undef]

v1 = rand(Int32, 5)
println(v1) #> Int32[1735276173,972339632,1303377282,1493859467,-788555652]

b = collect(1:7)
c = [100,200,300]
append!(b,c) # Now b is [1, 2, 3, 4, 5, 6, 7, 100, 200, 300]
pop!(b) #> 300, b is now [1, 2, 3, 4, 5, 6, 7, 100, 200]
push!(b, 42) # b is now [1, 2, 3, 4, 5, 6, 7, 100, 200, 42]
popfirst!(b) #> 1, b is now [2, 3, 4, 5, 6, 7, 100, 200, 42]
pushfirst!(b, 42) # b is now [1, 2, 3, 4, 5, 6, 7, 100, 200, 42]
splice!(b,8) #> 100, b is now [42, 2, 3, 4, 5, 6, 7, 200, 42]
in(42, b) #> true
in(43, b) #> false
println()
println("sorting:")
sort(b) #> [2,3,4,5,6,7,42,42,200], b is not changed
println(b) #> [42,2,3,4,5,6,7,200,42]
sort!(b) #> b is now changed to [2,3,4,5,6,7,42,42,200]
println(b) #> [2,3,4,5,6,7,42,42,200]
println()
arr = [1, 5, 3]
# looping
for e in arr
    print("$e ")
end # prints 1 5 3

arr = [1, 2, 3]
arr .+ 2 #> [3, 7, 5]
arr * 2 #> [2, 10, 6]
a1 = [1, 2, 3]
a2 = [4, 5, 6]
a1 .* a2 #> [4, 10, 18]

using LinearAlgebra
LinearAlgebra.dot(a1, a2) #> 32
sum(a1 .* a2)

repeat([1, 2, 3], inner = [2]) #> [1,1,2,2,3,3]

a = [1,2,4,6]
a1 = a
show(a1)
a[4] = 0
show(a) #> [1,2,4,0]
show(a1) #> [1,2,4,0]
b = copy(a)
b = deepcopy(a)

a = [1,2,3]
function change_array(arr)
	arr[2] = 25
end
change_array(a)
println(a) #>[ 1, 25, 3]

# the splat operator:
arr = ['a', 'b', 'c']
show(join(arr)) #> "abc"
show(string(arr)) #> "['a', 'b', 'c']"
show(string(arr...)) #> "abc"


instructions = string.(split(strip(read("input.txt", String)), "\n"))

struct Bitmask
    andmask::UInt64
    ormask::UInt64
end
Bitmask(raw::String) = Bitmask([parse(Int, replace(raw[8:end], "X" => l), base=2) for l in "10"]...)
mask(number::UInt64, bitmask::Bitmask) = (number & bitmask.andmask) | bitmask.ormask

struct Mutator
    address::UInt64
    value::UInt64
end
Mutator(raw::String) = Mutator(parse.(Int, match(r"\[(\d+)\] = (\d+)", raw).captures)...)

function run(instructions::Array{String})::UInt64
    program = [(startswith(instruction, "mask") ? Bitmask : Mutator)(instruction) for instruction in instructions]
    memory = zeros(UInt64, maximum(m -> m.value, filter(x -> isa(x, Mutator), program)))
    bitmask = nothing

    for instruction in program
        if instruction isa Bitmask
            bitmask = instruction
        elseif instruction isa Mutator
            memory[instruction.address] = mask(instruction.value, bitmask)
        end
    end

    sum(memory)
end

run(program)

function maskedaddresses(address::String, mask::String)
    result = []
    for maskno in 0:(1 << count("X", mask))
        b = string(maskno, base=2)
        push!(result, parse(Int64, address) | b)
    end
    result
end

using SparseArrays

function otherrun(instructions::Array{String})::UInt64
    memory = spzeros(UInt64, 2^36)
    bitmask = nothing

    for instruction in instructions
        if startswith("mask", instruction)
            bitmask = instruction[8:end]
        else
            address, value = parse.(Int, match(r"\[(\d+)\] = (\d+)", raw).captures)
            for maskedaddress in maskedaddresses(address, bitmask)
                memory[maskedaddress] = value
            end
        end
    end

    sum(memory)
end


using OptiMo
using LinearAlgebra: norm
using Printf

"""
    restore( prob, [x0, with_indicator, verbose, ...] )
"""
function restore(
    prob::TP;
    x0::Tx = prob.meta.x0,
    with_indicator::Bool = false,
    verbose::Bool = false,
    max_out_iter::Int = 100,
    max_sub_iter::Int = 1000,
    ifactor::Real = 0.1,
    wfactor::Real = 0.1,
    wmin::Real = 1e-16,
    tol_cviol::Real = 1e-8,
    tol_inf_cviol::Real = 1e-5,
    tol_infst::Real = 1e-10,
    tol_optim::Real = 1e-14,
    loss::Symbol = :l2sq,
    huber_rho::Real = 1.0,
    huber_mu::Real = 1.0,
) where {TP<:AbstractOptiModel,Tx<:AbstractVector}

    @assert 0 < tol_optim
    @assert 0 < tol_cviol
    @assert 0 < tol_inf_cviol
    @assert 0 < tol_infst
    @assert 0 < wmin
    @assert 0 < wfactor < 1
    @assert 0 < ifactor
    @assert 0 < max_sub_iter
    @assert 0 < max_out_iter

    solver = Bazinga.ZEROFPR( verbose = false,
                              max_iter = max_sub_iter,
                              tol_optim = tol_optim)

    # initial point = base point
    x = copy(x0)
    R = eltype(x)
    # allocation
    s = similar(x)
    infst = -1.0
    if with_indicator
        bigeps = 1e-10
    end
    # diagonal scaling
    d0 = min.(1.0, 1.0 ./ abs.(x0))
    # initial infeasibility
    fprob = FeasOptiModel(
        prob,
        x0 = x,
        with_indicator = with_indicator,
        loss = loss,
        huber_rho = huber_rho,
        huber_mu = huber_mu,
    )
    infea = infeasibility(fprob, x)
    cviol = cviolation(fprob, x)
    # initial weight estimate
    w = max(wmin, ifactor * infea)

    if verbose
        @printf(
            "%4s | %7s | %4s  %7s | %7s  %7s | %7s | \n",
            "o.it",
            "w",
            "i.it",
            "optim",
            "infea",
            "pdist",
            "cviol"
        )
    end

    niter = 0
    while true

        fprob = FeasOptiModel(
            prob,
            x0 = x,
            with_indicator = with_indicator,
            xprox = x0,
            wprox = w,
            dprox = d0,
            loss = loss,
            huber_rho = huber_rho,
            huber_mu = huber_mu,
        )
        out = solver(fprob)
        copyto!(x, out.x)
        infea = infeasibility(fprob, x)
        cviol = cviolation(fprob, x)
        w = max(wmin, min(wfactor * w, ifactor * infea))

        # infeasibility stationarity
        infeasibilitygrad!(fprob, x, s)
        if with_indicator
            s .= x .- bigeps .* s
            prox!(fprob, s, bigeps, s)
            s .= (x .- s) ./ bigeps
        end
        infst = norm(s, Inf)

        if verbose
            proxd = proxdistance(fprob, x)
            @printf(
                "%4d | %.1e | %4d  %.1e | %.1e  %.1e | %.1e | %.1e \n",
                niter,
                w,
                out.iterations,
                out.optimality,
                infea,
                proxd,
                cviol,
                infst
            )
        end

        niter += 1
        if cviol <= tol_cviol
            if verbose
                @printf "cviolation satisfied"
            end
            break
        elseif cviol >= tol_inf_cviol && infst <= tol_infst
            if verbose
                @printf "local infeasibility detected"
            end
            break
        elseif w <= wmin && out.iterations == 1
            if verbose
                @printf "tired restoration"
            end
            break
        elseif niter >= max_out_iter
            if verbose
                @printf "max iterations reached"
            end
            break
        end

    end
    return x, cviol, infst
end