#    Problem : GROUP A
#    *********
#    The second penalty function
#
#    This is a nonlinear least-squares problem with M=2*N groups.
#    Group 1 is linear.
#    Groups 2 to N use 2 nonlinear elements.
#    Groups N+1 to M-1 use 1 nonlinear element.
#    Group M uses N nonlinear elements.
#    The Hessian matrix is dense.
#
#    Origonal SIF Source: Problem 24 in
#    J.J. More', B.S. Garbow and K.E. Hillstrom,
#    "Testing Unconstrained Optimization Software",
#    ACM Transactions on Mathematical Software, vol. 7(1), pp. 17-41, 1981.
#
#    See also Buckley#112 (p. 80)
#
#    PENALTY2.SIF classification SUR2-AN-V-0
#
#    Number of variables is variable (n>3)
#
# Daniel Henderson, 08/2021

f = (x) -> begin
	@warn "PENALTY2.jl doesn't agree with CUTEst or a JuMP (all three conflict)"
	N  = lastindex(x)
	fx = (x[1]-0.2)^2
	a  = 1e-5
	y  = ones(2*N)

  	for i in 1:2*N
    	y[i] = exp(i / 10.0) + exp((i - 1) / 10.0)
  	end

	α  = N*x[1]^2-1

	for i in 2:N
		fx += a*(exp(x[i]/10)+exp(x[i-1]/10) - y[i])^2
		fx += a*(exp(x[i]/10)-exp(-1/10))^2
		α += (N-i+1)*x[i]^2 - 1
	end

	fx += α^2
	return fx
end

# check chain rule use on product of sum term
g! = (g, x) -> begin
    N  = lastindex(x)
	γ  = (1e-5)/5
	α  = N*x[1]^2-1

	g[1] = 2*(x[1]-0.2)*x[1]
	for i in 2:N
		g[i-1] += γ*exp(x[i-1]/10) * (exp(x[i]/10) + exp(x[i-1]/10) - exp((i+1)/10) - exp(i/10) )
		g[i]   -= γ*exp(x[i]/10) * (exp(1/10) + exp(i/10) + exp((i+1)/10) - exp(x[i-1]/10) - 2*exp(x[i]/10))
		α      += (N-i+1)*x[i]^2 - 1
	end
		α *= 2 # diff of outer function
	for i in 2:N
		g[i] += α * 2*(N-i+1)*x[i]
	end
	return g
end

fg! = (g, x) -> begin
    return error("PENALTY2 fg! not implmented")
end

init = (n::Int=1000) -> begin
	x0 = Vector(1.0:1.0:n)
    return n, x0
end

@warn "TODO: Debug PENALTY2"
#TestSet["PENALTY2"] = UncProgram("PENALTY2", f, g!, fg!, init)


using StatisticalRethinking, Optim
#gr(size=(600,600));

p_grid = range(0, step=0.001, stop=1)
prior = ones(length(p_grid))
likelihood = [pdf(Binomial(3, p), 3) for p in p_grid]
posterior = likelihood .* prior
posterior = posterior / sum(posterior)

N = 10000
samples = sample(p_grid, Weights(posterior), N);

chn = MCMCChains.Chains(reshape(samples, N, 1, 1), ["toss"]);

MCMCChains.describe(chn)

MCMCChains.hpd(chn, alpha=0.5)

mode(samples)

mean(samples)

median(samples)

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


# instantiate environment
using Pkg; Pkg.instantiate()

using Gadfly
using DataFrames
using Printf
using CSV

# intersection between two circles
function intersection(c₁, r₁, c₂, r₂)
  a = r₁ * r₁
  b = r₂ * r₂
  d = sqrt(2 * (c₁ - c₂) * (c₁ - c₂))

  # trivial cases
  d ≥ r₁ + r₂      && return 0.0
  d ≤ abs(r₂ - r₁) && return π * min(a, b)

  # general case
  C1 = a * acos((d * d + a - b) / (2d*r₁))
  C2 = b * acos((d * d + b - a) / (2d*r₂))
  C3 = (1 / 2) * sqrt((-d + r₂ + r₁) * (d + r₂ - r₁) * (d - r₂ + r₁) * (d + r₂ + r₁))

  C1 + C2 - C3
end

# -----------------
# SHIFT FUNCTIONS
# -----------------
kldiv(δ, τ) = δ^2 + τ^2 - log(τ^4) - 1

function jaccard(c₁, r₁, c₂, r₂)
  I = intersection(c₁, r₁, c₂, r₂)
  U = π * r₁^2 + π * r₂^2 - I

  1 - I / U
end

jaccard(δ, τ) = jaccard(0, 3, 3*√2δ + 0, 3τ)

function novelty(c₁, r₁, c₂, r₂)
  I = intersection(c₁, r₁, c₂, r₂)
  O = π * r₂^2 - I
  B = π * r₂^2
  R = (O - I) / B

  (R + 1) / 2
end

novelty(δ, τ) = novelty(0, 3, 3*√2δ + 0, 3τ)

# covariate shift configuration
shiftconfig(δ, τ) = 2δ ≤ 1 - τ ?
                    "inside" :
                    2δ ≥ 1 + τ ?
                    "outside" :
                    "partial"

# -------------
# MAIN SCRIPT
# -------------

# load results table
df = CSV.read(joinpath(@__DIR__,"gaussian.csv"), missingstring="NaN")

# drop outliers due to numerical instability
df = filter(row -> row[:DRV] ≤ 0.5 && row[:ACTUAL] ≤ 0.5, df)
df = dropmissing(df)

# shift functions
df[!,:KLDivergence]    = kldiv.(df[!,:δ], df[!,:τ])
df[!,:JaccardDistance] = jaccard.(df[!,:δ], df[!,:τ])
df[!,:NoveltyFactor]   = novelty.(df[!,:δ], df[!,:τ])

# shift configuration
df[!,:config] = shiftconfig.(df[!,:δ], df[!,:τ])

# correlation length as factor
df[!,:rfactor] = map(df[!,:r]) do r
  @sprintf "r=%.1f" r
end

# filter "inside" configuration
ff = filter(row -> row[:config] == "inside", df)

# set plotting theme
theme = Gadfly.get_theme(Val(:default))
Gadfly.push_theme(theme)
theme = style(point_size=2.5px,
              key_position=:top,
              default_color=colorant"black")
colors = ("#1b9e77","#7570b3","#d95f02")

# generalization error vs. shift function
Gadfly.with_theme(theme) do
  xcols = (:KLDivergence,:JaccardDistance,:NoveltyFactor)
  set_default_plot_size(24cm, 14cm)
  p = plot(df, x=Col.value(xcols...), y=:ACTUAL,
       xgroup=Col.index(xcols...), ygroup=:MODEL,
       color=:config,
       Guide.xlabel("Covariate shift"),
       Guide.ylabel("Error by models"),
       Guide.title("Error vs. covariate shift"),
       Guide.colorkey(title="Configuration"),
       Scale.color_discrete_manual(colors...),
       Geom.subplot_grid(layer(Geom.point), free_x_axis=true,
                         layer(yintercept=[0.0], Geom.hline(color="gray", style=:dash)),
                         layer(yintercept=[0.5], Geom.hline(color="gray", style=:dash)),
                         Coord.cartesian(xmin=0.0,ymax=0.5),
                         Guide.ylabel(orientation=:vertical)))
  p |> SVG(joinpath(@__DIR__,"gaussian-plot1.svg"))
  p
end

# generalization error by different methods
Gadfly.with_theme(theme) do
  ycols = (:CV,:BCV,:DRV,:ACTUAL)
  set_default_plot_size(24cm, 18cm)
  p1 = plot(df, x=:NoveltyFactor, y=Col.value(ycols...),
       xgroup=Col.index(ycols...), color=:rfactor,
       Guide.xlabel("Covariate shift by methods"),
       Guide.ylabel("Error"),
       Guide.title("Error vs. covariate shift"),
       Guide.colorkey(title="Correlation length"),
       Scale.color_discrete_manual(colors...),
       Geom.subplot_grid(layer(Geom.point),
                         layer(yintercept=[0.0], Geom.hline(color="gray", style=:dash)),
                         layer(yintercept=[0.5], Geom.hline(color="gray", style=:dash)),
                         Coord.cartesian(xmin=0.0,xmax=1.0,ymax=0.5)))
  p2 = plot(ff, x=:rfactor, y=Col.value(ycols...),
            xgroup=Col.index(ycols...), color=:rfactor,
            Guide.xlabel("Correlation length by methods"),
            Guide.ylabel("Error"),
            Guide.title("Error vs. correlation length for inside configuration"),
            Guide.colorkey(title="Correlation length"),
            Scale.color_discrete_manual(colors...),
            Geom.subplot_grid(Geom.boxplot))
  p = vstack(p1, p2)
  p |> SVG(joinpath(@__DIR__,"gaussian-plot2.svg"))
  p
end

# generalization error by different correlation lengths
theme2 = style(point_size=1.8px,
               key_position=:top,
               default_color=colorant"black")
Gadfly.with_theme(theme2) do
  set_default_plot_size(24cm, 18cm)
  ycols = (:CV,:BCV,:DRV,:ACTUAL)
  p1 = plot(df, x=:NoveltyFactor, y=Col.value(ycols...),
       xgroup=:rfactor, color=Col.index(ycols...),
       Guide.xlabel("Covariate shift by correlation length"),
       Guide.ylabel("Error"),
       Guide.title("Error vs. covariate shift"),
       Guide.colorkey(title="Method"),
       Scale.color_discrete_manual(colors...),
       Geom.subplot_grid(layer(Geom.point), layer(Geom.line, Stat.smooth(smoothing=1.0)),
                         layer(yintercept=[0.0], Geom.hline(color="gray", style=:dash)),
                         layer(yintercept=[0.5], Geom.hline(color="gray", style=:dash)),
                         Coord.cartesian(xmin=0.0,xmax=1.0,ymax=0.5)))
  xcols = (:CV,:BCV,:DRV)
  p2 = plot(ff, x=Col.value(xcols...), y=:ACTUAL, xgroup=Col.index(xcols...),
            Guide.xlabel("Estimated Error"), Guide.ylabel("Actual Error"),
            Guide.title("Q-Q plot for inside configuration"),
            Geom.subplot_grid(layer(Geom.point,Stat.qq),
                              layer(Geom.abline(color="grey",style=:dot))))
  p = vstack(p1, p2)
  p |> SVG(joinpath(@__DIR__,"gaussian-plot3.svg"))
  p
end


function en(an,anm1,Ln)
	return sqrt((1-abs2(an))/((one(an)-abs2(anm1))*(one(an)-abs2(Ln))))
end

function ln(z::Vector{ComplexPair{BigFloat}},a::Vector{ComplexPair{BigFloat}},phinm1::Vector{ComplexPair{BigFloat}})
    if length(a)==1
        an = copy(a[end])
        anm1 = zero(ComplexPair{BigFloat})
        phinm1star = phinm1
        f1 = [((z[i]-anm1)/(one(ComplexPair{BigFloat})-conj(an)*z[i]))*conj(phinm1[i]) for i=1:length(z)]
        f2 = [((one(ComplexPair{BigFloat})-conj(anm1)*z[i])/(one(ComplexPair{BigFloat})-conj(an)*z[i]))*phinm1star[i] for i=1:length(z)]
        num = sum(f1)
        den = sum(f2)
        r = -num/den
    else
        an = copy(a[end])
        anm1 = copy(a[end-1])
        bProd = length(a)>=3 ? conjBlaschkeProd(z,[a[1:end-2]]) : ones(ComplexPair{BigFloat},length(z))
        f1 = [((z[i]-anm1)/(one(ComplexPair{BigFloat})-conj(an)*z[i]))*conj(phinm1[i]) for i=1:length(z)]
        f2 = [((z[i]-anm1)/(one(ComplexPair{BigFloat})-conj(an)*z[i]))*phinm1[i]*bProd[i] for i=1:length(z)]
        num = sum(f1)
        den = sum(f2)
        r = -num/den
    end
    return r
end

function phiNext(z::Vector{ComplexPair{BigFloat}},a::Vector{ComplexPair{BigFloat}},phinm1::Vector{ComplexPair{BigFloat}})
    phin = zeros(ComplexPair{BigFloat},length(phinm1))
    if length(a)==1
        an = copy(a[end])
        anm1 = zero(ComplexPair{BigFloat})
        L = ln(z,a,phinm1)
        c = en(an,anm1,L)
        for k=1:length(z)
            q = (z[k]-anm1)
            p = (one(ComplexPair{BigFloat})-conj(anm1)*z[k])
            C = c/(z[k]-an)
            PHI = p*phinm1[k]
            PHISTAR = q*phinm1[k]
            phin[k] = C*(PHI+conj(L)*PHISTAR)
        end
    else
        an = a[end]
        anm1 = a[end-1]
        L = ln(z,a,phinm1)
        c = en(an,anm1,L)
        bProd = length(a)>=3 ? blaschkeProd(z,[a[1:end-2]]) : ones(ComplexPair{BigFloat},length(z))
        for k=1:length(z)
            q = (z[k]-an)
            p = (one(ComplexPair{BigFloat})-conj(anm1)*z[k])
            C = c*p/q
            phin[k] = C*(phinm1[k]+conj(L)*bProd[k]*conj(phinm1[k]))
        end
    end
    return (phin, L)
end

function phiNext(z::Vector{ComplexPair{BigFloat}},a::Vector{ComplexPair{BigFloat}},phinm1::Vector{ComplexPair{BigFloat}},L::ComplexPair{BigFloat})
    phin = zeros(ComplexPair{BigFloat},length(phinm1))
    if length(a)==1
        an = copy(a[end])
        anm1 = zero(ComplexPair{BigFloat})
        c = en(an,anm1,L)
        for k=1:length(z)
            q = (z[k]-anm1)
            p = (one(ComplexPair{BigFloat})-conj(anm1)*z[k])
            C = c/(z[k]-an)
            PHI = p*phinm1[k]
            PHISTAR = q*phinm1[k]
            phin[k] = C*(PHI+conj(L)*PHISTAR)
        end
    else
        an = a[end]
        anm1 = a[end-1]
        c = en(an,anm1,L)
        bProd = length(a)>=3 ? blaschkeProd(z,[a[1:end-2]]) : ones(ComplexPair{BigFloat},length(z))
        for k=1:length(z)
            q = (z[k]-an)
            p = (one(ComplexPair{BigFloat})-conj(anm1)*z[k])
            C = c*p/q
            phin[k] = C*(phinm1[k]+conj(L)*bProd[k]*conj(phinm1[k]))
        end
    end
    return phin
end

function blaschkeProd(z::Vector{ComplexPair{BigFloat}},aj::Vector{ComplexPair{BigFloat}})
    r = ones(ComplexPair{BigFloat},length(z))
    for i=1:length(aj)
        t = [(one(ComplexPair{BigFloat})-conj(aj[i])*z[k])./(z[k]-aj[i]) for k=1:length(z)]
        r = r.*t
    end
    return r
end

function conjBlaschkeProd(z::Vector{ComplexPair{BigFloat}},aj::Vector{ComplexPair{BigFloat}})
    r = ones(ComplexPair{BigFloat},length(z))
    for i=1:length(aj)
        t = [(z[k]-aj[i])/(one(ComplexPair{BigFloat})-conj(aj[i])*z[k]) for k=1:length(z)]
        r = r.*t
    end
    return r
end

