# Computes coefficients for the I_v and J_v series expansions.

function Iv_series_coeff!(t::Transit_Struct{T}) where {T <: Real}
# Use series expansion to compute coefficients in the
# series solution of I_{v_max}, needed for k^2 < 1:
v = t.v_max
# Computing leading coefficient (n=0):
coeff = 2/(2v+one(T))
# Compute leading coefficient in I_v series:
t.Iv_coeff[1] = convert(T,coeff)
# Now, compute higher order terms until desired precision is reached:
for i = 1:t.nmax-1
  n = 2i
  coeff *= (n-1)*(n+2v-1)
  coeff /= n*(n+2v+1)
  t.Iv_coeff[i+1] = coeff
end
return
end

function dJvdk_series_coeff!(t::Transit_Struct{T}) where {T <: Real}
# Use series expansion to compute J_v and its derivative, dJ_v/dk.
# Computing leading coefficient (n=0).  Need to compute for
# both v_max and v_max-1.  Also, need to keep track of k^2 < 1
# and k^2 >= 1.  So, Jv_coeff is 3-dimensional array:
# Jv_coeff[2,2,nmax], where first dimension is for k^2 <1 (1)
# and k^2 >= 1 (2); second dimension is for v_max (1) and v_max-1 (2),
# while third dimension is for the series coefficients.
# Computing leading coefficient (n=0):
for k2 = 0:1
coeff = zero(T)
# Loop over v_max and v_max-1:
for j=1:2
  # v is v_max or v_max-1 (we need both since downward recursion requires two):
  v = t.v_max-(j-1)
  if k2 < 1
    coeff = 0.75*convert(T,pi)/exp(lfactorial(v+2))
    # multiply by (2v-1)!!
    @inbounds for i=2:2:2v
      coeff *= 0.5*(i-1)
    end
    # Add leading term to J_v:
    t.Jv_coeff[1,j,1] = coeff
    if t.grad
      t.dJvdk_coeff[1,j,1] = coeff*(2v+1)
    end
    # Now, compute higher order terms until desired precision is reached:
    @inbounds for i=1:t.nmax-1
      n = 2i
      coeff *= (n-1)*(n+2v-1)
      coeff /= n*(n+2v+4)
      t.Jv_coeff[1,j,i+1] = coeff
      if t.grad
        t.dJvdk_coeff[1,j,i+1] = coeff*(n+2v+1)
      end
    end
  else # k^2 >= 1
    coeff = convert(T,pi)
    # Compute (2v-1)!!/(2^v v!):
    @inbounds for i=2:2:2v
      coeff *= (i-1)/i
    end
    # Store leading terms:
    t.Jv_coeff[2,j,1] = coeff
    if t.grad
      t.dJvdk_coeff[2,j,1] = zero(T)
    end
    # Loop over higher order terms:
    for i = 1:t.nmax-1
      n = 2i
      coeff *= (n-5)*(n+2v-1)
      coeff /= (n*(n+2v))
      t.Jv_coeff[2,j,i+1] = coeff
      if t.grad
        t.dJvdk_coeff[2,j,i+1] = -n*coeff
      end
    end
  end
end
end
end


type Convergents
  q::PartialQuotients
  accuracy::Int
end

function Convergents(q::PartialQuotients; accuracy = 10)
  if q.accuracy < accuracy && q.periodlength == 0
    # To do : add a warn
    accuracy = q.accuracy
  end
  return Convergents(q, accuracy)
end

function pos2index(period, pos)
  (period == 0) && return pos
  (pos <= period + 1) && return pos
  return 2 + (pos - 2) % period
end

# state = (n, h_{n}, h_{n-1}, k_{n}, k_{n-1})
function Base.start(c::Convergents)
  return (0, big(1), big(0), big(0), big(1))
end

function Base.done(c::Convergents, state)
  return state[1] == c.accuracy
end

function Base.next(c::Convergents, state)
  n, h1, h2, k1, k2 = state
  n += 1
  i = pos2index(c.q.periodlength, n)
  h = c.q.partialquotients[i] * h1 + h2
  k = c.q.partialquotients[i] * k1 + k2
  return h // k, (n, h, h1, k, k1)
end


export cg, cg!

cg(A, b, Pl=1; kwargs...) =  cg!(zerox(A,b), A, b, Pl; kwargs...)

function cg!(x, A, b, Pl=1; tol::Real=size(A,2)*eps(), maxiter::Int=size(A,2))
    K = KrylovSubspace(A, length(b), 1, Vector{Adivtype(A,b)}[])
    init!(K, x)
    cg!(x, K, b, Pl; tol=tol, maxiter=maxiter)
end

function cg!(x, K::KrylovSubspace, b, Pl=1;
        tol::Real=size(K.A,2)*eps(), maxiter::Integer=size(K.A,2))
    resnorms = zeros(maxiter)

    tol = tol * norm(b)
    r = b - nextvec(K)
    p = z = isa(Pl, Function) ? Pl(r) : Pl\r
    γ = dot(r, z)
    for iter=1:maxiter
        append!(K, p)
        q = nextvec(K)
        α = γ/dot(p, q)
        # α>=0 || throw(PosSemidefException("α=$α"))
        update!(x, α, p)
        r -= α*q
        resnorms[iter] = norm(r)
        if resnorms[iter] < tol #Converged?
            resnorms = resnorms[1:iter]
            break
        end
        z = isa(Pl, Function) ? Pl(r) : Pl\r
        oldγ = γ
        γ = dot(r, z)
        β = γ/oldγ
        p = z + β*p
      end
    x, ConvergenceHistory(0<resnorms[end]<tol, tol, K.mvps, resnorms)
end


using RollingFunctions, Distributions, DataFrames
function EARS_C1(cases; alpha_val = 0.05, epsilon = 0, previous_days = 7, initial_date = 1)
    #Author: rodrigo.zepeda@imss.gob.mx
    #EXAMPLE:
    #cases = [1,2,1,0,0,1,2,3,1,1,0,0,0,1,20,40,1,1,1,0,0,0,1,1]
    #detected = EARS_C1(cases, initial_date = 1)
    #detected = EARS_C1([0,0,100], previous_days = 2) #Alerta en 100
    #detected = EARS_C1([0,0,100], previous_days = 3) #No sirve
    if (previous_days < length(cases))
        Ybar        = rollmean(cases, previous_days)
        sd          = map((x) -> maximum([x,epsilon]), rollstd(cases, previous_days))
        Y           = cases[(previous_days + 1):length(cases)]
        upper_bound = Ybar + quantile.(Normal(0, 1), 1 - alpha_val) * sd
        upper_bound = upper_bound[1:(length(upper_bound) - 1)]
        alert       = Y .> upper_bound
        idate       = (initial_date + previous_days):(initial_date + length(cases) - 1)
    else
        error("Vector cases too small. Need cases > previous_days for rollingmean")
    end
    return DataFrame(Alert = alert, U = upper_bound, Date = idate)
end


"""
    ck45()

Returns coefficients rka,rkb,rkc for the 4th order 5-stage low storage Carpenter/Kennedy
Runge Kutta method. Coefficients evolve the residual, solution, and local time, e.g.,

# Example
```julia
res = rk4a[i]*res + dt*rhs # i = RK stage
@. u += rk4b[i]*res
```
"""
function ck45()
    rk4a = [            0.0 ...
    -567301805773.0/1357537059087.0 ...
    -2404267990393.0/2016746695238.0 ...
    -3550918686646.0/2091501179385.0  ...
    -1275806237668.0/842570457699.0];

    rk4b = [ 1432997174477.0/9575080441755.0 ...
    5161836677717.0/13612068292357.0 ...
    1720146321549.0/2090206949498.0  ...
    3134564353537.0/4481467310338.0  ...
    2277821191437.0/14882151754819.0]

    rk4c = [ 0.0  ...
    1432997174477.0/9575080441755.0 ...
    2526269341429.0/6820363962896.0 ...
    2006345519317.0/3224310063776.0 ...
    2802321613138.0/2924317926251.0 ...
    1.0];
    return rk4a, rk4b, rk4c
end


using OmegaGrav
using MatterPower
using PyCall
using Dierckx
using OrdinaryDiffEq
using Tables, CSV
# %% Call the python wrapper for CLASS, `classy`, via PyCall
classy = pyimport("classy")
# Create an instance of the CLASS wrapper
cosmo = classy.Class()
# Create a dictionary of the cosmological parameters
deg_ncdm = true # 3 massive neutrinos with degenerate mass?
params = Dict(
   "output" => "mPk",
   "P_k_max_h/Mpc" => 100,
   "z_pk" => 10,
   "non linear" => "halofit",
   "A_s" => 2.097e-9,
   "n_s" => 0.9652,
   "k_pivot" => 0.05,
   "h" => 0.6737,
   "omega_b" => 0.02233,
   "omega_cdm" => 0.1198,
   "N_ncdm" => 1,
)
if deg_ncdm
   push!(params, "m_ncdm" => 0.02, "deg_ncdm" => 3, "N_ur" => 0.00641)
   mν = params["m_ncdm"] * params["deg_ncdm"]
else
   push!(params, "m_ncdm" => 0.06, "N_ur" => 2.0328)
   mν = params["m_ncdm"]
end
h0 = params["h"]
Ωcb = (params["omega_b"] + params["omega_cdm"]) / h0^2
Ων = mν / 93.14 / h0^2 #  Ωνh^2 = ∑mν / (93.14 eV)
Ωm = Ωcb + Ων
# Set the parameters and run the CLASS code
cosmo.set(params)
cosmo.compute()

# %% Compute the analytical linear theory result for K/W
# Reference: Equation (12) of Davis, Miller, White, ApJ, 490, 63 (1997).
zini = 9
a = 1/(1+zini):0.01:1.0 # scale factor
nred = length(a)
sol_δθ = setup_growth(Ωm, 1 - Ωm)
KovW = zeros(nred)
for i = 1:nred
   KovW[i] = -2 / 3 * (sol_δθ(a[i])[2] / sol_δθ(a[i])[1])^2 * a[i] / Ωm
end

# %% Compute `Ωgrav = Ωm * W / 2` at various redshifts
Ωgpk = zeros(nred)
Ωglin = zeros(nred)
for ired = 1:nred
   z = 1 / a[ired] - 1
   # Define functions to return linear and non-linear power spectra
   # Note: The CLASS code takes wavenumbers in units of 1/Mpc (no h) and
   # return power spectra in units of Mpc^3 (no 1/h^3).
   pknl_class(kovh) = cosmo.pk(kovh * h0, z) * h0^3
   pklin_class(kovh) = cosmo.pk_lin(kovh * h0, z) * h0^3
   # Spline interpolate in log(k)
   lnk = log(1e-4):0.05:log(100)
   pknl = Spline1D(lnk, pknl_class.(exp.(lnk)))
   pklin = Spline1D(lnk, pklin_class.(exp.(lnk)))
   # %% Compute Ωgrav from non-linear total matter P(k)
   Ωgpk[ired] = ograv_pk(x -> pknl(log(x)), z, Ωm)
   # %% Compute Ωgrav from linear total matter P(k)
   Ωglin[ired] = ograv_pk(x -> pklin(log(x)), z, Ωm)
end
Ωgrav_pklin = Spline1D(a, Ωglin)
Ωgrav_pknl = Spline1D(a, Ωgpk)

# %% Solve the Layzer-Irvine equation
tspan = (1 / (1 + zini), 1.0)
t0 = tspan[1]
## Non-linear P(k)
W(x) = 2 * Ωgrav_pknl(x)
dWda(x) = 2 * derivative(Ωgrav_pknl, x)
f(u, p, t) = -(2 * u + W(t)) / t - dWda(t)
u0 = KovW[1] * W(t0)
prob = ODEProblem(f, u0, tspan)
sol_pknl = solve(prob, Tsit5())
## Linear P(k): This result must agree with the analytical linear theory result
W(x) = 2 * Ωgrav_pklin(x)
dWda(x) = 2 * derivative(Ωgrav_pklin, x)
f(u, p, t) = -(2 * u + W(t)) / t - dWda(t)
u0 = KovW[1] * W(t0)
prob = ODEProblem(f, u0, tspan)
sol_pklin = solve(prob, Tsit5())

#%% Save table data to table2.csv
redshift = [0, 0.3, 0.5, 0.7, 1, 1.3, 1.5]
a = 1 ./ (1 .+ redshift)
t = Tables.table([redshift sol_pklin.(a) sol_pknl.(a)])
header = ["z", "Omega_K_pklin", "Omega_K_pknl"]
CSV.write("table2.csv", t, header = header)

# %% Clean CLASS (the equivalent of the struct_free() in the `main`
# of CLASS. This step is primordial when running in a loop over different
# cosmologies, as you will saturate your memory very fast if you ommit
# it.
cosmo.struct_cleanup()
# If you want to change completely the cosmology, you should also
# clean the arguments, otherwise, if you are simply running on a loop
# of different values for the same parameters, this step is not needed
cosmo.empty()


using MultivariateStats
using LinearAlgebra
using Test
using StableRNGs
import Statistics: mean, cov, var

