using SymbolicUtils: PolyForm, Term, symtype
using Test, SymbolicUtils

include("utils.jl")

@testset "div and polyform" begin
    @syms x y z
    @test repr(PolyForm(x-y)) == "x - y"
    @test repr(x/y*x/z) == "(x^2) / (y*z)"
    @test repr(simplify_fractions(((x-y+z)*(x+4z+1)) /
                                  (y*(2x - 3y + 3z) +
                                   x*(x + z)))) == repr(simplify_fractions((1 + x + 4z) / (x + 3.0y)))
    @test simplify_fractions(x/(x+3) + 3/(x+3)) == 1
    @test repr(simplify(simplify_fractions(cos(x)/sin(x) + sin(x)/cos(x)))) == "2 / sin(2x)"
end

@testset "expand" begin
    @syms a b c d

    @test expand(term(*, 0, a)) == 0
    @test expand(a * (b + -1 * c) + -1 * (b * a + -1 * c * a)) == 0
    @eqtest simplify(expand(sin((a+b)^2)^2)) == simplify(sin(a^2+2*(b*a)+b^2)^2)
    @test simplify(expand(sin((a+b)^2)^2 + cos((a+b)^2)^2)) == 1
    @syms x1::Real f(::Real)::Real

    # issue 193
    #@test isequal(expand(f(2.0 + x1)), f(2.0 + x1))
    @test symtype(expand(f(x1 + 2.0))) == Real

    # cleanup rules
    # I don't think we need this anymore?
   #@test expand(Term{Number}(identity, 0)) == 0
   #@test expand(Term{Number}(one, 0)) == 1
   #@test expand(Term{Number}(zero, 0)) == 0
   #@test expand(identity(a * b) - b * a) == 0
    @test expand(a * b - b * a) == 0

    @syms A::Vector{Real}
    # test that the following works
    expand(Term{Real}(getindex, [A, 3]) - 3)
end