function genZArb(N::Int)
	return exp(convert(Vector{ComplexPair{BigFloat}},im*linspace2(0.0,2*pi,N)))
end

function genZArcArb(N::Int,Div::Int)
    return exp(convert(Vector{ComplexPair{BigFloat}},im*linspace(0.0,2*pi/Div,N)))
end

function genZArbUpSample(N::Int,Mult::Int)
    return exp(convert(Vector{ComplexPair{BigFloat}},im*linspace2upsample(0.0,2*pi,N,Mult)))
end

function genBasis(z::Vector,a::Vector{ComplexPair{BigFloat}})
    P = zeros(ComplexPair{BigFloat},(length(z),length(a)+1))
    P[:,1] = ones(ComplexPair{BigFloat},length(z))/sqrt(length(z));
    L = zeros(ComplexPair{BigFloat},length(a))
    for k=1:length(a)
        (P[:,k+1],L[k]) = phiNext(z,[a[1:k]],P[:,k])
    end
    return (P,L)
end

function genBasisGivenL(z::Vector{ComplexPair{BigFloat}},a::Vector{ComplexPair{BigFloat}},L::Vector{ComplexPair{BigFloat}},N::Int)
    P = zeros(ComplexPair{BigFloat},(length(z),length(a)+1))
    P[:,1] = ones(ComplexPair{BigFloat},length(z))/sqrt(N);
    for k=1:length(a)
        P[:,k+1] = phiNext(z,[a[1:k]],P[:,k],L[k])
    end
    return P
end

function unitDiskGrid(n::Int,delta::Float64)
    N = nextpow2(n)
    M = N>>1
    v = linspace(delta,one(Float64)-delta,M)
    Zjk = zeros(ComplexPair{BigFloat},N,N)
    for j=1:(M)
            #bottom, running counter-clockwise
            Zjk[(M+j),((M-(j-1)):(M+1+(j-1)))] = v[j]*exp(im*linspace(5*pi/4,7*pi/4,2*j))
            #top, running clockwise
            Zjk[(M-j+1),((M-(j-1)):(M+1+(j-1)))] = v[j]*exp(im*linspace(3*pi/4,1*pi/4,2*j))
    end
    for j=1:(M)
            #right, running clockwise
            Zjk[((M-(j-1)):(M+1+(j-1))),M+j] = v[j]*exp(im*linspace(pi/4,-pi/4,2*j)) #right
            #left, running counter-clockwise
            Zjk[((M-(j-1)):(M+1+(j-1))),M-j+1] = v[j]*exp(im*linspace(3*pi/4,5*pi/4,2*j)) #left
    end
    return (reshape(Zjk,N^2))
end


###
# Distribution only based on its laplace transform function `f` and exponential tilting `c²`.
# The sampling is from Ridout, M. S. (2009). Generating random numbers from a distribution specified by its laplace transform. Statistics and Computing, 19(4):439.
# And the inverse laplace operation is done via Grassmann, W. K. (Ed.). (2000). Computational Probability. International Series in #Operations Research & Management Science. doi:10.1007/978-1-4757-4828-4
struct LaplaceTransformDistribution{T,TAlg} <:
       Distributions.ContinuousUnivariateDistribution
    f::Function # Laplace transform of the pdf
    c²::T # Exponential tilting parameter
    alg::TAlg # Algorithm to compute the inverse Laplace transform
    function LaplaceTransformDistribution{T,TAlg}(
        f::Function, c²::T, alg::TAlg
    ) where {T<:Real,TAlg}
        _check_f(f) || error("The function passed is not valid") # Do a series of checks on f
        c² >= 0 || error("c² has to a be non-negative real")
        return new{T,TAlg}(f, c², alg)
    end
end

function LaplaceTransformDistribution(
    f::Function, c²::T=0.0, alg::TAlg=BromwichInverseLaplace()
) where {T,TAlg}
    return LaplaceTransformDistribution{T,TAlg}(f, c², alg)
end

function _check_f(f)
    return true # TODO Add tests for complete monotonicity / PDR
end
_gradf(d::LaplaceTransformDistribution, x::Real) = first(ForwardDiff.gradient(d.f, [x]))
function _gradlogf(d::LaplaceTransformDistribution, x::Real)
    return first(ForwardDiff.gradient(log ∘ d.f, [x]))
end
function _hessianlogf(d::LaplaceTransformDistribution, x::Real)
    return first(ForwardDiff.hessian(log ∘ d.f, [x]))
end

Distributions.pdf(dist::LaplaceTransformDistribution, x::Real) = apply_f(dist, x)
function Distributions.mean(dist::LaplaceTransformDistribution)
    return _gradf(dist, dist.c²) / dist.f(dist.c²)
end
function Distributions.var(dist::LaplaceTransformDistribution)
    return _hessianlogf(dist, dist.c²) / dist.f(dist.c²) - mean(dist)^2
end

function Random.rand(dist::LaplaceTransformDistribution)
    return first(rand(dist, 1))
end

# Sampling from Ridout (09)
function Random.rand(dist::LaplaceTransformDistribution, n::Int)
    nTries = 100
    i = 0
    while i < nTries
        try
            if i > 10
                @warn "This is taking more than 10 tries! Trie $i/$nTries!"
            end
            return laptrans(dist; n=n)
        catch e
            if e isa AssertionError
                @warn "$(e.msg), Trying again"
                i += 1
                continue
            else
                rethrow(e)
            end
        end
    end
    @error "Sampler failed to converge"
end

struct BromwichInverseLaplace{T}
    A::Int
    l::Int
    m::Int
    n::Int
    expA2l::T
    ijπl::Vector{Complex{T}}
    expijπl::Vector{Complex{T}}
    coefs::Vector{T}
    altern::Vector{Int}
    b::Vector{T}
    s::Vector{T}
    function BromwichInverseLaplace(A::Int=19, l::Int=1, m::Int=11, n::Int=28)
        @assert A > 0 "A should be a positive integer"
        @assert l > 0 "l should be a positive integer"
        @assert m > 0 "m should be a positive integer"
        @assert n > 0 "n should be a positive integer"
        T = typeof(exp(A / (2 * l)))
        ijπl = 1im .* (1:l) .* π / l
        return new{T}(
            A,
            l,
            m,
            n,
            exp(A / (2 * l)),
            ijπl,
            exp.(ijπl),
            binomial.(m, 0:m) ./ (2^m),
            (-1) .^ (0:(n + m)),
            zeros(n + m + 1),
            zeros(n + m + 1),
        )
    end
end

@inline function apply_f(dist::LaplaceTransformDistribution, t)
    if t == 0
        return zero(t)
    else
        invlaplace(dist.f, t, dist.alg)
    end
end

@inline function apply_F(dist::LaplaceTransformDistribution, t)
    if t == 0
        return zero(t)
    else
        invlaplace(s -> dist.f(s) / s, t, dist.alg)
    end
end

function invlaplace(f::Function, t::Real, A::Int=19, l::Int=1, m::Int=11, n::Int=28)
    # n is the burnin, m is the number of components, A and l aare discretization parameters
    alg = BromwichInverseLaplace(A, l, m, n)
    return invlaplace(f, t, alg)
end

function invlaplace(f::Function, t::Real, alg::BromwichInverseLaplace)
    scaled_t = 1 / (2 * alg.l * t)
    alg.b[1] =
        f(alg.A * scaled_t) +
        2 * sum(real.(f.(alg.A * scaled_t .+ alg.ijπl ./ t) .* alg.expijπl))
    @inbounds for k in 1:(alg.n + alg.m)
        alg.b[k + 1] =
            2 * sum(
                real.(
                    f.(alg.A * scaled_t .+ alg.ijπl ./ t .+ 1im * π * k / t) .* alg.expijπl
                ),
            )
    end
    acoeff = alg.expA2l * scaled_t
    cumsum!(alg.s, acoeff .* alg.b .* alg.altern)
    return dot(view(alg.s, (alg.n + 1):length(alg.s)), alg.coefs)
end

# Laplace transform implemented from "Computational Probability Grassmann 2000"
function laptrans(
    dist::LaplaceTransformDistribution;
    n::Int=10,
    jmax::Int=500,
    kmax::Int=1000,
    b::Real=2.0,
    τ::Real=1e-7,
)
    # Step 1
    u = sort!(rand(n))
    # Step 2
    xmax = 0.1
    j = 0
    while apply_F(dist, xmax) < u[end] && j < jmax
        xmax = b * xmax
        j += 1
    end
    @assert j != jmax "Xmax search failed to converge"
    x = zeros(n)
    # Step 3
    for i in 1:n
        x_L = i == 1 ? 0.0 : x[i - 1]
        x_U = xmax
        k = 0
        t = x_L
        while abs(apply_F(dist, t) - u[i]) > τ && k < kmax
            k += 1
            t = t - (apply_F(dist, t) - u[i]) / apply_f(dist, t)
            if t < x_L || t > x_U
                t = 0.5 * (x_L + x_U)
            end
            if apply_F(dist, t) <= u[i]
                x_L = t
            else
                x_U = t
            end
        end
        @assert k != kmax "Modified Newton-Raphson procedure failed to converge for i=$i, x[$i]=$t"
        x[i] = t
    end
    # Step 4
    return shuffle!(x)
end


# PREAMBLE

# PKG_SETUP

# ## Setup

using LinearAlgebra
using Plots
using DirectTrajectoryOptimization

# ## horizon
T = 101

# ## acrobot
num_state = 4
num_action = 1
num_parameter = 0

function acrobot(x, u, w)
    mass1 = 1.0
    inertia1 = 0.33
    length1 = 1.0
    lengthcom1 = 0.5

    mass2 = 1.0
    inertia2 = 0.33
    length2 = 1.0
    lengthcom2 = 0.5

    gravity = 9.81
    friction1 = 0.1
    friction2 = 0.1

    function M(x, w)
        a = (inertia1 + inertia2 + mass2 * length1 * length1
            + 2.0 * mass2 * length1 * lengthcom2 * cos(x[2]))

        b = inertia2 + mass2 * length1 * lengthcom2 * cos(x[2])

        c = inertia2

       return [a b; b c]
    end

    function Minv(x, w)
        a = (inertia1 + inertia2 + mass2 * length1 * length1
            + 2.0 * mass2 * length1 * lengthcom2 * cos(x[2]))

        b = inertia2 + mass2 * length1 * lengthcom2 * cos(x[2])

        c = inertia2

        return 1.0 / (a * c - b * b) * [c -b; -b a]
    end

    function τ(x, w)
        a = (-1.0 * mass1 * gravity * lengthcom1 * sin(x[1])
            - mass2 * gravity * (length1 * sin(x[1])
            + lengthcom2 * sin(x[1] + x[2])))

        b = -1.0 * mass2 * gravity * lengthcom2 * sin(x[1] + x[2])

        return [a; b]
    end

    function C(x, w)
        a = -2.0 * mass2 * length1 * lengthcom2 * sin(x[2]) * x[4]
        b = -1.0 * mass2 * length1 * lengthcom2 * sin(x[2]) * x[4]
        c = mass2 * length1 * lengthcom2 * sin(x[2]) * x[3]
        d = 0.0

        return [a b; c d]
    end

    function B(x, w)
        [0.0; 1.0]
    end

    q = view(x, 1:2)
    v = view(x, 3:4)

    qdd = Minv(q, w) * (-1.0 * C(x, w) * v
            + τ(q, w) + B(q, w) * u[1] - [friction1; friction2] .* v)

    return [x[3]; x[4]; qdd[1]; qdd[2]]
end

function midpoint_implicit(y, x, u, w)
    h = 0.05 # timestep
    y - (x + h * acrobot(0.5 * (x + y), u, w))
end

# ## model
dt = Dynamics(midpoint_implicit, num_state, num_state, num_action, num_parameter=num_parameter, evaluate_hessian=true)
dynamics = [dt for t = 1:T-1]

# ## initialization
x1 = [0.0; 0.0; 0.0; 0.0]
xT = [π; 0.0; 0.0; 0.0]

# ## objective
ot = (x, u, w) -> 0.1 * dot(x[3:4], x[3:4]) + 0.1 * dot(u, u)
oT = (x, u, w) -> 0.1 * dot(x[3:4], x[3:4])
ct = Cost(ot, num_state, num_action, num_parameter=num_parameter)
cT = Cost(oT, num_state, 0, num_parameter=num_parameter)
objective = [[ct for t = 1:T-1]..., cT]

# ## constraints
bnd1 = Bound(num_state, num_action)
bndt = Bound(num_state, num_action)
bndT = Bound(num_state, 0)
bounds = [bnd1, [bndt for t = 2:T-1]..., bndT]

constraints = [
        Constraint((x, u, w) -> x - x1, num_state, num_action),
        [Constraint() for t = 2:T-1]...,
        Constraint((x, u, w) -> x - xT, num_state, 0)
       ]

# ## problem
p = Solver(dynamics, objective, constraints, bounds,
    options=Options{Float64}())

# ## initialize
x_interpolation = linear_interpolation(x1, xT, T)
u_guess = [1.0 * randn(num_action) for t = 1:T-1]

initialize_states!(p, x_interpolation)
initialize_controls!(p, u_guess)

# ## solve
@time solve!(p)

# ## solution
x_sol, u_sol = get_trajectory(p)

@show x_sol[1]
@show x_sol[T]