@testset "Factor Analysis" begin

    rng = StableRNG(34568)

    ## FA with zero mean

    X = randn(rng, 5, 10)
    Y = randn(rng, 3, 10)

    W = qr(randn(rng, 5, 5)).Q[:, 1:3]
    Ψ = fill(0.1, 5)
    M = FactorAnalysis(Float64[], W, Ψ)

    @test size(M) == (5, 3)
    @test mean(M) == zeros(5)
    @test loadings(M) == W
    @test var(M) == Ψ

    T = inv(I+W'*diagm(0 => 1 ./ var(M))*W)*W'*diagm(0 => 1 ./ var(M))
    @test predict(M, X[:,1]) ≈ T * X[:,1]
    @test predict(M, X) ≈ T * X

    R = cov(M)*W*inv(W'W)
    @test reconstruct(M, Y[:,1]) ≈ R * Y[:,1]
    @test reconstruct(M, Y) ≈ R * Y

    ## PCA with non-zero mean

    mval = rand(rng, 5)
    M = FactorAnalysis(mval, W, Ψ)

    @test size(M) == (5, 3)
    @test mean(M) == mval
    @test loadings(M) == W
    @test var(M) == Ψ

    @test predict(M, X[:,1]) ≈ T * (X[:,1] .- mval)
    @test predict(M, X) ≈ T * (X .- mval)

    @test reconstruct(M, Y[:,1]) ≈ R * Y[:,1] .+ mval
    @test reconstruct(M, Y) ≈ R * Y .+ mval

    ## prepare training data

    d = 5
    n = 1000

    R = collect(qr(randn(rng, d, d)).Q)
    @test R'R ≈ Matrix(I, 5, 5)
    rmul!(R, Diagonal(sqrt.([0.5, 0.3, 0.1, 0.05, 0.05])))

    X = R'randn(rng, 5, n) .+ randn(rng, 5)
    mval = vec(mean(X, dims=2))
    Z = X .- mval

    # facm (default) & faem
    fa_methods = [:cm, :em]

    fas = FactorAnalysis[]
    for method in fa_methods
        let M = fit(FactorAnalysis, X, method=method, maxiter=5000),
            P = projection(M),
            W = loadings(M)
        push!(fas,M)

        @test size(M) == (5, 4)
        @test mean(M) == mval
        @test P'P ≈ Matrix(I, 4, 4)
        @test all(isapprox.(cov(M), cov(X, dims=2), atol=1e-3))

        M = fit(FactorAnalysis, X; mean=mval, method=method)
        @test loadings(M) ≈ W

        M = fit(FactorAnalysis, Z; mean=0, method=method)
        @test loadings(M) ≈ W

        M = fit(FactorAnalysis, X; maxoutdim=3, method=method)
        P = projection(M)

        @test size(M) == (5, 3)
        @test P'P ≈ Matrix(I, 3, 3)
        end
    end

    # compare two algorithms
    M1, M2 = fas
    @test all(isapprox.(cov(M1), cov(M2), atol=1e-3)) # noise
    LL(m, x) = (-size(x,2)/2)*(size(x,1)*log(2π) + log(det(cov(m))) + tr(inv(cov(m))*cov(x, dims=2)))
    @test LL(M1, X) ≈ LL(M2, X) # log likelihood

    # test with different types
    XX = convert.(Float32, X)
    YY = convert.(Float32, Y)

    for method in (:cm, :em)
        MM = fit(FactorAnalysis, XX; method=method, maxoutdim=3)

        # mixing types
        predict(M, XX)
        predict(MM, X)
        reconstruct(M, YY)
        reconstruct(MM, Y)

        # type stability
        for func in (mean, projection, cov, var, loadings)
            @test eltype(func(M)) == Float64
            @test eltype(func(MM)) == Float32
        end
    end

    # views
    M = fit(FactorAnalysis, view(XX, :, 1:100), method=:cm, maxoutdim=3)
    M = fit(FactorAnalysis, view(XX, :, 1:100), method=:em, maxoutdim=3)
end


include("imageio.jl")
import imageio

# Split out mandel as a function so we can call it if need be.
function mandel(xmin=-2.0, xmax=2.0, ymin=-2.0, ymax=2.0, max_iter=256, xres=2048, yres=2048, filename="mandel.pgm")

  rmax = 2.0
  pixels = zeros(Int64, xres, yres)

  for px = 1:xres
    x0 = (Float64(px)/Float64(xres)) * Float64(xmax - xmin) + xmin

    for py = 1:xres
      y0 = (Float64(py)/Float64(yres)) * Float64(ymax - ymin) + ymin

      iter = 0

      x = 0.0
      y = 0.0

      while (((x*x + y*y) < (rmax * rmax)) && (iter < max_iter))
        xt = (x*x) - (y*y) + x0
        y = 2.0 * x * y + y0
        x = xt
        iter = iter + 1
      end

      pixels[px,py] = max_iter - iter
    end

  end

  imageio.writepgm(pixels, max_iter - 1, "mandel.pgm")

end

mandel()


immutable Beta <: ContinuousUnivariateDistribution
    alpha::Float64
    beta::Float64
    function Beta(a::Real, b::Real)
        (a > zero(a) && b > zero(b)) || error("alpha and beta must be positive")
        new(float64(a), float64(b))
    end
end

Beta(a::Real) = Beta(a, a) # symmetric in [0, 1]
Beta() = Beta(1.0) # uniform

@_jl_dist_2p Beta beta

function entropy(d::Beta)
    o = lbeta(d.alpha, d.beta)
    o -= (d.alpha - 1.0) * digamma(d.alpha)
    o -= (d.beta - 1.0) * digamma(d.beta)
    o += (d.alpha + d.beta - 2.0) * digamma(d.alpha + d.beta)
    o
end

function kurtosis(d::Beta)
    α, β = d.alpha, d.beta
    num = 6.0 * ((α - β)^2 * (α + β + 1.0) - α * β * (α + β + 2.0))
    den = α * β * (α + β + 2.0) * (α + β + 3.0)
    num / den
end

mean(d::Beta) = d.alpha / (d.alpha + d.beta)

median(d::Beta) = quantile(d, 0.5)

function mode(d::Beta)
    α, β = d.alpha, d.beta
    if α >= 1.0
        if β > 1.0
            (α - 1.0) / (α + β - 2.0)
        elseif α == 1.0 && β == 1.0
            # Uniform[0,1]: what should be returned?
            0.5
        else
            1.0
        end
    else
        if β >= 1.0
            0.0
        else
            # not unique: return largest
            α > β ? 1.0 : 0.0
        end
    end
end

function modes(d::Beta)
    α, β = d.alpha, d.beta
    if α >= 1.0
        if β > 1.0
            [(α - 1.0) / (α + β - 2.0)]
        elseif α == 1.0 && β == 1.0
            # Uniform[0,1]: what should be returned?
            Float64[]
        else
            [1.0]
        end
    else
        if β >= 1.0
            [0.0]
        else
            [0.0,1.0]
        end
    end
end

function rand(d::Beta)
    u = randg(d.alpha)
    u / (u + randg(d.beta))
end

function rand!(d::Beta, A::Array{Float64})
    sa = GammaSampler(d.alpha)
    sb = GammaSampler(d.beta)

    for i = 1:length(A)
        u = rand(sa)
        v = rand(sb)
        @inbounds A[i] = u / (u + v)
    end

    return A
end

function skewness(d::Beta)
    num = 2.0 * (d.beta - d.alpha) * sqrt(d.alpha + d.beta + 1.0)
    den = (d.alpha + d.beta + 2.0) * sqrt(d.alpha * d.beta)
    num / den
end

function var(d::Beta)
    ab = d.alpha + d.beta
    d.alpha * d.beta / (ab * ab * (ab + 1.0))
end

function gradloglik(d::Beta, x::Float64)
  insupport(Beta, x) ? (d.alpha - 1.0) / x - (d.beta - 1.0) / (1 - x) : 0.0
end

### handling support

@continuous_distr_support Beta 0.0 1.0

## Fit model

# TODO: add MLE method (should be similar to Dirichlet)

# This is a moment-matching method (not MLE)
#
function fit(::Type{Beta}, x::Array)
    for xi in x
        insupport(Beta, xi) || error("Beta observations must be in [0,1]")
    end
    x_bar = mean(x)
    v_bar = varm(x, x_bar)
    α = x_bar * (((x_bar * (1.0 - x_bar)) / v_bar) - 1.0)
    β = (1.0 - x_bar) * (((x_bar * (1.0 - x_bar)) / v_bar) - 1.0)
    Beta(α, β)
end


using Polynomials
using Jacobi

include("testjacobi.jl")
include("testgauss.jl")


# Data container
immutable BiCGStabData{T}
    v::Vector{T}
    t::Vector{T}
    p::Vector{T}
    phat::Vector{T}
    s::Vector{T}
    shat::Vector{T}
    r::Vector{T}
    rtld::Vector{T}
end

BiCGStabData(n::Int, T::Type) = BiCGStabData(zeros(T, n), zeros(T, n),
                                             zeros(T, n), zeros(T, n),
                                             zeros(T, n), zeros(T, n),
                                             zeros(T, n), zeros(T, n))

# Main solver
function bicgstab!{T<:Real}(A, b::Vector{T}, x::Vector{T}; tol::Float64=1e-6, maxIter::Int64=100,
                                                           tolRho::Float64=1e-40, precon=copy!,
                                                           data=BiCGStabData(length(b), T))
    n = length(b)
    out_int = 0
    n_iter = 0
    converge = false
    bnrm2 = genblas_nrm2(b)
    if bnrm2 == 0.0
        x .= 0.0
        return 1, 0
    end
    A(data.r, x)
    genblas_scal!(-one(T), data.r)
    genblas_axpy!(one(T), b, data.r)
    residual = genblas_nrm2(data.r)/bnrm2
    if residual == 0.0
        return 2, 0
    end
    alpha = 1.0
    omega = 1.0
    data.rtld .= data.r
    rho1 = 0.0
    data.p .= data.r
    for iter = 1:maxIter
        rho = genblas_dot(data.rtld, data.r)
        if abs(rho) <= tolRho
            return -11, iter
        end
        if iter > 1
            beta  = (rho/rho1)*(alpha/omega)
            # p = r + beta*(p - omega*v)
            genblas_scal!(beta, data.p)
            genblas_axpy!(-(beta*omega), data.v, data.p)
            genblas_axpy!(one(T), data.r, data.p)
        end
        precon(data.phat, data.p)
        A(data.v, data.phat)
        alpha = rho/genblas_dot(data.rtld, data.v)
        # s = r - alpha*v
        data.s .= data.r
        genblas_axpy!(-alpha, data.v, data.s)
        # x = x + alpha*phat
        genblas_axpy!(alpha, data.phat, x)
        residual = genblas_nrm2(data.s)/bnrm2
        if residual <= tol
            return 31, iter
        end
        precon(data.shat, data.s)
        A(data.t, data.shat)
        omega = genblas_dot(data.t, data.s)/genblas_dot(data.t, data.t)
        # x = x + omega*shat
        genblas_axpy!(omega, data.shat, x)
        # r = s .- omega.*t
        data.r .= data.s
        genblas_axpy!(-omega, data.t, data.r)
        residual = genblas_nrm2(data.r)/bnrm2
        if residual <= tol
            return 32, iter
        end
        if abs(omega) < 1e-16
            return -12, iter
        end
        rho1 = rho
    end
    return -2, maxIter
end

# API

function bicgstab{T<:Real}(A, b::Vector{T}; tol::Float64=1e-6, maxIter::Int64=100,
                                            tolRho::Float64=1e-40, precon=copy!,
                                            data=BiCGStabData(length(b), T))
    x = zeros(b)
    exit_code, num_iters = bicgstab!(A, b, x, tol=tol, maxIter=maxIter, tolRho=tolRho, precon=precon, data=data)
    return x, exit_code, num_iters
end

export BiCGStabData, bicgstab!, bicgstab


#=

  Curious set of integers in Julia ConstraintSolver.jl

  Martin Gardner (February 1967):
  """
  The integers 1,3,8, and 120 form a set with a remarkable property:
  the product of any two integers is one less than a perfect square.
  Find a fifth number that can be added to the set without destroying
  this property.
  """

  Solution: The number is 0.

  There are however other sets of five numbers with this property.
  Here are the one in the range of 0..10000:

  [0, 1, 3, 8, 120]
  [0, 1, 3, 120, 1680]
  [0, 1, 8, 15, 528]
  [0, 1, 8, 120, 4095]
  [0, 1, 15, 24, 1520]
  [0, 1, 24, 35, 3480]
  [0, 1, 35, 48, 6888]
  [0, 2, 4, 12, 420]
  [0, 2, 12, 24, 2380]
  [0, 2, 24, 40, 7812]
  [0, 3, 5, 16, 1008]
  [0, 3, 8, 21, 2080]
  [0, 3, 16, 33, 6440]
  [0, 4, 6, 20, 1980]
  [0, 4, 12, 30, 5852]
  [0, 5, 7, 24, 3432]
  [0, 6, 8, 28, 5460]
  [0, 7, 9, 32, 8160]

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function curious_set_of_integers(print_solutions=true,all_solutions=false,timeout=6)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>timeout,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    n = 5
    dom_max = 120
    @variable(model, 0 <= x[1:n] <= dom_max, Int)

    @constraint(model, x in CS.AllDifferent())
    increasing_strict(model, x)

    # It's faster without these hints
    # @variable(model, 1 <= ixs[1:n] <= n, Int) # indices in X
    # @constraint(model, ixs in CS.AllDifferent())
    # my_element(model,ixs[1],x,1)
    # my_element(model,ixs[2],x,3)
    # my_element(model,ixs[3],x,8)
    # my_element(model,ixs[4],x,120)

    table = resize_matrix([ [i,j,i * j] for i in 0:dom_max, j in 0:dom_max])
    ps = []
    for i in 1:n, j in 1:i-1
        p = @variable(model, integer=true, lower_bound=0, upper_bound=dom_max)
        # @constraint(model, p*p -1 == x[i]*x[j]) # non-linear constraints!

        p_sq = @variable(model, integer=true, lower_bound=0, upper_bound=dom_max^2)
        @constraint(model, [p, p, p_sq] in CS.TableSet(table))

        xij_sq = @variable(model, integer=true, lower_bound=0, upper_bound=dom_max^2)
        @constraint(model, [x[i], x[j], xij_sq] in CS.TableSet(table))

        @constraint(model,p_sq-1 == xij_sq)

        push!(ps,p)
    end

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                println("x:$x_val")

            end
        end
    else
        println("status:$status")
    end

    return status
end

@time curious_set_of_integers(true,false)


using MomentClosure
using MomentClosure: define_μ
using ModelingToolkit: value
using Catalyst
using Test

@parameters t, σ_b, σ_u, ρ_b, ρ_u, σ_b_LMA
@variables g(t)

# simple feedback loop
rn_nonlinear = @reaction_network begin
      σ_b, g + p → 0
      σ_u*(1-g), 0 ⇒ g + p
      ρ_u, g → g + p
      ρ_b*(1-g), 0 ⇒ p
      1, p → 0
end σ_b σ_u ρ_b ρ_u

rn_linear = @reaction_network begin
      σ_b_LMA, g → 0
      σ_u*(1-g), 0 ⇒ g
      ρ_u, g → g+p
      (ρ_b*(1-g)), 0 ⇒ p
      1, p → 0
end σ_b_LMA σ_u ρ_b ρ_u

binary_vars = [speciesmap(rn_nonlinear)[g]]
LMA_eqs, _ = linear_mapping_approximation(rn_nonlinear, rn_linear, binary_vars, combinatoric_ratelaw=false)

μ = define_μ(2,3)
expr1 = LMA_eqs.odes.eqs[1].rhs
expr2 = σ_u - σ_b*μ[1,1] - σ_u*μ[1,0]
@test isequal(expr1, value.(expr2))
expr1 = LMA_eqs.odes.eqs[2].rhs
expr2 = ρ_b + ρ_u*μ[1,0] - ρ_b*μ[1,0] - μ[0,1]
@test isequal(expr1, value.(expr2))
expr1 = LMA_eqs.odes.eqs[3].rhs
expr2 = ρ_u*μ[1,0] + σ_u*μ[0,1] - μ[1,1] - σ_u*μ[1,1] - σ_b*μ[1,1]^2/μ[1,0]
@test isequal(expr1, value.(expr2))

# cooperativity cp=2
rn_nonlinear = @reaction_network begin
      σ_b, g + 2*p → 0
      σ_u*(1-g), 0 ⇒ g + 2*p
      ρ_u, g → g + p
      ρ_b*(1-g), 0 ⇒ p
      1, p → 0
end σ_b σ_u ρ_b ρ_u

binary_vars = [speciesmap(rn_nonlinear)[g]]
_, effective_params = linear_mapping_approximation(rn_nonlinear, rn_linear, binary_vars, combinatoric_ratelaw=false)
expr1 = effective_params[value(σ_b_LMA)]
expr2 = (σ_b*μ[1,2] - σ_b*μ[1,1]) / μ[1,0]
@test isequal(expr1, value.(expr2))


"""
    ex4_1_1(; formulation=1)
``trig'' is an MINLPLib test instance with one variable and power terms
(see http://www.minlplib.org/ex4_1_1.html). The global objective value for
this problem is -7.48731237. This code illustrates the strength of the
relaxation obtained by introducing one artificial variable for each non-linear
term as against one variable for the the entire univariate function.
Variables: -2 <= x <= 11; start = 10
Objective: Minimize 0.1 + x^6 - 2.08 * x^5 + 0.4875 * x^4 + 7.1 * x^3 - 3.95 * x^2 - x

The reformulation using one additional variable is as follows:
Variables:
    -2 <= x <= 11; start = 10
    y
Constraints: y = 0.1 + x^6 - 2.08 * x^5 + 0.4875 * x^4 + 7.1 * x^3 - 3.95 * x^2 - x
Objective: Minimize y

The reformulation using multiple additional variables is as follows:
Variables:
    -2 <= x <= 11; start = 10
    y[1:5]
Constraints:
    y[1] = x^6
    y[2] = -2.08 * x^5
    y[3] = 0.4875 * x^4
    y[4] = 7.1 * x^3
    y[5] = -3.95 * x^2
Objective: Minimize 0.1 + y[1] + y[2] + y[3] + y[4] + y[5] - x
"""

function ex4_1_1(; formulation = 1)
    if formulation == 1
        ex4_1_1_single()
    else
        ex4_1_1_multiple()
    end
end

function ex4_1_1_single()
    best_known_objective = -7.48731237
    cplex_optimizer =
        JuMP.optimizer_with_attributes(CPLEX.Optimizer, "CPX_PARAM_SCRIND" => 0)
    error_tolerances = [NaN64, 1e-1, 1e-2]
    base_partition = [-2.0, -1.7747, 0.1753, 11.0]
    milp_statistics = zeros(length(error_tolerances), 5)
    lp_statistics = zeros(length(error_tolerances), 5)
    func = x -> 0.1 + x^6 - 2.08 * x^5 + 0.4875 * x^4 + 7.1 * x^3 - 3.95 * x^2 - x
    for i = 1:length(error_tolerances)
        error_tolerance = error_tolerances[i]
        # construct MILP relaxations
        milp = Model(cplex_optimizer)
        milp_var, x_index, y_index =
            add_milp_relaxation(func, base_partition, error_tolerance, milp)
        @objective(milp, Min, milp_var[y_index])
        _, solve_time, solve_bytes_alloc, sec_in_gc = @timed optimize!(milp)
        relaxation_objective = objective_value(milp)
        @assert relaxation_objective ≤ best_known_objective
        relative_gap =
            abs(best_known_objective - relaxation_objective) / abs(relaxation_objective)
        milp_statistics[i, :] = [
            error_tolerance,
            best_known_objective,
            relaxation_objective,
            relative_gap,
            round(solve_time, digits = 2),
        ]
    end

    for i = 1:length(error_tolerances)
        error_tolerance = error_tolerances[i]
        # construct LP relaxations
        lp = Model(cplex_optimizer)
        lp_var, x_index, y_index =
            add_lp_relaxation(func, base_partition, error_tolerance, lp)
        @objective(lp, Min, lp_var[y_index])
        _, solve_time, solve_bytes_alloc, sec_in_gc = @timed optimize!(lp)
        relaxation_objective = objective_value(lp)
        @assert relaxation_objective ≤ best_known_objective
        relative_gap =
            abs(best_known_objective - relaxation_objective) / abs(relaxation_objective)
        lp_statistics[i, :] = [
            error_tolerance,
            best_known_objective,
            relaxation_objective,
            relative_gap,
            round(solve_time, digits = 2),
        ]
    end

    open("csv/single_ex4_1_1_milp.csv", "w") do io
        writedlm(io, ["error_tol" "g_opt" "relaxation" "relative_gap" "time"], ',')
        writedlm(io, milp_statistics, ",")
    end

    open("csv/single_ex4_1_1_lp.csv", "w") do io
        writedlm(io, ["error_tol" "g_opt" "relaxation" "relative_gap" "time"], ',')
        writedlm(io, lp_statistics, ",")
    end
end

function ex4_1_1_multiple()
    best_known_objective = -7.48731237
    cplex_optimizer =
        JuMP.optimizer_with_attributes(CPLEX.Optimizer, "CPX_PARAM_SCRIND" => 0)
    error_tolerances = [NaN64, 1e-1, 1e-2]
    base_partition = [-2.0, 0.0, 11.0]
    milp_statistics = zeros(length(error_tolerances), 5)
    lp_statistics = zeros(length(error_tolerances), 5)
    functions = Dict{Int,Function}()
    functions[1] = x -> x^6
    functions[2] = x -> -2.08 * x^5
    functions[3] = x -> 0.4875 * x^4
    functions[4] = x -> 7.1 * x^3
    functions[5] = x -> -3.95 * x^2
    for i = 1:length(error_tolerances)
        error_tolerance = error_tolerances[i]
        milp = Model(cplex_optimizer)
        milp_var = Dict{Int,Any}()
        x_indexes = Dict{Int,Int}()
        y_indexes = Dict{Int,Int}()
        # construct MILP relaxations
        for j = 1:5
            var, x_index, y_index = add_milp_relaxation(
                functions[j],
                base_partition,
                error_tolerance,
                milp,
                name = string(j),
            )
            milp_var[j] = var
            x_indexes[j] = x_index
            y_indexes[j] = y_index
            (j == 1) && (continue)
            @constraint(milp, var[x_index] == milp_var[1][x_indexes[1]])
        end
        @objective(
            milp,
            Min,
            0.1 + sum(milp_var[j][y_indexes[j]] for j = 1:5) - milp_var[1][x_indexes[1]]
        )
        _, solve_time, solve_bytes_alloc, sec_in_gc = @timed optimize!(milp)
        relaxation_objective = objective_value(milp)
        @assert relaxation_objective ≤ best_known_objective
        relative_gap =
            abs(best_known_objective - relaxation_objective) / abs(relaxation_objective)
        milp_statistics[i, :] = [
            error_tolerance,
            best_known_objective,
            relaxation_objective,
            relative_gap,
            round(solve_time, digits = 2),
        ]
    end

    for i = 1:length(error_tolerances)
        error_tolerance = error_tolerances[i]
        lp = Model(cplex_optimizer)
        lp_var = Dict{Int,Any}()
        x_indexes = Dict{Int,Int}()
        y_indexes = Dict{Int,Int}()
        # construct LP relaxations
        for j = 1:5
            var, x_index, y_index = add_lp_relaxation(
                functions[j],
                base_partition,
                error_tolerance,
                lp,
                name = string(j),
            )
            lp_var[j] = var
            x_indexes[j] = x_index
            y_indexes[j] = y_index
            (j == 1) && (continue)
            @constraint(lp, var[x_index] == lp_var[1][x_indexes[1]])
        end
        @objective(
            lp,
            Min,
            0.1 + sum(lp_var[j][y_indexes[j]] for j = 1:5) - lp_var[1][x_indexes[1]]
        )
        _, solve_time, solve_bytes_alloc, sec_in_gc = @timed optimize!(lp)
        relaxation_objective = objective_value(lp)
        @assert relaxation_objective ≤ best_known_objective
        relative_gap =
            abs(best_known_objective - relaxation_objective) / abs(relaxation_objective)
        lp_statistics[i, :] = [
            error_tolerance,
            best_known_objective,
            relaxation_objective,
            relative_gap,
            round(solve_time, digits = 2),
        ]
    end

    open("csv/multiple_ex4_1_1_milp.csv", "w") do io
        writedlm(io, ["error_tol" "g_opt" "relaxation" "relative_gap" "time"], ',')
        writedlm(io, milp_statistics, ",")
    end

    open("csv/multiple_ex4_1_1_lp.csv", "w") do io
        writedlm(io, ["error_tol" "g_opt" "relaxation" "relative_gap" "time"], ',')
        writedlm(io, lp_statistics, ",")
    end
end


# The Computer Language Benchmarks Game
# https://salsa.debian.org/benchmarksgame-team/benchmarksgame/

# based on the Javascript program
# optimizations by Kristoffer Carlsson

using Printf

A(i,j) = 1.0 / ( (((i+j)*(i+j+1)) >> 1) + i+1)

function Au!(w, u)
    n = length(u)
    Threads.@threads for i = 1:n
        w[i] = 0
        z = 0.0
        @simd for j = 1:n
           @inbounds z += A(i-1, j-1) * u[j]
        end
        w[i] = z
    end
end

function Atu!(v, w)
    n = length(w)
    Threads.@threads for i = 1:n
        z = 0.0
        @simd for j = 1:n
           @inbounds z += A(j-1,i-1) * w[j]
        end
        v[i] = z
    end
end

function AtAu!(w, v, u)
    Au!(w, u)
    Atu!(v, w)
end

function perf_spectralnorm(n::Int=100)
    u = ones(Float64, n)
    v = zeros(Float64 ,n)
    w = zeros(Float64, n)
    vv = vBv = 0
    for i = 1:10
        AtAu!(w, v, u)
        AtAu!(w, u, v)
    end
    for i = 1:n
        vBv += u[i]*v[i]
        vv += v[i]*v[i]
    end
    return sqrt(vBv/vv)
end

n = parse(Int,ARGS[1])
@printf("%.9f\n", perf_spectralnorm(n))


function clipdb(s,cutoff)
  # [clipped] = clipdb(s,cutoff)
  # Clip magnitude of s at its maximum + cutoff in dB.
  # Example: clip(s,-100) makes sure the minimum magnitude
  # of s is not more than 100dB below its maximum magnitude.
  # If s is zero, nothing is done.

  as = abs(s);
  mas = maximum(as);
  if mas==0; return; end
  if cutoff >= 0; return; end

  thresh = mas*10^(cutoff/20); # db to linear
  clipped = float(s);
  clipped[ as .< thresh ] = thresh;
  return clipped;
end


using Plots, Latexify, LaTeXStrings

e = [2.15, 3.11, 4.81, 7.06, 10, 13.1, 17.1, 22.1, 25.1, 28.4, 31.6, 34.7, 36.8, 41.7, 46.6, 51.7, 56]

b₀ = [21280, 12260, 9312, 6937, 6219, 5105, 5363, 3960, 3063, 2967, 3113, 2449, 2237, 1979, 2266, 2266, 2266]

b₁ = [1231, 973.3, 776.8, 697.9, 589.8, 572.8, 415.7, 362.1, 366, 333.3, 281.1, 302, 314, 325.3, 302, 302, 302]

scatter(e,b₀,legend=true,
 xlabel="Thickness", ylabel = "Unit pressure drop", labels = "b₀", title = "NC 2733 foam", xlim=(0,60),ylim=(0,25000))
plot!(e, b₁, label = "b₁", seriestype = :scatter)

A = [e[i]^j for i in 1:17, j in -1:0]

ê₀ = A \ b₀

b̂₀(e) = ê₀' * [e^-1,1]

ê₁ = A \ b₁

b̂₁(e) = ê₁' * [e^-1,1]

plot!(b̂₀, label = "b₀(e)")
plot!(b̂₁, label = "b₁(e)")

E = [5.12, 10, 13.2, 15.6, 20.4, 21.7, 25.3, 30.7, 37.5, 41.8, 44.9, 47.2, 50.8, 55, 59.9, 62.1]

b₂ = [5665, 2689, 2144, 2236, 1592, 1872, 1151, 1196, 1334, 947, 1076, 1060, 1093, 888.7, 817.1, 949.2]

b₃ = [256.2, 271.7, 235.9, 208.6, 171.3, 139.5, 155.8, 127.5, 93.9, 104.2, 91.3, 86.1, 79.9, 87.6, 86.7, 82.4]

scatter(E,b₂,legend=true,
 xlabel="Thickness", ylabel = "Unit pressure drop", label = "b₂", xlim=(0,70), ylim=(0,6000))
plot!(E, b₃, label = "b₃",seriestype = :scatter)

B = [E[i]^j for i in 1:16, j in -1:0]

ê₂ = B \ b₂

b̂₂(E) = ê₂' * [E^-1,1]

ê₃ = B \ b₃

b̂₃(E) = ê₃' * [E^-1,1]

plot!(b̂₂, label = "b₂(e)")
plot!(b̂₃, label = "b₃(e)")

compₐ = []
erroₐ = []
a_C = [1623.500, 433.890, 39415.000, 25798.000, 284.200, 83.501, 2177.600, 1230.200]
a_AR = [ê₀[2], ê₂[2], ê₀[1], ê₂[1], ê₁[2], ê₃[2], ê₁[1], ê₃[1]]

for i in 1:8
    append!(compₐ, abs(a_AR[i]-a_C[i]))
    append!(erroₐ, round(compₐ[i]/a_C[i]*100, digits=2))
end

NC2733_AR = []
NC2733_C = []
NC2733_erro = []
NCX1116_AR = []
NCX1116_C = []
NCX1116_erro = []

for i in 1:4
    imp = 2i - 1
    par = 2i
    append!(NC2733_AR, a_AR[imp])
    append!(NC2733_C, a_C[imp])
    append!(NC2733_erro, erroₐ[imp])
    append!(NCX1116_AR, a_AR[par])
    append!(NCX1116_C, a_C[par])
    append!(NCX1116_erro, erroₐ[par])

end

using DataFrames

df = DataFrame(Coeficientes=["a₀", "a₁", "a₂", "a₃"],
    NC2733_Artigo=NC2733_AR, NC2733_Experimento=NC2733_C, NC2733_Erro=NC2733_erro,
    NCX1116_Artigo=NCX1116_AR, NCX1116_Experimento=NCX1116_C, NCX1116_Erro=NCX1116_erro)

show(df, allcols=true)
print("\n\n(Erros já em porcentagem)")
display(df, )

using PyCall
import CoolProp.CoolProp

CP = pyimport("CoolProp.CoolProp")

T = 90 + 273.15
P = 100

stt = CP.State("Air", {"T":T, "P":P})


module DoubleIntegrator

using MDPs, PGFPlots
using DelimitedFiles

export state_space, action_space, transition, reward, simulate,
       viz_policy, viz_trajectory
export XMIN, XMAX, VMIN, VMAX

# mdp
const MX = 50
const MV = 50
const N = 1000

# state
const XMIN = -3.
const XMAX = 3.
const VMIN = -3.
const VMAX = 3.

# action
const AMIN = -1.0
const AMAX = 1.0

# transition
const DT = 0.1

# reward
const K = -.5

# test
const T = 200

function state_space(mx::Int64=MX, mv::Int64=MV)
    xs = range(XMIN, stop=XMAX, length=mx)
    vs = range(VMIN, stop=VMAX, length=mv)
    return RectangleGrid(xs, vs)
end

function action_space(n::Int64=N)
    return range(AMIN, stop=AMAX, length=n)
end

function transition(state::Vector{Float64}, action::Float64)
    snext = zeros(length(state))
    snext[2] = state[2] + action * DT
    snext[2] = clip(snext[2], VMIN, VMAX)
    snext[1] = clip(state[1] + snext[2] * DT, XMIN, XMAX)
    return snext
end

function clip(val::Float64, minval::Float64, maxval::Float64)
    return min(maxval, max(minval, val))
end

function reward(state::Vector{Float64}, action::Float64)
    return K * (state[1]^2 + state[2]^2)
end

function simulate(mdp::MDP, policy::Policy, state::Vector{Float64})
    trajectory = zeros(T, dimensions(mdp.S))
    actions = zeros(T)
    for t = 1:T
        trajectory[t, :] = state
        if state[1] == 0.0
            trajectory = trajectory[1:t, :]
            actions = actions[1:t - 1]
            break
        end
        action = policy!(policy, get_belief(mdp, state))
        state = mdp.transition(state, action)
        actions[t] = action
    end
    return trajectory, actions
end

function viz_policy(mdp::MDP, policy::Policy, title::String="Policy heatmap", saveplot::Bool=false, name::String="")
    function getmap(x::Float64, v::Float64)
        return policy!(policy, get_belief(mdp, [x, v]))
    end

    p = Axis([Plots.Image(getmap, (XMIN, XMAX), (VMIN, VMAX),
                      xbins = 250, ybins = 250,
                      colormap = ColorMaps.Named("Blues"))],
          width="14cm", height="14cm",
          xlabel="position", ylabel="speed",
          title=title)
    p
    if saveplot
      PGFPlots.save(name, p)
    end
end

function viz_trajectory(trajectory::Matrix{Float64}, actions::Vector{Float64},
                        title1::String="position over time", title2::String="acceleration over time",
                        saveplot::Bool=false, name::String="")
    g = GroupPlot(2, 1, groupStyle="horizontal sep=1.5cm")
    push!(g, Axis([Plots.Linear(trajectory[:, 1])],
                  width="12cm", height="12cm",
                  xlabel="time", ylabel="position",
                  title=title1))
    push!(g, Axis([Plots.Linear(actions)],
                  width="12cm", height="12cm",
                  xlabel="time", ylabel="acceleration",
                  title=title2))
    g
    if saveplot
      PGFPlots.save(name, g)
    end
end

end


# This is the first Stan model in Statistical Rethinking Edition 1 (page 249).
# In Edition 2 (page 242)

# \toc

# ## Data

import CSV
import Random

using DataFrames
using Turing
using TuringModels

Random.seed!(1)

data_path = joinpath(TuringModels.project_root, "data", "rugged.csv")
df = CSV.read(data_path, DataFrame)

df.log_gdp = log.(df.rgdppc_2000)
dropmissing!(df)

df = select(df, :log_gdp, :rugged, :cont_africa);

# DataFrame `df` is shown Section [df](#df).

# ## Model

@model function model_fn(y, x₁, x₂)
  σ ~ truncated(Cauchy(0, 2), 0, Inf)
  βAR ~ Normal(0, 10)
  βR ~ Normal(0, 10)
  βA ~ Normal(0, 10)
  α ~ Normal(0, 100)

  μ = α .+ βR * x₁ .+ βA * x₂ .+ βAR * x₁ .* x₂
  y ~ MvNormal(μ, σ)
end

model = model_fn(df.log_gdp, df.rugged, df.cont_africa);

# ## Output

chns = sample(model, NUTS(), 1000)

# \defaultoutput{}

# ## Original output

"""
Iterations = 1:1000
Thinning interval = 1
Chains = 1,2,3,4
Samples per chain = 1000

Empirical Posterior Estimates:
          Mean         SD       Naive SE       MCSE      ESS
    a  9.22360053 0.139119116 0.0021996664 0.0034632816 1000
   bR -0.20196346 0.076106388 0.0012033477 0.0018370185 1000
   bA -1.94430980 0.227080488 0.0035904578 0.0057840746 1000
  bAR  0.39071684 0.131889143 0.0020853505 0.0032749642 1000
sigma  0.95036370 0.052161768 0.0008247500 0.0009204073 1000

Quantiles:
          2.5%       25.0%       50.0%      75.0%        97.5%
    a  8.95307475  9.12719750  9.2237750  9.31974000  9.490234250
   bR -0.35217930 -0.25334425 -0.2012855 -0.15124725 -0.054216855
   bA -2.39010825 -2.09894500 -1.9432550 -1.78643000 -1.513974250
  bAR  0.13496995  0.30095575  0.3916590  0.47887625  0.650244475
sigma  0.85376115  0.91363250  0.9484920  0.98405750  1.058573750
""";

# ## df

df


module Weights
using RecipesBase
using ControlSystems, MonteCarloMeasurements
export makeweight, neglected_delay, gain_and_delay_uncertainty, neglected_lag
export fit_complex_perturbations

"""
    makeweight(low, [mid,] high)

Create a weighting function that goes from gain `low` at zero frequency, through gain `mid` to gain `high` at ∞

# Arguments:
- `low`: A number specifying the DC gain
- `mid`: A number specifying the frequency at which the gain is 1, or a tuple `(freq, gain)`. If `mid` is not specified, the geometric mean of `high` and `low` is used.
- `high`: A number specifying the gain at ∞
"""
function makeweight(low, mid::Number, high)
    makeweight(low, (mid, high < 1 ? sqrt(high*low) : 1), high)
end

makeweight(low, high) = makeweight(low, √(high*low), high)

function makeweight(low, mid, high)
    freq, mag = mid
    p = freq * √(((high/mag)^2-1)/(1-(low/mag)^2))
    W = tf([high, low*p], [1, p])
    W = ss(W)
end

"""
    neglected_delay(Lmax)

Return a multiplicative weight to represent the uncertainty coming from neglecting the dynamics `exp(-s*L)`
where `L ≤ Lmax`.
"Multivariable Feedback Control: Analysis and Design" Ch 7.4.5
"""
function neglected_delay(Lmax)
    gain_and_delay_uncertainty(1, 1, Lmax)
end

"""
    gain_and_delay_uncertainty(kmin, kmax, Lmax)

Return a multiplicative weight to represent the uncertainty coming from neglecting the dynamics `k*exp(-s*L)`
where `k ∈ [kmin, kmax]` and `L ≤ Lmax`.
This weight is slightly optimistic, an expression for a more exact weight appears in eq (7.27), "Multivariable Feedback Control: Analysis and Design"
"""
function gain_and_delay_uncertainty(kmin, kmax, Lmax)
    kb = (kmin+kmax)/2
    rk = (kmax-kmin)/2/kb
    tf([(1+rk/2)*Lmax, rk], [Lmax/2, 1])
end

"""
    neglected_lag(τmax)

Return a multiplicative weight to represent the uncertainty coming from neglecting the dynamics `1/(s*τ + 1)`
where `τ ≤ τmax`.
"Multivariable Feedback Control: Analysis and Design" Ch 7.4.5
"""
function neglected_lag(τmax)
    tf([τmax, 0], [τmax, 1])
end

"""
    centers, radii = fit_complex_perturbations(P, w; relative=true, nominal=:mean)

For each frequency in `w`, fit a circle in the complex plane that contains all models in the model set `P`, represented as an `LTISystem` with `Particles` coefficients. Note, the resulting radii can be quite unstable if the number of particles is small, in particular if the particles are normally distributed instead of uniformly.

If `realtive = true`, circles encompassing `|(P - Pn)/Pn|` will be returned (multiplicative/relative uncertainty).
If `realtive = false`, circles encompassing `|P - Pn|` will be returned (additive uncertainty).

If `nominal = :mean`, the mean of `P` will be used as nominal model. If `nominal = :first`, the first particle will be used. See `MonteCarloMeasurements.with_nominal` to set the nominal value in the first particle.

See also [`nyquistcircles`](@ref) to plot circles (only if relative=false).
"""
function fit_complex_perturbations(P, w; relative=true, nominal=:mean)
    ControlSystems.issiso(P) || throw(ArgumentError("This function only works for SISO systems."))
    r = freqresp(P, w)
    centers = Vector{Complex{eltype(w)}}(undef, length(w))
    radii = Vector{eltype(w)}(undef, length(w))
    for i in axes(r,1)
        ri = r[i,1,1] # NOTE: assumes SISO
        c = nominal === :mean ? pmean(ri) : MonteCarloMeasurements.nominal(ri)
        rad = abs(relative ? pmaximum((ri - c)/c) : pmaximum(ri - c))
        centers[i] = c
        radii[i] = rad
    end
    centers, radii
end

"""
    nyquistcircles(w, centers, radii)

Plot the nyquist curve with circles. It only makes sense to call this function if the circles represent additive uncertainty, i.e., were calculated with `relative=false`.

See also [`fit_complex_perturbations`](@ref)
"""
nyquistcircles
@userplot nyquistcircles
@recipe function nyquistcircles(h::nyquistcircles)
    w, centers, radii = h.args
    t = LinRange(0, 2pi, 100)
    re, im = cos.(t), sin.(t)
    @series begin
        lalbe --> "Center"
        centers
    end
    for i in eachindex(radii)
        @series begin
            alpha --> 0.3
            color --> 1
            fillcolor --> 1
            label --> ""
            re .* radii[i] .+ centers[i].re, im .* radii[i] .+ centers[i].im
        end
    end
end

end


"""
  generality(N::UnipartiteNetwork)

This function mesure the generality. The generality represents the link density (L/S) divided by the number of predators. The number of predators: (prop. of int + prop. of top)xS
"""
function generality(N::UnipartiteNetwork)
    out_degree = collect(values(degree(N; dims=1)))
    in_degree = collect(values(degree(N; dims=2)))
    degrees = DataFrame(out=out_degree, in=in_degree)
    B = nrow(degrees[degrees.out .== 0, :]) / nrow(degrees)
    T = nrow(degrees[degrees.in .== 0, :]) / nrow(degrees)
    I = nrow(degrees[(degrees.out .!= 0) .& (degrees.in .!= 0), :]) / nrow(degrees)
    L = links(N)
    Generality = L ./ (S.^2 .* (T .+ I))
    return Generality
end


using BenchmarkTools, FinancialMonteCarlo, DualNumbers, CUDA;
CUDA.allowscalar(false)
S0 = 100.0;
K = 100.0;
r = 0.02;
T = 1.0;
d = 0.01;
D = 90.0;

Nsim = 10000;
Nstep = 30;
sigma = 0.2;
sigma_zero = 0.2;
kappa = 0.01;
theta = 0.03;
lambda = 0.01;
rho = 0.0;
mc = MonteCarloConfiguration(Nsim, Nstep);

mc_2 = MonteCarloConfiguration(Nsim, Nstep, FinancialMonteCarlo.CudaMode());
toll = 0.8;

rfCurve = ZeroRate(r);

FwdData = Forward(T)
EUData = EuropeanOption(T, K)
AMData = AmericanOption(T, K)
BarrierData = BarrierOptionDownOut(T, K, D)
AsianFloatingStrikeData = AsianFloatingStrikeOption(T)
AsianFixedStrikeData = AsianFixedStrikeOption(T, K)
Model = HestonProcess(sigma, sigma_zero, lambda, kappa, rho, theta, Underlying(S0, d));

@show "STD fwd"
@btime FwdPrice = pricer(Model, rfCurve, mc, FwdData);

@show "CUDA_2 fwd"
@btime FwdPrice = pricer(Model, rfCurve, mc_2, FwdData);
@show "std eu"
@btime EuPrice = pricer(Model, rfCurve, mc, EUData);

@show "CUDA_2 eu"
@btime EuPrice = pricer(Model, rfCurve, mc_2, EUData);
@show "std am"
@btime AmPrice = pricer(Model, rfCurve, mc, AMData);

@show "CUDA_2 am"
@btime AmPrice = pricer(Model, rfCurve, mc_2, AMData);


using JuMP
using CPLEX
#using GLPK

function createLSP(params, nodes, demands, costs)

	#model = Model(GLPK.Optimizer)
	model = Model(CPLEX.Optimizer)

	#= Paramètres (Constantes) du problème

		demands[i][t] : d_{i,t}= demande du revendeur i au pas de temps t | i in [1,n], t in [1,l]=T
		nodes[i]["h"] : h_i = coût de stockage d'un produit pour le revendeur/fournisseur i | i in [0,n]
		node[i]["L"] : L_i le stock maximal du revendeur/fournisseur i | i in [0,n]
		params["u"] : u = coût de production d'un produit pour le fournisseur
		params["f"] : f = coût de set up (pour chaque période où le fournisseur décide de produire, il paye un coût f)
		params["C"] : C = Capacité de production max du fournisseur (=M dans le sujet)

	=#

	#= Variables de décision

		p[t] : p_t  (Variable de production): quantité produite à la période t
		y[t] : y_t  (Variable de lancement): variable binaire qui vaut 1 si une production est lancée à la période t, et 0 sinon

		I[i,j] : I_{i,t} (Variable de stockage): quantité en stock à la fin de la période t pour i.

		q[i,j] : q_{i,t} (Variable d’approvisionnement): quantité produite pour le revendeur i à la période t.

	=#

	#Variables Réel positive
	@variable(model, p[1:params["l"]] >= 0)
	@variable(model, I[0:params["n"], 0:params["l"]] >= 0)
	@variable(model, q[1:params["n"], 1:params["l"]] >= 0)

	#Variable binaire
	@variable(model, y[1:params["l"]], Bin)

	for i in 0:params["n"]
		@constraint(model, I[i, 0] == nodes[i]["L0"])
	end

	###Fonction objectif

	obj = 0
	for t in 1:params["l"]

		#On ajoute le cout de production de touts les produits au temps t
		obj += params["u"] * p[t]

		#Si on produit au temps t, on ajoute le cout de setup
		obj += params["f"] * y[t]

		#??? On ne prend pas le cout de stockage de fournisseur ?
		for i in 1:params["n"]
			obj += nodes[i]["h"] * I[i, t]
		end

	end
	@objective(model, Min, obj)

	###Contraintes
	for t in 1:params["l"]

		@constraint(model, I[0, t-1] + p[t] == I[0, t] + sum(q[i, t] for i in 1:params["n"])) #contraintes 1

		for i in 1:params["n"]
			@constraint(model, I[i, t-1] + q[i, t] == demands[i, t] + I[i, t]) #contraintes 2
			@constraint(model, I[i, t-1] + q[i, t] <= nodes[i]["L"]) #contraintes 5
		end

		@constraint(model, p[t] <= params["C"] * y[t]) #constraintes 3
		@constraint(model, I[0, t-1] <= nodes[0]["L"]) #contraintes 4

	end

	return model

end


"""
    hsvd(fid, rank)

Use the Hankel singular value decomposition method for fitting a sum of
exponential signals of the form

    xₙ = Σₖ Aₖ zₖⁿ

to compute the ampltudes Aₖ and complex-valued poles zₖ. Any poles with |zₖ| >
1 are discarded to avoid numerical instability during the fit. A tuple of poles,
amplitudes and the exponential basis (z,A,basis) is returned.

The discrete-time and continuous time signals are related by

    sₙ =  Aₖ zₖⁿ  =  Aₖ exp((-bₖ + iωₖ) * tₙ)

or

    zₖ  =  exp((-bₖ + iωₖ) * Δt)

where ωₖ is an angular frequency and bₖ the decay rate in Hz.

Ref. to NMR literature: Barkhuijsen et al., "Improved Algorithm for Noniterative
Time-Domain Model Fitting to Exponentially Damped Magnetic Resonance Signals",
J. Mag. Res., (1987) https://doi.org/10.1016/0022-2364(87)90023-0

More broadly these techniques seem to come under the title "Singular Spectrum
analysis".
"""
function hsvd(fid, rank)
    # In the code below we attempt to use the same notation as the paper.
    # Construct the Hankel matrix
    N = length(fid)
    L = N÷2
    M = N-L
    X = zeros(ComplexF64, L, M)
    for i = 1:M
        X[:, i] .= fid[i:i+L-1]
    end

    # Form a low-rank approximation to the signal basis by truncating the SVD U
    # factor to the given rank.
    hsvd = svd(X)
    Xl = hsvd.U[:, 1:rank]

    # As described in the paper, each row of Xl is the row below multiplied by
    # some matrix Zp, so we can compute Zp with a linear by solving the
    # overdetermined system `Ub*Zp = Ut` as a least squares problem:
    Ub = Xl[1:end-1, :]
    Ut = Xl[2:end, :]
    Zp = Ub \ Ut

    # The eigen decomposition of Zp now gives the eigenvalues as the poles
    Zeig = eigen(Zp)
    Zeig.values
end

"""
Use HSVD to suppress a spectral region.

This is rather a WIP implementation. TODO
* Make all this work with AxisArrays
* Proper testing...
"""
function hsvd_water_suppression(fid; dt = 0.0008, water_center_freq=0.0,
                                bandwidth=100.0, hsvd_rank=50)
    zs0 = hsvd(fid, hsvd_rank)

    # Filter out any exponentially growing components. These can occur if a
    # naive truncation-based frequency domain downsampling of the time domain
    # signal is used. (For example this seems to be the default on Siemens
    # scanners when using the "remove oversampling" option.)
    zs = zs0[abs.(zs0) .< 1]

    # Generate a basis from the components and fit to it to determine the
    # amplitudes of the various signal components.
    basis = [zs[i]^n for n=0:length(fid)-1, i=1:length(zs)]
    A = basis \ fid

    # Compute MR frequencies of components
    frequencies = imag(log.(zs))/dt

    # Assume all signal within some spectral bandwidth of water centre frequency
    water_components = abs.(frequencies .- water_center_freq) .< bandwidth

    water = basis[:,water_components] * A[water_components]

    fid .- water
end


using IntervalOptimisation, TaylorModels, Test, RangeEnclosures, SumOfSquares, SDPA

using DynamicPolynomials: @polyvar

available_solvers = (NaturalEnclosure(), TaylorModelsEnclosure(),
                     BranchAndBoundEnclosure(), MooreSkelboeEnclosure())

include("univariate.jl")
include("multivariate.jl")


function crandn(dims...)
    sqrt(0.5) * ( randn(dims) + 1im*randn(dims) )
end
"""
    scm_channel(angles,weights,nAntennas::Integer;nCoherence=1,AS=2.0)

Generate `nCoherence` channel realization according to angular power density (spectrum) `f`
where `f` is given as a superposition of weighted Laplace kernels located at `angles` and
with standard deviation `AS`
The spectrum `f` defines a covariance matrix of the channel `h`.
The first row of this covariance matrix is returned as second output.
"""
function scm_channel(angles,weights,nAntennas::Integer;nCoherence=1,AS=2.0)
    f(x)  = spectrum(x,angles,weights;AS=AS)
    (h,t) = chan_from_spectrum(f,nAntennas,nCoherence=nCoherence)
end

"""
    chan_from_spectrum(f,nAntennas::Integer;nCoherence=1)

Generate `nCoherence` channel realization according to angular power density (spectrum) `f`
(which is a cont. function).
The spectrum `f` defines a covariance matrix of the channel `h`.
The first row of this covariance matrix is returned as second output.
"""
function chan_from_spectrum(f,nAntennas::Integer;nCoherence=1)
    OF = 20 # oversampling factor (ideally, would use continuous freq. spectrum...)
    nFreqSamples = OF*nAntennas
    fs = Vector{Float64}([f(i/nFreqSamples*2*pi) for i=0:nFreqSamples-1])

    # avoid instabilities due to almost infinite energy at some frequencies
    # (this should only happen at "endfire" of a uniform linear array where --
    # because of the arcsin-transform -- the angular psd grows to infinity
    almostInfThreshold = max(1,nFreqSamples) # use nFreqSamples as threshold value...
    almostInfFreqs = abs.(fs) .> almostInfThreshold
    fs[almostInfFreqs] = almostInfThreshold * angle.(fs[almostInfFreqs])

    if sum(fs) > 0
        fs = fs./sum(fs)*nFreqSamples # normalize energy
    end

    x = crandn(nFreqSamples,nCoherence);
    h = ifft(sqrt.(fs).*x,1).*sqrt(nFreqSamples)
    h = h[1:nAntennas,:]

    # t is the first row of the covariance matrix of h (which is Toeplitz and Hermitian)
    t = fft(fs)./nFreqSamples
    t = t[1:nAntennas]
    (h,t)
end

"""
    spectrum(x_rad,angles,weights,AS=2.0)

Generate spectrum at values `asin(x_rad/pi) [deg]` corresponding
to superposition of weighted Laplace kernels with means `angles`
and standard deviations `AS`.
This functions takes care of the folding that happens when transforming
from [-180,180] deg. to [rad]-space...
"""
function spectrum(x_rad,angles,weights;AS=2.0)
    x_rad = mod(x_rad + pi,2*pi)-pi
    x_deg = asin(x_rad/pi)./pi*180

    v =     laplace(    x_deg,angles,weights,AS=AS)
    v = v + laplace(180-x_deg,angles,weights,AS=AS)
    v = v .* (180/pi*2*pi./sqrt(pi^2 - x_rad.^2))
end

"""
    laplace(x_deg::Float64,angles,weights;AS=2.0)

Value at `x_deg` (in deg) of sum of weighted Laplace kernels
(kernels are wrapped to [-180,180] degrees).
"""
function laplace(x_deg::Float64,angles,weights;AS=2.0)
    v = 0.0;
    for i in 1:length(weights)
        xshifted = x_deg - angles[i]
        xshifted = mod(xshifted + 180,360)-180;
        v += weights[i]/(2*AS)*exp(-abs(xshifted)/AS)
    end
    v
end

"""
    laplace(x_deg::Vector{Float64},angles,weights;AS=2.0)

Value at `x_deg` (in deg) of sum of weighted Laplace kernels
(kernels are wrapped to [-180,180] degrees).
"""
function laplace(x_deg::Vector{Float64},angles,weights;AS=2.0)
    v = zeros(x_deg)
    for j in 1:length(x_deg)
        vj = 0.0;
        for i in 1:length(weights)
            xshifted = x_deg[j] - angles[i]
            xshifted = mod(xshifted + 180,360)-180;
            vj += weights[i]/(2*AS)*exp(-abs.(xshifted)/AS)
        end
        v[j] = vj;
    end
    v
end


using Mimi
using OptiMimi
using Test

# Basic model to optimize:
# Each period, you have 1 unit to divide between consuming and staving off disaster.
# If disaster happens, you get 0 consumption for the remainder of the game.
# Utility is defined as sqrt(c).  The probability of disaster is 1-c in each period.
# Assume 5% discounting and 10 periods.

## Analytical solution

## V_t = max_c u(c) + gamma (1 - c) V_t+1
## u'(c) - gamma V_t+1 = 0
## Let u(c) = sqrt(c), u'(c) = .5 c^(-.5)
## c = 1/(2 gamma V_t+1)^2
gamma = exp(-.05) # 5% discount rate

# consume everything in last period
VV_reverse = [1.]
consumption_reverse = [1.]

for tt in 10:-1:1
    consumption = 1 / (2 * gamma * VV_reverse[end])^2
    VV = sqrt(consumption) + gamma * (1 - consumption) * VV_reverse[end]
    push!(consumption_reverse, consumption)
    push!(VV_reverse, VV)
end

println(reverse(consumption_reverse))
println(reverse(VV_reverse))

# OptiMimi solution

# Create simple model
@defcomp Bellmano begin
    # The x-value to evaluate the quadratic
    consumption = Parameter(index=[time])

    # The y-value of the quadratic at the x-value
    utility = Variable(index=[time])
    disaster::Bool = Variable(index=[time])
    bonus = Variable()

    function run_timestep(p, v, d, t)
        v.bonus = 0
        if is_first(tt) || !v.disaster[tt-1]
            v.utility[tt] = sqrt(p.consumption[tt])
            v.disaster[tt] = rand() < p.consumption[tt]
	    if tt == 10 && !v.disaster[tt]
	        v.bonus = 1
	    end
        else
            v.utility[tt] = 0
            v.disaster[tt] = true
        end
    end
end

m = Model()
set_dimension!(m, :time, collect(1:10))

bellmano = add_comp!(m, Bellmano)
bellmano[:consumption] = repmat([.25], 10)

run(m)
m[:Bellmano, :utility]

import OptiMimi.uncertainproblem
using DataFrames

df = DataFrame(algorithm=[], mcperlife=[], cons1=[], cons2=[], cons3=[], cons4=[], cons5=[], cons6=[], cons7=[], cons8=[], cons9=[], cons10=[], cons1serr=[], cons2serr=[], cons3serr=[], cons4serr=[], cons5serr=[], cons6serr=[], cons7serr=[], cons8serr=[], cons9serr=[], cons10serr=[], time=[])
push!(df, [:exact; Inf; reverse(consumption_reverse)[1:10]; repmat([0], 10); 0])

for samplemult in [1, 2, 4]
    for mcperlife in [5, 20, 50]
        println(mcperlife)

        prob = uncertainproblem(m, [:Bellmano], [:consumption], [0.], [1.], m -> sum(m[:Bellmano, :utility] .* exp(-(0:9) * .05)) + m[:Bellmano, :bonus] * exp(-10 * .05), (model) -> nothing)

        tic()
        soln = solution(prob, () -> repmat([.25], 10), :social, mcperlife, samplemult)
        time = toc()
        push!(df, [:social; mcperlife; soln.xmean; soln.xserr; time])

        println(df)

        tic()
        soln = solution(prob, () -> repmat([.25], 10), :biological, mcperlife, 110000*samplemult)
        time = toc()
        push!(df, [:biological; mcperlife; soln.xmean; soln.xserr; time])

        println(df)

        tic()
        soln = solution(prob, () -> repmat([.25], 10), :sampled, mcperlife, 300*samplemult)
        time = toc()
        push!(df, [:sampled; mcperlife; soln.xmean; soln.xserr; time])

        println(df)
    end
    writetable("attempts-compare.csv", df)
end


#largeeig.png
#Algorithm 1.4 of Random Eigenvalues by Alan Edelman

#Experiment:  Largest eigenvalue of random Hermitian matrices
#Plot:        Histogram of the normalized largest eigenvalues
#Theory:      Tracy-Widom as n->infinity

include("tracywidom.jl")

## Parameters
n=100        # matrix size
t=5000       # trials
dx=.2        # binsize
function largeeig_experiment(n,t,dx)
    ## Experiment
    v=Float64[] # eigenvalue samples
    for i=1:t
        a=randn(n,n)+im*randn(n,n) # random nxn complex matrix
        s=(a+a')/2                 # Hermitian matrix
        push!(v,max(eig(s)[1]))    # Largest Eigenvalue
    end
    v=n^(1/6)*(v-2*sqrt(n))        # normalized eigenvalues

    grid=-5:dx:2
    count=hist(v,grid)[2]/(t*dx)
    (grid, count)
end
(grid, count) = largeeig_experiment(n,t,dx)

## Theory
(t, f) = tracywidom(5., -8.)

## Plot
using Winston
p = FramedPlot()
h = Histogram(count, step(grid))
h.x0 = first(grid)
add(p, h)
add(p, Curve(t, f, "linewidth", 2, "color", "blue"))
if isinteractive()
    Winston.display(p)
else
    file(p, "largeeig.png")
end


using Distributions
using PyPlot
using Random
using Statistics

function plot_p_sigmoid(; lb, ub, alpha, filename=nothing)
    chisq_dist = Chisq(3.0)
    xs = [x for x in lb:0.05:ub]

    figure(figsize=(9.0, 6.0))
    pdfs = pdf.(chisq_dist, xs)
    plt.plot(xs, 5.0 * pdfs, label="cost", lw=1.0)

    alpha_quantile = quantile(chisq_dist, 1.0 - alpha)
    ys = [y for y in 0.0:0.05:1.2]
    plt.plot(fill(alpha_quantile, length(ys)), ys, label="α-quantile", linestyle="--")

    sigmoid_lw = 0.7
    sigmoid_1_0 = (x) -> 1 / (1 + exp.(-(x - alpha_quantile)))
    plt.plot(xs, sigmoid_1_0.(xs), label="sigmoid1", lw=sigmoid_lw)
    sigmoid_0_5 = (x) -> 1 / (1 + exp.(-0.5 * (x - alpha_quantile)))
    plt.plot(xs, sigmoid_0_5.(xs), label="sigmoid2", lw=sigmoid_lw)
    sigmoid_1_5 = (x) -> 1 / (1 + exp.(-2.0 * (x - alpha_quantile)))
    plt.plot(xs, sigmoid_1_5.(xs), label="sigmoid3", lw=sigmoid_lw)

    plt.legend()

    if !isnothing(filename)
        plt.savefig(filename, dpi=500)
    end
end

plot_p_sigmoid(lb=0.0, ub=15.0, alpha=0.05, filename="/Users/kykim/Desktop/sigmoid.png")


println("\nTesting linear expressions\n")

#### * ####
n, m1, m2, k = 3, 4, 5, 6
x1 = Variable(randn(m1))
x2 = Variable(randn(m2))
A1 = randn(n, m1)
A2 = randn(n, m2)
b  = randn(n)
b0  = pi
opA1 = MatrixOp(A1)
opA2 = MatrixOp(A2)
# multiply with Variable
ex1 = opA1*x1
@test variables(ex1) == (x1,)
@test operator(ex1)*(~x1) == A1*(~x1)
B1 = randn(k, n)
opB1 = MatrixOp(B1)
# multiply with Expression
ex2 = opB1*ex1
@test variables(ex2) == (x1,)
@test norm(operator(ex2)*(~x1) - B1*A1*(~x1)) < 1e-12
# multiply with Expression with multiple variables
ex3 = opB1*(opA1*x1+opA2*x2)
@test variables(ex3) == (x1,x2)
@test norm(operator(ex3)*(~x1,~x2) - B1*(A1*(~x1)+A2*(~x2))) < 1e-12
# multiply with displacemented Array Expression with multiple variables
ex3 = opB1*(opA1*x1+opA2*x2+b)
@test variables(ex3) == (x1,x2)
@test norm(displacement(ex3) - B1*b) < 1e-12
# multiply with displacemented scalar Expression with multiple variables
ex3 = opB1*(opA1*x1+opA2*x2+b0)
@test variables(ex3) == (x1,x2)
@test norm(displacement(ex3) - B1*(ones(size(B1,2))*b0)) < 1e-12
@test_throws ArgumentError MatrixOp(randn(n,m1+1))*x1
@test_throws ArgumentError MatrixOp(randn(n,m1))*Variable(Complex{Float64},m1)

n, m1, m2, k = 3, 4, 5, 6
A1 = randn(n, m1)
A2 = randn(n, m2)
b1 = randn(n,n)
b2 = randn(n,n)
opA1 = MatrixOp(A1,n)
opA2 = MatrixOp(A2,n)
x1, x2  = Variable(randn(m1,n)), Variable(randn(m2,n))
# multiply Expressions (NonLinearCompose)
ex = (opA1*x1)*(opA2*x2)
@test variables(ex) == (x1,x2)
@test norm(affine(ex)*(~x1,~x2) - (A1*(~x1))*(A2*(~x2))) < 1e-12
ex = (opA1*x1-b1)*(opA2*x2)
@test variables(ex) == (x1,x2)
@test norm(affine(ex)*(~x1,~x2) - (A1*(~x1)-b1)*(A2*(~x2))) < 1e-12
ex = (opA1*x1)*(opA2*x2+b2)
@test variables(ex) == (x1,x2)
@test norm(affine(ex)*(~x1,~x2) - (A1*(~x1))*(A2*(~x2)+b2)) < 1e-12
ex = (opA1*x1+b1)*(opA2*x2+b2)
@test variables(ex) == (x1,x2)
@test norm(affine(ex)*(~x1,~x2) - (A1*(~x1)+b1)*(A2*(~x2)+b2)) < 1e-12
@test_throws ErrorException (opA1*x1)*(opA1*x1)

n, m1, m2, k = 3, 4, 5, 6
A1 = randn(n, m1)
A2 = randn(n, m2)
b1 = randn(n,n)
b2 = randn(n,n)
opA1 = MatrixOp(A1,n)
opA2 = MatrixOp(A2,n)
x1, x2  = Variable(randn(m1,n)), Variable(randn(m2,n))
## multiply Expressions elementwise (Hadamard)
ex = (opA1*x1).*(opA2*x2)
@test variables(ex) == (x1,x2)
@test norm(affine(ex)*(~x1,~x2) - (A1*(~x1)).*(A2*(~x2))) < 1e-12
ex = (opA1*x1-b1).*(opA2*x2)
@test variables(ex) == (x1,x2)
@test norm(affine(ex)*(~x1,~x2) - (A1*(~x1)-b1).*(A2*(~x2))) < 1e-12
ex = (opA1*x1).*(opA2*x2+b2)
@test variables(ex) == (x1,x2)
@test norm(affine(ex)*(~x1,~x2) - (A1*(~x1)).*(A2*(~x2)+b2)) < 1e-12
ex = (opA1*x1+b1).*(opA2*x2+b2)
@test variables(ex) == (x1,x2)
@test norm(affine(ex)*(~x1,~x2) - (A1*(~x1)+b1).*(A2*(~x2)+b2)) < 1e-12
@test_throws ErrorException (opA1*x1)*(opA1*x1)

##### reshape ####
m,n = 8,10
A = randn(n,m)
x = Variable(randn(m))
b = randn(n)
B = reshape(b,2,5)

ex = reshape(x,4,2)
@test norm(operator(ex)*~x - reshape(~x,4,2)) < 1e-12

ex2 = reshape(A*x,2,5)
@test norm(operator(ex2)*~x - reshape(A*~x,2,5)) < 1e-12

ex3 = reshape(A*x,2,5)+B
@test norm(operator(ex2)*~x+displacement(ex3)- reshape(A*~x,2,5)-B) < 1e-12

ex4 = reshape(A*x-b,2,5)
@test norm(operator(ex4)*~x+displacement(ex4)- reshape(A*~x-b,2,5)) < 1e-12

#### + ####

# sum same variable
n, m = 3,4
x = Variable(randn(m))
A = randn(n, m)
B = randn(n, m)
opA = MatrixOp(A)
opB = MatrixOp(B)
ex1 = opA*x+opB*x
@test variables(ex1) == (x,)
@test norm(operator(ex1)*~x - (opA+opB)*~x) < 1e-12

# sum different variables
n, m1, m2 = 3, 4, 5
xa = Variable(randn(m1))
xb = Variable(randn(m2))
xc = Variable(randn(n))
xd = Variable(randn(n))
A = randn(n, m1)
A2 = randn(n, m1)
B = randn(n, m2)
b = randn(n)
b0 = pi
opA = MatrixOp(A)
opA2 = MatrixOp(A2)
opB = MatrixOp(B)
opI = Eye(n)

# (+) sum different variables no HCAT
ex1 = opA*xa+opB*xb
@test variables(ex1) == (xa,xb)
@test norm(operator(ex1)*(~xa,~xb) - hcat(opA,opB)*(~xa,~xb)) <1e-12

# (+) sum of same variables
ex2 = opA2*xa
exs1 = ex1+ex2
exs2 = ex2+ex1
@test variables(exs1) == (xa,xb)
@test norm(operator(exs1)*(~xa,~xb) - hcat(opA+opA2,opB)*(~xa,~xb)) <1e-12
@test variables(exs2) == (xa,xb)
@test norm(operator(exs2)*(~xa,~xb) - hcat(opA+opA2,opB)*(~xa,~xb)) <1e-12

# (+) sum of different variables with HCAT
exs3 = exs1+exs2
@test variables(exs3) == (xa,xb)
@test norm(operator(exs3)*(~xa,~xb) - hcat(2*(opA+opA2),2*opB)*(~xa,~xb)) <1e-12
# (+) sum of different variables with HCAT
exs4 = exs1+(xc+xd)
@test variables(exs4) == (xa,xb,xc,xd)
@test norm(operator(exs4)*(~xa,~xb,~xc,~xd) - hcat(opA+opA2,opB,opI,opI)*(~xa,~xb,~xc,~xd)) <1e-12

# (+) sum Array
ex1 = xd+b
@test norm(displacement(ex1) - b) == 0.

# (+) sum scalar
ex2 = opB*xb+b0
@test (displacement(ex2) - b0) == 0.

# (+) sum displacemented expressions
ex3 = ex1+ex2
@test norm(displacement(ex3) - (b.+b0)) == 0.

#### (.+) sum
n = 3
b = randn(n)

x1 = Variable(randn(1))
x2 = Variable(randn(n))
ex1 = x1.+x2
@test norm(operator(ex1)*(~x1,~x2)-((~x1).+(~x2))) < 1e-9

x1 = Variable(randn(1))
x2 = Variable(randn(n))
ex1 = x1.+(x2+2)
@test norm(operator(ex1)*(~x1,~x2)-((~x1).+(~x2))) < 1e-9
@test displacement(ex1) == 2

x1 = Variable(randn(1))
x2 = Variable(randn(n))
ex1 = (x1+2).+(x2+b)
@test norm(operator(ex1)*(~x1,~x2)-((~x1).+(~x2))) < 1e-9
@test displacement(ex1) == (b.+2)

x1 = Variable(randn(n))
x2 = Variable(randn(1))
ex1 = x1.+x2
@test norm(operator(ex1)*(~x1,~x2)-((~x1).+(~x2))) < 1e-9

x1 = Variable(randn(n))
x2 = Variable(randn(1))
ex1 = x1.+(x2+2)
@test norm(operator(ex1)*(~x1,~x2)-((~x1).+(~x2))) < 1e-9
@test displacement(ex1) == 2

x1 = Variable(randn(n))
x2 = Variable(randn(1))
ex1 = (x1+b).+(x2+2)
@test norm(operator(ex1)*(~x1,~x2)-((~x1).+(~x2))) < 1e-9
@test displacement(ex1) == (b.+2)

n,m =2,4
x1 = Variable(randn(n,m))
x2 = Variable(randn(1,m))
ex1 = x1.+x2+6
@test norm(operator(ex1)*(~x1,~x2)-((~x1).+(~x2))) < 1e-9
@test displacement(ex1) == 6

# #### (.-) sum

n = 3
b = randn(n)

x1 = Variable(randn(1))
x2 = Variable(randn(n))
ex1 = x1.-x2
@test norm(operator(ex1)*(~x1,~x2)-((~x1).-(~x2))) < 1e-9

x1 = Variable(randn(1))
x2 = Variable(randn(n))
ex1 = x1.-(x2+2)
@test norm(operator(ex1)*(~x1,~x2)-((~x1).-(~x2))) < 1e-9
@test displacement(ex1) == -2

x1 = Variable(randn(1))
x2 = Variable(randn(n))
ex1 = (x1+2).-(x2+b)
@test norm(operator(ex1)*(~x1,~x2)-((~x1).-(~x2))) < 1e-9
@test displacement(ex1) == (2 .-b)

x1 = Variable(randn(n))
x2 = Variable(randn(1))
ex1 = x1.-x2
@test norm(operator(ex1)*(~x1,~x2)-((~x1).-(~x2))) < 1e-9

x1 = Variable(randn(n))
x2 = Variable(randn(1))
ex1 = x1.-(x2+2)
@test norm(operator(ex1)*(~x1,~x2)-((~x1).-(~x2))) < 1e-9
@test displacement(ex1) == -2

x1 = Variable(randn(n))
x2 = Variable(randn(1))
ex1 = (x1+b).-(x2+2)
@test norm(operator(ex1)*(~x1,~x2)-((~x1).-(~x2))) < 1e-9
@test displacement(ex1) == (b.-2)

n,m =2,4
x1 = Variable(randn(n,m))
x2 = Variable(randn(1,m))
ex1 = x1.-x2+6
@test norm(operator(ex1)*(~x1,~x2)-((~x1).-(~x2))) < 1e-9
@test displacement(ex1) == 6

# (-) sum different variables no HCAT
ex1 = opA*xa-opB*xb
@test variables(ex1) == (xa,xb)
@test norm(operator(ex1)*(~xa,~xb) - hcat(opA,-opB)*(~xa,~xb)) <1e-12

# (-) sum of same variables
ex2 = opA2*xa
exs1 = ex1-ex2
exs2 = ex2-ex1
@test variables(exs1) == (xa,xb)
@test norm(operator(exs1)*(~xa,~xb) - hcat(opA-opA2,-opB)*(~xa,~xb)) <1e-12
@test variables(exs2) == (xa,xb)
@test norm(operator(exs2)*(~xa,~xb) - hcat(-opA+opA2,+opB)*(~xa,~xb)) <1e-12

# (-) sum of same variables with HCAT
exs3 = exs1-exs2
@test variables(exs3) == (xa,xb)
@test norm(operator(exs3)*(~xa,~xb) - hcat(2*(opA-opA2),-2*opB)*(~xa,~xb)) < 1e-12
# (-) sum of different variables with HCAT
exs4 = exs1-(xc-xd)
@test variables(exs4) == (xa,xb,xc,xd)
@test norm(operator(exs4)*(~xa,~xb,~xc,~xd) - hcat(opA-opA2,-opB,-opI,opI)*(~xa,~xb,~xc,~xd)) <1e-12

# (-) sum Array
ex1 = xd-b
@test norm(displacement(ex1) + b) == 0.

# (-) sum scalar
ex2 = opB*xb-b0
@test (displacement(ex2) + b0) == 0.

# (+) sum displacemented expressions
ex3 = ex1-ex2
@test norm(displacement(ex3) - (-b.+b0)) == 0.

@test_throws DimensionMismatch MatrixOp(randn(10,20))*Variable(20)+randn(11)
@test_throws ErrorException MatrixOp(randn(10,20))*Variable(20)+(3+im)


#!/usr/bin/env julia

using Luxor, Test, Random
Random.seed!(42)

function testcircleintersections(fname)
    # completely overlapping circles
    c1 = (O, 100)
    c2 = (O, 100)
    cia = intersection2circles(c1..., c2...)
    @test isapprox(cia, 10000π, atol=0.1)

    # nonoverlapping circles
    c1 = (O, 100)
    c2 = (O + (200, 200), 100)
    cia = intersection2circles(c1..., c2...)
    @test isapprox(cia, 0.0, atol=0.1)

    # test for transition between inside and outside
    c1 = (O, 100)
    for v in -200:200
        c2 = (O + (100 * sqrt(2), 100 * sqrt(2)), 100 + (v * eps(Float64)))
        cia = intersection2circles(c1..., c2...)
    end
    @test 0.0 < cia < 0.000000001

    pagewidth, pageheight = 1200, 1400
    Drawing(pagewidth, pageheight, fname)
    background("grey30")
    origin() # move 0/0 to center

    @layer  begin
        scale(2)
        setopacity(0.3)
        sethue("cyan")
        circle(c1..., :fill)
        circle(c2..., :fill)
        sethue("red")
        circle(c1..., :clip)
        circle(c2..., :fill)
        clipreset()
    end

    occultations  = []
    for (pt, n) in Table(10, 10, 120, 120)
        c1 = (O, rand(5:40))
        c2 = (O + (rand(-20:20, 2)...), rand(10:50))
        cia = intersection2circles(c1..., c2...)
        push!(occultations, (c1, c2, cia))
    end

    sort!(occultations, by = (x) -> last(x))

    for (pt, n) in Table(10, 10, 120, 120)
        c1, c2, cia = occultations[n]
        @layer begin
            translate(pt)
            randomhue()
            circle(c1..., :fill)
            circle(c2..., :fill)
            sethue("white")
            circle(c1..., :clip)
            circle(c2..., :fill)
            clipreset()
            text(string(round(cia)), O + (0, 50), halign=:center)
        end
    end
    fontsize(30)
    sethue("white")
    text("assorted occultations sorted by area", boxtopcenter(BoundingBox() * 0.9), halign=:center)
    @test finish() == true
    println("...finished circletest, saved in $(fname)")
end

fname = "circle-intersection-test.png"

testcircleintersections(fname)


"""
For greedy algorithm, it will the element based on the optimal value in the array at each loop in the function

This does 0-1 (each item can be chosen only once) knapsack :
pack capacity = capacity
weight of each item = weights
value of each item = values

Each loop the function will find the highest value in the array and check if the capacity is enough to store it, if enough then
the value will be added into the total_max_value.
After that the highest current value will be deleted.

```
julia> zero_one_pack!(20,[1,3,11],[2,5,30])
37
```
"""
function ga_zero_one_pack!(capacity::N, weights::V, values::V
) where {N <: Number,V <: AbstractVector}
    total_max_value = 0
    while capacity > 0 && (length(weights)!=0 || length(values)!=0)
        max_value, index = findmax(values) # get the highest value and index in the vector
        if capacity >= weights[index]
            total_max_value += max_value
            capacity -= weights[index]
        end
        splice!(values, index)
        splice!(weights, index)
    end
    return total_max_value
end

"""
This does complete/infinite (each item can be chosen infinite times) knapsack :
pack capacity = capacity
weight of each item = weights
value of each item = values

Each loop the function will find the highest value in the array and check if the capacity is enough to store it, if enough then
the value will be added into the total_max_value until the capacity cannot hold the weight of the highest current value.
After that the highest current value will be deleted.

```
julia> complete_pack!(20,[1,2,9],[1,3,20])
43
```
"""
function ga_complete_pack!(capacity::N, weights::V, values::V
) where {N <: Number,V <: AbstractVector}
    total_max_value = 0
    while capacity > 0 && (length(weights)!=0 || length(values)!=0)
        max_value, index = findmax(values)
        while capacity >= weights[index]
            total_max_value += max_value
            capacity -= weights[index]
        end
        splice!(values, index)
        splice!(weights, index)
    end
    return total_max_value
end


## Make an animation of a random walk
using SparseArrays
using DelimitedFiles
## Load the graph
ei,ej = copy.(eachcol(Int.(readdlm("undir-graph-14.edges"))))
xy = readdlm("undir-graph-14.xy")'
##
A = sparse(ei,ej,1,14,14) # make the sparse matrix
A = A+A' # make it symmetric (we only store half)
n = size(A,1)
##
d = vec(sum(A,dims=1)) # get the degrees
Di = Diagonal(1.0./d)
##
x = zeros(n)
x[1] = 1
A'*Di*x
##
P = Matrix(A'*Di)
vals,vecs = eigen(P)
vals
vecs[:,1]/sum(vecs[:,1]) # normalize to prob.
##
d/sum(d)


# State: 4 continuous variables
#   cart position:      ∈ [-4.8, 4.8]
#   cart speed:         ∈ [-Inf, Inf]
#   pole angle:         ∈ [-24°, 24°]
#   pole rotation rate: ∈ [-Inf, Inf]

# Action: Push cart right (1) or left (2)

@with_kw struct CartPole
    γ::Float64 = 1.0
    force_magnitude::Float64 = 10.0
    pole_length::Float64 = 1.0  # Distance from pole to mass at end of pole
    mass_cart::Float64 = 1.0
    mass_pole::Float64 = 0.1
    gravity::Float64 = 9.8
    Δt::Float64 = 0.02  # Time to advance by with each step
    Δx_max::Float64 = 4.8  # Max position deviation
    Δθ_max::Float64 = deg2rad(24)  # Max angular deviation
end

# const CART_POLE = CartPole(1.0, 10.0, 1.0, 1.0, 0.1, 9.8, 0.02, 4.8, deg2rad(24))

struct CartPoleState
    x::Float64 # cart pos [m]
    v::Float64 # cart speed [m/s]
    θ::Float64 # pole angle [rad]
    ω::Float64 # pole rotation range [rad/s]
end

n_actions(mdp::CartPole) = 2
discount(mdp::CartPole) = mdp.γ
ordered_actions(mdp::CartPole) = [1,2]

Base.vec(s::CartPoleState) = [s.x, s.v, s.θ, s.ω]

function generate_start_state(mdp::CartPole)
    U = Uniform(-0.05,0.05)
    return CartPoleState(rand(U), rand(U), rand(U), rand(U))
end

function is_terminal(mdp::CartPole, s::CartPoleState)
    return abs(s.x) > mdp.Δx_max || abs(s.θ) > mdp.Δθ_max
end

function cart_pole_transition(mdp::CartPole, s::CartPoleState, a::Int)
    @assert(1 <= a <= 2)

    if is_terminal(mdp, s)
        return s
    end

    x, v, θ, ω = s.x, s.v, s.θ, s.ω
    force = a == 1 ? mdp.force_magnitude : -mdp.force_magnitude

    cθ, sθ = cos(θ), sin(θ)
    half_length = mdp.pole_length / 2
    mass_tot = mdp.mass_cart + mdp.mass_pole
    temp = (force + half_length*sθ*ω^2) / mass_tot

    # Pole acceleration
    pole_α = (mdp.gravity*sθ - temp*cθ) / (half_length * (4.0/3.0 - mdp.mass_pole * cθ^2 / mass_tot))

    # Cart acceleration
    cart_a = temp - half_length * mdp.mass_pole * pole_α * cθ / mass_tot

    # Euler integration
    Δt = mdp.Δt
    x += Δt*v
    v += Δt*cart_a
    θ += Δt*ω
    ω += Δt*pole_α

    return CartPoleState(x,v,θ,ω)
end

function reward(mdp::CartPole, s::CartPoleState, a::Int)
    @assert(1 <= a <= 2)
    if is_terminal(mdp, s)
        return 0.0
    else
        return 1.0
    end
end

function MDP(mdp::CartPole; γ::Float64=mdp.γ)
    return MDP(
            γ,
            nothing, # no ordered states
            ordered_actions(mdp),
            nothing, # no probabilistic transition function
            (s,a) -> reward(mdp, s, a),
            (s, a)->begin
                s′ = cart_pole_transition(mdp, s, a)
                r = reward(mdp, s, a)
                return (s′, r)
            end
        )
end


# Diagonal Approximation for trace of inverse of matrix
# Values from Extrapolation Methods for Estimating Trace of the Matrix Inverse by P. Fika

export diagapp

using LinearAlgebra

struct Diagappspace{T, TM <: AbstractMatrix{T}, TV <: AbstractVector{T}, F<:Function}
    A::TM
    x::TV
    y::TV
    randfunc::F
end

function Diagappspace(A::AbstractMatrix, randfunc::Function)
    x = randfunc(eltype(A) <: Integer ? Float64 : eltype(A), size(A, 1))
    y = similar(x)
    return Diagappspace(A, x, y, randfunc)
end

# Extrapolation for c-1 and calculating value of v0
function v0(w::Diagappspace)

    copyto!(w.x, w.randfunc(-1:2:1, size(w.A, 1)))

    c0 = dot(w.x, w.x)
    mul!(w.y, w.A, w.x)
    c1 = dot(w.y, w.x)
    c2 = dot(w.y, w.y)
    mul!(w.x, w.A, w.y)
    c3 = dot(w.y, w.x)

    p = (c0 * c2)/(c1^2)

    v0 = log10((c1^2)/(c0 * c2)) / log10((c1 * c3)/(c2^2))
end

# Calculating the Approximation for ith value of inverse diagonal
function dfun(A, i, v)
    isum = zero(eltype(A))
    s = size(A, 1)
    for k in 1:s
        isum = isum + A[k, i]^2
    end
    #v = v0(A)
    p = isum / (A[i, i]^2)
    d = 1 / (p^v * A[i, i])
end

# Calculating Diagonal Approximation for Matrix Inverse
# This works a SPD Matrix with low condition number
"""
    diagapp(A::AbstractMatrix; randfunc=Base.rand)

Diagonal Approximation algorithm for inverse of matrix.
# Arguments
 - `A` : Symmetric Positive Definite Matrix with Low Condtion Number (k < 500)
 - `randfunc` : Random function for the particular type of matrix A
                (An example can be found in test/diagapp.jl)
"""
function diagapp(A; randfunc::Function=Base.rand)
    tr = zero(eltype(A))
    s = size(A, 1)
    w = Diagappspace(A, randfunc)
    v = v0(w)
    for i in 1:s
        tr = tr + dfun(w.A, i, v)
    end
    tr
end


"""
    discrete_tail_estimation(sempp::SEMPPExpKern, q::Real = 0.95; to_time::Union{DiscreteTimeTypes, Nothing} = nothing, from_time::Union{DiscreteTimeTypes, Nothing} = nothing)::Real

Compute the quantile at level *q* of the marks distribution between from_time and to_time, knowing history until from_time, with a discrete time approach.

See Li2020 for the mathematics behind this method.

from_time defaults to the last timestamp in the time series, and to_time defaults to from_time + 1.
Note that if from_time differs from the last timestamp in the time series, it will be considered that no event happened between this last timestamp and from_time.
"""
function discrete_tail_estimation(sempp::SEMPPExpKern, q::Real = 0.95; to_time::Union{DiscreteTimeTypes, Nothing} = nothing, from_time::Union{DiscreteTimeTypes, Nothing} = nothing, history::Union{MarkedTimeSeries, Nothing} = nothing)::Real
    mts = isnothing(history) ? MarkedTimeSeries([], []) : history

    if start_time(mts) isa TimeType

        isnothing(from_time) && (from_time = end_time(mts) + Day(1))

        isnothing(to_time) && (to_time = from_time)

        if !(from_time isa TimeType) || !(to_time isa TimeType)
            error("times and from_date and to_date must be either all Integers or all TimeTypes")
        end

        when = from_time:Dates.Day(1):to_time
    else

        isnothing(from_time) && (from_time = end_time(mts) + 1)

        isnothing(to_time) && (to_time = from_time)

        if !(from_time isa Integer) || !(to_time isa Integer)
            error("times and from_date and to_date must be either all Integers or all TimeTypes")
        end

        when = from_time:to_time
    end

    μ = sempp.μ
    ϕ = sempp.ϕ
    γ = sempp.γ
    δ = sempp.δ

    vol = volfunc(when, mts, γ, δ)
    intens = sum(μ .+ ϕ .* vol)
    prob = 1 - exp(-intens )

    prob <= 1-q && (error("Probability of an event below $(1-q), cannot estimate tail."))

    markdens = sempp.markdens
    ξ = sempp.ξ
    α = sempp.α
    β = sempp.β
    κ = sempp.κ

    σ = β + α * vol[1]

    if markdens == Distributions.GeneralizedPareto
        exceed_q = σ / ξ * ((prob / (1 - q))^ξ - 1)
    else
        exceed_q = σ / ξ * ((1 - (1 - (1 - q) / prob)^(1/κ))^(-ξ) - 1)
    end

    return exceed_q
end

"""
    tail_estimation(sempp::SEMPPExpKern, q::Real = 0.95; to_time::Union{DiscreteTimeTypes, Nothing} = nothing, from_time::Union{DiscreteTimeTypes, Nothing} = nothing)::Real

Compute the quantile at level *q* of the marks distribution between from_time and to_time, knowing history until from_time, with a continuous time approach.

See Li2020 for the mathematics behind this method.

from_time defaults to the last timestamp in the time series, and to_time defaults to from_time + 1.
Note that if from_time differs from the last timestamp in the time series, it will be considered that no event happened between this last timestamp and from_time.
"""
function tail_estimation(sempp::SEMPPExpKern, q::Real = 0.95; to_time::Union{DiscreteTimeTypes, Nothing} = nothing, from_time::Union{DiscreteTimeTypes, Nothing} = nothing, history::Union{MarkedTimeSeries, Nothing} = nothing)::Real
    mts = isnothing(history) ? MarkedTimeSeries([], []) : history

    times = mts.times
    marks = mts.marks

    isnothing(from_time) && (from_time = last(times))

    if first(times) isa TimeType

        isnothing(to_time) && (to_time = from_time + Day(1))

        if !(from_time isa TimeType) || !(to_time isa TimeType)
            error("times and from_date and to_date must be either all Real or all TimeTypes")
        end

        Dates.value(Date(to_time - from_time))

    else

        isnothing(to_time) && (to_time = from_time + 1)

        if (from_time isa TimeType) || (to_time isa TimeType)
            error("times and from_date and to_date must be either all Real or all TimeTypes")
        end

        T = to_time - from_time

    end

    μ = sempp.μ
    ϕ = sempp.ϕ
    γ = sempp.γ
    δ = sempp.δ

    intens = μ * T + ϕ/γ * sum((1 .+ δ .* marks) .* (1 .- exp.(-γ .* (to_time .- times))))
    prob = 1 - exp(-intens )

    prob <= 1-q && (error("Probability of an event below $(1-q), cannot estimate tail."))

    markdens = sempp.markdens
    ξ = sempp.ξ
    α = sempp.α
    β = sempp.β
    κ = sempp.κ

    σ = β + α * vol[1]

    if markdens == Distributions.GeneralizedPareto
        exceed_q = σ / ξ * ((prob / (1 - q))^ξ - 1)
    else
        exceed_q = σ / ξ * ((1 - (1 - (1 - q) / prob)^(1/κ))^(-ξ) - 1)
    end

    return exceed_q
end


using Test

@testset "Nesterov accel. ($R)" for R in [Float32, Float64]

using LinearAlgebra
using ProximalAlgorithms: NesterovAcceleration, update!

n, k = 20, 10

Q = randn(R, n, k)
H = Q*Q' + I
l = randn(R, n)

f(x) = dot(x, H*x)/2 + dot(x, l)

x_star = -H\l
f_star = f(x_star)
norm_x_star = norm(x_star)

Lip = opnorm(H)
gamma = 1/Lip
x = zeros(n)
res = zeros(n)
x_prev = zeros(n)
res_prev = zeros(n)

acc = NesterovAcceleration(x)

mul!(res, H, x)
res .+= l
res .*= gamma

d = acc * res

for it in 1:500
    # store iterate and residual for the operator update later
    res_prev .= res
    x_prev .= x

    # compute accelerated direction
    mul!(d, acc, res)

    # update iterate
    x .-= d

    # compute new residual
    mul!(res, H, x)
    res .+= l
    res .*= gamma

    # update operator
    update!(acc, x - x_prev, res - res_prev)

    # test that iterates satisfy Thm 4.4 from Beck, Teboulle (2009)
    @test f(x) - f_star <= 2/(gamma * (it+1)^2) * norm_x_star^2
end

end


using Flux, Plots
using Base.Iterators: repeated
using Flux: @epochs
using LinearAlgebra

gridsize = 150;
tstart=0.0
tend=20.0
myfunc(x) = sin(x);
x = collect(range(tstart,stop=tend,length=gridsize));
y= myfunc.(x);

xtest=collect(tend:0.1:2*tend)
xpred=collect(tstart:0.1:2*tend)

#Build the input data
data = []
for i in 1:length(x)
    push!(data, ([x[i]], y[i])) #have to input a list
end

Q = 20;
ann = Chain(Dense(1,Q,tanh),Dense(Q,Q,tanh),Dense(Q,1));

#Simple Mean-square loss2
function loss(x, y)
    pred=ann(x)
    loss=Flux.mse(ann(x), y)
    return loss
end

opt = ADAM()

ps=params(ann)
@epochs 5000 Flux.train!(loss,ps, data, opt)

plot(xpred, ann(xpred')', color=:red, lw=2.0, label="")
scatter!(x,myfunc.(x), color=:blue, legend=false)
scatter!(xtest,myfunc.(xtest), color=:green, lw=2.0, label="", markershape=:cross)
ylims!((-1.5,1.5))
xlims!((tstart,2*tend))
title!("Curve Fitting with Neural Networks")
xlabel!("Input")
ylabel!("Output")

savefig("sin_poor.png")


"""
    r_x(r, vl, vr)

Transfer coordinate r -> x from standard to real elements

"""
r_x(r, vl, vr) = ((1.0 - r) / 2.0) * vl + ((1.0 + r) / 2.0) * vr

"""
Transfer coordinate r -> x from standard to real elements

Triangle elements:

    rs_xy(r, s, v1, v2, v3)
    rs_xy(v, v1, v2, v3)

"""
function rs_xy(
    r,
    s,
    v1::AbstractVector{T},
    v2::AbstractVector{T},
    v3::AbstractVector{T},
) where {T<:Real}
    return @. -(r + s) / 2 * v1 + (r + 1) / 2 * v2 + (s + 1) / 2 * v3
end

function rs_xy(
    v::AbstractVector{T},
    v1::AbstractVector{T},
    v2::AbstractVector{T},
    v3::AbstractVector{T},
) where {T<:Real}
    r, s = v
    return rs_xy(r, s, v1, v2, v3)
end

function rs_xy(
    r,
    s,
    v1::AbstractVector{T},
    v2::AbstractVector{T},
    v3::AbstractVector{T},
    v4::AbstractVector{T},
) where {T<:Real}
    return @. (r - 1.0) * (s - 1.0) / 4.0 * v1 +
        (r + 1.0) * (1.0 - s) / 4.0 * v2 +
        (r + 1.0) * (s + 1.0) / 4.0 * v3 +
        (1.0 - r) * (s + 1.0) / 4.0 * v4
end

"""
    xy_rs(x, y)
    xy_rs(coords)

Transfer coordinates (x, y) -> (r,s) from equilateral to right triangle

"""
function xy_rs(x::T, y::T) where {T<:Real}
    L1 = (sqrt(3.0) * y + 1.0) / 3.0
    L2 = (-3.0 * x - sqrt(3.0) * y + 2.0) / 6.0
    L3 = (3.0 * x - sqrt(3.0) * y + 2.0) / 6.0
    r = -L2 + L3 - L1
    s = -L2 - L3 + L1

    return r, s
end

function xy_rs(x::AbstractVector{T}, y::AbstractVector{T}) where {T<:Real}
    Np = length(x)
    r = zeros(Np)
    s = zeros(Np)

    for n = 1:Np
        r[n], s[n] = xy_rs(x[n], y[n])
    end

    return r, s
end

function xy_rs(coords::AbstractMatrix{T}) where {T<:Real}
    return xy_rs(coords[:, 1], coords[:, 2])
end

"""
    rs_ab(r, s)
    rs_ab(coords)

Transfer coordinates (r,s) -> (a,b) in a triangle

"""
function rs_ab(r::T, s::T) where {T<:Real}
    a = ifelse(s != 1.0, 2.0 * (1.0 + r) / (1.0 - s) - 1.0, -1.0)
    b = 1.0 * s

    return a, b
end

function rs_ab(r::AbstractVector{T}, s::AbstractVector{T}) where {T<:Real}
    a = zero(r)
    b = zero(s)

    for n in eachindex(a)
        a[n], b[n] = rs_ab(r[n], s[n])
    end

    return a, b
end

function rs_ab(coords::AbstractMatrix{T}) where {T<:Real}
    return rs_ab(coords[:, 1], coords[:, 2])
end


"""
Calculates vertical functional principal component analysis on aligned data

    vert_fPCA(warp_data, qn; no=1)
    :param warp_data: fdawarp type from srsf_align of aligned data
    :param no: number of components to extract (default = 1)

    Returns vfpca type containing
    :return q_pca: srsf principal directions
    :return f_pca: functional principal directions
    :return latent: latent values
    :return coef: scores
    :return U: eigenvectors
    :return id: point used for f(0)
"""
function vert_fPCA(warp_data::fdawarp; no=1)
    qn = warp_data.qn
    fn = warp_data.fn
    timet = warp_data.time
    coef = collect(-2:3);
    Nstd = length(coef);

    # fPCA
    mq_new = mean(qn, dims=2);
    N = length(mq_new);
    mididx = round(Integer, length(timet)/2);
    m_new = sign.(fn[mididx, :]) .* sqrt.(abs.(fn[mididx, :]));
    mqn = [mq_new; mean(m_new)];
    qn2 = vcat(qn, m_new');
    K = Statistics.covm(qn2, mean(qn2,dims=2), 2);

    U, s, V = svd(K);
    stdS = sqrt.(s);

    # compute the PCA in the q domain
    q_pca = Array{Float64}(undef, (N+1, Nstd, no));
    for k in 1:no
        for l in 1:Nstd
            q_pca[:, l, k] = mqn + coef[l] * stdS[k] * U[:, k];
        end
    end

    # compute the correspondence in the f domain
    f_pca = Array{Float64}(undef, (N, Nstd, no));
    for k in 1:no
        for l in 1:Nstd
            f_pca[:, l, k] = cumtrapzmid(timet, q_pca[1:N,l,k].*
                                         abs.(q_pca[1:N,l,k]),
                                         sign.(q_pca[N+1,l,k]*q_pca[N,l,k]^2),
                                         mididx);
        end
        fbar = mean(fn,dims=2)
        fsbar = mean(f_pca[:, :, k],dims=2)
        err = repeat(fbar-fsbar,1,Nstd)
        f_pca[:, :, k] = f_pca[:, :, k] + err
    end

    N2 = size(qn,2);
    c = zeros(N2, no);
    for k in 1:no
        for l in 1:N2
            c[l,k] = sum(vec((qn2[:,l] - mqn).* U[:,k]));
        end
    end

    out = vfpca(q_pca, f_pca, s[1:no], c, U[:,1:no], mididx, mqn[:], timet);

    return out
end

"""
Calculates horizontal functional principal component analysis on aligned data

    horiz_fPCA(warp_data; no=1)
    :param warp_data: fdawarp type from srsf_align of aligned data
    :param no: number of components to extract (default = 1)

    Returns hfpca containing
    :return gam_pca: warping principal directions
    :return psi_pca: srsf functional principal directions
    :return latent: latent values
    :return U: eigenvectors
    :return coef: scores
    :return vec: shooting vectors
    :return gam_mu: mean warping function
"""
function horiz_fPCA(warp_data::fdawarp; no=1)
    gam = warp_data.gam;
    mu, gam_mu, psi, vec1 = sqrt_mean(gam);
    tau = collect(1:5);
    m = length(mu);
    n = length(tau);
    TT = size(vec1,1)

    # TFPCA
    K = Statistics.covm(vec1, mean(vec1,dims=2), 2);

    U, s, V = svd(K);
    vm = mean(vec1, dims=2);

    gam_pca = Array{Float64}(undef, (m, n, no));
    psi_pca = Array{Float64}(undef, (m, n, no));
    for j in 1:no
        for k in tau
            v = (k-3)*sqrt(s[j]).*U[:,j];
            vn = norm(v) / sqrt(TT);
            if vn < 0.0001
                psi_pca[:, k, j] = mu;
            else
                psi_pca[:, k, j] = cos(vn).*mu + sin(vn).*v/vn;
            end

            gam0 = cumtrapz(collect(LinRange(0,1,TT)), psi_pca[:,k,j].*psi_pca[:,k,j]);
            gam_pca[:,k,j] = norm_gam(gam0)
        end
    end

    N2 = size(gam,2);
    c = zeros(N2, no);
    for k in 1:no
        for l in 1:N2
            c[l,k] = sum((vec1[:,l]-vm).* U[:,k]);
        end
    end

    out = hfpca(gam_pca, psi_pca, s[1:no], U[:,1:no], c, vec1, gam_mu)
    return out
end


# Vortex detection and imprinting

"""
    P, Q = poles(u)

Return the Wirtinger derivatives of u, assuming a unit grid step
"""
function poles(u)
    N = size(u, 1)
    @assert N == size(u, 2)
    u = complex.(u)
    rs = (-1:1) .+ 1im * (-1:1)'
    rs /= sum(abs2.(rs))
    conv(u, rs) = [rs .* u[j:j+2, k:k+2] |> sum for j = 1:N-2, k = 1:N-2]
    P = zero(u)
    P[2:end-1, 2:end-1] .= conv(u, conj(rs))
    Q = zero(u)
    Q[2:end-1, 2:end-1] .= conv(u, rs)
    P, Q
end

"""
    unroll(θs)

Return θs adjusted mod 2π to be as continuous as possible
"""
function unroll(θ)
    Θ = similar(θ)
    Θ[1] = θ[1]
    poff = 0.0
    for j = 2:length(θ)
        jump = θ[j] - θ[j-1]
        if jump > π
            poff -= 2π
        elseif jump < -π
            poff += 2π
        end
        Θ[j] = θ[j] + poff
    end
    Θ
end

"""
    find_vortex(d, u)

Return the centre of the sole vortex in u
"""
function find_vortex(d::Discretisation, u, ixs = nothing)
    P, Q = poles(u)
    if isnothing(ixs)
        ixs = abs.(P) .> 0.5maximum(abs.(P))
    end
    regress_core(d, u, ixs)
end

find_vortex(u) = find_vortex(default(:discretisation), u)

function find_vortices(d::Discretisation, u)
    P, Q = poles(u)
    ixs = abs.(P) .> 0.5maximum(abs.(P))
    clusters = cluster_adjacent(adjacent_index, keys(ixs)[ixs])
    [find_vortex(d, u, [j ∈ C for j in keys(u)]) for C in clusters]
end

function find_moat(d::Discretisation, u)
    P, Q = poles(u)
    v = @. (R-w/2 < r < R+w/2)*abs(P+conj(Q))/abs(u)
    ixs = v .> 0.5maximum(v)
    regress_core(d, u, ixs)
end

function regress_core(d, u, ixs)
    z = argand(d)
    ixs = expand_box(box(ixs))[:]
    a, b, c = [z[ixs] conj(z[ixs]) ones(size(z[ixs]))] \ u[ixs]
    (b * conj(c) - conj(a) * c) / (abs2(a) - abs2(b))
end

function expand_box(bb::CartesianIndices{2})
    e1, e2 = ((1, 0), (0, 1)) .|> CartesianIndex
    # TODO raise issue on CartesianIndex half-pregnancy
    if length(bb) == 1
        bb = bb[]
        bb = [
            bb-e1-e2 bb-e1 bb-e1+e2
            bb-e2 bb bb+e2
            bb+e1-e2 bb+e1 bb+e1+e2
        ]
    elseif size(bb, 1) == 1
        bb = [bb .- e1; bb; bb .+ e1]
    elseif size(bb, 2) == 1
        bb = [bb .- e2 bb bb .+ e2]
    end
    bb
end

"Bounding box of cartesian indices"
function box(cixs::AbstractVector{CartesianIndex{2}})
    j1 = k1 = typemax(Int)
    j2 = k2 = typemin(Int)
    for c in cixs
        j1 = min(j1, c[1])
        j2 = max(j2, c[1])
        k1 = min(k1, c[2])
        k2 = max(k2, c[2])
    end
    CartesianIndices((j1:j2, k1:k2))
end

box(ixs::AbstractMatrix{Bool}) = box(keys(ixs)[ixs])

function cluster_adjacent(f, ixs)
    # function f determines adjacency
    clusters = Set()
    for i in ixs
        out = Set()
        ins = Set()
        for C in clusters
            if any(j -> f(i, j), C)
                push!(ins, C)
            else
                push!(out, C)
            end
        end
        push!(out, union(Set([i]), ins...))
        clusters = out
    end
    clusters
end

adjacent_index(j, k) = -1 ≤ j[1] - k[1] ≤ 1 && -1 ≤ j[2] - k[2] ≤ 1

"""
    PinnedVortices([d], rvs, as) <: Optim.Manifold

Constrain a field to have vortices centred at the rvs

TODO document the new point-estimates and Gaussians approach
"""
struct PinnedVortices <: Manifold
    # V'*ψ[:] is a vector of interpolated ψ.(rvs)
    # (1-U*V')*ψ projects the singularities back to the rvs
    # TODO make these Tensars
    # TODO solve a linear system in the u(rvs) to make the projection exact
    U::Matrix{Complex{Float64}}
    V::Matrix{Complex{Float64}}

    function PinnedVortices(
        d::Sampling{2},
        rvs::Vector{Complex{Float64}},
        as::Vector{Float64}
    )
        U = Array{Complex{Float64}}(undef, d.n^2, length(rvs))
        V = similar(U)
        for j in eachindex(rvs)
            U[:, j] = finterp(d, rvs[j], as[j])[:]
            V[:, j] = interpolant(d, rvs[j])[:]
        end
        new(U, V)
    end
end

"""
    finterp(d, r, a)

Return interpolant(d, r) convoluted with a Gaussian of width a.

Subtracting this is a stable way to shift a vortex by further than one pixel.
"""
function finterp end

# use coordinates instead of this

# PinnedVortices(d::Discretisation, rvs::Vararg{Number}; kwargs...) =
#     PinnedVortices(d, [convert(Complex{Float64}, r) for r in rvs]...; kwargs...)
# PinnedVortices(rvs::Vararg{Number}; kwargs...) =
#     PinnedVortices(default(:discretisation), rvs...; kwargs...)

function prjct!(M, q)
    w = q[:]
    q .-= reshape(M.U * (M.V' * w), size(q))
end

# The "vortex at R" space is invariant under normalisation
Optim.retract!(M::PinnedVortices, q) = Optim.retract!(Sphere(), prjct!(M, q))
Optim.project_tangent!(M::PinnedVortices, dq, q) =
    Optim.project_tangent!(Sphere(), prjct!(M, dq), q)