@testset "simplify_fractions with quick-cancel" begin
    @syms x y
    @test simplify_fractions(x/2x) == 1//2
    @test simplify_fractions(2x//x) == 2
    @eqtest simplify_fractions((x+y) * (x-y) / (x+y)) == (x-y)
    @eqtest simplify_fractions(x^3 * y / x) == y*x^2
    @eqtest simplify_fractions(2x^3 * y / x) == 2y*x^2
    @eqtest simplify_fractions(x / (3(x^3)*y)) == simplify_fractions(1/(3*(y*x^2)))
    @eqtest simplify_fractions(2x / (3(x^3)*y)) == simplify_fractions(2/(3*(y*x^2)))
    @eqtest simplify_fractions(x^2 / (3(x^3)*y)) == simplify_fractions(1/(3*(y*x)))
    @eqtest simplify_fractions((3(x^3)*y) / x^2) == simplify_fractions(3*(y*x))
    @eqtest simplify_fractions(x^2/x^4) == (1/x^2)
    @eqtest simplify_fractions(x^2/x^3) == 1/x
    @eqtest simplify_fractions(x^3/x^2) == x
    @eqtest simplify_fractions(x^2/x^2) == 1
    @eqtest simplify_fractions(3x^2/x^3) == 3/x
    @eqtest simplify_fractions(3*(x^2)*(y^3)/(3*(x^3)*(y^2))) == y/x
    @eqtest simplify_fractions(3*(x^x)/x*y) == 3*(x^x)/x*y

    ##404#issuecomment-939404030
    a = 1 / (x - (2//1)) + ((-5//1) - x) / ((x - (2//1))^2)
    @test isequal(simplify_fractions(a), -7/(x-2)^2)
end

@testset "isone iszero" begin
    @syms a b c d e f g h i
    x = (f + ((((g*(c^2)*(e^2)) / d - e*h*(c^2)) / b + (-c*e*f*g) / d + c*e*i) /
              (i + ((c*e*g) / d - c*h) / b + (-f*g) / d) - c*e) / b +
         ((g*(f^2)) / d + ((-c*e*f*g) / d + c*f*h) / b - f*i) /
         (i + ((c*e*g) / d - c*h) / b + (-f*g) / d)) / d

    o = (d + (e*((c*(g + (-d*g) / d)) / (i + (-c*(h + (-e*g) / d)) / b + (-f*g) / d))) / b + (-f*(g + (-d*g) / d)) / (i + (-c*(h + (-e*g) / d)) / b + (-f*g) / d)) / d
    @test SymbolicUtils.fraction_iszero(x)
    @test !SymbolicUtils.fraction_isone(x)
    @test SymbolicUtils.fraction_isone(o)
end


# ---
# title: 552. Student Attendance Record II
# id: problem552
# author: Indigo
# date: 2021-08-18
# difficulty: Hard
# categories: Dynamic Programming
# link: <https://leetcode.com/problems/student-attendance-record-ii/description/>
# hidden: true
# ---
#
# Given a positive integer **n** , return the number of all possible attendance
# records with length n, which will be regarded as rewardable. The answer may be
# very large, return it after mod 109 \+ 7.
#
# A student attendance record is a string that only contains the following three
# characters:
#
#   1. **'A'** : Absent.
#   2. **'L'** : Late.
#   3. **'P'** : Present.
#
# A record is regarded as rewardable if it doesn't contain **more than one 'A'
# (absent)** or **more than two continuous 'L' (late)**.
#
# **Example 1:**
#
#
#
#     Input: n = 2
#     Output: 8
#     Explanation:
#     There are 8 records with length 2 will be regarded as rewardable:
#     "PP" , "AP", "PA", "LP", "PL", "AL", "LA", "LL"
#     Only "AA" won't be regarded as rewardable owing to more than one absent times.
#
#
# **Note:** The value of **n** won't exceed 100,000.
#
#
## @lc code=start
using LeetCode

function check_record(n::Int)
    P = [1 1 0 1 0 0
        1 0 1 1 0 0
        1 0 0 1 0 0
        0 0 0 1 1 0
        0 0 0 1 0 1
        0 0 0 1 0 0]
    res = mat_fast_mul(P, n, Int(1e9 + 7))
    sum(@view res[1, :]) % Int(1e9 + 7)
end
## @lc code=end


# Load the CytofResearch environment
import Pkg; Pkg.activate("../../")

println("Loading libraries. Expect this to take a minute or two ...")
using CytofResearch
using Random
using BSON
println("Finished loading libraries ...")

flipbits(x; prob) = [prob > rand() ? 1 - xi : xi for xi in x]

function createSimData(;seed::Int, nfac::Int,
                       K=5, L=Dict(0=>3, 1=>3), J=20, eps=zeros(3))

  N = [8, 1, 2] * nfac
  I = length(N)

  # Set rng seed for reproducibility.
  Random.seed!(seed)

  # Initialize true Z matrix.
  Z = zeros(Int, J, K)

  # Regenerate Z unitl it is valid.
  while !CytofResearch.Model.isValidZ(Z)
    Z = CytofResearch.Model.genZ(J, K, .5)
    Z[:, 2] .= flipbits(Z[:, 1], prob=.1)
  end

  # Create true mu*
  mus = Dict(0 => -[1.0, 2.3, 3.5],
             1 => +[1.0, 2.0, 3.0])

  # Create true a_W
  a_W = rand(K) * 10

  # Create true a_eta
  a_eta = Dict(z => rand(L[z]) * 10 for z in 0:1)

  # Make simulated data
  simdat = CytofResearch.Model.genData(J=J, N=N, K=K, L=L, Z=Z,
                                beta=[-9.2, -2.3],
                                sig2=[0.2, 0.1, 0.3],
                                mus=mus,
                                a_W=a_W,
                                a_eta=a_eta,
                                sortLambda=false, propMissingScale=0.7,
                                eps=eps)
  return simdat
end

### Main ###
settings = [Dict(:nfac => 500, :K => 5, :seed => 90),
            Dict(:nfac => 5000, :K => 10, :seed => 1)]

for setting in settings
  nfac = setting[:nfac]
  seed = setting[:seed]
  K = setting[:K]

  println("Creating simulated data with: $(setting)")
  simdat = createSimData(seed=seed, nfac=nfac,
                         K=K, L=Dict(0=>3, 1=>3), J=20, eps=fill(.005, 3))

  # Save simulated data
  path_to_simdat = "../data/simdat-nfac$(nfac).bson"
  BSON.bson(path_to_simdat, simdat)
end

# NOTE: The ".bson" file is cross-platform, but requires the same BSON version.
#       To load the simulated data:
# simdat_small = BSON.load("../data/simdat-nfac500.bson")
# simdat_large = BSON.load("../data/simdat-nfac5000.bson")


# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module SpigotPi

export Pi, V000796

"""
Computes the first n decimal digits of Pi, uses a variant of the spigot algorithm valid as long as the number of digits <= 54900. Based on ideas of A. Sale (1968). Algorithm due to D. Saada (1988) and S. Rabinowitz (1991). Proof due to [Rabinowitz and S. Wagon](https://www.maa.org/sites/default/files/pdf/pubs/amm_supplements/Monthly_Reference_12.pdf) (1995).

* Pi, V000796
"""
const ModuleSpigotPi = ""

"""

Return the decimal expansion of π up to ``n`` digits as a string.
Note the feature that the precision is increased only in steps of 4, i.e. only if ``n`` is divisible by 4. If ``n`` is not divisible by 4 then the number of digits in the output is rounded down to the largest number divisible by 4 smaller than ``n``.
# Examples
```julia-repl
julia> Pi(8)
31415926
julia> Pi(10)
31415926
julia> Pi(12)
314159265358
```
"""
function Pi(digits::Int)
    digits > 54900 && throw(ValueError("digits must be <= 54900"))

    buf = IOBuffer()
    zero = [ "0", "00", "000" ]
    c = (digits >> 2 + 1) * 14
    a = fill(20000000, c)
    d, f = 0, 10000

    while (b = c -= 14) > 0
        d = e = d % f
        while (b -= 1) > 0
            d = d * b + a[b]
            g = (b << 1) - 1
            a[b] = (d % g) * f
            d = div(d, g)
        end
        r = e + div(d, f)
        r < 1000 && print(buf, zero[r > 99 ? 1 : r > 9 ? 2 : 3])
        print(buf, r)
    end
    return String(take!(buf))
end

"""

Return the decimal expansion of π up to ``n`` digits as a string. (Note what the docstring of 'Pi'' says about the number of digits in the output.)
# Examples
```julia-repl
julia π(100)[100]
7
```
"""
π(digits) = Pi(digits)

"""

Return the decimal expansion of π up to ``n`` digits as a string. (Note what the docstring of 'Pi'' says about the number of digits in the output.)
"""
V000796(digits) = Pi(digits)

#START-TEST-########################################################

using Test

function test()
    @testset "Pi" begin
        @test Pi(12) == "314159265358"
        @test Pi(100)[100] == '7'
    end
end

function demo()
    π(1600) |> println
end

function perf()
    @time Pi(1600)
end

function main()
    test()
    demo()
    perf()
end

main()

end # module


abstract AbstractStateSpaceModel
typealias AbstractSSM AbstractStateSpaceModel
abstract AbstractGaussianSSM <: AbstractStateSpaceModel

# LinearGaussianSSM     LinearKalmanFilter
#                       NonlinearKalmanFilter
#                       NonlinearFilter
#
# NonlinearGaussianSSM  NonlinearKalmanFilter
#                       NonlinearFilter
#
# NonlinearSSM          NonlinearFilter

###########################################################################
# Linear Gaussian
###########################################################################

immutable LinearGaussianSSM <: AbstractGaussianSSM
    # Process transition matrix, control matrix, and noise covariance
    F::Function
    B::Function
    V::Function

    # Observation matrix and noise covariance
    G::Function
    W::Function

    # Model dimensions
    nx::Int
    ny::Int
    nu::Int

    function LinearGaussianSSM(F::Function, B::Function, V::Function,
                                G::Function, W::Function)
        @assert ispossemidef(V(1))
        @assert ispossemidef(W(1))

        nx, ny, nu = confirm_matrix_sizes(F(1), B(1), V(1), G(1), W(1))
        new(F, B, V, G, W, nx, ny, nu)
    end
end

# Time-dependent constructor
LinearGaussianSSM{T}(F::Function, V::Function, G::Function, W::Function;
                          B::Function=_->zeros(T, size(V(1), 1), 1)) =
	  LinearGaussianSSM(F, B, V, G, W)

# Time-independent constructor
LinearGaussianSSM{T}(F::Matrix{T}, V::Matrix{T}, G::Matrix{T}, W::Matrix{T};
                          B::Matrix{T}=zeros(T, size(F, 1), 1)) =
	  LinearGaussianSSM(_->F, _->B, _->V, _->G, _->W)

function show(io::IO, mod::LinearGaussianSSM)
    dx, dy = mod.nx, mod.ny
    println("LinearGaussianSSM, $dx-D process x $dy-D observations")
    println("Process evolution matrix F:")
    show(mod.F(1))
    println("\n\nGontrol input matrix B:")
    show(mod.B(1))
    println("\n\nProcess error covariance V:")
    show(mod.V(1))
    println("\n\nObservation matrix G:")
    show(mod.G(1))
    println("\n\nObseration error covariance W:")
    show(mod.W(1))
end

## Core methods
process_matrix(m::LinearGaussianSSM, state::Vector, t::Real=0.0) = m.F(t)
process_matrix(m::LinearGaussianSSM, state::AbstractMvNormal, t::Real=0.0) = m.F(t)

observation_matrix(m::LinearGaussianSSM, state::Vector, t::Real=0.0) = m.G(t)
observation_matrix(m::LinearGaussianSSM, state::AbstractMvNormal, t::Real=0.0) = m.G(t)

control_input(m::LinearGaussianSSM, u, t::Real=0.0) = m.B(t) * u

###########################################################################
# Nonlinear Gaussian
###########################################################################

immutable NonlinearGaussianSSM <: AbstractGaussianSSM
    # Process transition function, jacobian, and noise covariance matrix
    f::Function
    V::Function

    # Control input function
    b::Function

    # Observation function,
    g::Function # actual observation function
    W::Function

    nx::Int
    ny::Int
    nu::Int

    function NonlinearGaussianSSM(f::Function, V::Function,
            b::Function, g::Function, W::Function, nx::Int, ny::Int, nu::Int)
        @assert ispossemidef(V(1))
        @assert ispossemidef(W(1))
        @assert (nx, nx) == size(V(1))
        @assert (ny, ny) == size(W(1))

        new(f, V, b, g, W, nx, ny, nu)
    end
end

function NonlinearGaussianSSM{T}(f::Function, V::Matrix{T}, g::Function,
        W::Matrix{T}; b::Function=_ -> 0, nu::Int=0)
    nx = size(V, 1)
    ny = size(W, 1)

    return NonlinearGaussianSSM(f, _-> V, b, g, _-> W, nx, ny, nu)
end

## Core methods
process_matrix(m::NonlinearGaussianSSM, x::Vector, t::Real=0.0) = ForwardDiff.jacobian(m.f, x)
function process_matrix(m::NonlinearGaussianSSM, x::AbstractMvNormal, t::Real=0.0)
    return process_matrix(m, mean(x), t)
end

observation_matrix(m::NonlinearGaussianSSM, x::Vector, t::Real=0.0) = ForwardDiff.jacobian(m.g, x)
function observation_matrix(m::NonlinearGaussianSSM, x::AbstractMvNormal, t::Real=0.0)
    return process_matrix(m, mean(x), t)
end

control_input(m::NonlinearGaussianSSM, u, t::Real=0.0) = m.b(u)


module MatrixCompletion

import Optim
import NLopt

export complete, cross_complete

function nuclear_norm{T <: AbstractFloat}(m::Matrix{T})
    f = svdfact(m)
    sum(f[:S])
end

function d_nuclear_norm{T <: AbstractFloat}(m::Matrix{T})
    # Alt (http://math.stackexchange.com/users/123685/alt), Derivative
    # of nuclear norm, URL (version: 2015-12-29):
    # http://math.stackexchange.com/q/701104
    f = svdfact(m)
    f[:U] * f[:Vt]
end

type CompletionResults{T <: AbstractFloat}
    method::Symbol
    initial::Matrix{T}
    minimum::Matrix{T}
    f_minimum::Float64
    xtol::T
    ftol::T
end

function complete{T <: AbstractFloat}(m::Matrix{T}, lambda::Real;
                                      method = :LD_LBFGS, tol = eps(T),
                                      x0 = 0.5 * m + mean(m),
                                      maxeval = Inf, maxtime = Inf)
    sz = size(m)

    mones = zeros(sz)
    for i = eachindex(m)
        if m[i] != 0.
            mones[i] = 1.
        end
    end

    function f(vx::Vector, grad::Vector)
        x = reshape(vx, sz)

        fn = 0.
        for i = eachindex(m)
            s = mones[i] * (x[i] - m[i])
            fn += s * s
        end
        fn = sqrt(fn)

        f = svdfact(x)
        sum(f[:S]) + lambda * fn
    end

    function g!(vx::Vector, grad::Vector)
        x = reshape(vx, sz)
        z = reshape(grad, sz)

        fn = 0.
        for i = eachindex(m)
            s = mones[i] * (x[i] - m[i])
            z[i] = s
            fn += s * s
        end
        fn = sqrt(fn)

        f = svdfact(x)
        BLAS.gemm!('N', 'N', 1., f[:U], f[:Vt], lambda / fn, z)
    end

    function fg!(vx::Vector, grad::Vector)
        x = reshape(vx, sz)
        z = reshape(grad, sz)

        # z = mones .* (x - m)
        # fn = vecnorm(z)
        # grad[:] = vec(d_nuclear_norm(x) + (lambda/fn) * z)
        # nuclear_norm(x) + lambda * fn

        fn = 0.
        for i = eachindex(m)
            s = mones[i] * (x[i] - m[i])
            z[i] = s
            fn += s * s
        end
        fn = sqrt(fn)

        f = svdfact(x)
        BLAS.gemm!('N', 'N', 1., f[:U], f[:Vt], lambda / fn, z)
        x = sum(f[:S]) + lambda * fn

        x
    end

    d4 = Optim.DifferentiableFunction(f, g!, fg!)
    res = Optim.optimize(d4, vec(x0))
    optf = res.f_minimum
    optx = res.minimum

    # opt = NLopt.Opt(method, length(m))
    # NLopt.min_objective!(opt, fg!)
    # NLopt.xtol_rel!(opt, tol)
    # NLopt.ftol_rel!(opt, sqrt(tol))

    # if maxeval < Inf
    #     NLopt.maxeval!(opt, maxeval)
    # end
    # if maxtime < Inf
    #     NLopt.maxtime!(opt, maxeval)
    # end

    # (optf,optx,ret) = NLopt.optimize(opt, vec(x0))

    CompletionResults(method, m, reshape(optx, sz), optf, tol, sqrt(tol))
end

function roundn(v, n::Int)
    round(v * 10. ^ n) * 10. ^ (-n)
end

function cross_complete(m; n=5, js=1:n,
                        lambda_min = 0, lambda_max = 20,
                        trace = false)
    subset = rand(1:n, size(m))

    insample = map(j -> subset .!= j, js)
    outsample = map(j -> ! insample[j], js)

    m_in  = map(j -> m .* insample[j], js)

    if trace
        cnt = countnz(m)
        for j = js
            cnt_in = countnz(insample[j] .* m)
            cnt_out = countnz(outsample[j] .* m)

            println("j = ", j)
            println("insample count: ", cnt_in, " / ", cnt, " = ", roundn(cnt_in / cnt, 2))
            println("outsample count: ", cnt_out, " / ", cnt, " = ", roundn(cnt_out / cnt, 2))
            println()
        end
    end

    function x_utility(lambda)
        function j_utility(j)
            # @time r = complete(m_in, lambda, x0=res.minimum)
            # r = complete(m_in[j], lambda, x0=res[j].minimum)
            @time r = complete(m_in[j], lambda)
            u = vecnorm((m - r.minimum) .* outsample[j])

            x = reshape(r.minimum, size(m))
            f = svdfact(x)
            println(j, " : ", @sprintf("%.2f", lambda), " : ", round(Int, f[:S]))

            u
        end

        us = map(j_utility, js)
        u = mean(us)

        if trace
            println(@sprintf("%.2f", u), " ", roundn(us, 2))
            println()
        end

        u
    end

    Optim.optimize(x_utility, lambda_min, lambda_max, rel_tol = 1e-4)
end

end # module


@testset "synthesis" begin
    N  = 1024; t = 1:N; scale=100
    X  = sin.(2π*t/scale)
    wav = CauchyWavelet(10)
    wt = ContinuousWaveletTransform(wav, LinRange(2, 200, 64))
    wX = transform(wt, X)
    Y  = synthesis(wX)
    edge = Int32(ceil(scale*ContinuousWavelet.cutoff_time(wav)))
    @test maximum(abs.(X[edge:end-edge]-Y[edge:end-edge])) < 1e-1
end


"""
# Hearthstone

Functions to calculate the probabilities of different events in [Hearthstone: Heroes of Warcraft](http://us.battle.net/hearthstone/en/).

"""
module Hearthstone

export card_in_opening_hand_no_coin,
       card_in_opening_hand_coin,
       card_in_next_N_turns,
       missile_lethal
"""
    hypergeom_pmf(N::Integer, K::Integer, n::Integer, k::Integer)

Calculate the hypergeometric probability mass function.

# Arguments
* `N::Integer`: the population size.
* `K::Integer`: the number of success states in the population.
* `n::Integer`: the number of draws.
* `k::Integer`: the number of observed successes.
"""
function hypergeom_pmf(N::Integer, K::Integer, n::Integer, k::Integer)
  @assert(k >= max(0, n+K-N))
  @assert(k <= min(n, K))
  binomial(K, k) * binomial(N-K, n-k) / binomial(N, n)
end

"""
    hypergeom_cmf(N::Integer, K::Integer, n::Integer, ks::UnitRange{Integer})

Calculate the hypergeometric cumulative mass function.

# Arguments
* `N::Integer`: the population size.
* `K::Integer`: the number of success states in the population.
* `n::Integer`: the number of draws.
* `ks::UnitRange{Integer}`: the range of observed successes to accumulate over.
"""
function hypergeom_cmf{T<:Integer}(N::Integer, K::Integer, n::Integer, ks::UnitRange{T})
  f(a,b) = a + binomial(K, b) * binomial(N-K, n-b)
  num =  foldl(f, 0, ks)
  num / binomial(N,n)
end

"""
    card_in_opening_hand_coin(cardcount::Integer)

Returns the probability to obtain at least one card of cardcount total in the opening hand of the player with the coin.

# Arguments
* `cardcount::Integer`: the total number of cards of interest in the deck.
"""
function card_in_opening_hand_coin(cardcount::Integer)
  """
  Define events:
  A = Card not drawn in opening hand
  B = Card not drawn in initial 4 cards
  C = Card not drawn after replacing initial 4
  D = Card not drawn on first turn
  Then P(A) = P(B AND C AND D) = P(B) * P(C) * P(D)
  Calculate 1.0 - P(A)
  """
  not_in_initial = hypergeom_pmf(30, cardcount, 4, 0)
  not_in_replacement = hypergeom_pmf(26, cardcount, 4, 0)
  not_in_firstturn = hypergeom_pmf(26, cardcount, 1, 0)
  1.0 - (not_in_initial * not_in_replacement * not_in_firstturn)
end

"""
    card_in_opening_hand_no_coin(cardcount::Integer)

Returns the probability to obtain at least one card of cardcount total in the opening hand of the player without the coin.

# Arguments
* `cardcount::Integer`: the total number of cards of interest in the deck.
"""
function card_in_opening_hand_no_coin(cardcount::Integer)
  """
  Define events:
  A = Card not drawn in opening hand
  B = Card not drawn in initial 3 cards
  C = Card not drawn after replacing initial 3
  D = Card not drawn on first turn
  Then P(A) = P(B AND C AND D) = P(B) * P(C) * P(D)
  Calculate 1.0 - P(A)
  """
  not_in_initial = hypergeom_pmf(30, cardcount, 3, 0)
  not_in_replacement = hypergeom_pmf(27, cardcount, 3, 0)
  not_in_firstturn = hypergeom_pmf(27, cardcount, 1, 0)
  1.0 - (not_in_initial * not_in_replacement * not_in_firstturn)
end

"""
    card_in_next_N_turns(cardcount::Integer, deckcount::Integer, N::Integer)

Returns the probability to obtain at least one card of cardcount toal in the next N turns given
deckcount cards remaining in the deck.

# Arguments
* `cardcount::Integer`: the total number of cards of interest in the deck.
* `deckcount::Integer`: the total number of cards remaining in the deck.
* `N::Integer`: the number of turns.
"""
function card_in_next_N_turns(cardcount::Integer, deckcount::Integer, N::Integer)
  "1.0 - P(Do not draw desired card in N turns)"
  1.0 - hypergeom_pmf(deckcount, cardcount, N, 0)
end

"""
    missile_lethal(missilecount::Integer, enemyhealth::Integer, healthonboard::Integer)

Returns the probability to reduce the health of a character (minion or hero) with enemyhealth health to 0 health given missilecount missiles fired and healthonboard total health on the board, not including enemyhealth.

To account for friendly fire, such as playing Mad Bomber, include the health of your board and hero in healthonboard.

# Arguments
* `missilecount::Integer`: the total number of missiles to fire.
* `enemyhealth::Integer`: the total health of the character (minion or hero) to kill.
* `healthonboard::Integer`: the total health of all other characters on the board.
"""
function missile_lethal(missilecount::Integer, enemyhealth::Integer, healthonboard::Integer)
  hypergeom_pmf(enemyhealth + healthonboard, enemyhealth, missilecount, enemyhealth)
end

end # module


mutable struct Moon
    pos
    vel
end

function Moon(pos)
    Moon(pos, zeros(Int64, 3))
end

mutable struct System
    moons
end

function apply_gravity_pair!(moon1, moon2)
end

function step!(moons)
    for i=1:length(moons)
        for j=i+1:length(moons)
            for d=1:3
                @inbounds if moons[j].pos[d] > moons[i].pos[d]
                    @inbounds moons[i].vel[d] += 1
                    @inbounds moons[j].vel[d] -= 1
                elseif moons[j].pos[d] < moons[i].pos[d]
                    @inbounds moons[i].vel[d] -= 1
                    @inbounds moons[j].vel[d] += 1
                end
            end
        end
    end
    for j=1:length(moons)
        for d=1:3
            @inbounds moons[j].pos[d] += moons[j].vel[d]
        end
    end
end

function run!(system, steps)
    for i=1:steps
        step!(system.moons)
    end
end

function total_energy(system::System)
    sum(total_energy(moon) for moon in system.moons)
end

function potential_energy(moon::Moon)
    sum(abs.(moon.pos))
end

function kinetic_energy(moon::Moon)
    sum(abs.(moon.vel))
end

function total_energy(moon::Moon)
    potential_energy(moon) * kinetic_energy(moon)
end

test_moons = [Moon([-1,0,2]), Moon([2,-10,-7]), Moon([4,-8,8]), Moon([3,5,1])]

# input for day 12:
input_system = System([
    Moon([15, -2, -6]),
    Moon([-5, -4, -11]),
    Moon([0, -6, 0]),
    Moon([5, 9, 6]),
])


struct SARIMAOrder
    p::Int
    d::Int
    q::Int
    P::Int
    D::Int
    Q::Int
    s::Int
    r::Int # Model order
    n_states_diff::Int
    n_states::Int
    function SARIMAOrder(p::Int, d::Int, q::Int, P::Int, D::Int, Q::Int, s::Int)
        @assert p >= 0
        @assert d >= 0
        @assert q >= 0
        @assert P >= 0
        @assert D >= 0
        @assert Q >= 0
        @assert s >= 0
        r = max(p + s * P, q + s * Q + 1)
        n_states_diff = d + D * s
        n_states = d + D * s + r
        return new(p, d, q, P, D, Q, s, r, n_states_diff, n_states)
    end
end

# This is used to some optimizations such as
# defining the polynomials beforehand
# and avoiding string interpolations such as "ar_L$i"
# for "ar_L1"
mutable struct SARIMAHyperParametersAuxiliary{Fl <: AbstractFloat}
    ar_poly::Vector{Fl}
    ar_names::Vector{String}
    seasonal_ar_poly::Vector{Fl}
    seasonal_ar_names::Vector{String}
    ma_poly::Vector{Fl}
    ma_names::Vector{String}
    seasonal_ma_poly::Vector{Fl}
    seasonal_ma_names::Vector{String}
    function SARIMAHyperParametersAuxiliary{Fl}(order::SARIMAOrder) where Fl
        ar_poly = create_poly(order.p, Fl)
        seasonal_ar_poly = create_seasonal_poly(order.P, order.s, Fl)
        ma_poly = create_poly(order.q, Fl)
        seasonal_ma_poly = create_seasonal_poly(order.Q, order.s, Fl)
        return new{Fl}(
            ar_poly,
            create_names("ar_L", order.p),
            seasonal_ar_poly,
            create_names("s_ar_L", order.P, order.s),
            ma_poly,
            create_names("ma_L", order.q),
            seasonal_ma_poly,
            create_names("s_ma_L", order.Q, order.s)
        )
    end
end

@doc raw"""
    SARIMA(
        y::Vector{Fl};
        order::Tuple{Int,Int,Int} = (1, 0, 0),
        seasonal_order::Tuple{Int, Int, Int, Int} = (0, 0, 0, 0),
        include_mean::Bool = false
    ) where Fl

A SARIMA model (Seasonal AutoRegressive Integrated Moving Average) implemented within the state-space
framework.

The SARIMA model is specified ``(p, d, q) \times (P, D, Q, s)``. We can also consider a polynomial ``A(t)`` to
model a trend, here we only allow to add a constante term with the `include_mean` keyword argument.
```math
\begin{gather*}
    \begin{aligned}
        \phi_p (L) \tilde \phi_P (L^s) \Delta^d \Delta_s^D y_t = A(t) + \theta_q (L) \tilde \theta_Q (L^s) \zeta_t
    \end{aligned}
\end{gather*}
```
In terms of a univariate structural model, this can be represented as
```math
\begin{gather*}
    \begin{aligned}
    y_t & = u_t + \eta_t \\
    \phi_p (L) \tilde \phi_P (L^s) \Delta^d \Delta_s^D u_t & = A(t) + \theta_q (L) \tilde \theta_Q (L^s) \zeta_t
    \end{aligned}
\end{gather*}
```

# Example
```jldoctest
julia> model = SARIMA(rand(100); order=(1,1,1), seasonal_order=(1,2,3,12))
SARIMA(1, 1, 1)x(1, 2, 3, 12) model
```

See more on [Airline passengers](@ref)

# References
 * Durbin, James, & Siem Jan Koopman.
   Time Series Analysis by State Space Methods: Second Edition.
   Oxford University Press, 2012
"""
mutable struct SARIMA <: StateSpaceModel
    order::SARIMAOrder
    hyperparameters_auxiliary::SARIMAHyperParametersAuxiliary
    hyperparameters::HyperParameters
    system::LinearUnivariateTimeInvariant
    results::Results
    include_mean::Bool

    function SARIMA(y::Vector{Fl};
                    order::Tuple{Int,Int,Int} = (1, 0, 0),
                    seasonal_order::Tuple{Int, Int, Int, Int} = (0, 0, 0, 0),
                    include_mean::Bool = false) where Fl
        or = SARIMAOrder(order[1], order[2], order[3],
                        seasonal_order[1], seasonal_order[2], seasonal_order[3], seasonal_order[4])
        hyperparameters_auxiliary = SARIMAHyperParametersAuxiliary{Fl}(or)

        Z = build_Z(or, Fl)
        T = build_T(or, Fl)
        R = build_R(or, Fl)
        d_ss = zero(Fl)
        c = build_c(or, Fl)
        H = zero(Fl)
        Q_ss = zeros(Fl, 1, 1)

        system = LinearUnivariateTimeInvariant{Fl}(y, Z, T, R, d_ss, c, H, Q_ss)

        num_hyperparameters = or.p + or.q + or.P + or.Q + 1
        if include_mean
            num_hyperparameters += 1
        end
        names = Vector{String}(undef, num_hyperparameters)
        for i in 1:(or.p)
            names[i] = "ar_L$i"
        end
        for j in 1:(or.q)
            names[or.p + j] = "ma_L$j"
        end
        for is in 1:(or.P)
            names[or.p + or.q + is] = "s_ar_L$(or.s * is)"
        end
        for js in 1:(or.Q)
            names[or.p + or.q + or.P + js] = "s_ma_L$(or.s * js)"
        end
        if include_mean
            names[end-1] = "mean"
        end
        names[end] = "sigma2_η"

        hyperparameters = HyperParameters{Fl}(names)

        return new(or, hyperparameters_auxiliary, hyperparameters, system, Results{Fl}(), include_mean)
    end
end

# Auxiliary functions
function build_Z(or, Fl)
    seasonal_entries = or.s - 1 >= 0 ? or.s - 1 : 0
    return vcat(
            fill(one(Fl), or.d),
            repeat(vcat(zeros(Fl, seasonal_entries), one(Fl)), or.D),
            one(Fl),
            fill(zero(Fl), or.r - 1)
        )
end

triu_indices(k::Int) = findall(isone, triu(ones(k, k)))
function build_T(or, Fl)
    # zeros matrix in the correct dimension
    T = zeros(Fl, or.n_states, or.n_states)
    if or.r > 0
        T[end - or.r + 1:end, end - or.r + 1:end] = diagm(1 => ones(or.r - 1))
    end
    if or.D > 0
        seasonal_companion = diagm(-1 => ones(or.s - 1))
        seasonal_companion[1, end] = 1
        for l in 1:or.D
            start_idx = or.d + (l - 1) * or.s
            end_idx = or.d + l * or.s
            T[start_idx + 1:end_idx, start_idx + 1:end_idx] .= seasonal_companion
            if l + 1 <= or.D
                T[start_idx + 1, end_idx + or.s] = 1
            end
            T[start_idx + 1, or.n_states_diff + 1] = 1
        end
    end
    if or.d > 0
        idx = triu_indices(or.d)
        T[idx] .= 1
        if or.s > 0
            start_idx = or.d
            end_idx = or.n_states_diff
            T[1:or.d, start_idx+1:end_idx] .= repeat(hcat(zeros(Fl, 1, or.s - 1), one(Fl)), 1, or.D)
        end
        T[1:or.d, or.n_states_diff + 1] .= 1
    end
    return T
end
function build_R(or, Fl)
    return vcat(
        fill(zero(Fl), or.n_states_diff, 1),
        fill(one(Fl), 1, 1),
        fill(zero(Fl), or.r - 1, 1)
    )
end
function build_c(or, Fl)
    return zeros(Fl, or.n_states)
end

function create_names(str::String, dim::Int, s::Int = 1)
    names = Vector{String}(undef, dim)
    for i in 1:dim
        names[i] = str*"$(s * i)"
    end
    return names
end
function create_poly(n::Int, Fl::DataType)
    return fill(one(Fl), n + 1)
end
function create_seasonal_poly(n::Int, s::Int, Fl::DataType)
    seasonal_poly = fill(zero(Fl), (s * n) + 1)
    seasonal_poly[1] = one(Fl)
    for i in 1:n
        seasonal_poly[(i * s) + 1] = one(Fl)
    end
    return seasonal_poly
end
function poly_mul(p1::Vector{Fl}, p2::Vector{Fl}) where Fl
    n, m = length(p1), length(p2)
    if n > 1 && m > 1
        c = zeros(Fl, m + n - 1)
        for i in 1:n, j in 1:m
            c[i + j - 1] += p1[i] * p2[j]
        end
        return c
    elseif n <= 1
        return p2
    else
        return p1
    end
end
get_ar_name(model::SARIMA, i::Int) = model.hyperparameters_auxiliary.ar_names[i]
get_ma_name(model::SARIMA, i::Int) = model.hyperparameters_auxiliary.ma_names[i]
get_seasonal_ar_name(model::SARIMA, i::Int) = model.hyperparameters_auxiliary.seasonal_ar_names[i]
get_seasonal_ma_name(model::SARIMA, i::Int) = model.hyperparameters_auxiliary.seasonal_ma_names[i]

function update_T_terms!(model::SARIMA)
    update_ar_poly!(model)
    update_seasonal_ar_poly!(model)
    # TODO increaase performance
    # We could preallocate this polynomial p3
    p3 = poly_mul(model.hyperparameters_auxiliary.ar_poly, model.hyperparameters_auxiliary.seasonal_ar_poly)
    for i in model.order.n_states_diff + 1:model.order.n_states_diff + model.order.p + model.order.s * model.order.P
        model.system.T[i, model.order.n_states_diff + 1] = p3[i - model.order.n_states_diff + 1]
    end
    return nothing
end

function update_R_terms!(model::SARIMA)
    update_ma_poly!(model)
    update_seasonal_ma_poly!(model)
    # TODO increaase performance
    # We could preallocate this polynomial p3
    p3 = poly_mul(model.hyperparameters_auxiliary.ma_poly, model.hyperparameters_auxiliary.seasonal_ma_poly)
    for i in model.order.n_states_diff + 1:model.order.n_states_diff + model.order.q + model.order.s * model.order.Q
        model.system.R[i + 1] = p3[i - model.order.n_states_diff + 1]
    end
    return nothing
end

function update_ar_poly!(model::SARIMA)
    if model.order.p > 0
        for i in 1:(model.order.p)
            model.hyperparameters_auxiliary.ar_poly[i + 1] = get_constrained_value(
                model, get_ar_name(model, i)
            )
        end
    end
end
function update_seasonal_ar_poly!(model::SARIMA)
    if model.order.P > 0
        for i in 1:(model.order.P)
            model.hyperparameters_auxiliary.seasonal_ar_poly[(i * model.order.s) + 1] = get_constrained_value(
                model, get_seasonal_ar_name(model, i)
            )
        end
    end
end
function update_ma_poly!(model::SARIMA)
    if model.order.q > 0
        for i in 1:(model.order.q)
            model.hyperparameters_auxiliary.ma_poly[i + 1] = get_constrained_value(
                model, get_ma_name(model, i)
            )
        end
    end
end
function update_seasonal_ma_poly!(model::SARIMA)
    if model.order.Q > 0
        for i in 1:(model.order.Q)
            model.hyperparameters_auxiliary.seasonal_ma_poly[(i * model.order.s) + 1] = get_constrained_value(
                model, get_seasonal_ma_name(model, i)
            )
        end
    end
end

function impose_unit_root_constraint(hyperparameter_values::Vector{Fl}) where Fl
    # TODO improve performance possibly preallocating y
    #     Monahan, John F. 1984.
    #    "A Note on Enforcing Stationarity in
    #    Autoregressive-moving Average Models."
    #    Biometrika 71 (2) (August 1): 403-404.
    n = length(hyperparameter_values)
    y = fill(zero(Fl), n, n)
    r = @. hyperparameter_values / sqrt((1 + hyperparameter_values^2))
    if n == 1
        return -r
    end
    @inbounds for k in 1:n
        for i in 1:(k - 1)
            y[k, i] = y[k - 1, i] + r[k] * y[k - 1, k - i]
        end
        y[k, k] = r[k]
    end
    return -@view y[n, :]
end

function relax_unit_root_constraint(hyperparameter_values::Vector{Fl}) where Fl
    #     Monahan, John F. 1984.
    #    "A Note on Enforcing Stationarity in
    #    Autoregressive-moving Average Models."
    #    Biometrika 71 (2) (August 1): 403-404.
    n = length(hyperparameter_values)
    y = fill(zero(Fl), n, n)
    y[n, :] = -hyperparameter_values
    for k in n:-1:2, i in 1:(k - 1)
        y[k - 1, i] = (y[k, i] - y[k, k] * y[k, k - i]) / (1 - y[k, k]^2)
    end
    r = diag(y)
    x = r ./ sqrt.(1 .- r .^ 2)
    return x
end

function SARIMA_exact_initialization!(filter::KalmanFilter, model::SARIMA)
    SARIMA_exact_initialization!(filter.kalman_state, model.system, model.order.r)
    return nothing
end

function SARIMA_exact_initialization!(kalman_state,
                                      system,
                                      num_arma_states::Int)
    # TODO we could preallocate these arrays
    arma_T = system.T[(end - num_arma_states + 1):end, (end - num_arma_states + 1):end]
    arma_R = system.R[(end - num_arma_states + 1):end]
    # Calculate exact initial P1
    arma_P1 = lyapd(arma_T, arma_R * arma_R')
    # Fill P1 in the arma states
    kalman_state.P[(end - num_arma_states + 1):end, (end - num_arma_states + 1):end] .=
        system.Q[1] .* arma_P1
    return nothing
end

function concatenate_on_bottom(X1::Matrix{Fl}, X2::Matrix{Fl}) where Fl
    n = min(size(X1, 1), size(X2, 1))
    return hcat(X1[end-n+1:end, :], X2[end-n+1:end, :])
end

function diff_sarima(y::Vector{Fl}, d::Int, D::Int, s::Int) where Fl
    # Seasonal differencing
    for _ in 1:D
        y = y[s:end] - y[1:end-s+1]
    end
    # Simple differencing
    for _ in 1:d
        y = diff(y)
    end
    return y
end

function constrain_ar!(model::SARIMA, Fl::DataType)
    if model.order.p > 0
        unconstrained_ar = Vector{Fl}(undef, model.order.p)
        for i in 1:(model.order.p)
            unconstrained_ar[i] = get_unconstrained_value(model, get_ar_name(model, i))
        end
        constrained_ar = impose_unit_root_constraint(unconstrained_ar)
        for i in 1:(model.order.p)
            update_constrained_value!(model, get_ar_name(model, i), constrained_ar[i])
        end
    end
    return nothing
end
function constrain_ma!(model::SARIMA, Fl::DataType)
    if model.order.q > 0
        unconstrained_ma = Vector{Fl}(undef, model.order.q)
        for j in 1:(model.order.q)
            unconstrained_ma[j] = get_unconstrained_value(model, get_ma_name(model, j))
        end
        constrained_ma = -impose_unit_root_constraint(unconstrained_ma)
        for j in 1:(model.order.q)
            update_constrained_value!(model, get_ma_name(model, j), constrained_ma[j])
        end
    end
    return nothing
end
function constrain_seasonal_ar!(model::SARIMA, Fl::DataType)
    if model.order.P > 0
        unconstrained_seasonal_ar = Vector{Fl}(undef, model.order.P)
        for i in 1:(model.order.P)
            unconstrained_seasonal_ar[i] = get_unconstrained_value(model, get_seasonal_ar_name(model, i))
        end
        constrained_seasonal_ar = impose_unit_root_constraint(unconstrained_seasonal_ar)
        for i in 1:(model.order.P)
            update_constrained_value!(model, get_seasonal_ar_name(model, i), constrained_seasonal_ar[i])
        end
    end
    return nothing
end
function constrain_seasonal_ma!(model::SARIMA, Fl::DataType)
    if model.order.Q > 0
        unconstrained_seasonal_ma = Vector{Fl}(undef, model.order.Q)
        for j in 1:(model.order.Q)
            unconstrained_seasonal_ma[j] = get_unconstrained_value(model, get_seasonal_ma_name(model, j))
        end
        constrained_seasonal_ma = -impose_unit_root_constraint(unconstrained_seasonal_ma)
        for j in 1:(model.order.Q)
            update_constrained_value!(model, get_seasonal_ma_name(model, j), constrained_seasonal_ma[j])
        end
    end
    return nothing
end
function constrain_mean!(model::SARIMA)
    if model.include_mean
        constrain_identity!(model, "mean")
    end
    return nothing
end

function unconstrain_ar!(model::SARIMA, Fl::DataType)
    if model.order.p > 0
        constrained_ar = Vector{Fl}(undef, model.order.p)
        for i in 1:(model.order.p)
            constrained_ar[i] = get_constrained_value(model, get_ar_name(model, i))
        end
        unconstrained_ar = relax_unit_root_constraint(constrained_ar)
        for i in 1:(model.order.p)
            update_unconstrained_value!(model, get_ar_name(model, i), unconstrained_ar[i])
        end
    end
    return nothing
end
function unconstrain_ma!(model::SARIMA, Fl::DataType)
    if model.order.q > 0
        constrained_ma = Vector{Fl}(undef, model.order.q)
        for j in 1:(model.order.q)
            constrained_ma[j] = get_constrained_value(model, get_ma_name(model, j))
        end
        unconstrained_ma = -relax_unit_root_constraint(constrained_ma)
        for j in 1:(model.order.q)
            update_unconstrained_value!(model, get_ma_name(model, j), unconstrained_ma[j])
        end
    end
    return nothing
end
function unconstrain_seasonal_ar!(model::SARIMA, Fl::DataType)
    if model.order.P > 0
        constrained_seasonal_ar = Vector{Fl}(undef, model.order.P)
        for i in 1:(model.order.P)
            constrained_seasonal_ar[i] = get_constrained_value(model, get_seasonal_ar_name(model, i))
        end
        unconstrained_seasonal_ar = relax_unit_root_constraint(constrained_seasonal_ar)
        for i in 1:(model.order.P)
            update_unconstrained_value!(model, get_seasonal_ar_name(model, i), unconstrained_seasonal_ar[i])
        end
    end
    return nothing
end
function unconstrain_seasonal_ma!(model::SARIMA, Fl::DataType)
    if model.order.Q > 0
        constrained_seasonal_ma = Vector{Fl}(undef, model.order.Q)
        for j in 1:(model.order.Q)
            constrained_seasonal_ma[j] = get_constrained_value(model, get_seasonal_ma_name(model, j))
        end
        unconstrained_seasonal_ma = -relax_unit_root_constraint(constrained_seasonal_ma)
        for j in 1:(model.order.Q)
            update_unconstrained_value!(model, get_seasonal_ma_name(model, j), unconstrained_seasonal_ma[j])
        end
    end
    return nothing
end
function unconstrain_mean!(model::SARIMA)
    if model.include_mean
        unconstrain_identity!(model, "mean")
    end
    return nothing
end

function ar_polinomial(p::Vector{Fl}) where Fl
    return Polynomial([one(Fl); -p])
end

function ma_polinomial(q::Vector{Fl}) where Fl
    return Polynomial([one(Fl); q])
end

function roots_of_inverse_polinomial(poly::Polynomial)
    return roots(poly).^-1
end

function assert_stationarity(p::Vector{Fl}) where Fl
    poly = ar_polinomial(p)
    return all(abs.(roots_of_inverse_polinomial(poly)) .< 1)
end

function assert_invertibility(q::Vector{Fl}) where Fl
    poly = ma_polinomial(q)
    return all(abs.(roots_of_inverse_polinomial(poly)) .< 1)
end

function conditional_sum_of_squares(y_diff::Vector{Fl}, k_ar::Int, k_ma::Int) where Fl
    if (k_ar == 0) && (k_ma == 0)
        return (Fl[], Fl[])
    end
    k = 2 * k_ma
    r = max(k + k_ma, k_ar)
    residuals = nothing
    X = nothing
    if k_ar + k_ma > 0
        # If we have MA terms, get residuals from an AR(k) model to use
        # as data for conditional sum of squares estimates of the MA
        # parameters
        if k_ma > 0
            y = y_diff[k:end]
            X = lagmat(y_diff, k)
            params_ar = X \ y[end - size(X, 1) + 1:end]
            residuals = y[end - size(X, 1) + 1:end] - X * params_ar
        end
        # Run an ARMA(p,q) model using the just computed residuals as
        # data
        y = y_diff[r:end]
        X = Matrix{Fl}(undef, length(y), 0)
        X = concatenate_on_bottom(X, lagmat(y_diff, k_ar))
        if k_ma > 0
            X = concatenate_on_bottom(X, lagmat(residuals, k_ma))
        end
    end

    initial_params = X \ y_diff[end - size(X, 1) + 1:end]
    params_ar = Fl[]
    params_ma = Fl[]
    offset = 1
    if k_ar > 0
        params_ar = initial_params[offset:offset + k_ar - 1]
        offset += k_ar
    end
    if k_ma > 0
        params_ma = initial_params[offset:offset + k_ma - 1]
    end
    return params_ar, params_ma
end

# Obligatory functions
function default_filter(model::SARIMA)
    Fl = typeof_model_elements(model)
    a1 = zeros(Fl, model.order.n_states)
    P1 = Fl(1e6) .* Matrix{Fl}(I, model.order.n_states, model.order.n_states)
    skip_llk_instants = model.order.d + model.order.s * model.order.D
    steadystate_tol = Fl(1e-5)
    return UnivariateKalmanFilter(a1, P1, skip_llk_instants, steadystate_tol)
end

function initial_hyperparameters!(model::SARIMA)
    Fl = typeof_model_elements(model)
    initial_hyperparameters = Dict{String,Fl}()
    # Heuristic inspired in statsmodels from python
    # TODO find a reference to this heuristic
    # conditional sum of squares
    y_diff = diff_sarima(model.system.y, model.order.d, model.order.D, model.order.s)
    y_diff = filter(!isnan, y_diff)
    # Non-seasonal ARMA
    (initial_ar, initial_ma) = conditional_sum_of_squares(y_diff, model.order.p, model.order.q)
    if !assert_stationarity(initial_ar)
        @warn("Conditional sum of squares estimated initial_ar out of the unit circle, using zero as starting params")
        initial_ar .= zero(Fl)
    end
    if !assert_invertibility(initial_ma)
        @warn("Conditional sum of squares estimated initial_ma out of the unit circle, using zero as starting params")
        initial_ma .= zero(Fl)
    end
    # Seasonal ARMA
    (initial_seasonal_ar,
    initial_seasonal_ma) = conditional_sum_of_squares(y_diff, model.order.P, model.order.Q)
    if !assert_stationarity(initial_seasonal_ar)
        @warn("Conditional sum of squares estimated initial_seasonal_ar out of the unit circle, using zero as starting params")
        initial_seasonal_ar .= zero(Fl)
    end
    if !assert_invertibility(initial_seasonal_ma)
        @warn("Conditional sum of squares estimated initial_seasonal_ma out of the unit circle, using zero as starting params")
        initial_seasonal_ma .= zero(Fl)
    end
    for i in 1:(model.order.p)
        initial_hyperparameters[get_ar_name(model, i)] = initial_ar[i]
    end
    for j in 1:(model.order.q)
        initial_hyperparameters[get_ma_name(model, j)] = initial_ma[j]
    end
    for i in 1:(model.order.P)
        initial_hyperparameters[get_seasonal_ar_name(model, i)] = initial_seasonal_ar[i]
    end
    for j in 1:(model.order.Q)
        initial_hyperparameters[get_seasonal_ma_name(model, j)] = initial_seasonal_ma[j]
    end
    if model.include_mean
        initial_hyperparameters["mean"] = mean(y_diff)
    end

    initial_hyperparameters["sigma2_η"] = var(y_diff)
    set_initial_hyperparameters!(model, initial_hyperparameters)
    return nothing
end

function constrain_hyperparameters!(model::SARIMA)
    Fl = typeof_model_elements(model)
    constrain_ar!(model, Fl)
    constrain_ma!(model, Fl)
    constrain_seasonal_ar!(model, Fl)
    constrain_seasonal_ma!(model, Fl)
    constrain_mean!(model)
    constrain_variance!(model, "sigma2_η")
    return nothing
end

function unconstrain_hyperparameters!(model::SARIMA)
    Fl = typeof_model_elements(model)
    unconstrain_ar!(model, Fl)
    unconstrain_ma!(model, Fl)
    unconstrain_seasonal_ar!(model, Fl)
    unconstrain_seasonal_ma!(model, Fl)
    unconstrain_mean!(model)
    unconstrain_variance!(model, "sigma2_η")
    return nothing
end

function fill_model_system!(model::SARIMA)
    update_T_terms!(model)
    update_R_terms!(model)
    if model.include_mean
        model.system.d = get_constrained_value(model, "mean")
    end
    model.system.Q[1] = get_constrained_value(model, "sigma2_η")
    return nothing
end

function fill_model_filter!(filter::KalmanFilter, model::SARIMA)
    SARIMA_exact_initialization!(filter, model)
    return nothing
end

function reinstantiate(model::SARIMA, y::Vector{Fl}) where Fl
    return SARIMA(y;
                order = (model.order.p, model.order.d, model.order.q),
                seasonal_order = (model.order.P, model.order.D, model.order.Q, model.order.s),
                include_mean = model.include_mean
            )
end

has_exogenous(::SARIMA) = false

function model_name(model::SARIMA)
    p, d, q = model.order.p, model.order.d, model.order.q
    P, D, Q, s = model.order.P, model.order.D, model.order.Q, model.order.s
    str = "SARIMA($p, $d, $q)x($P, $D, $Q, $s)"
    if model.include_mean
        str *= " with non-zero mean"
    end
    return str
end


"""
returns Core-type relperms for gas and liquid as a
function of the liquid saturation and their derivatives
with respect to the liquid saturation; it also returns
the stars model for the gas relperm in presence of foam.
inputs [unit] \\(default):
   + swc [-] \\(0.1)
   + sgr [-] \\(0.05)
   + krg0 [-] \\(0.9)
   + ng [-] \\(1.8)
   + krw0 [-] \\(0.2)
   + nw [-] \\(4.2)
   + fmmob [-] \\(25000)
   + epdry [-] \\(10000)
   + fmdry [-] \\(0.29)
Outputs:
   + krw
   + krg
   + dkrwdsw
   + dkrgdsw
   + krgf
   + dkrgfdsw
Usage:
   + (krw, krg, dkrwdsw, dkrgdsw, krgf, dkrgfdsw)=corey_rel_perms()
   + (krw, krg, dkrwdsw, dkrgdsw, krgf, dkrgfdsw)=corey_rel_perms(swc,
sgr, krg0, ng, krw0, nw, fmmob, epdry, fmdry)
"""
function corey_rel_perms(;swc=0.1, sgr=0.05, krg0=0.9, ng=1.8, krw0=0.2,
nw=4.2, fmmob=25000.0, epdry=10000.0, fmdry=0.29)
  sws(sw::Real)=((sw>swc)*(sw<1-sgr)*(sw-swc)/(1-sgr-swc)+(sw>=1-sgr)*1.0)
  sws(sw::Array{Float64})=((sw.>swc).*(sw.<1-sgr).*(sw-swc)/(1-sgr-swc)+(sw.>=1-sgr).*ones(size(sw)))
  krg(sw)=((sw.>=swc).*krg0.*(1-sws(sw)).^ng+(sw.<swc).*(1+(krg0-1)/swc*sw))
  krw(sw)=((sw.<=1-sgr).*krw0.*sws(sw).^nw+(sw.>1-sgr).*(-(1-krw0)/sgr.*(1.0-sw)+1.0))
  fm(sw)=(1+fmmob*(0.5+atan(epdry.*(sw-fmdry))/π))
  krgf(sw)=(krg(sw)./fm(sw))
  dkrwdsw(sw)=(nw*krw0*(1/(1-sgr-swc))*sws(sw).^(nw-1))
  dkrgdsw(sw)=((-krg0*ng*(1-sws(sw)).^(ng-1))/(-swc-sgr+1))
  dfmdsw(sw)=(((epdry*fmmob)./(π*(epdry^2*(sw-fmdry).^2+1))))
  dkrgfdsw(sw)=((dkrgdsw(sw).*fm(sw)-dfmdsw(sw).*krg(sw))./fm(sw).^2)

  return krw, krg, dkrwdsw, dkrgdsw, krgf, dkrgfdsw
end

"""
fractional flow formulation and solution of a WAG process
for Corey-type gas and water relative permeabilities on a
Cartesian coordinate (not radial)
inputs [unit] \\(default):
   + muw [Pa.s] \\(1e-3)
   + mug [Pa.s] \\(2e-5)
   + ut [m/s] \\(1e-5)
   + phi [-] \\(0.2)
   + k [m^2] \\(1e-12)
   + swc [-] \\(0.1)
   + sgr [-] \\(0.05)
   + krg0 [-] \\(0.9)
   + ng [-] \\(1.8)
   + krw0 [-] \\(0.2)
   + nw [-] \\(4.2)
   + sw0 [-] \\(1.0)
   + sw_inj [-] \\(0.0)
   + L [m] \\(domain length = 1.0)
   + pv_inj [-] \\(injected pore volumes = 5)
Outputs:
   + xt_shock (x/t shock position)
   + sw_shock
   + xt_prf
   + sw_prf
   + t [s] \\time
   + p_inj [Pa] \\(injection pressure history)
Usage:
   using roots, Dierckx
   + (xt_shock, sw_shock, xt_prf, sw_prf, t, p_inj)=frac_flow_wag()
   + (xt_shock, sw_shock, xt_prf, sw_prf, t, p_inj)=frac_flow_wag(muw, mug, ut, phi,
k, swc, sgr, krg0, ng, krw0, nw, sw0, sw_inj, L, pv_inj)
"""
function frac_flow_wag(;muw=1e-3, mug=2e-5, ut=1e-5, phi=0.2,
  k=1e-12, swc=0.1, sgr=0.05, krg0=0.9, ng=1.8, krw0=0.2,
  nw=4.2, sw0=1.0, sw_inj=0.0, L=1, pv_inj=5)

  sws(sw::Real)=((sw>swc)*(sw<1-sgr)*(sw-swc)/(1-sgr-swc)+(sw>=1-sgr)*1.0)
  sws(sw::Array{Float64})=((sw.>swc).*(sw.<1-sgr).*(sw-swc)/(1-sgr-swc)+(sw.>=1-sgr).*ones(size(sw)))
  kr(sw)=((sw.>=swc).*krg0.*(1-sws(sw)).^ng+(sw.<swc).*(1+(krg0-1)/swc*sw))
  krw(sw)=((sw.<=1-sgr).*krw0.*sws(sw).^nw+(sw.>1-sgr).*(-(1-krw0)/sgr.*(1.0-sw)+1.0))
  krg(sw)=kr(sw)
  dkrwdsw(sw)=(nw*krw0*(1/(1-sgr-swc))*sws(sw).^(nw-1))
  dkrdsw(sw)=((-krg0*ng*(1-sws(sw)).^(ng-1))/(-swc-sgr+1))
  dkrgdsw(sw)=dkrdsw(sw)
  fw(sw)=((krw(sw)/muw)./(krw(sw)/muw+krg(sw)/mug))
  dfw(sw)=((dkrwdsw(sw)/muw.*(krw(sw)/muw+krg(sw)/mug)-
      (dkrwdsw(sw)/muw+dkrgdsw(sw)/mug).*krw(sw)/muw)./
      (krg(sw)/mug+krw(sw)/muw).^2)
    eps1=1e-3
    dfw_num(sw) = (fw(sw+eps1)-fw(sw-eps1))/(2*eps1)
    ftot(sw)=krg(sw)/mug+krw(sw)/muw

# solve the nl equation to find the shock front saturation
    f_shock(sw)=(dfw(sw)-(fw(sw0)-fw(sw))./(sw0-sw))
    sw_shock = fzero(f_shock, [swc+eps(),1-sgr-eps()])
    s=collect(linspace(0.0,1.0,100))
    s1 = collect(linspace(sw_inj, sw_shock, 500))
    xt_s1 = ut/phi*dfw(s1)
    xt_s = ut/phi*dfw(s)
    xt_shock = ut/phi*dfw(sw_shock)
    xt_prf=[xt_s1; xt_shock; xt_shock+eps(); 2*xt_shock]
    sw_prf=[s1; sw_shock; sw0; sw0]

# Go through the data first
    i=1
    while(true)
      if (i+1)==length(xt_prf)
        break
      elseif xt_prf[i]==xt_prf[i+1]
        deleteat!(xt_prf, i+1)
        deleteat!(sw_prf, i+1)
      else
        i+=1
      end
    end

# find the injection pressure history
    x = collect(linspace(0,L,1000))
    sw_int = Spline1D([xt_prf; L/eps()], [sw_prf; sw0], k=1);
    t_inj=pv_inj*phi*L/ut
    t = collect(linspace(eps(),t_inj, 100)) # [s] time
    p_inj = zeros(length(t))
    p_inj[1]=trapz(x, ut./(k*(krg(sw0*ones(size(x)))/mug+krw(sw0*ones(size(x)))/muw)))
    for i in 2:length(t)
        xt = x/t[i]
        p_inj[i] = trapz(x, ut./(k*(krg(sw_int(xt))/mug+krw(sw_int(xt))/muw)))
    end

# Return the results
    return xt_shock, sw_shock, xt_prf, sw_prf, t, p_inj
end

"""
Foam flow in porous media: a BL formulation of SAG process.
CMG Stars foam model is used for the gas mobility in
presence of foam, and Corey-type relperms are used for
the gas and water. The equations are solved on a
Cartesian coordinate in one dimension.
inputs [unit] \\(default):
   + fmmob [-] \\(25000)
   + epdry [-] \\(10000)
   + fmdry [-] \\(0.29)
   + muw [Pa.s] \\(=1e-3)
   + mug [Pa.s] \\(2e-5)
   + ut [m/s] \\(1e-5)
   + phi [-] \\(0.2)
   + k [m^2] \\(1e-12)
   + swc [-] \\(0.1)
   + sgr [-] \\(0.05)
   + krg0 [-] \\(0.9)
   + ng [-] \\(1.8)
   + krw0 [-] \\(0.2)
   + nw [-] \\(4.2)
   + sw0 [-] \\(1.0)
   + sw_inj [-] \\(0.0)
Outputs:
   + xt_shock (x/t shock position)
   + sw_shock
   + xt_prf
   + sw_prf
Usage:
   using roots
   + (xt_shock, sw_shock, xt_prf, sw_prf)=frac_flow_sag()
   + (xt_shock, sw_shock, xt_prf, sw_prf)=frac_flow_sag(fmmob, epdry, fmdry, muw, mug, ut, phi,
k, swc, sgr, krg0, ng, krw0, nw, sw0, sw_inj)
"""
function frac_flow_sag(;fmmob=25000.0, epdry=10000.0, fmdry=0.29,
  muw=1e-3, mug=2e-5, ut=1e-5, phi=0.2,
  k=1e-12, swc=0.1, sgr=0.05, krg0=0.9, ng=1.8, krw0=0.2,
  nw=4.2, sw0=1.0, sw_inj=0.0, L=1, pv_inj=5)

    sws(sw::Real)=((sw>swc)*(sw<1-sgr)*(sw-swc)/(1-sgr-swc)+(sw>=1-sgr)*1.0)
    sws(sw::Array{Float64})=((sw.>swc).*(sw.<1-sgr).*(sw-swc)/(1-sgr-swc)+(sw.>=1-sgr).*ones(size(sw)))
    kr(sw)=((sw.>=swc).*krg0.*(1-sws(sw)).^ng+(sw.<swc).*(1+(krg0-1)/swc*sw))
    krw(sw)=((sw.<=1-sgr).*krw0.*sws(sw).^nw+(sw.>1-sgr).*(-(1-krw0)/sgr.*(1.0-sw)+1.0))
    fm(sw)=(1+fmmob*(0.5+atan(epdry.*(sw-fmdry))/π))
    krg(sw)=(kr(sw)./fm(sw))
    dkrwdsw(sw)=(nw*krw0*(1/(1-sgr-swc))*sws(sw).^(nw-1))
    dkrdsw(sw)=((-krg0*ng*(1-sws(sw)).^(ng-1))/(-swc-sgr+1))
    dfmdsw(sw)=(((epdry*fmmob)./(π*(epdry^2*(sw-fmdry).^2+1))))
    dkrgdsw(sw)=((dkrdsw(sw).*fm(sw)-dfmdsw(sw).*kr(sw))./fm(sw).^2)
    fw(sw)=((krw(sw)/muw)./(krw(sw)/muw+krg(sw)/mug))
    dfw(sw)=((dkrwdsw(sw)/muw.*(krw(sw)/muw+krg(sw)/mug)-
      (dkrwdsw(sw)/muw+dkrgdsw(sw)/mug).*krw(sw)/muw)./
      (krg(sw)/mug+krw(sw)/muw).^2)
    eps1=1e-3
    dfw_num(sw) = (fw(sw+eps1)-fw(sw-eps1))/(2*eps1)
    ftot(sw)=krg(sw)/mug+krw(sw)/muw

# solve the nl equation to find the shock front saturation
    f_shock(sw)=(dfw(sw)-(fw(sw0)-fw(sw))./(sw0-sw))
    sw_shock = fzero(f_shock, [swc+eps(),1-sgr-eps()])
    s=collect(linspace(0.0,1.0,100))
    s1 = collect([linspace(sw_inj, swc, 100); linspace(swc+eps(), sw_shock, 1000)])
    xt_s1 = ut/phi*dfw(s1)
    xt_s = ut/phi*dfw(s)
    xt_shock = ut/phi*dfw(sw_shock)
    xt_prf=[xt_s1; xt_shock; xt_shock+eps(); 2*xt_shock]
    sw_prf=[s1; sw_shock; sw0; sw0]

# Go through the data first
    i=1
    while(true)
      if (i+1)==length(xt_prf)
        break
      elseif xt_prf[i]==xt_prf[i+1]
        deleteat!(xt_prf, i+1)
        deleteat!(sw_prf, i+1)
      else
        i+=1
      end
    end

# find the injection pressure history
    x = collect(linspace(0,L,1000))
    sw_int = Spline1D([xt_prf; L/eps()], [sw_prf; sw0], k=1);
    t_inj=pv_inj*phi*L/ut
    t = collect(linspace(eps(),t_inj, 100)) # [s] time
    p_inj = zeros(length(t))
    p_inj[1]=trapz(x, ut./(k*(krg(sw0*ones(size(x)))/mug+krw(sw0*ones(size(x)))/muw)))
    for i in 2:length(t)
        xt = x/t[i]
        p_inj[i] = trapz(x, ut./(k*(krg(sw_int(xt))/mug+krw(sw_int(xt))/muw)))
    end

    return xt_shock, sw_shock, xt_prf, sw_prf, t, p_inj
end

function trapz(x,y)
# copied from https://github.com/hwborchers/NumericalMath.jl/blob/master/src/integrate.jl
  n=length(x)
  if (length(y) != n)
    error("Vectors 'x', 'y' must be of same length")
  end
  r=0
  for i in 2:n
    r+=(x[i] - x[i-1])*(y[i] + y[i-1])
  end
  return r/2.0
end


    using DifferentialEquations
    using OrdinaryDiffEq
    using DiffEqBase
    using Sundials
    using ODEInterfaceDiffEq
    using Turing
    using MCMCChains

        # du -> Derivative
        # u -> State at time t
        # p-> paramter vector
        # t-> time as tuple (init, end)

        function PLacExampleODE!(du,u,p,t)
            Cit_mrna, Cit_foldedP, Cit_fluo, Cit_AU = u;

            alpha1, Vm1, h1, Km1, d1, alpha2, d2, Kf, sc_molec, IPTG = p;

        du[1] = dCit_mrna=alpha1+Vm1*(IPTG^h1/(Km1^h1+IPTG^h1))-d1*Cit_mrna;
        du[2] = dCit_foldedP=alpha2*Cit_mrna-(d2+Kf)*Cit_foldedP;
        du[3] = dCit_fluo=Kf*Cit_foldedP-d2*Cit_fluo;
        du[4] = dCit_AU = sc_molec*dCit_fluo;

        end

        # p -> parameter vector (plus inducer vector at the end)
        # I -> initial Y0 vector

        function PLacExampleSteadyState(p,I)

            expCit_mrna, expCit_foldedP, expCit_fluo, expCit_AU = I;
        alpha1, Vm1, h1, Km1, d1, alpha2, d2, Kf, sc_molec, IPTG = p;

         alp = zeros(4);

        alp[1] = Cit_mrna = (alpha1 + Vm1*(IPTG^h1/(Km1^h1+IPTG^h1)))/d1;;
        alp[2] = Cit_foldedP = (alpha2*Cit_mrna)/(Kf+d2);;
        alp[3] = Cit_fluo = (Kf*Cit_foldedP)/d2;;
        alp[4] = Cit_AU= sc_molec*Cit_fluo;;
        return(alp)

        end

        # ts -> time vector going from t=0 to t=end every 1
        # p -> parameter vector
        # sp -> vector with switching times (t=0 and t=end needs to be included). Length of vector will be length(inducer steps)+1
        # inputs -> Vecor with the inducers (if more than 1 inducer in the system, the shape will be a1,b1,c1,..., a2,b2,c2,... aend,bend,cend,... where each leter represents a different inducer)
        # ivss -> Initial Y0 vector
        # pre -> Vector with the inducers in the ON. It can be empty

        @model function PLacExample_TuringModel(data,prob1)

            p1 ~ Truncated(Distributions.Normal(0.2475194, 0.1237403), 3.88e-5, 0.495)
            p2 ~ Truncated(Distributions.Normal(0.2669, 0.11405), 0.0388, 0.495)
            p3 ~ Truncated(Distributions.Normal(2.7, 1.1), 0.5, 4.9)
            p4 ~ Truncated(Distributions.Normal(6.0, 2.0), 2.0, 10.0)
            p5 ~ Truncated(Distributions.Normal(0.11885000000000001, 0.055575), 0.0077, 0.23)
            p6 ~ Truncated(Distributions.Normal(3.525, 1.6408500000000001), 0.2433, 6.8067)
            p7 ~ Truncated(Distributions.Normal(0.1224799, 0.06121005), 5.98e-5, 0.2449)
            p8 ~ Truncated(Distributions.Normal(0.01685, 0.002425), 0.012, 0.0217)
            p9 ~ Truncated(Distributions.Normal(5.0005, 2.49975), 0.001, 10.0)

            p = [p1,p2,p3,p4,p5,p6,p7,p8,p9,];

            Neq = 4;

            for exp in 1:2

                ts=collect(0.:1:data["finalTime"][exp])
                sp = data["switchT"][exp];
                inputs = data["uInd"][exp];
                ivss = data["y0"][exp];
                pre=data["preInd"][exp];

                maxtime = length(ts);
                Nsp = length(sp);
                Nevents = length(sp)-1;

                if typeof(p) == Array{Float64, 1}
                    if pre != []
                        pSS = vcat(p, pre);
                    else
                        pSS = p;
                    end
                else
                    if pre != []
                        pSS = vcat([p[i].value for i in 1:9], pre);
                    else
                        pSS = [p[i].value for i in 1:9];
                    end
                end

                final = zeros(maxtime, Neq);
                finalBay2 = Array{Any,2}(undef, 1, maxtime)

                y_al = PLacExampleSteadyState(pSS,ivss) # Calculation of initial guesses for steady state
                y0 = y_al;
                initialV = y0;
                i = 1;

                for q in collect(1:Nevents)

                    lts = length(ts[(sp[q]+1):sp[q+1]+1]);  # General way to define the number of elements in each event series

                    Tevent = ts[(sp[q]+1):sp[q+1]+1];  # General way to extract the times of each event
                    I = inputs[i:(i+(0))];
                    pSte = vcat(p, I);

                    if q == 1
                        if typeof(p) == Array{Float64, 1}
                            prob = ODEProblem(PLacExampleODE!,initialV,(ts[(sp[q]+1)],ts[sp[q+1]+1]),pSte);
                            part1 = DifferentialEquations.solve(prob, DifferentialEquations.Tsit5(),reltol=1.0e-9,abstol=1.0e-9,saveat=1);
                        else
                            prob = remake(prob1, p=pSte)
                            prob = remake(prob, u0=initialV)
                            prob = remake(prob, tspan = convert.(Int, (ts[(sp[q]+1)],ts[sp[q+1]+1])))
                            part1 = DifferentialEquations.solve(prob,DifferentialEquations.Tsit5(),reltol=1.0e-9,abstol=1.0e-9,saveat=1)
                        end
                    else
                        if typeof(p) == Array{Float64, 1}
                            prob = ODEProblem(PLacExampleODE!,initialV,(ts[(sp[q]+1)],ts[sp[q+1]+1]),pSte);
                            part1 = DifferentialEquations.solve(prob, DifferentialEquations.Tsit5(),reltol=1.0e-9,abstol=1.0e-9,saveat=1);
                        else
                            prob = remake(prob1, p=pSte)
                            prob = remake(prob, u0=initialV)
                            prob = remake(prob, tspan = (ts[(sp[q]+1)],ts[sp[q+1]+1]))
                            part1 = DifferentialEquations.solve(prob,DifferentialEquations.Tsit5(),reltol=1.0e-9,abstol=1.0e-9,saveat=1)
                        end
                    end

                    if typeof(p) == Array{Float64, 1}
                        initialV = part1[lts];
                    else
                        initialV = [part1[:,lts][h].value for h in 1:4];
                    end

                    i+=1;

                    if typeof(p) == Array{Float64, 1}
                        for d in collect((sp[q]+1):(sp[q]+lts))
                            final[d,:] = part1[d-sp[q]];
                        end
                    else
                        for d in collect((sp[q]+1):(sp[q]+lts))
                            finalBay2[1,d] = 3 .* part1[4, d-sp[q]];

                        end
                    end

                end

                obs = Array{Any,1}(undef, 1)
                if typeof(p) == Array{Float64, 1}
                    obs[1] = 3 .* final[convert.(Int,round(data["tsamps"][exp])).+1, 4];

                else
                    obs[1] = finalBay2[1, convert.(Int,round(data["tsamps"][exp])).+1];
                end

                for ob in 1:1
                    for dat in 1:length(obs)
                        data["DataMean"][ob][dat] ~ Normal(obs[ob][dat], data["DataError"][ob][1][dat])
                    end
                end

            end

        end


# using Revise
using TopOpt, LinearAlgebra, StatsFuns
using StatsFuns: logsumexp

E = 1.0 # Young’s modulus
v = 0.3 # Poisson’s ratio
f = 1.0 # downward force
rmin = 3.0

# ### Define the problem
problem = PointLoadCantilever(Val{:Linear}, (160, 40), (1.0, 1.0), E, v, f)

# ### Parameter settings
V = 0.5 # volume fraction
xmin = 0.0001 # minimum density
steps = 40 # maximum number of penalty steps, delta_p0 = 0.1

# ### Continuation SIMP
x0 = fill(1.0, 160 * 40) # initial design
x = copy(x0)
for p in [1.0, 2.0, 3.0]
    global penalty, stress, filter, result, stress, x
    penalty = TopOpt.PowerPenalty(p)
    solver = FEASolver(Direct, problem; xmin=xmin, penalty=penalty)
    stress = TopOpt.MicroVonMisesStress(solver)
    filter = DensityFilter(solver; rmin=rmin)
    volfrac = TopOpt.Volume(problem, solver)

    obj = x -> volfrac(filter(x)) - V
    thr = 10 # stress threshold
    constr = x -> begin
        s = stress(filter(x))
        vcat((s .- thr) / 100, logsumexp(s) - log(length(s)) - thr)
    end
    alg = PercivalAlg()
    options = PercivalOptions()
    optimizer = Optimizer(obj, constr, x, alg; options=options)
    simp = SIMP(optimizer, solver, p)
    result = simp(x)
    x = result.topology
end

maximum(stress(filter(x0))) # 0.51
maximum(stress(filter(x))) # 10.01

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

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


# function NSDEBase.cost(problem::AbstractInitialValueProblem, solver::MovingWindow.AbstractMovingWindowSolver, solution::MovingWindow.AbstractMovingWindowSolution)
#     @↓ (t0, tN) ← tspan = problem
#     @↓ τ, Δτ, parallelsolver, budget = solver
#     @↓ finesolver, coarsolver,  P = parallelsolver
#     @↓ hF ← h = finesolver.stepsize
#     @↓ hG ← h = coarsolver.stepsize
#     ξ = hG / hF
#     @↓ B ← budget, a₁, a₂, a₃ = budget
#     M = length(solution)
#     r = solution.restarts
#     α = τ / Δτ
#     return τ / abs(tN - t0) * (M*(M-1)/(2α*ξ) + (P+ξ)/(2P*ξ)*sum(1 + (B+1)*r[m] + length(solution[m].errors) for m=1:M))
# end
# cost(problem, mowa, mowasolution) * Cₛ


"""
**FUNÇÃO ÁREA TOTAL**
Sintaxe:
areatotal(raio, altura_cilindro)
"""

function areatotal(raio, altura)
    println("Área total:")
    return arealateral(raio, altura) + 2*areaBC(raio)
end


"""
    s = inpoly2(vert, node, [edge=Nothing, atol=0.0, rtol=eps(), outformat=InOnOut{1,0,-1}])

Check all points defined by `vert` are inside, outside or on bounds of
polygon defined by `node` and `edge`.

`vert` and `node` are matrices with the x-coordinates in the first
and the y-coordinates in the second column.

`edge` is an integer matrix of same size as `node`. The first column contains
the index of the starting node, the second column the index of the ending node.
The polygon needs to be closed. It may contain unconnected cycles.

A point is considered on boundary, if its Euclidian distance to any edge is less
than or equal `max(atol, rtol * span)`, where span is the maximum extension of the
polygon in either direction.

If a point is considered on bounds, no further check is made to determine it it is inside or
outside numerically.

Expected effort of the algorithm is `(N+M)*log(M)`, where `M` is the number of points
and `N` is the number of polygon edges.
"""
function inpoly2(vert, node, edge=zeros(Int); atol::T=0.0, rtol::T=NaN, iyperm=nothing, outformat=InOnBit) where T<:AbstractFloat

    rtol = !isnan(rtol) ? rtol : iszero(atol) ? eps(T)^0.85 : zero(T)
    poly = PolygonMesh(node, edge)
    points = PointsInbound(vert)
    npoints = length(points)

    vmin = minimum(points)
    vmax = maximum(points)
    pmin = minimum(poly)
    pmax = maximum(poly)

    lbar = sum(pmax - pmin)
    tol = max(abs(rtol * lbar), abs(atol))

    ac = areacount(poly)
    stat = ac > 1 ? falses(npoints,2,ac) : falses(npoints,2)
    # flip coordinates so expected efford is minimal

    dvert = vmax - vmin
    ix = dvert[1] < dvert[2] ? 1 : 2
    if isnothing(iyperm)
        iyperm = sortperm(points, 3 - ix)
    end

    inpoly2!(points, iyperm, poly, ix, tol, stat)

    convertout(outformat, InOnBit, stat)
end

"""
    inpoly2_mat(vert, node, edge, fTolx, ftoly, stats)

INPOLY2_MAT the local m-code version of the crossing-number
test. Loop over edges; do a binary-search for the first ve-
rtex that intersects with the edge y-range; do crossing-nu-
mber comparisons; break when the local y-range is exceeded.
"""
function inpoly2!(points, iyperm, poly, ix::Integer, veps::T, stat::S) where {N,T<:AbstractFloat,S<:AbstractArray{Bool,N}}

    nvrt = length(points)   # number of points to be checked
    nedg = edgecount(poly)  # number of edges of the polygon mesh
    vepsx = vepsy = veps
    iy = 3 - ix

    #----------------------------------- loop over polygon edges
    for epos = 1:nedg

        inod = edgeindex(poly, epos, 1)  # from
        jnod = edgeindex(poly, epos, 2)  # to
        # swap order of vertices
        if vertex(poly, inod, iy) > vertex(poly, jnod, iy)
            inod, jnod = jnod, inod
        end

        #------------------------------- calc. edge bounding-box
        xone = vertex(poly, inod, ix)
        yone = vertex(poly, inod, iy)
        xtwo = vertex(poly, jnod, ix)
        ytwo = vertex(poly, jnod, iy)

        xmin0 = min(xone, xtwo)
        xmax0 = max(xone, xtwo)
        xmin = xmin0 - vepsx
        xmax = xmax0 + vepsx
        ymin = yone - vepsy
        ymax = ytwo + vepsy

        ydel = ytwo - yone
        xdel = xtwo - xone
        xysq = xdel^2 + ydel^2
        feps = sqrt(xysq) * veps

        # find top points[:,iy] < ymin by binary search
        ilow = searchfirst(points, iy, iyperm, ymin)
        #------------------------------- calc. edge-intersection
        # loop over all points with y ∈ [ymin,ymax)
        for jpos = ilow:nvrt
            jorig = iyperm[jpos]
            ypos = vertex(points, jorig, iy)
            ypos > ymax && break
            xpos = vertex(points, jorig, ix)

            if xpos >= xmin
                if xpos <= xmax
                    #--------- inside extended bounding box of edge
                    mul1 = ydel * (xpos - xone)
                    mul2 = xdel * (ypos - yone)
                    if abs(mul2 - mul1) <= feps
                        #------- distance from line through edge less veps
                        mul3 = xdel * (2xpos-xone-xtwo) + ydel * (2ypos-yone-ytwo)
                        if abs(mul3) <= xysq ||
                           hypot(xpos- xone, ypos - yone) <= veps ||
                           hypot(xpos- xtwo, ypos - ytwo) <= veps
                            # ---- round boundaries around endpoints of edge
                            setonbounds!(poly, stat, jorig, epos)
                        end
                        if mul1 < mul2 && yone <= ypos < ytwo
                            #----- left of line && ypos exact to avoid multiple counting
                            flipio!(poly, stat, jorig, epos)
                        end
                    elseif mul1 < mul2 && yone <= ypos < ytwo
                        #----- left of line && ypos exact to avoid multiple counting
                        flipio!(poly, stat, jorig, epos)
                    end
                end
            else # xpos < xmin - left of bounding box
                if yone <= ypos <  ytwo
                    #----- ypos exact to avoid multiple counting
                    flipio!(poly, stat, jorig, epos)
                end
            end
        end
    end
    stat
end

"""
    search lowest iy coordinate >= ymin
"""
function searchfirst(points::PointsInbound, iy::Integer, iyperm, ymin)
    ilow = 0
    iupp = length(points) + 1
    @inbounds while ilow < iupp - 1
        imid = ilow + (iupp-ilow) >>> 0x01
        if vertex(points, iyperm[imid], iy) < ymin
            ilow = imid
        else
            iupp = imid
        end
    end
    iupp
end

"""
    flipio!(poly, stat, p, epos)

Flip boolean value of in-out-status of point `p` for all areas associated with edge `epos`.
"""
function flipio!(poly::PolygonMesh, stat::AbstractArray{Bool}, i::Integer, j::Integer)
    statop!((stat,j,a) -> begin stat[j,1,a] = !stat[j,1,a]; end, poly, stat, i, j)
end

"""
    setonbounds!(poly, stat, p, epos)

Set boolean value of on-bounds-status of point `p` for all areas associated with edge `epos`.
"""
function setonbounds!(poly::PolygonMesh, stat::AbstractArray{Bool}, i::Integer, j::Integer)
    statop!((stat, j,a) -> begin stat[j,2,a] = true; end, poly, stat, i, j)
end
function statop!(f!::Function, poly::PolygonMesh{A}, stat::AbstractArray{Bool,N}, p::Integer, ed::Integer) where {A,N}
    for k in 1:max(A, 1)
        area = N == 3 ? areaindex(poly, ed, k) : 1
        if area > 0
            f!(stat, p, area)
        end
    end
end


abstract type AbstractTableauSplitting{T <: Real} <: AbstractTableau{T} end

function get_splitting_coefficients(r, a::Vector{T}, b::Vector{T}) where {T}
    @assert length(a) == length(b)

    s = length(a)
    f = zeros(Int, 2r*s)
    c = zeros(T,   2r*s)

    for i in 1:s
        for j in 1:r
            f[(2i-2)*r+j] = j
            c[(2i-2)*r+j] = a[i]
        end
        for j in 1:r
            f[(2i-1)*r+j] = r-j+1
            c[(2i-1)*r+j] = b[i]
        end
    end

    return f, c
end

@doc raw"""
Tableau for general splitting methods for vector fields with two components A and B.

Integrator:
```math
\varphi_{\tau} = \varphi_{b_s \tau}^{B} \circ \varphi_{a_s \tau}^{A} \circ \dotsc \circ \varphi_{b_1 \tau}^{B} \circ \varphi_{a_1 \tau}^{A}
```
"""
struct TableauSplitting{T} <: AbstractTableauSplitting{T}
    @HeaderTableau

    a::Vector{T}
    b::Vector{T}

    function TableauSplitting{T}(name, o, s, a, b) where {T}
        @assert s == length(a) == length(b)
        new(name, o, s, a, b)
    end
end

function TableauSplitting(name, o, a::Vector{T}, b::Vector{T}) where {T}
    TableauSplitting{T}(name, o, length(a), a, b)
end

function get_splitting_coefficients(nequs, tableau::TableauSplitting{T}) where {T}
    @assert nequs == 2
    @assert length(tableau.a) == length(tableau.b)

    s = length(tableau.a)
    f = zeros(Int, 2s)
    c = zeros(T,   2s)

    for i in eachindex(tableau.a, tableau.b)
        f[2*(i-1)+1] = 1
        c[2*(i-1)+1] = tableau.a[i]
        f[2*(i-1)+2] = 2
        c[2*(i-1)+2] = tableau.b[i]
    end

    # select all entries with non-vanishing coefficients
    y = c .!= 0

    return f[y], c[y]
end

@doc raw"""
Tableau for non-symmetric splitting methods.
    See McLachlan, Quispel, 2003, Equ. (4.10).
    The methods A and B are the composition of all vector fields in the SODE
    and its adjoint, respectively.

Basic method: Lie composition
```math
\begin{aligned}
\varphi_{\tau}^{A} &= \varphi_{\tau}^{v_1} \circ \varphi_{\tau}^{v_2} \circ \dotsc \circ \varphi_{\tau}^{v_{r-1}} \circ \varphi_{\tau}^{v_r} \\
\varphi_{\tau}^{B} &= \varphi_{\tau}^{v_r} \circ \varphi_{\tau}^{v_{r-1}} \circ \dotsc \circ \varphi_{\tau}^{v_2} \circ \varphi_{\tau}^{v_1}
\end{aligned}
```

Integrator:
```math
\varphi_{\tau}^{NS} = \varphi_{b_s \tau}^{B} \circ \varphi_{a_s \tau}^{A} \circ \dotsc \circ \varphi_{b_1 \tau}^{B} \circ \varphi_{a_1 \tau}^{A}
```
"""
struct TableauSplittingNS{T} <: AbstractTableauSplitting{T}
    @HeaderTableau

    a::Vector{T}
    b::Vector{T}

    function TableauSplittingNS{T}(name, o, s, a, b) where {T}
        @assert s == length(a) == length(b)
        new(name, o, s, a, b)
    end
end

function TableauSplittingNS(name, o, a::Vector{T}, b::Vector{T}) where {T}
    TableauSplittingNS{T}(name, o, length(a), a, b)
end

function get_splitting_coefficients(nequs, tableau::TableauSplittingNS)
    # R = length(equation.v)
    # S = tableau.s
    #
    # f = zeros(Int, 2R*S)
    # c = zeros(TT,  2R*S)
    #
    # for i in 1:S
    #     for j in 1:R
    #         f[(2i-2)*R+j] = j
    #         c[(2i-2)*R+j] = tableau.a[i]
    #     end
    #     for j in 1:R
    #         f[(2i-1)*R+j] = R-j+1
    #         c[(2i-1)*R+j] = tableau.b[i]
    #     end
    # end

    get_splitting_coefficients(nequs, tableau.a, tableau.b)
end

@doc raw"""
Tableau for symmetric splitting methods with general stages.
    See McLachlan, Quispel, 2003, Equ. (4.11).

Basic method: Lie composition
```math
\begin{aligned}
\varphi_{\tau}^{A} &= \varphi_{\tau}^{v_1} \circ \varphi_{\tau}^{v_2} \circ \dotsc \circ \varphi_{\tau}^{v_{r-1}} \circ \varphi_{\tau}^{v_r} \\
\varphi_{\tau}^{B} &= \varphi_{\tau}^{v_r} \circ \varphi_{\tau}^{v_{r-1}} \circ \dotsc \circ \varphi_{\tau}^{v_2} \circ \varphi_{\tau}^{v_1}
\end{aligned}
```

Integrator:
```math
\varphi_{\tau}^{GS} = \varphi_{a_1 \tau}^{A} \circ \varphi_{b_1 \tau}^{B} \circ \dotsc \circ \varphi_{b_1 \tau}^{B} \circ \varphi_{a_1 \tau}^{A}
```
"""
struct TableauSplittingGS{T} <: AbstractTableauSplitting{T}
    @HeaderTableau

    a::Vector{T}
    b::Vector{T}

    function TableauSplittingGS{T}(name, o, s, a, b) where {T}
        @assert s == length(a) == length(b)
        new(name, o, s, a, b)
    end
end

function TableauSplittingGS(name, o, a::Vector{T}, b::Vector{T}) where {T}
    TableauSplittingGS{T}(name, o, length(a), a, b)
end

function get_splitting_coefficients(nequs, tableau::TableauSplittingGS)
    # R = nequs
    # S = tableau.s
    #
    # f = zeros(Int, 2R*S)
    # c = zeros(TT,  2R*S)
    #
    # for i in 1:S
    #     for j in 1:R
    #         f[(2i-2)*R+j] = j
    #         c[(2i-2)*R+j] = tableau.a[i]
    #     end
    #     for j in R:-1:1
    #         f[(2i-1)*R+j] = R-j+1
    #         c[(2i-1)*R+j] = tableau.b[i]
    #     end
    # end

    f, c = get_splitting_coefficients(nequs, tableau.a, tableau.b)
    vcat(f, f[end:-1:1]), vcat(c, c[end:-1:1])
end

@doc raw"""
Tableau for symmetric splitting methods with symmetric stages.
    See McLachlan, Quispel, 2003, Equ. (4.6).

Basic method: symmetric Strang composition
```math
\varphi_{\tau}^{A} = \varphi_{\tau/2}^{v_1} \circ \varphi_{\tau/2}^{v_2} \circ \dotsc \circ \varphi_{\tau/2}^{v_{r-1}} \circ \varphi_{\tau/2}^{v_r} \circ \varphi_{\tau/2}^{v_r} \circ \varphi_{\tau/2}^{v_{r-1}} \circ \dotsc \circ \varphi_{\tau/2}^{v_2} \circ \varphi_{\tau/2}^{v_1}
```

Integrator:
```math
\varphi_{\tau}^{SS} = \varphi_{a_1 \tau}^{A} \circ \varphi_{a_2 \tau}^{A} \circ \dotsc \circ \varphi_{a_s \tau}^{A} \circ \dotsc \circ \varphi_{a_2 \tau}^{A} \circ \varphi_{a_1 \tau}^{A}
```
"""
struct TableauSplittingSS{T} <: AbstractTableauSplitting{T}
    @HeaderTableau

    a::Vector{T}

    function TableauSplittingSS{T}(name, o, s, a) where {T}
        @assert s == length(a)
        new(name, o, s, a)
    end
end

function TableauSplittingSS(name, o, a::Vector{T}) where {T}
    TableauSplittingSS{T}(name, o, length(a), a)
end

function get_splitting_coefficients(nequs, tableau::TableauSplittingSS{T}) where {T}
    r = nequs
    a = vcat(tableau.a, tableau.a[end-1:-1:1]) ./ 2
    s = length(a)

    f = zeros(Int, 2r*s)
    c = zeros(T,   2r*s)

    for i in 1:s
        for j in 1:r
            f[(2i-2)*r+j] = j
            c[(2i-2)*r+j] = a[i]
            f[(2i-1)*r+j] = r-j+1
            c[(2i-1)*r+j] = a[i]
        end
    end

    f, c
end


const lines = readlines()

parse_dot(l) = CartesianIndex((1 .+ parse.(Int, split(l, ",")))...)
function parse_fold(l)
    axis, v = split(l, "=")
    axis = Symbol(axis[end])
    v = parse(Int, v) + 1
    axis, v
end

const dots = parse_dot.(lines[begin:793])
const folds = parse_fold.(lines[795:end])

const paper = begin
    p = falses(maximum(dots).I)
    p[dots] .= true
    p
end

function fold(m, ::Val{:x}, pos)
    xsize = size(m, 1)
    n = copy(m[begin:pos-1, :])
    for I in CartesianIndices(n)
        xx, yy = I.I
        if (2pos - xx) ≤ xsize
            n[I] |= m[2pos-xx, yy]
        end
    end
    n
end

function fold(m, ::Val{:y}, pos)
    ysize = size(m, 2)
    n = copy(m[:, begin:pos-1])
    for I in CartesianIndices(n)
        xx, yy = I.I
        if (2pos - yy) ≤ ysize
            n[I] |= m[xx, 2pos-yy]
        end
    end
    n
end

fold(m, s::Symbol, pos) = fold(m, Val{s}(), pos)

# Part One
fold(paper, folds[1]...) |> count |> println

# Part Two
let pp = paper
    for f in folds
        pp = fold(pp, f...)
    end
    for r in eachcol(pp)
        println(join([rr ? "#" : " " for rr in r]))
    end
end


"""
    sinkhorn_unbalanced(μ, ν, C, λ1::Real, λ2::Real, ε; kwargs...)

Compute the optimal transport plan for the unbalanced entropically regularized optimal
transport problem with source and target marginals `μ` and `ν`, cost matrix `C` of size
`(length(μ), length(ν))`, entropic regularization parameter `ε`, and marginal relaxation
terms `λ1` and `λ2`.

The optimal transport plan `γ` is of the same size as `C` and solves
```math
\\inf_{\\gamma} \\langle \\gamma, C \\rangle
+ \\varepsilon \\Omega(\\gamma)
+ \\lambda_1 \\operatorname{KL}(\\gamma 1 | \\mu)
+ \\lambda_2 \\operatorname{KL}(\\gamma^{\\mathsf{T}} 1 | \\nu),
```
where ``\\Omega(\\gamma) = \\sum_{i,j} \\gamma_{i,j} \\log \\gamma_{i,j}`` is the entropic
regularization term and ``\\operatorname{KL}`` is the Kullback-Leibler divergence.

The keyword arguments supported here are the same as those in the `sinkhorn_unbalanced`
for unbalanced optimal transport problems with general soft marginal constraints.
"""
function sinkhorn_unbalanced(
    μ, ν, C, λ1::Real, λ2::Real, ε; proxdiv_F1=nothing, proxdiv_F2=nothing, kwargs...
)
    if proxdiv_F1 !== nothing && proxdiv_F2 !== nothing
        Base.depwarn(
            "keyword arguments `proxdiv_F1` and `proxdiv_F2` are deprecated",
            :sinkhorn_unbalanced,
        )

        # have to wrap the "proxdiv" functions since the signature changed
        # ε was fixed in the function, so we ignore it
        proxdiv_F1_wrapper(s, p, _) = copyto!(s, proxdiv_F1(s, p))
        proxdiv_F2_wrapper(s, p, _) = copyto!(s, proxdiv_F2(s, p))

        return sinkhorn_unbalanced(
            μ, ν, C, proxdiv_F1_wrapper, proxdiv_F2_wrapper, ε; kwargs...
        )
    end

    # define "proxdiv" functions for the unbalanced OT problem
    proxdivF!(s, p, ε, λ) = (s .= (p ./ s) .^ (λ / (ε + λ)))
    proxdivF1!(s, p, ε) = proxdivF!(s, p, ε, λ1)
    proxdivF2!(s, p, ε) = proxdivF!(s, p, ε, λ2)

    return sinkhorn_unbalanced(μ, ν, C, proxdivF1!, proxdivF2!, ε; kwargs...)
end

"""
    sinkhorn_unbalanced(
        μ, ν, C, proxdivF1!, proxdivF2!, ε;
        atol=0, rtol=atol > 0 ? 0 : √eps, check_convergence=10, maxiter=1_000,
    )

Compute the optimal transport plan for the unbalanced entropically regularized optimal
transport problem with source and target marginals `μ` and `ν`, cost matrix `C` of size
`(length(μ), length(ν))`, entropic regularization parameter `ε`, and soft marginal
constraints ``F_1`` and ``F_2`` with "proxdiv" functions `proxdivF!` and `proxdivG!`.

The optimal transport plan `γ` is of the same size as `C` and solves
```math
\\inf_{\\gamma} \\langle \\gamma, C \\rangle
+ \\varepsilon \\Omega(\\gamma)
+ F_1(\\gamma 1, \\mu)
+ F_2(\\gamma^{\\mathsf{T}} 1, \\nu),
```
where ``\\Omega(\\gamma) = \\sum_{i,j} \\gamma_{i,j} \\log \\gamma_{i,j}`` is the entropic
regularization term and ``F_1(\\cdot, \\mu)`` and ``F_2(\\cdot, \\nu)`` are soft marginal
constraints for the source and target marginals.

The functions `proxdivF1!(s, p, ε)` and `proxdivF2!(s, p, ε)` evaluate the "proxdiv"
functions of ``F_1(\\cdot, p)`` and ``F_2(\\cdot, p)`` at ``s`` for the entropic
regularization parameter ``\\varepsilon``. They have to be mutating and overwrite the first
argument `s` with the result of their computations.

Mathematically, the "proxdiv" functions are defined as
```math
\\operatorname{proxdiv}_{F_i}(s, p, \\varepsilon)
= \\operatorname{prox}^{\\operatorname{KL}}_{F_i(\\cdot, p)/\\varepsilon}(s) \\oslash s
```
where ``\\oslash`` denotes element-wise division and
``\\operatorname{prox}_{F_i(\\cdot, p)/\\varepsilon}^{\\operatorname{KL}}`` is the proximal
operator of ``F_i(\\cdot, p)/\\varepsilon`` for the Kullback-Leibler
(``\\operatorname{KL}``) divergence.  It is defined as
```math
\\operatorname{prox}_{F}^{\\operatorname{KL}}(x)
= \\operatorname{argmin}_{y} F(y) + \\operatorname{KL}(y|x)
```
and can be computed in closed-form for specific choices of ``F``. For instance, if
``F(\\cdot, p) = \\lambda \\operatorname{KL}(\\cdot | p)`` (``\\lambda > 0``), then
```math
\\operatorname{prox}_{F(\\cdot, p)/\\varepsilon}^{\\operatorname{KL}}(x)
= x^{\\frac{\\varepsilon}{\\varepsilon + \\lambda}} p^{\\frac{\\lambda}{\\varepsilon + \\lambda}},
```
where all operators are acting pointwise.[^CPSV18]

Every `check_convergence` steps it is assessed if the algorithm is converged by checking if
the iterates of the scaling factor in the current and previous iteration satisfy
`isapprox(vcat(a, b), vcat(aprev, bprev); atol=atol, rtol=rtol)` where `a` and `b` are the
current iterates and `aprev` and `bprev` the previous ones. The default `rtol` depends on
the types of `μ`, `ν`, and `C`. After `maxiter` iterations, the computation is stopped.

[^CPSV18]: Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F.-X. (2018). [Scaling algorithms for unbalanced optimal transport problems](https://doi.org/10.1090/mcom/3303). Mathematics of Computation, 87(314), 2563–2609.

See also: [`sinkhorn_unbalanced2`](@ref)
"""
function sinkhorn_unbalanced(
    μ,
    ν,
    C,
    proxdivF1!,
    proxdivF2!,
    ε;
    tol=nothing,
    atol=tol,
    rtol=nothing,
    max_iter=nothing,
    maxiter=max_iter,
    check_convergence::Int=10,
)
    # deprecations
    if tol !== nothing
        Base.depwarn(
            "keyword argument `tol` is deprecated, please use `atol` and `rtol`",
            :sinkhorn_unbalanced,
        )
    end
    if max_iter !== nothing
        Base.depwarn(
            "keyword argument `max_iter` is deprecated, please use `maxiter`",
            :sinkhorn_unbalanced,
        )
    end

    # compute Gibbs kernel
    K = @. exp(-C / ε)

    # set default values of squared tolerances
    T = float(Base.promote_eltype(μ, ν, K))
    sqatol = atol === nothing ? 0 : atol^2
    sqrtol = rtol === nothing ? (sqatol > zero(sqatol) ? zero(T) : eps(T)) : rtol^2

    # initialize iterates
    a = similar(μ, T)
    sum!(a, K)
    proxdivF1!(a, μ, ε)
    b = similar(ν, T)
    mul!(b, K', a)
    proxdivF2!(b, ν, ε)

    # caches for convergence checks
    a_old = similar(a)
    b_old = similar(b)

    isconverged = false
    _maxiter = maxiter === nothing ? 1_000 : maxiter
    for iter in 1:_maxiter
        # update cache if necessary
        ischeck = iter % check_convergence == 0
        if ischeck
            copyto!(a_old, a)
            copyto!(b_old, b)
        end

        # compute next iterates
        mul!(a, K, b)
        proxdivF1!(a, μ, ε)
        mul!(b, K', a)
        proxdivF2!(b, ν, ε)

        # check convergence of the scaling factors
        if ischeck
            # compute norm of current and previous scaling factors and their difference
            sqnorm_a_b = sum(abs2, a) + sum(abs2, b)
            sqnorm_a_b_old = sum(abs2, a_old) + sum(abs2, b_old)
            a_old .-= a
            b_old .-= b
            sqeuclidean_a_b = sum(abs2, a_old) + sum(abs2, b_old)
            @debug "Sinkhorn algorithm (" *
                   string(iter) *
                   "/" *
                   string(_maxiter) *
                   ": squared Euclidean distance of iterates = " *
                   string(sqeuclidean_a_b)

            # check convergence of `a`
            if sqeuclidean_a_b < max(sqatol, sqrtol * max(sqnorm_a_b, sqnorm_a_b_old))
                @debug "Sinkhorn algorithm ($iter/$_maxiter): converged"
                isconverged = true
                break
            end
        end
    end

    if !isconverged
        @warn "Sinkhorn algorithm ($_maxiter/$_maxiter): not converged"
    end

    return K .* a .* b'
end

"""
    sinkhorn_unbalanced2(μ, ν, C, λ1, λ2, ε; plan=nothing, kwargs...)
    sinkhorn_unbalanced2(μ, ν, C, proxdivF1!, proxdivF2!, ε; plan=nothing, kwargs...)

Compute the optimal transport plan for the unbalanced entropically regularized optimal
transport problem with source and target marginals `μ` and `ν`, cost matrix `C` of size
`(length(μ), length(ν))`, entropic regularization parameter `ε`, and marginal relaxation
terms `λ1` and `λ2` or soft marginal constraints with "proxdiv" functions `proxdivF1!` and
`proxdivF2!`.

A pre-computed optimal transport `plan` may be provided. The other keyword arguments
supported here are the same as those in the [`sinkhorn_unbalanced`](@ref) for unbalanced
optimal transport problems with general soft marginal constraints.

See also: [`sinkhorn_unbalanced`](@ref)
"""
function sinkhorn_unbalanced2(
    μ, ν, C, λ1_or_proxdivF1, λ2_or_proxdivF2, ε; plan=nothing, kwargs...
)
    γ = if plan === nothing
        sinkhorn_unbalanced(μ, ν, C, λ1_or_proxdivF1, λ2_or_proxdivF2, ε; kwargs...)
    else
        # check dimensions
        size(C) == (length(μ), length(ν)) ||
            error("cost matrix `C` must be of size `(length(μ), length(ν))`")
        size(plan) == size(C) || error(
            "optimal transport plan `plan` and cost matrix `C` must be of the same size",
        )
        plan
    end
    return dot(γ, C)
end


push!(LOAD_PATH, "..")

using LinearAlgebra
using JuMP, Optim
using LogisticOptTools
using Xpress, GLPK

const LOT = LogisticOptTools

SOLVER = optimizer_with_attributes(Xpress.Optimizer, "PRESOLVE" => 0,
                                   "MIPPRESOLVE" => 0)
SOLVER = () -> Xpress.Optimizer(
        PRESOLVE = 0,
        CUTSTRATEGY = 0,
        HEURSTRATEGY = 0,
        SYMMETRY = 0,
        OUTPUTLOG = 3
    )
SOLVER = GLPK.Optimizer

sparsity(s::Vector; tol=0.5) = findall(x -> x > tol, s)

function oracle(X, y, γ, s)
    pattern = sparsity(s)
    Xs = X[:, pattern]
    glm = LOT.DualLogisticRegressor(Xs, y, penalty=LOT.L2Penalty(γ))
    f, gradient!, _, _ = LOT.generate_callbacks(glm)
    algo = LBFGS()
    options = Optim.Options(iterations=250, g_tol=1e-5, show_trace=false)
    lower = LOT.lowerbound(glm.data)
    upper = LOT.upperbound(glm.data)
    x0 = 0.5 * (lower .+ upper)
    res = Optim.optimize(f, gradient!, lower, upper,
                         x0, Fminbox(algo), options)
    c = res.minimum
    α♯ = res.minimizer
    ∇c = zeros(length(s))

    @inbounds for i in eachindex(s)
        atmp = dot(X[:, i], α♯)
        ∇c[i] = - γ * atmp^2
    end
    return -c, ∇c
end

# TODO: for some reason lazy constraints are not working with Xpress yet
function benders(X, y, s0, γ, k; solver=SOLVER)
    n, p = size(X)
    c0, ∇c0 = oracle(X, y, γ, s0)

    benders_model = Model(solver)

    # Variables s and t
    @variable(benders_model, 0 <= s[i=1:p] <= 1, Bin, start=s0[i])
    @variable(benders_model, 0 <= t , start=c0)
    # Global constraints
    @constraint(benders_model, t >= c0 + dot(∇c0, s - s0))
    @constraint(benders_model, sum(s) <= k)
    # Objective
    @objective(benders_model, Min, t)

    function lazy_constraint_callback(cb_data)
        s_current = [callback_value(cb_data, s[j]) for j in eachindex(s)]
        println(s_current)
        ck, ∇ck = oracle(X, y, γ, s_current)
        new_optimality_cons = @build_constraint(ck + dot(∇ck, s - s_current) <= t)
        MOI.submit(benders_model, MOI.LazyConstraint(cb_data), new_optimality_cons)
    end

    MOI.set(benders_model, MOI.LazyConstraintCallback(), lazy_constraint_callback)
    optimize!(benders_model)
    return benders_model
end

function benders_manual(X, y, s0, γ, k; solver=SOLVER, niter=100)
    n, p = size(X)
    c0, ∇c0 = oracle(X, y, γ, s0)

    benders_model = Model(solver)

    # Variables s and t
    @variable(benders_model, 0 <= s[i=1:p] <= 1, Bin, start=s0[i])
    @variable(benders_model, 0 <= t , start=c0)
    # Global constraints
    @constraint(benders_model, t >= c0 + dot(∇c0, s - s0))
    @constraint(benders_model, sum(s) <= k)
    # Objective
    @objective(benders_model, Min, t)

    for i in 1:niter
        optimize!(benders_model)
        sk = JuMP.value.(s)
        ck, ∇ck = oracle(X, y, γ, sk)
        @constraint(benders_model, t >= ck + dot(∇ck, s - sk))
        lb = JuMP.value(t)
        ub = ck
        println("LB: ", lb,  " UB ", ub)

        if (ub - lb) / max(1.0, abs(lb)) <= 1e-6
            break
        end
    end

    return benders_model
end

γ = 1e-2
# Set number of non-null component
k = 2
DATASET = "diabetes"
#= DATASET = "covtype.libsvm.binary.bz2" =#
#= DATASET = "colon-cancer.bz2" =#
dataset = LOT.LogitData(DATASET, scale_data=true)
X, y = dataset.X, dataset.y

n, p = size(X)
# Scale penalty for dual problem
dγ = 1.0 / (4.0 * γ * n)

s0 = zeros(p)
s0[1:k] .= 1

m = benders_manual(X, y, s0, dγ, k)