# ## state
plot(hcat(x_sol...)')

# ## control
plot(hcat(u_sol[1:end-1]..., u_sol[end-1])', linetype = :steppost)


get_input() = map(s -> parse(Int, s), readlines(joinpath(INPUT_DIR, "day01.txt")))

fuelForMass(mass::Int) = div(mass, 3) - 2

function solve(::Puzzle(2019, 1, 1), input)
    map(fuelForMass, get_input()) |> sum
end

recursiveFuel(mass::Int) = begin
    fuel = fuelForMass(mass)

    if fuel <= 0
        return 0
    end

    return fuel + recursiveFuel(fuel)
end

function solve(::Puzzle(2019, 1, 2), input)
    map(recursiveFuel, input) |> sum
end


using Cbc
using Test
using LinearAlgebra
solver = CbcSolver(logLevel=0, integerTolerance=1e-9, primalTolerance=1e-9, ratioGap=1e-8)

# using Gurobi
# solver = GurobiSolver(OutputFlag=0)

# using CPLEX
# solver = CplexSolver(CPX_PARAM_SCRIND=0)

using JuMP
using PiecewiseLinearOpt

methods_1D = (:CC,:MC,:Logarithmic,:LogarithmicIB,:ZigZag,:ZigZagInteger,:SOS2,:GeneralizedCelaya,:SymmetricCelaya,:Incremental,:DisaggLogarithmic)
methods_2D = (:CC,:MC,:Logarithmic,:LogarithmicIB,:ZigZag,:ZigZagInteger,:SOS2,:GeneralizedCelaya,:SymmetricCelaya,:DisaggLogarithmic)
patterns_2D = (:Upper,:Lower,:BestFit,:UnionJack,:K1,:Random) # :OptimalTriangleSelection and :Stencil not supported currently

println("\nunivariate tests")
@testset "1D: $method" for method in methods_1D
    model = Model(solver=solver)
    @variable(model, x)
    z = piecewiselinear(model, x, range(1,stop =2π, length =8), sin, method=method)
    @objective(model, Max, z)

    @test solve(model) == :Optimal
    @test getvalue(x) ≈ 1.75474 rtol=1e-4
    @test getvalue(z) ≈ 0.98313 rtol=1e-4

    @constraint(model, x ≤ 1.5z)

    @test solve(model) == :Optimal
    @test getvalue(x) ≈ 1.36495 rtol=1e-4
    @test getvalue(z) ≈ 0.90997 rtol=1e-4
end

println("\nbivariate tests")
@testset "2D: $method, $pattern" for method in methods_2D, pattern in patterns_2D
    model = Model(solver=solver)
    @variable(model, x)
    @variable(model, y)
    d = range(0,stop=1,length=8)
    f = (x,y) -> 2*(x-1/3)^2 + 3*(y-4/7)^4
    z = piecewiselinear(model, x, y, BivariatePWLFunction(d, d, f, pattern=pattern), method=method, subsolver=solver)
    @objective(model, Min, z)

    @test solve(model) == :Optimal
    @test getvalue(x) ≈ 0.285714 rtol=1e-4
    @test getvalue(y) ≈ 0.571429 rtol=1e-4
    @test getvalue(z) ≈ 0.004535 rtol=1e-3

    @constraint(model, x ≥ 0.6)

    @test solve(model) == :Optimal
    @test getvalue(x) ≈ 0.6 rtol=1e-4
    @test getvalue(y) ≈ 0.571428 rtol=1e-4
    @test getvalue(z) ≈ 0.148753 rtol=1e-4
end

println("\nbivariate optimal IB scheme tests")
@testset "2D: optimal IB, UnionJack" begin
    model = Model(solver=solver)
    @variable(model, x)
    @variable(model, y)
    d = range(0,stop=1,length=3)
    f = (x,y) -> 2*(x-1/3)^2 + 3*(y-4/7)^4
    z = piecewiselinear(model, x, y, BivariatePWLFunction(d, d, f, pattern=:UnionJack), method=:OptimalIB, subsolver=solver)
    @objective(model, Min, z)

    @test solve(model) == :Optimal
    @test getvalue(x) ≈ 0.5 rtol=1e-4
    @test getvalue(y) ≈ 0.5 rtol=1e-4
    @test getvalue(z) ≈ 0.055634 rtol=1e-3

    @constraint(model, x ≥ 0.6)

    @test solve(model) == :Optimal
    @test getvalue(x) ≈ 0.6 rtol=1e-4
    @test getvalue(y) ≈ 0.5 rtol=1e-4
    @test getvalue(z) ≈ 0.222300 rtol=1e-3
end


""" Média Ponderada
# Implemente um programa que lê quatro valores
# e calcule a média ponderada para pesos 1, 2, 3 e 4,
# respectivamente (multiplique cada nota pelo seu peso,
# some os produtos e divida o resultado pela soma dos pesos).
# Exemplo: 5*2 + 7*2 + 6*3 + 3*3
"""

println("Informe o primeiro produto")
avg = parse( Float16, readline() )
println("Informe o peso do primeiro produto")
weight = parse( Float16, readline() )
numerator = avg * weight
denominator = weight

println("Informe o segundo produto")
avg = parse( Float16, readline() )
println("Informe o peso do segundo produto")
weight = parse( Float16, readline() )
numerator += avg * weight
denominator += weight

println("Informe o terceiro produto")
avg = parse( Float16, readline() )
println("Informe o peso do terceiro produto")
weight = parse( Float16, readline() )
numerator += avg * weight
denominator += weight

println("Informe o quarto produto")
avg = parse( Float16, readline() )
println("Informe o peso do quarto produto")
weight = parse( Float16, readline() )
numerator += avg * weight
denominator += weight

avg = (numerator / denominator)
println( "A média ponderada = ",avg )


@testset "Fifth order" begin
    @testset "Bits $i" for i in 11:12
        x = 0:1 << i
        y_approx = fpcos.(Int32.(x), Val(i))
        y = cos.(x / 1 << i * π / 2) .* 1 << i
        @test all(abs.(y_approx .- y) .< 2.0)
    end
end


"""
    logseriescdf(p)

Returns the vector of samples from the discrete cdf of logarithmic distribution
"""
function logseriescdf(p)
  T = eltype(p)
  cdfs = T.([0.])
  for i in 1:100000000
    @inbounds push!(cdfs, cdfs[i]-(p^i)/(i*log(1-p)))
    # one can change atol here
    if isapprox(cdfs[i],  1.)
      return cdfs
    end
  end
  cdfs
end

"""
    levyel(θ, u1, u2)

u1 and u2 are supposed to be random numbers form [0,1]
An element from Levy stable distribution with parameters α = 1/θ, β = 1,
γ = (cos(pi/(2*θ)))^θ and δ = 0.
"""
function levyel(θ, u1, u2)
  ϕ = pi*u1-pi/2
  v = quantile.(Exponential(1.), u2)
  γ = (cos(pi/(2*θ)))^θ
  v = ((cos(pi/(2*θ)+(1/θ-1)*ϕ))/v)^(θ-1)
  γ*v*sin(1/θ*(pi/2+ϕ))*(cos(pi/(2*θ))*cos(ϕ))^(-θ)
end

"""
    tiltedlevygen(V0, α; rng::AbstractRNG)

Returns a sample from the expotencialy tilted Levy stable pdf
f(x; V0, α) = exp(-V0^α) g(x; α)/exp(-V0), where g(x; α) is a stable Levy pdf
with parameters α = 1/θ, β = 1, γ = (cos(pi/(2*θ)))^θ and δ = 0.
"""
function tiltedlevygen(V0, α; rng)
  T = eltype(V0)
  x = levyel(α, rand(rng, T), rand(rng, T))
  u = rand(rng, T)
  while exp(-V0^α*x)/(1500*exp(-V0)) < u
    x = levyel(α, rand(rng, T), rand(rng, T))
    u = rand(rng, T)
  end
  return x.*V0.^α
end

"""
    Ginv(y, α)

Returns Real, helper for the joe/frank nested copula generator
"""
Ginv(y, α) = ((1-y)*SpecialFunctions.gamma(1-α))^(-1/α)

"""
    InvlaJ(n::Int, α)

Returns the inverse Laplace transform of generator of Joe nested copula
parametered by α
"""
function InvlaJ(n, α)
  T = eltype(α)
  return 1-1/(n*beta(T(n), 1-α))
end

"""
  sampleInvlaJ(α, v)

Returns Int, a sample of inverce laplacea transform of generator of Joe nested copula,
given a parameter α and a random numver v ∈ [0,1]
"""
function sampleInvlaJ(α, v)
  if v <= α
    return 1
  else
    G = Ginv(v, α)
    if G > 2^62
      return 2^62
    else
      return (InvlaJ(floor(Int, G), α) < v) ? ceil(Int, G) : floor(Int, G)
    end
  end
end

"""
  elInvlaF(θ₁, θ₀, logseriescdf; rng::AbstractRNG)

Returns Int, a single sample of the inverse Laplace transform of the generator
of nested Frank copula
"""
function elInvlaF(θ₁, θ₀, logseriescdf; rng)
  T = eltype(θ₁)
  c1 = 1-exp(-θ₁)
  α = θ₀/θ₁
  if θ₀ <= 1
    v = rand(rng, T)
    X = findlast(logseriescdf .< rand(rng, T))[1]
    while v > 1/((X-α)*beta(X, 1-α))
      v = rand(rng, T)
      X = findlast(logseriescdf .< rand(rng, T))[1]
    end
    return X
  else
    v = rand(rng, T)
    X = sampleInvlaJ(α, rand(rng, T))
    while v > c1^(X-1)
      X = sampleInvlaJ(α, rand(rng, T))
      v = rand(rng, T)
    end
    return X
  end
end

"""
  nestedfrankgen(θ₁, θ₀, V0::Int, logseriescdf; rng::AbstractRNG)

Return int, sampled from the Inverse Laplace trensform of nested
Frank copula given parametes θ₁ θ₀ (child and parent)
and an element of V0 - vector of samples of invlaplace of the parents copula
"""
function nestedfrankgen(θ₁, θ₀, V0, logseriescdf; rng)
  T = eltype(θ₁)
  return sum([elInvlaF(θ₁, θ₀, logseriescdf; rng = rng) for j in 1:V0])
end


function simulate(b::Beta, n::Vector{<:Integer})
    @assert all(n .> 0) "All nᵢ must be strictly positive for i = 1,...,N"
    N = length(n)

    x = zeros(Integer, N)
    θ = zeros(N)
    for i in 1:N
        θ[i] = rand(b)
        x[i] = rand(Binomial(n[i], θ[i]))
    end
    return (x = x, n = n, θ = θ)
end

n   = rand(1:100, 100)
sim = simulate(Beta(10, 30), n)


println("\nTesting MCDTS complete tree, coupled Logistic, CCM:")
@time begin
@testset "MCDTS single rollout on univariate Logistic map data" begin

    Lval = 3500
    x = zeros(Lval)
    y = zeros(Lval)
    r = 3.8
    r2 = 3.5
    βxy = 0.02
    βyx = 0.1
    x[1]=0.4
    y[1]=0.2

    for i = 2:Lval
        x[i]=x[i-1]*(r-r*x[i-1]-βxy*y[i-1])
        y[i]=y[i-1]*(r2-r2*y[i-1]-βyx*x[i-1])
    end

    w1 = DelayEmbeddings.estimate_delay(x, "mi_min")
    w2 = DelayEmbeddings.estimate_delay(y, "mi_min")

    test1 = DelayEmbeddings.standardize(x)
    test2 = DelayEmbeddings.standardize(y)
    # try MCDTS with CCM
    taus1 = 0:10 # the possible delay vals
    trials = 20 # the sampling of the tree

    Random.seed!(1234)
    optmodel = MCDTS.MCDTSOptimGoal(MCDTS.CCM_ρ(test2), MCDTS.Range_function())
    tree = mcdts_embedding(Dataset(test1), optmodel, w1, taus1, trials)
    best_node = MCDTS.best_embedding(tree)
    τ_mcdts = best_node.τs
    ts_mcdts = best_node.ts
    Lval = L(best_node)

    # less fid points
    Random.seed!(1234)
    optmodel2 = MCDTS.MCDTSOptimGoal(MCDTS.CCM_ρ(test2,1,0.5), MCDTS.Range_function())
    tree2 = mcdts_embedding(Dataset(test1), optmodel2, w1, taus1, trials)
    best_node2 = MCDTS.best_embedding(tree2)
    τ_mcdts2 = best_node2.τs
    ts_mcdts2 = best_node2.ts
    L2 = L(best_node2)

    @test Lval < -0.95
    @test Lval - .01 < L2 < Lval + .01
    @test length(τ_mcdts) ==  3 == length(τ_mcdts2)
    @test τ_mcdts[1] == 0 == τ_mcdts2[1]
    @test τ_mcdts[2] == 2 == τ_mcdts2[2]
    @test τ_mcdts[3] == 1 == τ_mcdts2[3]

    Random.seed!(1234)
    optmodel2 = MCDTS.MCDTSOptimGoal(MCDTS.CCM_ρ(test1), MCDTS.Range_function())
    tree2 = mcdts_embedding(Dataset(test2), optmodel2, w2, taus1, trials)
    best_node = MCDTS.best_embedding(tree2)
    τ_mcdts = best_node.τs
    ts_mcdts = best_node.ts
    Lval = L(best_node)

    @test Lval < -0.997
    @test length(τ_mcdts) ==  3
    @test τ_mcdts[1] == 0
    @test τ_mcdts[2] == 1
    @test τ_mcdts[3] == 2

end
end
true


import LinearAlgebra.cross
cross(a::Pt{2,T}, b::Pt{2,T}) where {T} = a[1]*b[2] - a[2]*b[1]

"""
Compute whether two flat patches of the same dimension overlap or not
"""
function overlap(p1::Simplex{U,D,C,N,T}, p2::Simplex{U,D,C,N,T}) where {U,D,C,N,T}

    # Are the patches in the same D-plane?
    throw(ErrorException("Not implemented yet!"))

end

###
# Version for points: overlap in this case is taken to mean that points approx conincide
###
function overlap(p::Simplex{U,0,C,1,T}, q::Simplex{U,0,C,1,T}) where {T,U,C}

    tol = sqrt(eps(T))
    return norm(p[1]-q[1]) < tol
end

"""
Compute whether two segments in 3D space overlap
"""
function overlap(p::Simplex{U,1,C,2,T}, q::Simplex{U,1,C,2,T}) where {T,U,C}

    tol = sqrt(eps(T))

    # Are these edges on the same line?
    o = q.vertices[2]
    u = p.vertices[1] - o
    t = q.tangents[1]
    t2 = dot(t,t)
    e = u - dot(u,t) / t2 * t

    # if vertex 1 of p is not in the line defined by q, they do not overlap
    norm(e) > tol && return false

    # if the tangents are not linearly dependent return false
    norm(cross(p.tangents[1],q.tangents[1])) > tol && return false

    # we determined p and q are on the same line. Do standard collision testing
    a1 = dot(p.vertices[1]-o, t)
    b1 = dot(p.vertices[2]-o, t)

    a1 < b1 ? (x1 = a1; y1 = b1) : (x1 = b1; y1 = a1)

    y1 <= zero(T) && return false # p to the left of q
    t2 <= x1 && return false      # q to the left of p

    return true
end

"""
Compute whether two triangles in 3D space overlap
"""
function overlap(p::Simplex{3,2,1,3,T}, q::Simplex{3,2,1,3,T}) where T

  tol = sqrt(eps(T))

  # Are the patches in the same plane?
  u1 = q.tangents[1]
  u2 = q.tangents[2]
  v = p.vertices[1] - q.vertices[2]

  # if vertex 1 of p is not in q, return false
  norm(dot(cross(u1,u2),v)) > tol && return false

  # if the two triangles are not coplanar return false
  norm(cross(p.normals[1], q.normals[1])) > tol && return false

  n = p.normals[1]
  for i in 1:3
    a = p.vertices[mod1(i+1,3)]
    b = p.vertices[mod1(i+2,3)]
    c = p.vertices[i]
    t = b - a
    m = cross(t,n)

    sp = zeros(T,3); sp[i] = dot(c-a, m)
    sq = T[dot(q.vertices[j]-a, m) for j in 1:3]

    minp, maxp = extrema(sp)
    minq, maxq = extrema(sq)

    maxq <= minp + tol && return false
    maxp <= minq + tol && return false
  end

  return true
end


# check if vector is sorted and elements are unique
function _checksu(v::AbstractVector)
    isu = true
    for i in 2:length(v)
        if v[i] <= v[i-1]
            isu = false
            break
        end
    end
    isu
end

# a vector whose elements are sorted and unique
struct SUVector{T} <: AbstractVector{T}
    # elements are stored in a standard vector
    data::Vector{T}
    # constructor
    function SUVector(data::AbstractVector{T}) where T
        # check sorted and unique properties
        @assert _checksu(data)
        # NOTE: data is copied
        new{T}(copy(data))
    end
end

# empty constructor
SUVector{T}() where T = SUVector(T[])

# formatted printing
show(io::IO, v::SUVector) = show(io::IO, v.data)

# shape of the vector
size(v::SUVector) = size(v.data)

# indexed reading
getindex(v::SUVector, i::Int) = getindex(v.data, i)

# indexed writing
function setindex!(v::SUVector, vi, i::Int)
    # check to make sure order and uniqueness of elements is preserved
    l = length(v)
    if l > 1
        errmsg = "This doesn't preserve order or uniqueness of elements."
        if i == 1
            @assert vi < v[i+1] errmsg
        elseif i == l
            @assert v[i-1] < vi errmsg
        else
            @assert v[i-1] < vi < v[i+1] errmsg
        end
    end
    setindex!(v.data, vi, i)
end

# pushing an element onto the end of a sorted unique vector
function push!(v::SUVector{T}, el::T) where T
    if !isempty(v)
        @assert el > v[end] "Must be larger than current last element."
    end
    push!(v.data, el)
end

# appending elements onto the end of a sorted unique vector
function append!(v::SUVector{T}, els::AbstractVector{T}) where T
    if !isempty(v)
        errmsg = "Doesn't preserve order or uniqueness of elements."
        @assert _checksu(els) && els[1] > v[end] errmsg
    end
    append!(v.data, els)
end

# intersection of two sorted unique vectors
function intersect(v1::SUVector{T}, v2::SUVector{T}) where T
    s1 = length(v1)
    s2 = length(v2)
    # empty intersection to be added to
    int = SUVector{T}()
    # step through vectors simultaneously
    i1 = 1
    i2 = 1
    while i1 <= s1 && i2 <= s2
        if v1[i1] < v2[i2]
            # not an intersection, move towards one
            i1 += 1
        elseif v2[i2] < v1[i1]
            # not an intersection, move towards one
            i2 += 1
        else
            # an intersection
            push!(int, v1[i1])
            i1 += 1
            i2 += 1
        end
    end
    int
end

# symmetric difference of two sorted unique vectors
function symdiff(v1::SUVector{T}, v2::SUVector{T}) where T
    s1 = length(v1)
    s2 = length(v2)
    # empty symmetric difference to be added to
    sdiff = SUVector{T}()
    if s1 == 0
        # v1 is empty, so return all of v2
        append!(sdiff, v2)
    elseif s2 == 0
        # v2 is empty, so return all of v1
        append!(sdiff, v1)
    else
        # step through vectors simultaneously
        i1 = 1
        i2 = 1
        while i1 <= s1 && i2 <= s2
            if v1[i1] < v2[i2]
                # add to sdiff and move towards and intersection
                push!(sdiff, v1[i1])
                i1 += 1
            elseif v2[i2] < v1[i1]
                # add to sdiff and move towards and intersection
                push!(sdiff, v2[i2])
                i2 += 1
            else
                # intersection
                i1 += 1
                i2 += 1
            end
        end
        # once one vector is finished, add the rest of the other vector
        if i1 > s1 && i2 <= s2
            append!(sdiff, view(v2, i2:s2))
        elseif i2 > s2 && i1 <= s1
            append!(sdiff, view(v1, i1:s1))
        end
    end
    sdiff
end

# union of two sorted unique vectors
function union(v1::SUVector{T}, v2::SUVector{T}) where T
    s1 = length(v1)
    s2 = length(v2)
    # empty union to be added to
    uni = SUVector{T}()
    if s1 == 0
        # v1 is empty, return all of v2
        append!(uni, v2)
    elseif s2 == 0
        # v2 is empty, return all of v1
        append!(uni, v1)
    else
        # step through vectors simultaneously
        i1 = 1
        i2 = 1
        while i1 <= s1 && i2 <= s2
            if v1[i1] < v2[i2]
                # add to union and move towards an intersection
                push!(uni, v1[i1])
                i1 += 1
            elseif v2[i2] < v1[i1]
                # add to union and move towards and intersection
                push!(uni, v2[i2])
                i2 += 1
            else
                # intersection, only add one but increment both
                push!(uni, v1[i1])
                i1 += 1
                i2 += 1
            end
        end
        # when finished one of the vectors, add the rest of the other
        if i1 > s1 && i2 <= s2
            append!(uni, view(v2, i2:s2))
        elseif i2 > s2 && i1 <= s1
            append!(uni, view(v1, i1:s1))
        end
    end
    uni
end


function graphdiff(g1::Graph, g2::Graph)
    new_edges = Set{Pair{Int,Int}}()
    for u in nodes(g1)
        for v in neighbors(g1,u)
            if u<=v
                push!(new_edges, Pair(u,v))
            end
        end
    end
    for u in nodes(g2)
        for v in neighbors(g2,u)
            if u<=v
                delete!(new_edges, Pair(u,v))
            end
        end
    end
    new_edges
end

function graphdiff(graph1, graph2)
    edges_g1 = readdlm(graph1, Int)
    edges_g2 = readdlm(graph2, Int)
    new_edges = Set{Pair{Int,Int}}()
    for i=1:size(edges_g1, 1)
        u = edges_g1[i,1]
        v = edges_g1[i,2]
        if u>v
            u,v = v,u
        end
        push!(new_edges, Pair(u,v))
    end
    for i=1:size(edges_g2, 1)
        u = edges_g2[i,1]
        v = edges_g2[i,2]
        if u>v
            u,v = v,u
        end
        delete!(new_edges, Pair(u,v))
    end
    new_edges
end

function modularity(g, c)
    labels = unique(c)
    length(labels) == 1 && return 0.0
    m = ne(g)
    e = Dict{Int,Int}()
    a = Dict{Int,Int}()
    for u in nodes(g)
        for v in neighbors(g,u)
            if u<v
                c1 = c[u]
                c2 = c[v]
                if c1 == c2
                    e[c1] = get(e,c1,0) + 2
                end
                a[c1] = get(a,c1,0) + 1
                a[c2] = get(a,c2,0) + 1
            end
        end
    end
    modularity_value = 0.0
    if m > 0
        for i in labels
            if i > 0
                tmp = haskey(a,i) ? a[i]/2/m : 0
                modularity_value += haskey(e,i) ? e[i]/2/m : 0
                modularity_value -= tmp*tmp
            end
        end
    end
    modularity_value
end


abstract type AbstractIntegrator end

struct Leapfrog{T<:AbstractFloat} <: AbstractIntegrator
    ϵ   ::  T
end

# Create a `Leapfrog` with a new `ϵ`
function (::Leapfrog)(ϵ::AbstractFloat)
    return Leapfrog(ϵ)
end

function is_valid(v::AbstractVector{<:Real})
    if any(isnan, v) || any(isinf, v)
        return false
    end
    return true
end

function lf_momentum(ϵ::T, h::Hamiltonian, θ::AbstractVector{T}, r::AbstractVector{T}) where {T<:Real}
    _∂H∂θ = ∂H∂θ(h, θ)
    !is_valid(_∂H∂θ) && return r, false
    return r - ϵ * _∂H∂θ, true
end

function lf_position(ϵ::T, h::Hamiltonian, θ::AbstractVector{T}, r::AbstractVector{T}) where {T<:Real}
    return θ + ϵ * ∂H∂r(h, r)
end

function step(lf::Leapfrog{F}, h::Hamiltonian, θ::AbstractVector{T}, r::AbstractVector{T}) where {F<:AbstractFloat,T<:Real}
    r_new, _is_valid = lf_momentum(lf.ϵ / 2, h, θ, r)
    !_is_valid && return θ, r, false
    θ_new = lf_position(lf.ϵ, h, θ, r_new)
    r_new, _is_valid = lf_momentum(lf.ϵ / 2, h, θ_new, r_new)
    !_is_valid && return θ, r, false
    return θ_new, r_new, true
end

# TODO: double check the function below to see if it is type stable or not
function steps(lf::Leapfrog{F}, h::Hamiltonian, θ::AbstractVector{T}, r::AbstractVector{T}, n_steps::Int) where {F<:AbstractFloat,T<:Real}
    n_valid = 0
    r_new, _is_valid = lf_momentum(lf.ϵ / 2, h, θ, r)
    !_is_valid && return θ, r, n_valid
    r = r_new
    for i = 1:n_steps
        θ_new = lf_position(lf.ϵ, h, θ, r)
        r_new, _is_valid = lf_momentum(i == n_steps ? lf.ϵ / 2 : lf.ϵ, h, θ, r)
        if !_is_valid
            # The reverse function below is guarantee to be numerical safe.
            # This is because we know the previous step was valid.
            r, _ = lf_momentum(-lf.ϵ / 2, h, θ, r)
            return θ, r, n_valid
        end
        θ, r = θ_new, r_new
        n_valid = n_valid + 1
    end
    return θ, r, n_valid
end


using Plots
using LaTeXStrings
using Printf

function rotate_text(TextHeight::Float64, Vec::Vector{T}) where {T<:Real}
    modVec = hypot(Vec[1], Vec[2])
    Vec[1] >= 0 && (angVec = asind(Vec[2]/modVec))
    Vec[1] < 0 && (angVec = 180 - asind(Vec[2]/modVec))
    rotM = [cosd(angVec) sind(angVec)
           -sind(angVec) cosd(angVec)]
    TextPos = [modVec/2 TextHeight]
    return TextPos*rotM
end

default(grid = true, markersize = 8, framestyle = :none, linewidth = 3, legend = :false,
    aspect_ratio = :equal, gridalpha = 1)

#definindo número de vetores a serem somados
nvec = 8

#criando nvec vetores aleatórios
vecs = rand(-5:1:5, (nvec,2))

#resultante dos nvec vetores
res = [sum(vecs[:,1]), sum(vecs[:,2])]

#calculando ângulo que a resultante faz com a direção x
ang_res = atand(res[2]/res[1])

#string em latex com os valores de cada vetor.
vecstr = String("")
for i = 1:nvec
    global vecstr = string(vecstr, string("\\vec{a}_{", i,"} = (", vecs[i,1], ",", vecs[i,2], ")\\, \\newline"))
end

latexstring(vecstr)

#calculando limites na direção x e na direção y para a criação da imagem
xmax = max(vecs[1,1],0)
ymax = max(vecs[1,2],0)

xmin = min(vecs[1,1],0)
ymin = min(vecs[1,2],0)

#gerando vetor com os pontos de origem de cada um dos vetores que serão somados
origins = [0 0]

#calculando limites e origem dos vetores
for i = 2: nvec
    global origins = vcat(origins, [(origins[i-1,1] + vecs[i-1,1]) (origins[i-1,2] + vecs[i-1,2])])
    global xmax = max(xmax, origins[i,1], vecs[i,1])
    global ymax = max(ymax, origins[i,2], vecs[i,2])
    global xmin = min(xmin, origins[i,1], vecs[i,1])
    global ymin = min(ymin, origins[i,2], vecs[i,2])
end

xmax = max(xmax, res[1])
xmin = min(xmin, res[1])

ymax = max(ymax, res[2])
ymin = min(ymin, res[2])

xlimits = (xmin - 1, xmax + 1)
ylimits = (ymin - 1, ymax + 1)

#gerando texto que será inserido na imagem
res_pos = rotate_text(0.35, res)

resstr = latexstring("{\\rm Resultante:\\;} \\vec{S}  = (", res[1],",", res[2],")")

ann_res = [(res_pos[1,1], res_pos[1,2], text(resstr, halign = :center, valign = :center, pointsize = 24, rotation = ang_res))]

emailstr = "ghralima at cefetmg.br"
ann_autor = [((xmax +0.5),(ymax+ymin)/2, text(emailstr, halign = :center, valign = :center,
    pointsize = 18, rotation = 90, color = :dodgerblue))]

npoints = 30
a2_ani_x = vcat(range(0,origins[2,1], length = npoints), origins[2,1].+zeros(nvec*npoints))
a2_ani_y = vcat(range(0,origins[2,2], length = npoints), origins[2,2].+zeros(nvec*npoints))

origens_ani = [[a2_ani_x, a2_ani_y]]

line_somax = vcat(zeros((nvec-1)*npoints), range(0, res[1], length = npoints), res[1].+zeros(npoints))
line_somay = vcat(zeros((nvec-1)*npoints), range(0, res[2], length = npoints), res[2].+zeros(npoints))

for i = 3: nvec
    ai_ani_x = vcat(zeros((i-2)*npoints), range(0,origins[i,1], length = npoints), origins[i,1].+zeros((nvec+2-i)*npoints))
    ai_ani_y = vcat(zeros((i-2)*npoints), range(0,origins[i,2], length = npoints), origins[i,2].+zeros((nvec+2-i)*npoints))
    push!(origens_ani, [ai_ani_x, ai_ani_y])
end

anim = @animate for i = 1: (nvec+1)*npoints
    quiver([0],[0], quiver = ([vecs[1,1]],[vecs[1,2]]), lw = 2, alpha = 0.5, framestyle = :grid,
        xaxis = ("direção x", xlimits, xlimits[1]:1:xlimits[2]),
        yaxis = ("direção y", ylimits, ylimits[1]:1:ylimits[2]), ann = ann_autor, size = (800,800))
     for j = 2:nvec
        quiver!([origens_ani[j-1][1][i]],[origens_ani[j-1][2][i]], quiver = ([vecs[j,1]],[vecs[j,2]]),
            lw = 2, alpha = 0.5)
     end
     quiver!([0],[0], quiver = ([line_somax[i]],[line_somay[i]]), lw = 4)

     if i > nvec*30
        annotate!(ann_res)
     end
end

gif(anim, "soma_vetores_ani.gif", fps = 25)


using Push
using Base.Test

# Load the base configuration used by all tests.
cfg_path = joinpath(dirname(@__FILE__), "configuration/boolean.cfg")
cfg = Push.load_configuration(cfg_path)

# EQUALS.
s = Push.run("(FALSE FALSE BOOLEAN.=)", cfg)
@test s.boolean == [true]
s = Push.run("(FALSE TRUE BOOLEAN.=)", cfg)
@test s.boolean == [false]

# NOT.
s = Push.run("(FALSE BOOLEAN.NOT)", cfg)
@test s.boolean == [true]
s = Push.run("(TRUE BOOLEAN.NOT)", cfg)
@test s.boolean == [false]

# AND.
s = Push.run("(FALSE FALSE BOOLEAN.AND)", cfg)
@test s.boolean == [false]
s = Push.run("(FALSE TRUE BOOLEAN.AND)", cfg)
@test s.boolean == [false]
s = Push.run("(TRUE FALSE BOOLEAN.AND)", cfg)
@test s.boolean == [false]
s = Push.run("(TRUE TRUE BOOLEAN.AND)", cfg)
@test s.boolean == [true]

# OR.
s = Push.run("(FALSE FALSE BOOLEAN.OR)", cfg)
@test s.boolean == [false]
s = Push.run("(FALSE TRUE BOOLEAN.OR)", cfg)
@test s.boolean == [true]
s = Push.run("(TRUE FALSE BOOLEAN.OR)", cfg)
@test s.boolean == [true]
s = Push.run("(TRUE TRUE BOOLEAN.OR)", cfg)
@test s.boolean == [true]

# DUP.
s = Push.run("(FALSE BOOLEAN.DUP)", cfg)
@test length(s.boolean) == 2

# FLUSH.
s = Push.run("(FALSE TRUE FALSE FALSE BOOLEAN.FLUSH)", cfg)
@test isempty(s.boolean)

# FROM FLOAT.
s = Push.run("(0.0 BOOLEAN.FROMFLOAT)", cfg)
@test isempty(s.float) && s.boolean == [false]
s = Push.run("(1.871 BOOLEAN.FROMFLOAT)", cfg)
@test isempty(s.float) && s.boolean == [true]

# FROM INT.
s = Push.run("(0 BOOLEAN.FROMINT)", cfg)
@test isempty(s.integer) && s.boolean == [false]
s = Push.run("(32 BOOLEAN.FROMINT)", cfg)
@test isempty(s.integer) && s.boolean == [true]

# POP.
s = Push.run("(TRUE FALSE BOOLEAN.POP)", cfg)
@test s.boolean == [true]

# ROT.
s = Push.run("(FALSE TRUE TRUE BOOLEAN.ROT)", cfg)
@test s.boolean == [true, false, true]

# SWAP.
s = Push.run("(FALSE TRUE BOOLEAN.SWAP)", cfg)
@test s.boolean == [true, false]

# STACKDEPTH.
s = Push.run("(TRUE FALSE FALSE FALSE BOOLEAN.STACKDEPTH)", cfg)
@test length(s.boolean) == 4 && s.integer == [4]

# SHOVE.
s = Push.run("(3 FALSE FALSE FALSE TRUE BOOLEAN.SHOVE)", cfg)
@test isempty(s.integer) && s.boolean == [true, false, false, false]

# YANK.
s = Push.run("(3 TRUE FALSE FALSE FALSE BOOLEAN.YANK)", cfg)
@test isempty(s.integer) && s.boolean == [false, false, false, true]

# YANKDUP.
s = Push.run("(3 TRUE FALSE FALSE FALSE BOOLEAN.YANKDUP)", cfg)
@test isempty(s.integer) && s.boolean == [true, false, false, false, true]

# RAND.
s = Push.run("(BOOLEAN.RAND)", cfg)
@test length(s.boolean) == 1

# DEFINE.
s = Push.run("(X BOOLEAN.DEFINE X)", cfg)
@test isempty(s.boolean) && s.name == [:X, :X]
s = Push.run("(X TRUE BOOLEAN.DEFINE X)", cfg)
@test s.boolean == [true]


include("../../../opti/src/opti_module.jl")

module ntest
using ..opti

using ..noisydiff
using PyPlot
PyPlot.pygui(true)

function test()

    t = collect( 0:0.01:7 )
    y_theo = sin.(t) # abs.(t .- 3.5)# abs.(t .- 3.5)
    noise = 0.01 * sin.(100 .* t)
    y_noisy = y_theo .+ noise # 0.01 .* randn(size(y_theo))
    y_est = noisydiff.pt1_smooth(t, y_noisy; T = 0.1)

    # @time (x_sol, y_sol, iter, stopping_crit) = noisydiff.tvd(t, y_noisy, lamda = 1 * 0.2,
    #         tol_options = opti.ToleranceOptions(max_iter = 500, step_size_tol = 1e-20, f_abs_tol = 1e-15))
    # dy_tvd_fw = stat.x_sol
    # dy_tvd_fw = x_sol
    # show(stat)
    # y_tvd = y_noisy[1] .+ noisydiff.cumtrapezoidal(t, dy_tvd_fw)

    t_mids = 0.5 * (t[1:end-1] + t[2:end])
    dy_theo = cos.(t_mids)
    dy_straightfw = noisydiff.diff_fw(t, y_noisy)
    dy_pt1_fw = noisydiff.diff_fw(t, y_est)

    w1 = 50
    t_lsq_d1 = t[(w1+1):end]
    dy_lsq_d1 = noisydiff.diff_sliding_least_squares_d1(t, y_noisy, w = w1)

    w3 = w1# 3*w1
    t_lsq_d3 = t[(w3+1):end]
    dy_lsq_d3 = noisydiff.diff_sliding_least_squares_d3(t, y_noisy, w = w3)

    # dy_tvd_fw = noisydiff.diff_fw(t, y_tvd)

    # PyPlot.figure()
    # PyPlot.grid()
    # PyPlot.plot(t, y_noisy)
    # # PyPlot.plot(t, y_theo, linewidth = 2)
    # PyPlot.plot(t, y_est, linewidth = 2)
    # PyPlot.plot(t, y_tvd, linewidth = 2)

    # PyPlot.figure()

    # PyPlot.subplot(2,1,1)
    # PyPlot.grid()
    # PyPlot.plot(t, y_noisy)

    # PyPlot.subplot(2,1,2)
    # PyPlot.grid()
    # # PyPlot.plot(t_mids, dy_theo, linewidth = 2)
    # # PyPlot.plot(t_mids, dy_straightfw)
    # PyPlot.plot(t_mids, dy_pt1_fw, linewidth = 2)
    # PyPlot.plot(t, dy_tvd_fw, linewidth = 4)

    PyPlot.figure()
    PyPlot.grid()
    PyPlot.plot(t_mids, dy_pt1_fw, linewidth = 2)
    PyPlot.plot(t_lsq_d1, dy_lsq_d1, linewidth = 2)
    PyPlot.plot(t_lsq_d3, dy_lsq_d3, linewidth = 2)

end

end


using MarkovTransitionMatrices
using Test
using BenchmarkTools
using SharedArrays
using Distributed
using LinearAlgebra

using Distributions

isapprox_oneval(x::AbstractArray) = all( x .≈ x[1] )

include("FarmerTodaBmark.jl")

# setup for test
testvec = -3.0:0.5:0.0
ss = collect(Base.product(testvec, testvec))
n = length(testvec)
nn = length(ss)
mvd = MvNormal(zeros(2), Matrix{Float64}(I,2,2))

# test the 1-d version
fullp1 = Matrix(markov_transition((s) -> 0., (s) -> 1., 1e-8, testvec))
ratio1 = fullp1 ./ pdf.(Normal(0,1), testvec')
@test isapprox_oneval(ratio1)

# test the 2-d version
fullp2 = Matrix(markov_transition((s) -> zeros(2), (s) -> Matrix{Float64}(I,2,2), 1e-8, testvec, testvec))
ratio2 = reshape(fullp2[1,:], n, n) ./  [pdf(mvd, [s...]) for s in ss]
@test isapprox_oneval(ratio2)

# test the markov-switching version
μswitch(r::Real,s::NTuple{N,T}) where {N,T<:Real} = r==1 ? zeros(2) : [s...]
Σswitch(r::Real,s::NTuple{N,T}) where {N,T<:Real} = r==1 ? Matrix{Float64}(I,2,2)   : eps(1.0)*Matrix(I,2,2)
πswitch = Matrix{Float64}(I, 2, 2)
fullp3 = Matrix(markovswitching_transition(μswitch, Σswitch, πswitch, 0.0, testvec, testvec))

@test all(fullp3[1:nn, 1:nn] .== fullp2)
@test all(fullp3[nn+1:2*nn, 1:nn] .== 0.0)
@test all(fullp3[1:nn, nn+1:2*nn] .== 0.0)
@test all(fullp3[nn+1:2*nn, nn+1:2*nn] .== Matrix(I,nn,nn))

# ---------------------------------------------------------------

println("Testing moment-matching")

s = -3.0:0.25:3.0
ss = -3.0:0.125:3.0, -3.0:0.125:3.0

pids = addprocs()
@everywhere using MarkovTransitionMatrices

# simulate random walks w/ moment matching
p_P_match_1, p_mom, p_er = markov_transition_moment_matching_parallel((s) -> 0.0   , (s) -> 1.0   , 1e-8, s)
s_P_match_1, s_mom, s_er = markov_transition_moment_matching_serial(  (s) -> 0.0   , (s) -> 1.0   , 1e-8, s)
@test true == all(s_P_match_1 .== p_P_match_1) == all(s_mom .== p_mom) == all(s_er[isfinite.(s_er)] .== p_er[isfinite.(p_er)])
println("Single-var serial/parallel OK")

# p_P_match_2, p_mom, p_er = markov_transition_moment_matching_parallel((s) -> [s...], (s) -> [1.0 0.0; 0.0 1.0], 1e-8, ss...)
p_P_match_2, p_mom, p_er = markov_transition_moment_matching_parallel((s) -> [s...], (s) -> Matrix{Float64}(LinearAlgebra.I,2,2), 1e-8, ss...)
s_P_match_2, s_mom, s_er = markov_transition_moment_matching_serial(  (s) -> [s...], (s) -> Matrix{Float64}(LinearAlgebra.I,2,2), 1e-8, ss...)
@test true == all(s_P_match_2 .≈ p_P_match_2) == all(s_mom .≈ p_mom) == all(s_er[isfinite.(s_er)] .== p_er[isfinite.(p_er)])
println("Multi-var serial/parallel OK")

# share of states that get moment-matching
share = sum(s_mom .> 0) / length(s_mom) * 100
@show "$share percent of states get moment-matching"
@test share .> 0.5

# without matching
P_nomatch_2 = markov_transition((s) -> [s...], (s) -> Matrix{Float64}(I,2,2), 1e-8, ss...)
number_diff = sum(s_P_match_2 .!= P_nomatch_2)
sd_diff = sum((s_P_match_2 .- P_nomatch_2).^2)/prod(size(s_P_match_2))
abs_diff = sum(abs.(s_P_match_2 .- P_nomatch_2))/prod(size(s_P_match_2))
sharediff = number_diff / prod(size(s_P_match_2))
nm = norm(vec(s_P_match_2 .- P_nomatch_2), Inf)
@show "max difference between matching & no-matching is $nm"
@show "$sharediff percent of probabilities change"

# benchmarking
println("Benchmarking parallel...")
@show @benchmark markov_transition_moment_matching_parallel(  (s) -> [s...], (s) -> Matrix{Float64}(LinearAlgebra.I,2,2), 1e-8, ss...)
println("Benchmarking serial...")
@show @benchmark markov_transition_moment_matching_serial(    (s) -> [s...], (s) -> Matrix{Float64}(I,2,2), 1e-8, ss...)

rmprocs(pids)

# using Plots
# gr()
#
# heatmap(s, s, reshape(P_match_2[10,:]  , length.(ss)...))
# heatmap(s, s, reshape(P_nomatch_2[10,:], length.(ss)...))
# heatmap(s, s, reshape(P_match_2[20,:] .- P_nomatch_2[20,:], length.(ss)...))
# heatmap(s, s, reshape(mom.>0, length.(ss)...))
# histogram(vec(clamp.(P_match_2 .- P_nomatch_2, -0.05, 0.05)) )


@testset "Testing Point Constructors" begin
    @test Point(1,2) === Point(Vec(1.0,2.0))
    @test Point((1,2)) === Point(Vec(1.0,2.0))
    @test Point([1,2]) === Point(Vec(1.0,2.0))
end

@testset  "Adding a point to vector gives a point" begin
   a  =  Point(3, -2, 5)
   v =   Vec(-2, 3, 1)
   @test a + v == Point(1, 1, 6)
end
@testset  "Addition of point to vector is orderless" begin
    a  =  Point(3, -2, 5)
    v =   Vec(-2, 3, 1)
    @test a + v == v + a
 end

@testset "Subtracting two points gives a vector" begin
     p1  = Point(3, 2, 1)
     p2 =  Point(5, 6, 7)
    @test p1 - p2 === Vec(-2.0, -4.0, -6.0)
end

@testset "Subtracting a vector from a point gives a point" begin
     p = Point(3, 2, 1)
     v = Vec(5, 6, 7)
    @test p - v === Point(-2, -4, -6)
end

@testset "Subtracting a point from a vector gives a point" begin
    p = Point(3, 2, 1)
    v = Vec(5, 6, 7)
   @test v - p === Point(-2, -4, -6)
end


using BenchmarkTools, Random
using Mill, Flux

using Plots, LaTeXStrings; pgfplotsx()

Random.seed!(42)
params = Base.Iterators.product([100, 200, 300], [100, 200, 300]) |> collect |> vec
for (m, n) in params
    @show m, n
    # A = rand(Float32, m, n)
    A = PreImputingMatrix(rand(Float32, m, n))
    # B = rand(Float32, n, m)
    B = [rand() < 0.5 ? missing : rand() for _ in 1:n, _ in 1:m]
    @btime $A * $B
end

xticklabels = [latexstring("$a|$b") for (a,b) in params]

t1 = [22.859, 50.551, 95.680, 38.428, 82.887, 200.255, 54.975, 133.416, 232.263] ./ 1000
l1 = "matrix(m,n) * Full Matrix(n,m)"
t2 = [3.868, 17.397, 39.650, 8.290, 35.292, 78.747, 13.223, 53.783, 121.800]
l2 = "Matrix(m,n) * Half Full Matrix(n,m)"
t3 = [22.114, 59.371, 96.551, 29.781, 67.402, 208.026, 36.066, 127.561, 227.967] ./ 1000
l3 = "Preimputingmatrix(m,n) * Full Matrix(n,m)"
t4 = [55.889, 119.639, 196.984, 94.915, 204.092, 335.075, 136.566, 298.371, 487.540] ./ 1000
l4 = "PreImputingmatrix(m,n) * Half Full Matrix(n,m)"
groupedbar(xticklabels, hcat(t1, t2, t3, t4); ylabel=L"ms", xlabel="n | m", labels=[l1 l2 l3 l4], title="PreImputing (log y)", yscale=:log10)

savefig("pre_imputing.png")

Random.seed!(42)
params = Base.Iterators.product([100, 200, 300], [100, 200, 300]) |> collect |> vec
for (m, n) in params
    @show m, n
    # A = rand(Float32, m, n)
    A = PreImputingMatrix(rand(Float32, m, n))
    # B = rand(Float32, n, m)
    B = [rand() < 0.5 ? missing : rand() for _ in 1:n, _ in 1:m]
    @btime gradient((A, B) -> sum(A*B), $A, $B)
end

xticklabels = [latexstring("$a|$b") for (a,b) in params]

t1 = [901.895 / 1000, 3.553, 7.889, 1.883, 7.232, 16.162, 2.648, 10.688, 24.344]
l1 = "Matrix(m,n) * Full Matrix(n,m)"
t2 = [909.361 / 1000, 3.650, 7.865, 1.792, 7.260, 16.200, 2.782, 10.598, 24.071]
l2 = "PreImputingMatrix(m,n) * Full Matrix(n,m)"
t3 = [1.077, 4.042, 9.111, 2.135, 7.637, 17.257, 3.075, 12.057, 24.922]
l3 = "PreImputingMatrix(m,n) * Half Full Matrix(n,m)"
groupedbar(xticklabels, hcat(t1, t2, t3); ylabel=L"ms", xlabel="n | m", labels=[l1 l2 l3], title="PreImputing (grad)")

savefig("pre_imputing_grad.png")


# utilities.jl

using LinearAlgebra
using CSV
using DataFrames

"""
    setup(path::String)

Setup a first-passage time problem by specifying a path to a .csv file containing the
transition.

The .csv file should have the column names "condition", "from", and "to". Currently, the
transition rates are set to be equal to the number of states in each condition. This keeps
the processing times comparable between conditions, and when fitting the scaling parameter
τ, that τ represents the transition rate per jump.
"""
function setup(path::String)
    df = CSV.read(path, DataFrame, stringtype=String)
    @assert names(df) == ["from", "to"] "Incorrect column names in .csv file."

    statenames = []
    ctr = 1
    nametoidx = Dict()
    rows = []
    columns = []
    # Getting a mapping from state name to its index in the W matrix
    for row = 1:nrow(df)
        if df[row, :from] ∉ statenames
            push!(statenames, df[row, :from])
            nametoidx[df[row, :from]] = ctr
            ctr += 1
        end

        if df[row, :to] ∉ statenames
            push!(statenames, df[row, :to])
            nametoidx[df[row, :to]] = ctr
            ctr += 1
        end
        push!(rows, nametoidx[df[row, :to]])
        push!(columns, nametoidx[df[row, :from]])
    end

    # Setting transition rates
    nstates = length(statenames)
    W = zeros(nstates, nstates)
    for (col, row) in zip(columns, rows)
        W[row, col] = nstates
    end
    W = setdiagonal!(W)

    @assert 0 in W[diagind(W)] "No absorbing states provided."

    # Separating transient and absorbing submatrices
    (transient, absorbing) = getdims(W)
    T = W[transient, transient]
    A = W[absorbing, transient]

    return T, A
end

"""
    setdiagonal!(W::Matrix)

Set the diagonal correctly for W-matrices (van Kampen, 2007, Ch. V.2).
"""
function setdiagonal!(W::Matrix)
    W[diagind(W)] .= 0.0
    W[diagind(W)] = -sum(W, dims=1)
    return W
end

"""
    rescale!(W::Matrix, τ)

Rescale a matrix a matrix by a constant τ.
"""
function rescale!(W::Matrix{Float64}, τ)
    W *= τ
end

"""
    getdims(W::Matrix)

Takes a W-matrix and returns the indices of the transient and absorbing states.
"""
function getdims(W::Matrix)
    transient = []
    absorbing =[]
    for i = 1:size(W, 1)
        if W[i, i] == 0
            push!(absorbing, i)
        else
            push!(transient, i)
        end
    end
    return transient, absorbing
end


# # Production-Destruction
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/models/ProductionDestruction.ipynb)
#
#md # !!! note "Overview"
#md #     System type: rational function\
#md #     State dimension: 3\
#md #     Application domain: Chemical kinetics
#
# ## Model description
#
# A production-destruction system consists of an ordinary differential equation and two constraints:
# positivity and conservativity. It means that the system states are quantities which are always positive
# and the sum of these quantities is constant. This particular family of systems is often used to test
# the stability of integration schemes.
#
# As proposed by **Michaelis-Menten** theory [^KM18], a model for three quantities can be defined as follows:
#
# ```math
#    \left\{ \begin{array}{lcl}
#    \dot{x} &=& \dfrac{-xy}{1+x}\\ \\
#    \dot{y} &=& \dfrac{xy}{1+x} - a y\\ \\
#    \dot{z} &=& ay
#    \end{array} \right.
# ```
# with nominal values ``a = 0.3`` and an initial condition ```x(0) = 9.98``, ``y(0) = 0.01`` and
# ``z(0) = 0.01``. In this model, ``x`` is the nutrients, ``y`` the phytoplankton
# and ``z`` the detritus. In this notebook we analyze the behavior of the system
# subject to uncertainties in the initial condition and/or the model's parameter `a`
# using forward reachability analysis.
#
# The solution should satisfy the following two constraints:
#
# - **Positivity:** ``x(t)``, ``y(t)``, ``z(t)`` are positive, and
# - **Conservativity:** ``x(t) + y(t) + z(t) = 10`` for all ``t \ge 0``.

# In this notebook we are interested in computing the reachable tube until the time horizon ``T=100``,
# and to verify (or embed) the  constraints. Three setups are considered, depending on the source
# of bounded uncertainties which is taken into account:
#
# - **Case I**: ``x(0) \in [9.5, 10.0]``, i.e., uncertainty on the initial condition;
# - **Case P**: ``a \in [0.296, 0.304]``, i.e., uncertainty on the parameter;
# - **Case I & P**: ``x(0) \in [9.5, 10.0]`` and ``a \in [0.296, 0.304]``, i.e., both uncertainties are mixed.

# In terms of objectives, at ``T = 100``, the constraints ``10 \in x+y+z`` and
# ``x, y, z \geq 0`` have to be verified.

# The reachability settings considered above are a variation of those in [^ARCHCOMP20_NLN].
# Please note that variables $x$ and $y$ converge towards zero, so the final volume can be used
# as a quality measure of the overapproximation.

# In the next section we explain how to check whether the constraints are satisfied
# using `LazySets.jl`.

# ## Checking satisfiability with sets

# Two measures of quality of the approximation are considered: the volume of the box
# (``x \times y \times z``) enclosing the final state (at ``T = 100``) and the total
# time of computation for evolution and verification.
# All of these results are obtained for the three setups.

# Let us discuss the implementation of each constraint satisfaction problem.
# We begin by defining the positive orthant in three-dimensional space,
# `positive_orthant`, as an unbounded polyhedron in constraint representation.

using ReachabilityAnalysis, Symbolics, Plots

@variables x y z
const positive_orthant = HPolyhedron([x >= 0, y >= 0, z >= 0], [x, y, z])

# Given a set $X \subseteq \mathbb{R}^n$, to check whether the positivity constraint
# holds correspnods to checking wheter ``X`` is included in the positive orthant.
# This computation can be done efficiently using support functions, and it is available
# in `LazySets.jl`. Multiple dispatch takes care based on the types of the arguments
# in the call `X ⊆ positive_orthant` depending on the type of ``X``.

# On the other hand, given a set ``X``, a quick way to check the conservativity constraint
# is to first overapproximate ``X`` with a box, then represent such box as a
# product-of-intervals (`IntervalBox`) `B`, and finally take the Minkowski
# sum of each line segment, using `sum(B)`. This check is only sufficient; splitting
# ``X`` can be used to refine the check.

# The verification problem is summarized in the function `prod_dest_verif`.
# It receives the solution of a reachability problem represented with
# a Taylor model flowpipe and two optional arguments that specify the time horizon
# `T` and the conservativity condition given a target state `target`,
# which defaults to `10`. The function returns the tuple `(flag, vol)` where `flag`
# is true if and only if the positivity and conservativity conditions are satisfied at the time horizon
# `T`, which defaults to `100`. The second output, `vol`, corresponds to the volume
# of the box overapproximation of the final reach-set.

function prod_dest_verif(sol; T=100.0, target=10.0)

    ## convert to a zonotopic flowpipe
    solz = overapproximate(sol, Zonotope)

    ## project the final reach-set onto the space variables x, y, z
    X = project(solz(T), vars=(1, 2, 3))

    ## check that all variables are nonnegative
    nonnegative = X ⊆ positive_orthant

    ## compute the volume of the last reach-set
    H = overapproximate(X, Hyperrectangle)
    vol = volume(H)

    ## check that that target belongs to the minkowski sum of the reach-sets projected in each coordinate
    B = convert(IntervalBox, H) # get the product-of-intervals representation
    contains_target = target ∈ sum(B)

    return nonnegative && contains_target, vol
end

# ## Case I

# Case I corresponds to initial states uncertainty. We begin by writing the system
# of ODEs in the function `prod_dest_I!`.

@taylorize function prod_dest_I!(du, u, params, t)
    local a = 0.3
    x, y, z = u[1], u[2], u[3]

    du[1] = - (x * y) / (1 + x)
    du[2] = (x * y) / (1 + x) - a * y
    du[3] = a * y
    return du
end

# We define the initial states as the cartesian product of intervals.

X0 = (9.5 .. 10.0) × (0.01 .. 0.01) × (0.01 .. 0.01)
prob = @ivp(x'= prod_dest_I!(x), dim:3, x(0) ∈ X0)

solI = solve(prob, T=100.0, alg=TMJets(abstol=1e-11, orderT=7, orderQ=1));

# Verifying that the specification holds:
property, vol = prod_dest_verif(solI)

# Now we plot ``z`` (in the ``[0, 11]`` range) w.r.t. time (in the ``[0, 100]`` range).

plot(solI, vars=(0, 3), linecolor=:orange, color=:orange, alpha=0.3, lab="I")

# ## Case P

# For the case of an uncertain parameter, we add a new state variable that corresponds
# to `a`, with constant (zero) dynamics.

@taylorize function prod_dest_IP!(du, u, params, t)
    x, y, z, a = u[1], u[2], u[3], u[4]

    du[1] = - (x * y) / (1 + x)
    du[2] = (x * y) / (1 + x) - a * y
    du[3] = a * y
    du[4] = zero(x)
    return du
end

X0 = (9.98 .. 9.98) × (0.01 .. 0.01) × (0.01 .. 0.01) × (0.296 .. 0.304)
prob = @ivp(x'= prod_dest_IP!(x), dim:4, x(0) ∈ X0)

solP = solve(prob, T=100.0, alg=TMJets(abstol=1e-12, orderT=7, orderQ=1));

# Verifying that the specification holds:
property, vol = prod_dest_verif(solP)

# Now we plot ``z`` (in the ``[0, 11]`` range) w.r.t. time (in the ``[0, 100]`` range).

plot(solP, vars=(0, 3), linecolor=:blue, color=:blue, alpha=0.3, lab="P")

# ## Case I & P

# When uncertainty in both the intial states and the paramters are present, we can
# reuse the function `prod_dest_IP!`, but setting an uncertain initial condition
# and an uncertain parameter.
# Recall that we are interested in ``x(0) \in [9.5, 10.0]`` and ``a \in [0.296, 0.304]``.

X0 = (9.5 .. 10.0) × (0.01 .. 0.01) × (0.01 .. 0.01) × (0.296 .. 0.304)
prob = @ivp(x'= prod_dest_IP!(x), dim:4, x(0) ∈ X0)

solIP = solve(prob, T=100.0, alg=TMJets(abstol=1e-11, orderT=7, orderQ=1));

# Verifying that the specification holds:
property, vol = prod_dest_verif(solIP)

# Now we plot ``z`` (in the ``[0, 11]`` range) w.r.t. time (in the ``[0, 100]`` range).

plot(solIP, vars=(0, 3), linecolor=:red, color=:red, alpha=0.3, lab="I & P")

# ## Performance optimization

# We have observed that the specifications are satisfied in all cases. The question remains
# if we can tune the problem implementation such that we can verify these properties in less time.
# In this section we put into practice the ideas explained in [Some common gotchas](@ref).
# First, consider the function `prod_dest_I!` defined above. The term `du[1] = -x*y / (1+x)`
# has the numerator and denominator in common with respect to `du[2]`, so we can store
# those values in an auxiliary variable. Similar comments apply to the function
# with varying parameter `a`. The final functions are presented below.

@taylorize function prod_dest_I_optimized!(du, u, params, t)
    local a = 0.3
    x, y, z = u[1], u[2], u[3]

    num = x * y
    den = 1 + x
    aux = num/den
    aux2 = a * y
    du[1] = -aux
    du[2] = aux - aux2
    du[3] = aux2
    return du
end

@taylorize function prod_dest_IP_optimized!(du, u, params, t)
    x, y, z, a = u[1], u[2], u[3], u[4]

    num = x * y
    den = 1 + x
    aux = num/den
    aux2 = a * y
    du[1] = -aux
    du[2] = aux - aux2
    du[3] = aux2
    du[4] = zero(x)
    return du
end

# With respect to parameter tuning, we have observed that for all three cases
# a good parameter choice is ``n_T = 7``, ``n_Q = 1``, and an adaptive absolute
# tolerance with initial value ``10^{-11}`` (resp ``10^{-12}``) for ``I``
# and ``I`` and ``P`` (resp. ``P``). This choice offers a good balance between runtime
# and precision.

# ## References

# [^ARCHCOMP20_NLN]: *ARCH-COMP20 Category Report: Continuous and Hybrid Systems with Nonlinear Dynamics.* Luca Geretti, Julien Alexandre dit Sandretto, Matthias Althoff, Luis Benet, Alexandre Chapoutot, Xin Chen, Pieter Collins, Marcelo Forets, Daniel Freire, Fabian Immler, Niklas Kochdumper, David P. Sanders and Christian Schilling (2020) ARCH20. To appear in 7th International Workshop on Applied Verification of Continuous and Hybrid Systems. 7th International Workshop on Applied Verification of Continuous and Hybrid Systems (ARCH20), vol 74, pages 49--75. 10.29007/zkf6.

# [^KM18]: Kopecz, Stefan, and Andreas Meister. *On order conditions for modified Patankar–Runge–Kutta schemes.* Applied Numerical Mathematics 123 (2018): 159-179.


# todo : document and add some other nasty functions
using SparseArrays

"""
    Mesh1d

A type describing a mesh for a one dimensional domain.
"""
struct Mesh1d
    nodes::Array{Float64, 1}
    elements::Array{Int64, 2}
end

"""
    assemble_1d_FE_matrix(elem::Array{Float64, 2}, nbNodes::Int64;
      intNodes = 0, dof1 = 1, dof2 = 1)

Assemble a finite elements matrix corresponding to a 1 dimensional uniform mesh.

# Arguments
  * `elem` : elementary finite elements matrix
  * `nbNodes` : number of nodes in the mesh
  * `intNodes1` : number of interior nodes in the interior of each element for lhs
  * `intNodes2` : number of interior nodes in the interior of each element for rhs
  * `dof1`     : number of degrees of freedom for each node for lhs
  * `dof2`     : number of degrees of freedom for each node for rhs
"""
function assemble_1d_FE_matrix(elem::Array{Float64, 2}, nbNodes::Int64;
                               intNodes1 = 0, intNodes2 = 0, dof1 = 1, dof2 = 1)

    nT = Threads.nthreads()

    nbNodesTotal1 = (nbNodes + (nbNodes - 1) * intNodes1)
    nbNodesTotal2 = (nbNodes + (nbNodes - 1) * intNodes2)

    MM = Array{SparseMatrixCSC{Float64,Int64}}(undef, nT)
    for i = 1:nT
        MM[i] = spzeros(Float64, nbNodesTotal1 * dof1, nbNodesTotal2 * dof2)
    end

    l1 = zeros(UInt64, nT)
    r1 = zeros(UInt64, nT)
    l2 = zeros(UInt64, nT)
    r2 = zeros(UInt64, nT)

    Threads.@threads for i = 1:(nbNodes - 1)
        k = Threads.threadid()

        l1[k] = (i - 1 + (i - 1) * intNodes1) * dof1 + 1
        r1[k] = (i + 1 + i * intNodes1) * dof1
        l2[k] = (i - 1 + (i - 1) * intNodes2) * dof2 + 1
        r2[k] = (i + 1 + i * intNodes2) * dof2
        MM[k][l1[k]:r1[k], l2[k]:r2[k]] += elem
    end
    M = MM[1]
    for i=2:nT
        M += MM[i]
    end
    M
end

"""
    assemble_1d_nu_FE_matrix(elem::Matrix{SymPy.Sym}, nodes::Array{Float64, 1};
      intNodes = 0, dof1 = 1, dof2 = 1)

Assemble a finite elements matrix corresponding to a 1 dimensional non-uniform mesh.

# Arguments
  * `elem` : elementary finite elements matrix (with elements of type SymPy.Sym)
  * `nodes` : vector of nodes composing the non uniform-mesh
  * `intNodes1` : number of interior nodes in the interior of each element for lhs
  * `intNodes2` : number of interior nodes in the interior of each element for rhs
  * `dof1`     : number of degrees of freedom for each node for lhs
  * `dof2`     : number of degrees of freedom for each node for rhs
"""
function assemble_1d_nu_FE_matrix(elem::Matrix{SymPy.Sym}, nodes::Array{Float64, 1};
                                  intNodes1 = 0, intNodes2 = 0, dof1 = 1, dof2 = 1)

    nT = Threads.nthreads()
    global h
    nbNodes = length(nodes)
    nbNodesTotal1 = (nbNodes + (nbNodes - 1) * intNodes1)
    nbNodesTotal2 = (nbNodes + (nbNodes - 1) * intNodes2)

    MM = Array{SparseMatrixCSC{Float64,Int64}}(undef, nT)
    for i = 1:nT
        MM[i] = spzeros(Float64, nbNodesTotal1 * dof1, nbNodesTotal2 * dof2)
    end

    l1 = zeros(UInt64, nT)
    r1 = zeros(UInt64, nT)
    l2 = zeros(UInt64, nT)
    r2 = zeros(UInt64, nT)
    Threads.@threads for i = 1:nbNodes - 1
        k = Threads.threadid()
        l1[k] = (i - 1 + (i - 1) * intNodes1) * dof1 + 1
        r1[k] = (i + 1 + i * intNodes1) * dof1
        l2[k] = (i - 1 + (i - 1) * intNodes2) * dof2 + 1
        r2[k] = (i + 1 + i * intNodes2) * dof2
        elem_loc = elem.subs(h, nodes[i + 1] - nodes[i])
        MM[k][l1[k]:r1[k], l2[k]:r2[k]] += elem_loc
    end
    M = MM[1]
    for i=2:nT
        M += MM[i]
    end
    M
end


d0(Lmin) = 1.24 * (Lmin - 15)^(1/3) - 1.8

function similarityscore(x::AbstractMatrix, y::AbstractMatrix; metric=SqEuclidean())
    @assert size(x,1) == size(y,1)
    Lmin = min(size(x,2), size(y,2))
    d₀ = d0(Lmin)
    sqdist = pairwise(metric, x, y)
    return 1 ./ (1 .+ sqdist ./ (d₀^2))
end

function needlemanwunsch_align(x,y; gapopen=-0.6, gapextend=0.0)
    similarity = similarityscore(x,y)
    nw = DistanceNeedlemanWunsch(similarity, gapopen,gapextend)
    fill!(nw)
    return traceback(nw)[2]
end

function tm_align(x,y; kwargs...)
    return iterate_kabsch(x, y; correspondence=needlemanwunsch_align, kwargs...)
end


using JuMP
using GLPK
using StatsBase

# Initialize problem constants
const maturity   = [3, 4]
const yield      = [0.04, 0.03]
const risk_level = [2, 1]
const available_capital = 1e5
const max_avg_maturity = 3.6
const max_avg_risk_level = 1.5
N_bonds = length(yield)

# Initialize model
model = Model(GLPK.Optimizer)

@variable(model, x[1:N_bonds] >= 0)

@constraint(model, mean(x .* maturity) <= max_avg_maturity)

@constraint(model, mean(x .* risk_level) <= max_avg_risk_level)

@constraint(model, sum(x) <= available_capital)

@objective(model, Max, sum(x .* yield))

#--- Solve
optimize!(model)
@show objective_value(model)
@show termination_status(model)
@show solution_summary(model)
value.(x)


using Soss

simple_poisson = @model traps begin
    α ~ Normal(log(4), .1)
    β ~ Normal(-0.25, .1)

    complaints ~ For(eachindex(traps)) do i
        λ = exp(α + β * traps[i])
        Poisson(λ)
    end
end

alt_poisson = @model (N, mean_traps) begin
    traps ~ For(1:N) do i
        traps[i] ~ Poisson(mean_traps)
    end

    complaints ~ For(eachindex(traps)) do i
        λ = exp(α + β * traps[i])
        Poisson(λ)
    end
end


# TODO: Rotations.jl

# General axis
function rotation_matrix(axis::NTuple{3, <: Real}, α::Real; out=Matrix{Float64}(undef, 3, 3))
	# Arbitrary axis
	# Apply vector from the right
	(sinα, cosα) = sincos(α)
	one_minus_cosα = (1 - cosα)
	# First column
	out[1, 1] = cosα + axis[1]^2 * one_minus_cosα
	out[2, 1] = axis[1] * axis[2] * one_minus_cosα + axis[3] * sinα
	out[3, 1] = axis[1] * axis[3] * one_minus_cosα - axis[2] * sinα
	# Second column
	out[1, 2] = axis[1] * axis[2] * one_minus_cosα - axis[3] * sinα
	out[2, 2] = cosα + axis[2]^2 * one_minus_cosα
	out[3, 2] = axis[2] * axis[3] * one_minus_cosα + axis[1] * sinα
	# Third column
	out[1, 3] = axis[1] * axis[3] * one_minus_cosα + axis[2] * sinα
	out[2, 3] = axis[2] * axis[3] * one_minus_cosα - axis[1] * sinα
	out[3, 3] = cosα + axis[3]^2 * one_minus_cosα
	return out
end

# Around axis in plane alinged with coordinate system
function rotation_matrix(plane::Val{:xy}, axis::NTuple{2, <: Real}, α::Real; out=Matrix{Float64}(undef, 3, 3))
	# axis in xy plane, axis must be normalised!
	# Apply vector from the right
	(sinα, cosα) = sincos(α)
	one_minus_cosα = (1 - cosα)
	# First column
	out[1, 1] = cosα + axis[1]^2 * one_minus_cosα
	out[2, 1] = axis[1] * axis[2] * one_minus_cosα
	out[3, 1] = -axis[2] * sinα
	# Second column
	out[1, 2] = axis[1] * axis[2] * one_minus_cosα
	out[2, 2] = cosα + axis[2]^2 * one_minus_cosα
	out[3, 2] = axis[1] * sinα
	# Third column
	out[1, 3] = axis[2] * sinα
	out[2, 3] = -axis[1] * sinα
	out[3, 3] = cosα
	return out
end

# Around axis in plane alinged with coordinate system
function rotation_matrix(plane::Val{:xz}, axis::NTuple{2, <: Real}, α::Real; out=Matrix{Float64}(undef, 3, 3))
	# Arbitrary axis
	# Apply vector from the right
	(sinα, cosα) = sincos(α)
	one_minus_cosα = (1 - cosα)
	# First column
	out[1, 1] = cosα + axis[1]^2 * one_minus_cosα
	out[2, 1] = axis[2] * sinα
	out[3, 1] = axis[1] * axis[2] * one_minus_cosα
	# Second column
	out[1, 2] = -axis[2] * sinα
	out[2, 2] = cosα
	out[3, 2] = axis[1] * sinα
	# Third column
	out[1, 3] = axis[1] * axis[2] * one_minus_cosα
	out[2, 3] = -axis[1] * sinα
	out[3, 3] = cosα + axis[2]^2 * one_minus_cosα
	return out
end

# Axis defined by phase
function rotation_matrix(plane::Val{:xy}, θ::Real, α::Real; out=Matrix{Float64}(undef, 3, 3))
	# Rotate by α, around axis specified by θ
	# Right handed coordinate system, x,y,z
	# This is for B1 in xy plane
	# Apply vector from the right
	(sinα, cosα) = sincos(α)
	(sinθ, cosθ) = sincos(θ)
	# First column
	out[1, 1] = sinθ^2 + cosα * cosθ^2
	out[2, 1] = sinθ * cosθ - cosα * sinθ * cosθ
	out[3, 1] = -sinα * cosθ
	# Second column
	out[1, 2] = sinθ * cosθ - cosα * sinθ * cosθ
	out[2, 2] = cosθ^2 + cosα * sinθ^2
	out[3, 2] = sinα * sinθ
	# Third column
	out[1, 3] = sinα * cosθ
	out[2, 3] = -sinα * sinθ
	out[3, 3] = cosα
	return out
end

function rotation_matrix(plane::Val{:xz}, θ::Real, α::Real; out=Matrix{Float64}(undef, 3, 3))
	# Rotate by α, around axis specified by θ
	# Right handed coordinate system, x,y,z
	# This is for B1 in xz plane
	# Apply vector from the right
	(sinα, cosα) = sincos(α)
	(sinθ, cosθ) = sincos(θ)
	# First column
	out[1, 1] = sinθ^2 + cosα * cosθ^2
	out[2, 1] = -sinα * cosθ
	out[3, 1] = sinθ * cosθ - cosα * sinθ * cosθ
	# Second column
	out[1, 2] = sinα * cosθ
	out[2, 2] = cosα
	out[3, 2] = -sinα * sinθ
	# Third column
	out[1, 3] = sinθ * cosθ - cosα * sinθ * cosθ
	out[2, 3] = sinα * sinθ
	out[3, 3] = cosθ^2 + cosα * sinθ^2

	return out
end


# This script provides an example of simulating changepoint data,
# using PELT to find these changepoints, and plotting the results
#
# Requires Plots package

using Changepoints, Distributions
using Plots

#########################
# Simulate changepoints #
#########################

n = 1000           # Sample size
λ = 100            # Frequencey of changes

# Sigma parameter sampled from a Uniform(2,15) distribution for each changepoint
μ, σ = 1.0, Uniform(2.0, 15.0)
# Generate a sample with from Normal distribution with changing variance
y, cps = @changepoint_sampler n λ Normal(μ, σ)

########################
# Finding changepoints #
########################

pelt_cps, cost = @PELT y Normal(μ, :?)

println("True changepoints:", cps)
println("Changepoint found by PELT:", pelt_cps)

changepoint_plot(y, pelt_cps)
gui()


using JuMP

function test(n::Int)
    m = Model()
    a = rand(n,n)
    @variable(m,x[1:n,1:n])
    b = rand(n,n,n)
    @variable(m,y[1:n,1:n,1:n])
    c = rand(n)
    @variable(m,z[1:n])

    #initialize
    @constraint(m,sum{c[i]*z[i],i=1:n}<=0)

    #Vector
    tic()
    @constraint(m,sum{c[i]*z[i],i=1:n}<=1)
    println("Vector with sum{}: $(toq())")
    tic()
    @constraint(m,vecdot(c,z) <= 1)
    println("Vector with vecdot() : $(toq())")

    #2D Matrix
    tic()
    @constraint(m,sum{a[i,j]*x[i,j],i=1:n,j=1:n}<=1)
    println("2D Matrix with sum{}: $(toq())")
    tic()
    @constraint(m,vecdot(a,x)<=1)
    println("2D Matrix with bigvecdot(): $(toq())")

    #3D Matrix
    tic()
    @constraint(m,sum{b[i,j,k]*y[i,j,k],i=1:n,j=1:n,k=1:n}<=1)
    println("3D Matrix with sum{}: $(toq())")
    tic()
    @constraint(m,vecdot(b,y)<=1)
    println("3D Matrix with vecdot(): $(toq())")
    return 0
end

for n in [10, 50, 100, 200, 300, 400]
    println("n = $n")
    test(n)
end


using FractionalDiffEq, Plots, SpecialFunctions

α = 0.5

function testfun(t, y)
    return 40320/gamma(9-α)*t^(8-α) - 3*gamma(5+α/2)/gamma(5-α/2)*t^(4-α/2) + 9/4*gamma(α+1) + (3/2*t^α/2-t^4)^3-y^1.5
end

prob = FODEProblem(testfun, α, 0, 1, 0.001)

result = solve(prob, PECE())

tfun(t) = t^8-3t^4.25+9/4*t^0.5

tspan = collect(0:0.001:1)

plot(tspan, result)

plot!(tspan, tfun, lw=2)


adapters = open("adapters.txt") do f
    sort(parse.(Int, readlines(f)))
end

function find_contiguous_sets(numbers::Array{Int64,1}, n::Int64)::Array{Int64, 1}
    sets = []
    i = 1

    while i < length(numbers)
        if numbers[i] != n
            i += 1
            continue
        end

        j = i
        while j < length(numbers)
            if numbers[j+1] != n
                break
            end
            j += 1
        end

        if j > i
            push!(sets, length(numbers[i:j]))
        end

        i = j + 1
    end

    return sets
end

function permutations(x::Int64)::Int64
    sum = 0

    for i in (x ÷ 3):-1:0
        y = x - 3*i

        for j in (y ÷ 2):-1:0
            z = y - 2*j
            p = factorial(z+j+i) ÷ (factorial(z)*factorial(j)*factorial(i))
            sum += p
            # println("$(x): [$(z), $(j), $(i)]: $(p)")
        end
    end

    return sum
end

pushfirst!(adapters, 0)
diffs = push!(map(i -> adapters[i]-adapters[i-1], 2:length(adapters)), 3)
println(count(i -> i == 1, diffs) * count(i -> i == 3, diffs))
println(prod(map(x -> permutations(x), find_contiguous_sets(diffs, 1))))


@testset "Inner product" begin
    @test v1 ⋅ v1 == 1v
    @test v1 ⋅ v2 == 0v
    @test v12 ⋅ v == 0v
    @test A ⋅ (B + C) ≈ A ⋅ B + A ⋅ C
    @test (B + C) ⋅ A ≈ B ⋅ A + C ⋅ A

    for r ∈ 0:4
        Aᵣ = As[r+1]
        for s ∈ 0:4
            Bₛ = Bs[s+1]
            for t ∈ 0:4
                Cₜ = Cs[t+1]
                if r + s ≤ t && r > 0 && s > 0
                    @test Aᵣ ⋅ (Bₛ ⋅ Cₜ) ≈ (Aᵣ ∧ Bₛ) ⋅ Cₜ
                elseif r + t ≤ s
                    @test Aᵣ ⋅ (Bₛ ⋅ Cₜ) ≈ (Aᵣ ⋅ Bₛ) ⋅ Cₜ
                end
            end
        end
        @test a ⋅ Aᵣ ≈ 0.5 * (a * Aᵣ + (-1) ^ (r + 1) * Aᵣ * a)
    end
end


import LinearAlgebra: I, Diagonal
import FFTW: fftfreq, fft, ifft

"""
    compute_RCWA_method(η::Function, μ::Function, L::Float64, N::Int, k::Float64, i::Int)

Compute the fields using Luke's RCWA method.
The output fields are on a grid compatible with the BIE method.

Arguments:
η:: A periodic function returning the metasurface magnetic conductance
μ:: A periodic function returning the metasurface electric resistance
L:: the length of the unit cell (some multiple of the period of η, μ)
N:: the number of discretization points along the unit cell
k:: the wavenumber/frequency of the incident planewave (assumed TM polarization)
i:: the index of the incident planewave in the Fourier basis

Returns:
NamedTuple with fields:
x:: the positions of the grid points
kz₁:: the values of the wavevector in the z component for y<0
kz₂:: the values of the wavevector in the z component for y>0
Ĩ₁:: an array of the incident fields in the Fourier basis for y<0
Õ₁:: an array of the outgoing fields in the Fourier basis for y<0
Ĩ₂:: an array of the incident fields in the Fourier basis for y>0
Õ₂:: an array of the outgoing fields in the Fourier basis for y>0
"""
function compute_RCWA_method(η::Function, μ::Function, L::Float64, N::Int, k::Float64, i::Int)
    # discretization parameters
    dx = L/N
    x⃗ = range(0, step=dx, length=N)
    k⃗ = 2*pi*fftfreq(N, 1/dx)
    β⃗ = @. sqrt(Complex(k^2 - k⃗^2))
    # construct Fourier-basis operators
    β̂ = Diagonal(β⃗)
    η̃ = ifft(fft(Matrix(Diagonal(η.(x⃗))), 2), 1)
    μ̃ = ifft(fft(Matrix(Diagonal(μ.(x⃗))), 2), 1)
    # build GSTC matrices and scattering matrix
    A = [
        -β̂-k*η̃    -β̂-k*η̃;
        -β̂-k*μ̃     β̂+k*μ̃
    ]
    B = [
        -β̂+k*η̃    -β̂+k*η̃;
        -β̂+k*μ̃     β̂-k*μ̃
    ]
    S = A\B
    # solve
    Ĩ₁ = [n == i ? 1 : 0 for n in 1:N]
    Ĩ₂ = zeros(N)
    inc = [Ĩ₁; Ĩ₂]
    scat = S*inc
    Õ₁, Õ₂ = scat[1:N], scat[N+1:2N]
    (x=x⃗, kz₁=β⃗, kz₂=β⃗, Ĩ₁=Ĩ₁, Õ₁=Õ₁, Ĩ₂=Ĩ₂, Õ₂=Õ₂)
end


using AbstractPlotting
using Meshing: SignedDistanceField, HyperRectangle, Vec, HomogenousMesh, MarchingTetrahedra

# algebraic dual boundary of the elliptopes
f = x -> (x[1]^2*x[2]^2 + x[1]^2*x[3]^2 + x[2]^2*x[3]^2 - 2*prod(x))

sdf = SignedDistanceField(f, HyperRectangle(Vec(-1, -1, -1), Vec(2, 2, 2)),0.01)

m = HomogenousMesh(sdf, MarchingTetrahedra())

using MeshCat
vis = Visualizer()
open(vis)

setobject!(vis, m,
           MeshPhongMaterial(color=RGBA{Float32}(1, 0, 0, 0.5)))

# Makie
# mesh_scene = Scene()
# mesh!(mesh_scene, m, color=:blue)
# display(mesh_scene)


module UKrig

using SpecialFunctions: besselk, gamma
import DynamicPolynomials: @polyvar, monomials
using FixedPolynomials
using LinearAlgebra
using Statistics
using Random

export Mnu, Gnu, generate_Gnu_krigY, generate_Mnu_krigY

tν𝒦t(t,ν) = t^ν * besselk(ν, t)

function Mnu(t, ν)::Float64
	pt, pν, p0, p1 = promote(t, ν, Float64(0), Float64(1)) # note besselk always returns a Float64 apparently
	return (pt==p0) ? p1 : tν𝒦t(√(2pν)*pt,pν) * 2^(1-pν) / gamma(pν)
end

# the const on the principle irregular term
scν(ν)       = - (2ν)^ν * gamma(ν + 1//2) * gamma(1-ν) / gamma(2ν+1) / sqrt(π)
scν(ν::Int)  = - 2 * (-2ν)^ν * gamma(ν + 1//2) / gamma(ν) / gamma(2ν+3) / sqrt(π)

function Gnu(t::T, ν::Int) where T<:Real
	if t==0
		return T(0)
	end
	return scν(ν) * t^(2ν) * log(t)
end

# Gnu(t::T, ν) where T<:Real = scν(ν) * t^(2ν)
function Gnu(t::T, ν) where T<:Real
	return scν(ν) * t^(2ν)
end

"""
`generate_Gnu_krigY(;fdata, xdata, ν, σg, σε) -> (x::Array->krigY(x), x::Array->fp(x'), b, c)`
"""
function generate_Gnu_krigY(Ydata,Xdata1,Xdata2,Xdata3, ν)
	m   =  7
	n₁  = length(Xdata1)
	#kriging matrix F
	@polyvar x y z
	p=1+x+y+x^2+y^2+x*y+z
	F=[t(x=>vx,y=>vy,z=>vz) for (vx, vy,vz) in zip(Xdata1,Xdata2,Xdata3), t in p]
	#distance matrix and Gnu
 	dist_mat= sqrt.(norm.( (Xdata1.-Xdata1').^2  + (Xdata2.-Xdata2').^2) )
	Gnu_mat=Gnu.(dist_mat, ν) |>Symmetric
	#generate the nullspace of F
	Mᵀ = nullspace(F')
	M = transpose(Mᵀ)
	#pre-compute matrix
	#This one takes so much time
 	Mat₁=Symmetric(M*Gnu_mat*Mᵀ)
	Mat₂=Symmetric(M*Mᵀ)

	Ydata_M=M*Ydata

    function KrigY(σg,σε,X1,X2,X3)
		G₁₁ = (σg^2) .* Gnu_mat
		Ξ   = [
			G₁₁ .+ σε^2*I(n₁)    	 	F
			F'       zeros(m,m)
		]
		cb = Ξ \ vcat(Ydata, zeros(m))
		c  = cb[1:length(Ydata)]
		b  = cb[length(Ydata)+1:end]
		K= (σg^2) .* Gnu.( sqrt.(norm.( (X1.-Xdata1').^2  + (X2.-Xdata2').^2) )      , ν)
		FF = [t(x=>vx,y=>vy,z=>vz) for (vx,vy,vz) in zip(X1,X1,X3), t in p]
		return K*c.+FF*b
	end

	function loglike(σg,σε)

		Σ₁=σg^2*Mat₁+σε^2*Mat₂
		ch=cholesky(Σ₁)
 		return -1/2*sum( (ch.L\(Ydata_M)).^2)- sum(log.(diag(ch.L)))

	end
	return KrigY, loglike
end

end


VERSION >= v"0.4.0-dev+6521" && __precompile__()

module OnlineMoments

import StatsBase

#-----------------------------------------------------------------------------

immutable Mean
  N::Float64
  M::Float64
end

Base.zero(::Type{Mean}) = Mean(0.0, 0.0)

Base.mean(a::Mean) = a.M

StatsBase.nobs(a::Mean) = convert(Int64, a.N)

function Base.(:+)(a::Mean, x::Real)
  N = a.N + 1.0
  Δ = x - a.M
  δ = Δ / N
  M = a.M + δ
  Mean(N, M)
end

function Base.(:+)(a::Mean, b::Mean)
  N = a.N + b.N
  Δ = b.M - a.M
  δ = Δ / N
  M = a.M + b.N*δ
  Mean(N, M)
end

#-----------------------------------------------------------------------------

immutable Variance
  N::Float64
  M::Float64
  V::Float64
end

Base.zero(::Type{Variance}) = Variance(0.0, 0.0, 0.0)

Base.mean(a::Variance) = a.M

Base.var(a::Variance) = a.V / (a.N - 1.0)

Base.std(a::Variance) = sqrt(var(a))

StatsBase.sem(a::Variance) = sqrt(a.V / ((a.N - 1.0) * a.N))

StatsBase.nobs(a::Variance) = convert(Int64, a.N)

function Base.(:+)(a::Variance, x::Real)
  N = a.N + 1.0
  Δ = x - a.M
  δ = Δ / N
  ξ = Δ * a.N
  M = a.M + δ
  V = a.V + δ*ξ
  Variance(N, M, V)
end

function Base.(:+)(a::Variance, b::Variance)
  N = a.N + b.N
  Δ = b.M - a.M
  δ = Δ / N
  ξ = Δ * (a.N*b.N)
  M = a.M + b.N*δ
  V = a.V + b.V + δ*ξ
  Variance(N, M, V)
end

#-----------------------------------------------------------------------------

immutable Moments
  N::Float64
  M::Float64
  V::Float64
  S::Float64
  K::Float64
end

Base.zero(::Type{Moments}) = Moments(0.0, 0.0, 0.0, 0.0, 0.0)

Base.mean(a::Moments) = a.M

Base.var(a::Moments) = a.V / (a.N - 1.0)

Base.std(a::Moments) = sqrt(var(a))

StatsBase.sem(a::Moments) = sqrt(a.V / ((a.N - 1.0) * a.N))

StatsBase.skewness(a::Moments) = sqrt(a.N) * a.S / a.V^1.5

StatsBase.kurtosis(a::Moments) = a.N * a.K / (a.V * a.V) - 3.0

StatsBase.nobs(a::Moments) = convert(Int64, a.N)

function Base.(:+)(a::Moments, x::Real)
  N = a.N + 1.0
  Δ = x - a.M
  δ = Δ / N
  δ² = δ * δ
  δ³ = δ² * δ
  ξ = Δ * a.N
  M = a.M + δ
  V = a.V + δ*ξ
  S = a.S + δ²*ξ*(a.N - 1.0) - 3.0*δ*a.V
  K = a.K + δ³*ξ*(a.N*a.N - a.N + 1.0) + 6.0*δ²*a.V - 4.0*δ*a.S
  Moments(N, M, V, S, K)
end

function Base.(:+)(a::Moments, b::Moments)
  N = a.N + b.N
  Δ = b.M - a.M
  δ = Δ / N
  δ² = δ * δ
  δ³ = δ² * δ
  ξ = Δ * (a.N*b.N)
  M = a.M + b.N*δ
  V = a.V + b.V + δ*ξ
  S = a.S + b.S + δ²*ξ*(a.N - b.N) + 3.0*δ*(a.N*b.V - b.N*a.V)
  K = a.K + b.K + δ³*ξ*(a.N*a.N - a.N*b.N + b.N*b.N) + 6.0*δ²*(a.N*a.N*b.V + b.N*b.N*a.V) + 4.0*δ*(a.N*b.S - b.N*a.S)
  Moments(N, M, V, S, K)
end

#-----------------------------------------------------------------------------

end


module SIAMFANLEquations
using LinearAlgebra
using LinearAlgebra.BLAS
using SparseArrays
using SuiteSparse
using BandedMatrices
using Printf

export nsolsc
export ptcsolsc
export ptcsol
export ptcsoli
export nsol
export nsoli
export aasol
export nofact
export secant
export armijosc
export kl_gmres
export kl_bicgstab
export Katv
export kstore
export Orthogonalize!
export EvalF!
export solhistinit

include("Tools/armijo.jl")
include("Tools/PrintError.jl")
include("Tools/FunctionJacobianEvals.jl")
include("Tools/ManageStats.jl")
include("Tools/IterationInit.jl")
include("Tools/ErrorTest.jl")
include("Tools/NewtonKrylov_Tools.jl")
include("Tools/PTCTools.jl")
include("Tools/PTCToolsi.jl")
include("Tools/AA_Tools.jl")
include("Solvers/Chapter1/nsolsc.jl")
include("Solvers/Chapter1/ptcsolsc.jl")
include("Solvers/Chapter1/secant.jl")
include("Solvers/ptcsol.jl")
include("Solvers/ptcsoli.jl")
include("Solvers/nsol.jl")
include("Solvers/nsoli.jl")
include("Solvers/aasol.jl")
include("Solvers/aa_qr_update.jl")
include("Solvers/LinearSolvers/kl_gmres.jl")
include("Solvers/LinearSolvers/kl_bicgstab.jl")
include("Solvers/LinearSolvers/Orthogonalize!.jl")
include("PlotsTables/printhist.jl")

module TestProblems
using LinearAlgebra
using SparseArrays
using SuiteSparse
using BandedMatrices
using AbstractFFTs
using FFTW
using Printf

export
    #Functions
    fcos,
    fpatan,
    spitchfork,
    linatan,
    sptestp,
    sptest,
    ftanx,
    ftanxp,
    heqinit,
    heqf!,
    heqJ!,
    HeqFix!,
    simple!,
    jsimple!,
    JVsimple,
    heqbos!,
    setc!,
    chandprint,
    bvpinit,
    Fbvp!,
    Jbvp!,
    FBeam!,
    FBeamtd!,
    BeamJ!,
    BeamtdJ!,
    beaminit,
    ptctest,
    pdeF!,
    pdeJ!,
    Jvec2d,
    pdeinit,
    pdegminit,
    fishinit,
    fish2d,
    sintv,
    isintv,
    Pfish2d,
    Pvec2d,
    Lap2d,
    Dx2d,
    Dy2d,
    solexact,
    l2dexact,
    dxexact,
    dyexact,
    hardleft!,
    hardleftFix!

include("TestProblems/Scalars/fcos.jl")
include("TestProblems/Scalars/fpatan.jl")
include("TestProblems/Scalars/spitchfork.jl")
include("TestProblems/Scalars/linatan.jl")
include("TestProblems/Scalars/ftanx.jl")
include("TestProblems/Systems/simple!.jl")
include("TestProblems/Systems/Fbvp!.jl")
include("TestProblems/Systems/FBeam!.jl")
include("TestProblems/Systems/Hequation.jl")
include("TestProblems/Systems/EllipticPDE.jl")
include("TestProblems/Systems/PDE_Tools.jl")
end

module Examples
using SIAMFANLEquations
using SIAMFANLEquations.TestProblems
using LinearAlgebra
using BandedMatrices

export ptciBeam
export ptcBeam
export ivpBeam
export BVP_solve
export nsolheq
export NsolPDE
export NsoliPDE
export PDE_aa

include("Examples/ptciBeam.jl")
include("Examples/ptcBeam.jl")
include("Examples/ivpBeam.jl")
include("Examples/BVP_solve.jl")
include("Examples/NsolPDE.jl")
include("Examples/NsoliPDE.jl")
include("Examples/PDE_aa.jl")
include("Examples/Internal/nsolheq.jl")
end

end # module


"""
Simplex in 2D.
PyDEC: Software and Algorithms for Discretization of Exterior Calculus
by Bell and Hirani.
"""
struct Simplex2D
    # vertex locations
    v::Array{Float64,2}
    # vertices that form triangles
    s::Array{Int64,2}
end

function buildsimplex2d(v, s)
    # The simplex has to be normalized.
    # The hilbert order increases locality of the geometry.
    reorder = hilbert_order(v)
    v = v[:, reorder]
    remap = sortperm(reorder)
    for j in 1:size(s, 2)
        for i in 1:size(s, 1)
            s[i, j] = remap[s[i, j]]
        end
    end
    # We will rely on simplex vertices being sorted.
    # circular permuations perserve orientation.
    for shift_idx in 1:size(s, 2)
        verts = s[:, shiftidx]
        s[:, shiftidx] = circshift(verts, findmin(verts))
    end
    s = sort(s, dims = 1)
    Simplex2D(v, s)
end

"""
Assume s is a list of triangles.
Pydec paper page 13. But transposed b/c this is column first ordering.
"""
function boundary_faces(s)
    d, n = size(s)  # dimension (3) and number of simplices
    splus = zeros(Int64, d + 2, n)
    splus[1:d, :] = s
    # need orientation of simplex for edges.
    for orient_idx in 1:d
        orient = iseven(countedinsertionsort!(splus[1:d, orient_idx]))
        splus[d + 1, orient_idx] = orient ? 1 : -1
        splus[d + 2, orient_idx] = orient_idx
    end
    resultant = zeros(Int64, d + 1, d * n)
    for vert in 1:d
        for tri in 1:n
            for copy_idx in 1:vert
                cnt = 1
                if copy_idx != vert
                    resultant[cnt, tri] = splus[copy_idx, tri]
                    cnt += 1
                end
                resultant[d, tri] = splus[d + 1, tri] * (-1)^(vert - 1)
                resultant[d + 1, tri] = splus[d + 2, tri]
            end
        end
    end
    ordered_resultant = sort(resultant, dims = 1)

    faces_pad = zeros(Int64, d - 1, size(resultant, 2))
    face_cnt = 1
    faces_pad[:, 1] = resultant[1:(d - 1), 1]
    for fidx in 2:size(resultant, 2)
        if faces_pad[1:(d - 1), fidx] != faces_pad[:, face_cnt]
            face_cnt += 1
            faces_pad[:, face_cnt] = resultant[1:(d - 1), fidx]
        end
    end
    faces = faces_pad[:, 1:face_cnt]

    edges = SparseMatrixCSC{Int64,Int64}(
        n,  # row count is number of simplices
        size(resultant, 2),  # column count is number of edges
        1:size(resultant, 2),  # one column per edge
        resultant[d + 1, :],  # row indices are the simplex index
        resultant[d, :]  # values are the orientation of the edge
        )
    edges, faces
end