export arithmetic_mean, geometric_mean, harmonic_mean, generalized_means
export pref_attach3, common3, jaccard3, adamic_adar3
export common_nbr_set

"""
arithmetic_mean
---------------

Returns the arithmetic mean of the weights of the edges of a list of triangles.

arithmetic_mean(triangles::Vector{NTuple{3,Int64}}, B::SpIntMat)

Input parameters:
- triangles::Vector{NTuple{3,Int64}}: The vector of triangles upon which to compute scores.
- B::SpIntMat: Projected graph as a Sparse integer matrix, where B[i, j] is the number of times that i and j co-appear in a simplex.
"""
function arithmetic_mean(triangles::Vector{NTuple{3,Int64}}, B::SpIntMat)
    scores = zeros(Float64, length(triangles))
    Threads.@threads for ind = 1:length(triangles)
        i, j, k = triangles[ind]
        scores[ind] = (B[i, j] + B[j, k] + B[i, k]) / 3
    end
    return scores
end

"""
geometric_mean
---------------

Returns the geometric mean of the weights of the edges of a list of triangles.

geometric_mean(triangles::Vector{NTuple{3,Int64}}, B::SpIntMat)

Input parameters:
- triangles::Vector{NTuple{3,Int64}}: The vector of triangles upon which to compute scores.
- B::SpIntMat: Projected graph as a Sparse integer matrix, where B[i, j] is the number of times that i and j co-appear in a simplex.
"""
function geometric_mean(triangles::Vector{NTuple{3,Int64}}, B::SpIntMat)
    scores = zeros(Float64, length(triangles))
    Threads.@threads for ind = 1:length(triangles)
        i, j, k = triangles[ind]
        scores[ind] = (B[i, j] * B[j, k] * B[i, k])^(1 / 3)
    end
    return scores
end

"""
harmonic_mean
---------------

Returns the harmonic mean of the weights of the edges of a list of triangles.

harmonic_mean(triangles::Vector{NTuple{3,Int64}}, B::SpIntMat)

Input parameters:
- triangles::Vector{NTuple{3,Int64}}: The vector of triangles upon which to compute scores.
- B::SpIntMat: Projected graph as a Sparse integer matrix, where B[i, j] is the number of times that i and j co-appear in a simplex.
"""
function harmonic_mean(triangles::Vector{NTuple{3,Int64}}, B::SpIntMat)
    scores = zeros(Float64, length(triangles))
    Threads.@threads for ind = 1:length(triangles)
        i, j, k = triangles[ind]
        scores[ind] = 3.0 / (1.0 / B[i, j] + 1.0 / B[j, k] + 1.0 / B[i, k])
    end
    return scores
end

"""
generalized_means
----------------

Computes the generalized p-means of the weights of the edges of a list of triangles.
The generalized mean of 3 values is

```math
M_p(x, y, z) = ((x^p + y^p + z^p) / 3)^{1/p}
```

generalized_means(triangles::Vector{NTuple{3,Int64}}, B::SpIntMat, ps::Float64=[-Inf; collect(-4:0.25:4); Inf])

Input parameters:
- triangles::Vector{NTuple{3,Int64}}: The vector of triangles upon which to compute scores.
- B::SpIntMat: Projected graph as a Sparse integer matrix, where B[i, j] is the number of times that i and j co-appear in a simplex.
- ps::Vector{Float64}=[-Inf; collect(-4:0.25:4); Inf]: the values of p for which to compute the means

Returns a matrix of size length(triangles) x length(ps) of the scores for the various generalized means.
"""
function generalized_means(triangles::Vector{NTuple{3,Int64}}, B::SpIntMat, ps::Vector{Float64}=[-Inf; collect(-4:0.25:4); Inf])
    scores = zeros(Float64, length(triangles), length(ps))
    Threads.@threads for ind = 1:length(triangles)
        i, j, k = triangles[ind]
        Bij, Bjk, Bik = B[i, j], B[j, k], B[i, k]
        for (r, p) in enumerate(ps)
            if     p == -Inf; scores[ind, r] = min(Bij, Bjk, Bik)
            elseif p == Inf;  scores[ind, r] = max(Bij, Bjk, Bik)
            elseif p == 0;    scores[ind, r] = (Bij * Bjk * Bik)^(1.0 / 3)
            else              scores[ind, r] = ((Bij^p + Bjk^p + Bik^p) / 3)^(1.0 / p)
            end
        end
    end
    return scores
end

"""
pref_attach3
------------

Returns the preferential attachment score for a degree vector

pref_attach3(triangles::Vector{NTuple{3,Int64}}, degrees::Vector{Int64})

The score of triangle (i, j, k) is degrees[i] * degrees[j] * degrees[k]

Input parameters:
- triangles::Vector{NTuple{3,Int64}}: The vector of triangles upon which to compute scores.
- degrees::Vector{Int64}: the degree vector
"""
function pref_attach3(triangles::Vector{NTuple{3,Int64}}, degrees::Vector{Int64})
    scores = zeros(Float64, length(triangles))
    Threads.@threads for ind = 1:length(triangles)
        i, j, k = triangles[ind]
        scores[ind] = 1.0 * degrees[i] * degrees[j] * degrees[k]
    end
    return scores
end

""" Return common neighbors of two nodes u and v. """
function common_nbr_set(common_nbrs::NbrSetMap, u::Int64, v::Int64)
    min_node, max_node = min(u, v), max(u, v)
    if !haskey(common_nbrs, (min_node, max_node))
        return Set{Int64}()
    end
    return common_nbrs[(min_node, max_node)]
end

"""
common3
------------

Returns the number of common 4th neighbors of a list of triangles.

common3(triangles::Vector{NTuple{3,Int64}}, common_nbrs::NbrSetMap)

Input parameters:
- triangles::Vector{NTuple{3,Int64}}: The vector of triangles upon which to compute scores.
- common_nbrs::NbrSetMap: the common neighbors map attained from common_neighbors_map()
"""
function common3(triangles::Vector{NTuple{3,Int64}}, common_nbrs::NbrSetMap)
    scores = zeros(Float64, length(triangles))
    Threads.@threads for ind = 1:length(triangles)
        i, j, k = triangles[ind]
        common_ij = common_nbr_set(common_nbrs, i, j)
        common_ik = common_nbr_set(common_nbrs, i, k)
        common_jk = common_nbr_set(common_nbrs, j, k)
        common_ijk = intersect(common_ij, common_ik, common_jk)
        scores[ind] = length(common_ijk)
    end
    return scores
end

"""
jaccard3
------------

Returns the 3-way Jaccard index of the neighbor lists of nodes in triangle.

jaccard3(triangles::Vector{NTuple{3,Int64}}, common_nbrs::NbrSetMap, degrees::Vector{Int64})

Input parameters:
- triangles::Vector{NTuple{3,Int64}}: The vector of triangles upon which to compute scores.
- common_nbrs::NbrSetMap: the common neighbors map attained from common_neighbors_map()
- degrees::Vector{Int64}: degree of each node in the projected graph
"""
function jaccard3(triangles::Vector{NTuple{3,Int64}}, common_nbrs::NbrSetMap,
                  degrees::Vector{Int64})
    scores = zeros(Float64, length(triangles))
    Threads.@threads for ind = 1:length(triangles)
        i, j, k = triangles[ind]
        common_ij = common_nbr_set(common_nbrs, i, j)
        common_ik = common_nbr_set(common_nbrs, i, k)
        common_jk = common_nbr_set(common_nbrs, j, k)
        common_ijk = intersect(common_ij, common_ik, common_jk)
        nij, nik, njk = length(common_ij), length(common_ik), length(common_jk)
        nijk = length(common_ijk)
        di, dj, dk = degrees[[i, j, k]]
        scores[ind] = nijk / (di + dj + dk - nij - nik - njk + nijk)
    end
    return scores
end

"""
adamic_adar3
------------

Returns the 3-way Adamic-Adar score for a list of triangles.

adamic_adar3(triangles::Vector{NTuple{3,Int64}}, common_nbrs::NbrSetMap, degrees::Vector{Int64})

Input parameters:
- triangles::Vector{NTuple{3,Int64}}: The vector of triangles upon which to compute scores.
- common_nbrs::NbrSetMap: the common neighbors map attained from common_neighbors_map()
- degrees::Vector{Int64}: degree of each node in the projected graph
"""
function adamic_adar3(triangles::Vector{NTuple{3,Int64}}, common_nbrs::NbrSetMap,
                      degrees::Vector{Int64})
    scores = zeros(Float64, length(triangles))
    Threads.@threads for ind = 1:length(triangles)
        i, j, k = triangles[ind]
        common_ij = common_nbr_set(common_nbrs, i, j)
        common_ik = common_nbr_set(common_nbrs, i, k)
        common_jk = common_nbr_set(common_nbrs, j, k)
        common_ijk = intersect(common_ij, common_ik, common_jk)
        scores[ind] = sum([1.0 / log(degrees[z]) for z in common_ijk])
    end
    return scores
end


abstract type AbstractPoint end

function iselliptic(𝑥::Number,𝑦::Number,𝑎::Number,𝑏::Number)
    𝑦^2 == 𝑥^3 + 𝑎*𝑥 + 𝑏
end

POINTTYPES = Union{Integer,PrimeField}

"""
Represents a point with coordinates (𝑥,𝑦) on an elliptic curve where 𝑦² = 𝑥³ + 𝑎𝑥 + 𝑏
Optional parameter 𝑝 represents finite field 𝐹ₚ and will convert all other parameter to PrimeField
Point(∞,∞,𝑎,𝑏) represents point at infinity
Returns an error if elliptic curve equation isn't satisfied
"""
struct Point{T<:Number,S<:Number} <: AbstractPoint
    𝑥::T
    𝑦::T
    𝑎::S
    𝑏::S
    Point{T,S}(𝑥,𝑦,𝑎,𝑏) where {T<:Number,S<:Number} = new(𝑥,𝑦,𝑎,𝑏)
end

Point(𝑥::Infinity,𝑦::Infinity,𝑎::T,𝑏::T) where {T<:POINTTYPES} = Point{Infinity,T}(𝑥,𝑦,𝑎,𝑏)
Point(𝑥::T,𝑦::T,𝑎::T,𝑏::T) where {T<:POINTTYPES} = !iselliptic(𝑥,𝑦,𝑎,𝑏) ? throw(DomainError("Point is not on curve")) : Point{T,T}(𝑥,𝑦,𝑎,𝑏)
Point(𝑥::Infinity,𝑦::Infinity,𝑎::T,𝑏::T,𝑝::T) where {T<:Integer} = Point(𝑥,𝑦,FieldElement(𝑎,𝑝),FieldElement(𝑏,𝑝))
Point(𝑥::T,𝑦::T,𝑎::T,𝑏::T,𝑝::T) where {T<:Integer} = Point(FieldElement(𝑥,𝑝),FieldElement(𝑦,𝑝),FieldElement(𝑎,𝑝),FieldElement(𝑏,𝑝))

"Formats AbstractPoint as (𝑥, 𝑦) on 𝑦² = 𝑥³ + 𝑎𝑥 + 𝑏 (: 𝐹ₚ)"
function show(io::IO, z::AbstractPoint)
    if typeof(z.𝑥) <: PrimeField
        x, y = z.𝑥.𝑛, z.𝑦.𝑛
    else
        x, y = z.𝑥, z.𝑦
    end

    if typeof(z.𝑎) <: PrimeField
        a, b = z.𝑎.𝑛, z.𝑏.𝑛
        field = string(" : 𝐹", z.𝑎.𝑝)
    else
        a, b = z.𝑎, z.𝑏
        field = ""
    end
    print(io, "(", x, ", ", y, ") on 𝑦² = 𝑥³ + ", a, "𝑥 + ", b, field)
end

"""
Returns the point resulting from the intersection of the curve and the
straight line defined by the points P and Q
"""
function +(𝑃::AbstractPoint,𝑄::AbstractPoint)
    T = typeof(𝑃)
    S = typeof(𝑃.𝑎)
    if 𝑃.𝑎 != 𝑄.𝑎 || 𝑃.𝑏 != 𝑄.𝑏
        throw(DomainError("Points are not on the same curve"))

    # Case 0
    elseif 𝑃.𝑥 == ∞
        return 𝑄
    elseif 𝑄.𝑥 == ∞
        return 𝑃
    elseif 𝑃.𝑥 == 𝑄.𝑥 && 𝑃.𝑦 != 𝑄.𝑦
        # something more elegant should exist to return correct point type
        if T <: Point
            return Point{Infinity,S}(∞, ∞, 𝑃.𝑎, 𝑃.𝑏)
        elseif T <: S256Point
            return S256Point{Infinity}(∞, ∞, 𝑃.𝑎, 𝑃.𝑏)
        end

    # Case 1
    elseif 𝑃.𝑥 != 𝑄.𝑥
        λ = (𝑄.𝑦 - 𝑃.𝑦) ÷ (𝑄.𝑥 - 𝑃.𝑥)
        𝑥 = λ^2 - 𝑃.𝑥 - 𝑄.𝑥
    # Case 2
    else
        λ = (3 * 𝑃.𝑥^2 + 𝑃.𝑎) ÷ (2 * 𝑃.𝑦)
        𝑥 = λ^2 - 2 * 𝑃.𝑥
    end
    𝑦 = λ * (𝑃.𝑥 - 𝑥) - 𝑃.𝑦
    return T(S(𝑥), S(𝑦), 𝑃.𝑎, 𝑃.𝑏)
end

"Scalar multiplication of a Point"
function *(λ::Integer,𝑃::Point)
    𝑅 = Point(∞, ∞, 𝑃.𝑎, 𝑃.𝑏)
    while λ > 0
        𝑅 += 𝑃
        λ -= 1
    end
    return 𝑅
end


function algebry(rel::Parallelism, coords)
    A = coords[rel.l.l]
    B = coords[rel.l.r]
    C = coords[rel.r.l]
    D = coords[rel.r.r]
    return ((B.x - A.x)*(D.y - C.y) - (B.y - A.y)*(D.x - C.x),)
end

function algebry(rel::Perpendicularity, coords)
    A = coords[rel.l.l]
    B = coords[rel.l.r]
    C = coords[rel.r.l]
    D = coords[rel.r.r]
    return ((B.x - A.x)*(D.x - C.x) + (B.y - A.y)*(D.y - C.y),)
end

function algebry(rel::Congruence{Segment}, coords)
    A = coords[rel.l.l]
    B = coords[rel.l.r]
    C = coords[rel.r.l]
    D = coords[rel.r.r]
    return ((B.x-A.x)^2 + (B.y-A.y)^2 - (D.x-C.x)^2 - (D.y-C.y)^2,)
end

function algebry(rel::Parallelogram, coords)
    A = coords[rel.A]
    B = coords[rel.B]
    C = coords[rel.C]
    D = coords[rel.D]
    return (B.x + D.x - A.x - C.x, B.y + D.y - A.y - C.y)
end

function algebry(rel::Midpoint, coords)
    M = coords[rel.M]
    A = coords[rel.seg.l]
    B = coords[rel.seg.r]
    return (2*M.x-A.x-B.x, 2*M.y-A.y-B.y)
end

function algebry(rel::Appartenence{Segment}, coords)
    P = coords[rel.l]
    A = coords[rel.r.l]
    B = coords[rel.r.r]
    return ((B.x - A.x)*(P.y - A.y) - (B.y - A.y)*(P.x - A.x),)
end

function algebry(rel::Intersection{Segment, Segment}, coords)
    return (algebry(Appartenence(rel.P, rel.l), coords)[1], algebry(Appartenence(rel.P, rel.r), coords)[1])
end

function algebry(rel::Projection, coords)
    return (algebry(Appartenence(rel.H, rel.seg), coords)[1], algebry(Perpendicularity(Segment(rel.P, rel.H), rel.seg), coords)[1])
end

function algebry(rel::Appartenence{Circle}, coords)
    OP = Segment(rel.r.O, rel.l)
    OA = Segment(rel.r.O, rel.r.A)
    return algebry(Congruence(OA, OP), coords)
end

function algebry(rel::CircleThreePoints, coords)
    OA = Segment(rel.O, rel.A)
    OB = Segment(rel.O, rel.B)
    OC = Segment(rel.O, rel.C)
    return (algebry(Congruence(OA, OB), coords)[1], algebry(Congruence(OA, OC), coords)[1])
end


using DifferentialEquations
using Plots

function sir_ode!(du,u,p,t)
    (S,I,R) = u
    (β,c,γ) = p
    N = S+I+R
    @inbounds begin
        du[1] = -β*c*I/N*S
        du[2] = β*c*I/N*S - γ*I
        du[3] = γ*I
    end
    nothing
end;

δt = 0.1
tmax = 80.0
tspan = (0.0,tmax)
t = 0.0:δt:tmax;

u0 = [990.0,10.0,0.0]; # S,I.R

p = [0.05,10.0,0.25]; # β,c,γ

lockdown_times = [10.0, 20.0]
condition(u,t,integrator) = t ∈ lockdown_times
function affect!(integrator)
    if integrator.t < lockdown_times[2]
        integrator.p[1] = 0.01
    else
        integrator.p[1] = 0.05
    end
end
cb = PresetTimeCallback(lockdown_times, affect!);

prob_ode = ODEProblem(sir_ode!,u0,tspan,p)

sol_ode = solve(prob_ode, callback = cb);

plot(sol_ode, label = ["S" "I" "R"], title = "Lockdown in a SIR model")
vline!(lockdown_times, c = :red, w = 2, label = "")


# # [Basic Epidemiology Models](@id epidemiology_example)
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/examples/covid/epidemiology.ipynb)

using AlgebraicPetri
using AlgebraicPetri.Epidemiology

using LabelledArrays
using OrdinaryDiffEq
using Plots

using Catlab
using Catlab.Graphics
using Catlab.WiringDiagrams
using Catlab.CategoricalAlgebra
using Catlab.Programs.RelationalPrograms

display_uwd(ex) = to_graphviz(ex, box_labels=:name, junction_labels=:variable, edge_attrs=Dict(:len=>".75"));

# #### SIR Model:

# define model
sir = @relation (s,i,r) begin
    infection(s,i)
    recovery(i,r)
end
display_uwd(sir)
#-
p_sir = apex(oapply_epi(sir))
Graph(p_sir)

# define initial states and transition rates, then
# create, solve, and visualize ODE problem

u0 = LVector(S=10, I=1, R=0);
p = LVector(inf=0.4, rec=0.4);

# The C-Set representation has direct support for generating a DiffEq vector field

prob = ODEProblem(vectorfield(p_sir),u0,(0.0,7.5),p);
sol = solve(prob,Tsit5())

plot(sol)

# #### SEIR Model:

# define model
seir = @relation (s,e,i,r) begin
    exposure(s,i,e)
    illness(e,i)
    recovery(i,r)
end
display_uwd(seir)
#-
p_seir = apex(oapply_epi(seir))
Graph(p_seir)

# define initial states and transition rates, then
# create, solve, and visualize ODE problem

u0 = LVector(S=10, E=1, I=0, R=0);
p = LVector(exp=.9, ill=.2, rec=.5);

prob = ODEProblem(vectorfield(p_seir),u0,(0.0,15.0),p);
sol = solve(prob,Tsit5())

plot(sol)

# #### SEIRD Model:

# define model
seird = @relation (s,e,i,r,d) begin
    exposure(s,i,e)
    illness(e,i)
    recovery(i,r)
    death(i,d)
end
display_uwd(seird)
#-
p_seird = apex(oapply_epi(seird))
Graph(p_seird)

# define initial states and transition rates, then
# create, solve, and visualize ODE problem

u0 = LVector(S=10, E=1, I=0, R=0, D=0);
p = LVector(exp=0.9, ill=0.2, rec=0.5, death=0.1);

prob = ODEProblem(vectorfield(p_seird),u0,(0.0,15.0),p);
sol = solve(prob,Tsit5())

plot(sol)


using Plots, Measures; pyplot()

pDiscrete = [0.25, 0.25, 0.5]
xGridD = 0:2

pContinuous(x) = 3/4*(1 - x^2)
xGridC = -1:0.01:1

pContinuous2(x) = x < 0 ? x+1 : 1-x

p1 = plot(xGridD, line=:stem, pDiscrete, marker=:circle, c=:blue, ms=6, msw=0)
p2 = plot(xGridC, pContinuous.(xGridC), c=:blue)
p3 = plot(xGridC, pContinuous2.(xGridC), c=:blue)

plot(p1, p2, p3, layout=(1,3), legend=false, ylims=(0,1.1), xlabel="x",
	ylabel=["Probability" "Density" "Density"], size=(1200, 400), margin=5mm)


using Turing, Random, MCMCChains
using CSV, DataFrames, StatsBase
using ParetoSmooth
using StatisticalRethinking

#Random.seed!(129111)

df = CSV.read(sr_datadir("WaffleDivorce.csv"), DataFrame)
df.D = zscore(df.Divorce)
df.M = zscore(df.Marriage)
df.A = zscore(df.MedianAgeMarriage)
data = (D=df.D, A=df.A)

function lin(a, b, c, x...)
    result = @. a + b * c
    for i in 1:2:length(x)
        @. result += x[i] * x[i+1]
    end
    return result
end

@model function m5_1t(A, D)
    a ~ Normal(0, 0.2)
    bA ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    for i in eachindex(D)
        μ = lin(a, A[i], bA)
        D[i] ~ Normal(μ, σ)
    end
end

chn5_1t = sample(m5_1t(df.A, df.D), NUTS(1000, .9), MCMCThreads(), 1000, 4)

@model function m5_2t(M, D)
    a ~ Normal(0, 0.2)
    bM ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    for i in eachindex(D)
        μ = lin(a, M[i], bM)
        D[i] ~ Normal(μ, σ)
    end
end

chn5_2t = sample(m5_2t(df.M, df.D), NUTS(1000, .9), MCMCThreads(), 1000, 4)

@model function m5_3t(A, M, D)
    a ~ Normal(0, 0.2)
    bA ~ Normal(0, 0.5)
    bM ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    for i in eachindex(D)
        μ = a + M[i] * bM + A[i] * bA
        D[i] ~ Normal(μ, σ)
    end
end

chn5_3t = sample(m5_3t(df.A, df.M, df.D), NUTS(1000, .9), MCMCThreads(), 1000, 4)

pw_lls5_1t = pointwise_log_likelihoods(m5_1t(df.A, df.D), chn5_1t)
pw_lls5_2t = pointwise_log_likelihoods(m5_2t(df.M, df.D), chn5_2t)
pw_lls5_3t = pointwise_log_likelihoods(m5_3t(df.A, df.M, df.D), chn5_3t)

loo_comparison = ParetoSmooth.loo_compare([pw_lls5_1t, pw_lls5_2t, pw_lls5_3t];
    model_names=[:m5_1t, :m5_2t, :m5_3t])

for (i, psis) in enumerate(loo_comparison.psis)
    psis |> display
    pk_plot(psis.pointwise(:pareto_k))
    savefig(joinpath(@__DIR__, "m5.$(i)t.png"))
end
println()
Particles(chn5_1t[:,[:a, :bA, :σ],:]) |> display
Particles(chn5_2t[:,[:a, :bM, :σ],:]) |> display
Particles(chn5_3t[:,[:a, :bA, :bM, :σ],:]) |> display
loo_comparison |> display


export ortho_latin, check_ortho

"""
`A,B = ortho_latin(n)` returns a pair of orthogonal `n`-by-`n`
Latin squares.

`A,B = ortho_latin(n,true)` returns a pair of orthogonal Latin squares
that are transposes of each other.

`A,B = ortho_latin(n,r,s)` builds the Latin squares `latin(n,r)`
and `latin(n,s)` and, if they are orthogonal, returns them as the
answer. (Otherwise, throws an error.) See: `find_ortho_parameters`.
"""
function ortho_latin(n::Int, self::Bool=false)
    if !self
        try
            r,s = find_ortho_parameters(n)
            return ortho_latin(n,r,s)
        catch
        end
    end
    println("No quick solution. Using integer programming.")

    return ortho_latin_IP(n,self)
end

function ortho_latin_IP(n::Int,self::Bool=false)
    MOD = Model(get_solver())
    # Z[i,j,k,l] is an indicator that there is a k in A[i,j] and
    # an l in B[i,j]
    @variable(MOD,Z[1:n,1:n,1:n,1:n], Bin)

    # one entry per cell constraint
    for i=1:n
        for j=1:n
            @constraint(MOD, sum(Z[i,j,k,l] for k=1:n for l=1:n) == 1)
        end
    end

    # Top row 11 22 33 ... nn
    if !self
        for i=1:n
            @constraint(MOD, Z[1,i,i,i]==1)  # A[1,i] = B[1,i] = i
        end
    end

    # orthogonality constraint
    for k=1:n
        for l=1:n
            @constraint(MOD, sum(Z[i,j,k,l] for i=1:n for j=1:n) == 1)
        end
    end

    # Row constraints

    for i=1:n
        for k=1:n
            @constraint(MOD, sum(Z[i,j,k,l] for j=1:n for l=1:n) == 1)
        end
    end

    for i=1:n
        for l=1:n
            @constraint(MOD, sum(Z[i,j,k,l] for j=1:n for k=1:n) == 1)
        end
    end

    # Col constraints
    for j=1:n
        for k=1:n
            @constraint(MOD, sum(Z[i,j,k,l] for i=1:n for l=1:n) == 1)
        end
    end

    for j=1:n
        for l=1:n
            @constraint(MOD, sum(Z[i,j,k,l] for i=1:n for k=1:n) == 1)
        end
    end

    if self   # force A^T==B
        for i=1:n
            for j=1:n
                for k=1:n
                    for l=1:n
                        @constraint(MOD, Z[i,j,k,l] == Z[j,i,l,k])
                    end
                end
            end
        end
    end

    optimize!(MOD)
    status = Int(termination_status(MOD))

    if status != 1
        error("No pair of orthogonal Latin squares of order $n can be found.")
    end

    ZZ = value.(Z)
    A = zeros(Int,n,n)
    B = zeros(Int,n,n)

    for i=1:n
        for j=1:n
            for k=1:n
                for l=1:n
                    if ZZ[i,j,k,l]>0
                        A[i,j] = k
                        B[i,j] = l
                    end
                end
            end
        end
    end

    return A,B
end

function ortho_latin(n::Int, r::Int, s::Int)
    A = latin(n,r)
    B = latin(n,s)
    @assert check_ortho(A,B) "Parameters n=$n, r=$r, and s=$s do not generate a pair of orthogonal Latin squares"
    return A,B
end

"""
`find_ortho_parameters(n)` tries to find parameters `r` and `s`
so that `ortho_latin(n,r,s)` will succeed. Returns `(r,s)` if
successful or throws an error if not.
"""
function find_ortho_parameters(n::Int)
    for r=1:n-1
        for s=1:n-1
            if gcd(n,r)==1 && gcd(n,s)==1 && gcd(n,r-s)==1
                return r,s
            end
        end
    end
    error("No parameters for n=$n found")
end

"""
`check_ortho(A,B)` checks that matrices `A` and `B` are a pair of
orthogonal Latin squares.
"""
function check_ortho(A::Matrix{Int},B::Matrix{Int})::Bool
    if size(A) != size(B)
        return false
    end
    if !check_latin(A) || !check_latin(B)
        return false
    end
    n,r = size(A)

    vals = unique((n+1)*A + B)
    return length(vals) == n*n
end


# These functions compute the sizes of combinations of parameters.

using Combinatorics: combinations

"""
    total_combinations(arity, n_way)

Given an array of the number of values for each parameter and the
level of coverage, return the number of n-tuples required for complete
coverage.
"""
function total_combinations(arity, n_way)
  param_cnt = length(arity)
  sum(prod(arity[key_set]) for key_set in combinations(1:param_cnt, n_way))
end

"""
    next_multiplicative!(values, arity)

Given a set of values for parameters, this returns the next possible
value. When it reaches the end, it cycles back to the start. The first
value is all ones. The last value equals the arity.
"""
function next_multiplicative!(values, arity)
    carry = 1
    for slot_idx in length(values):-1:1
        values[slot_idx] += carry
        if arity[slot_idx] < values[slot_idx]
            values[slot_idx] = 1
            carry = 1
        else
            carry = 0
        end
    end
end

"""
    all_combinations(arity, n_way)

This represents possible coverage as a matrix, one row per parameter,
zero if not used. `arity` is a list of the number of values for each parameter,
and `n_way` is the order of the combinations, most commonly 2-way.
"""
function all_combinations(arity, n_way)
    v_cnt = length(arity)
    # This returns a list of lists, so it has length, not size.
    indices = collect(combinations(1:v_cnt, n_way))
    combinations_cnt = total_combinations(arity, n_way)

    coverage = zeros(eltype(arity), v_cnt, combinations_cnt)
    idx = 1
    for indices_idx in 1:length(indices)
        offset = indices[indices_idx]
        sub_arity = arity[offset]
        sub_cnt = prod(sub_arity)
        values = copy(sub_arity)
        for sub_idx in 1:sub_cnt
            next_multiplicative!(values, sub_arity)
            coverage[offset, idx] = values
            idx += 1
        end
    end
    coverage
end

"""
    one_parameter_combinations(arity, n_way)

Generates all combinations that are nonzero for the last parameter.
This is for in-parameter-order generation, where we need
only those tuples that end with this column being nonzero.

The construction method is to leave out the given parameter
and construct all `n_way` - 1 tuples. Then copy and paste that
once for each possible value of the given parameter.
"""
function one_parameter_combinations(arity, n_way)
    param_cnt = length(arity)
    if n_way > 1
        partial = all_combinations(arity[1:(param_cnt-1)], n_way - 1)
        one_set = size(partial, 2)
        comb = zeros(eltype(arity), param_cnt, one_set * arity[param_cnt])
        for vidx in 1:(arity[param_cnt])
            col_begin = (vidx - 1) * one_set + 1
            col_end = vidx * one_set
            comb[1:size(partial, 1), col_begin:col_end] .= partial
            comb[param_cnt, col_begin:col_end] .= vidx
        end
    else  # n_way == 1
        comb = zeros(eltype(arity), param_cnt, arity[param_cnt])
        comb[param_cnt, :] .= 1:(arity[param_cnt])
    end
    comb
end

combination_number(n, m) = prod(n:-1:(n-m+1)) ÷ factorial(m)

function pairs_in_entry(entry, n_way)
    n = length(entry)
    ans = zeros(Int, combination_number(n, n_way), n)
    col_set_idx = 1
    for param_idx in combinations(1:length(entry), n_way)
        for row_idx in 1:n_way
            ans[param_idx[row_idx], col_set_idx] = entry[param_idx[row_idx]]
        end
        col_set_idx += 1
    end
    ans
end


# Example from "Bounds on the reachable sets of nonlinear control systems" Constant Temperature Profile
id = :ScottB2013a

pL = [420.0]
pU = [480.0]
ScottB2013a_x0(p::Vector{T}) where T = [1.5*one(T); 0.5*one(T); zero(T)]
function ScottB2013a_f!(du::Vector{T}, u::Vector{T}, p::Vector{T}, t) where T

    R = 8.134
    A1 = 2400
    A2 = 8800
    E1 = 6.9*10^3
    E2 = 1.69*10^4

    du[1] = -A1*exp(E1/(R*p[1]))*u[1]
    du[2] = A1*exp(E1/(R*p[1]))*u[1] - A2*exp(E2/(R*p[1]))*u[2]
    du[3] = A2*exp(E2/(R*p[1]))*u[2]
    return
end
tspan = (0.0, 0.08)

prob = ODERelaxProb(ScottB2013a_f!, tspan, ScottB2013a_x0, pL, pU)
url = "https://www.sciencedirect.com/science/article/pii/S0005109812004839"
source = "ScottB2013"
desc = "3 Series Rxn Network"


using StanModels, CSV

df = CSV.read(joinpath(@__DIR__, "..", "..", "data", "milk.csv"), delim=';')
dcc = filter(row -> !(row[:neocortex_perc] == "NA"), df)
dcc[!, :kcal_per_g] = convert(Vector{Float64}, dcc[!, :kcal_per_g])
dcc[!, :neocortex_perc] = parse.(Float64, dcc[!, :neocortex_perc])

# Show first 5 rows

first(dcc, 5)

# Define the Stan language model

m5_5s = "
data{
    int N;
    vector[N] kcal_per_g;
    vector[N] neocortex_perc;
}
parameters{
    real a;
    real bn;
    real sigma;
}
model{
    vector[N] mu = a + bn * neocortex_perc;
    sigma ~ uniform( 0 , 1 );
    bn ~ normal( 0 , 1 );
    a ~ normal( 0 , 100 );
    kcal_per_g ~ normal( mu , sigma );
}
";

# Define the Stanmodel and set the output format to :mcmcchains.

sm = SampleModel("m5.5s", m5_5s);

# Input data for cmdstan

m5_5_data = Dict("N" => size(dcc, 1),
  "kcal_per_g" => dcc[!, :kcal_per_g],
  "neocortex_perc" => dcc[!, :neocortex_perc]);

# Sample using cmdstan

(sample_file, log_file) = stan_sample(sm, data=m5_5_data);

# Rethinking results

rethinking_results = "
       mean   sd  5.5% 94.5%
 a     0.04 0.15 -0.21  0.29
 bN    0.13 0.22 -0.22  0.49
 sigma 1.00 0.16  0.74  1.26
"

# Describe the draws
if !(sample_file == nothing)
  chn = read_samples(sm)
  describe(chn)
end


"A basis of Chebyshev polynomials of the second kind on the interval `[-1,1]`."
struct ChebyshevU{T} <: OPS{T}
    n			::	Int
end

ChebyshevU(n::Int) = ChebyshevU{Float64}(n)

similar(b::ChebyshevU, ::Type{T}, n::Int) where {T} = ChebyshevU{T}(n)

show(io::IO, d::ChebyshevU{Float64}) = print(io, "ChebyshevU($(length(d)))")
show(io::IO, d::ChebyshevU{T}) where T = print(io, "ChebyshevU{$(T)}($(length(d)))")

function unsafe_eval_element(b::ChebyshevU, idx::PolynomialDegree, x::Real)
    # Don't use the formula when |x|=1, because it will generate NaN's
    d = degree(idx)
    abs(x) < 1 ? sin((d+1)*acos(x))/sqrt(1-x^2) : recurrence_eval(b, idx, x)
end

first_moment(b::ChebyshevU{T}) where {T} = convert(T, pi)/2

interpolation_grid(b::ChebyshevU{T}) where {T} = ChebyshevUNodes{T}(b.n)

iscompatible(dict::ChebyshevU, grid::ChebyshevUNodes) = length(dict) == length(grid)
issymmetric(::ChebyshevU) = true
measure(dict::ChebyshevU{T}) where {T} = ChebyshevUWeight{T}()
hasmeasure(::ChebyshevU) = true

function innerproduct_native(b1::ChebyshevU, i::PolynomialDegree, b2::ChebyshevU, j::PolynomialDegree, m::ChebyshevUWeight;
			T = coefficienttype(b1), options...)
	if i == j
		convert(T, pi)/2
	else
		zero(T)
	end
end

# Parameters alpha and beta of the corresponding Jacobi polynomial
jacobi_α(b::ChebyshevU{T}) where {T} = one(T)/2
jacobi_β(b::ChebyshevU{T}) where {T} = one(T)/2

# See DLMF, Table 18.9.1
# http://dlmf.nist.gov/18.9#i
rec_An(b::ChebyshevU{T}, n::Int) where {T} = convert(T, 2)

rec_Bn(b::ChebyshevU{T}, n::Int) where {T} = zero(T)

rec_Cn(b::ChebyshevU{T}, n::Int) where {T} = one(T)

support(b::ChebyshevU{T}) where {T} = ChebyshevInterval{T}()

"A Chebyshev polynomial of the second kind"
struct ChebyshevUPolynomial{T} <: OrthogonalPolynomial{T}
    degree  ::  Int
end

ChebyshevUPolynomial{T}(p::ChebyshevUPolynomial) where {T} = ChebyshevUPolynomial{T}(p.degree)

show(io::IO, p::ChebyshevUPolynomial) = print(io, "U_$(degree(p))(x) (Chebyshev polynomial of the second kind)")

convert(::Type{TypedFunction{T,T}}, p::ChebyshevUPolynomial) where {T} = ChebyshevUPolynomial{T}(p.degree)

support(::ChebyshevUPolynomial{T}) where {T} = ChebyshevInterval{T}()

(p::ChebyshevUPolynomial{T})(x) where {T} = eval_element(ChebyshevU{T}(degree(p)+1), degree(p)+1, x)

basisfunction(dict::ChebyshevU, idx) = basisfunction(dict, native_index(dict, idx))
basisfunction(dict::ChebyshevU{T}, idx::PolynomialDegree) where {T} = ChebyshevUPolynomial{T}(degree(idx))

dictionary(p::ChebyshevUPolynomial{T}) where {T} = ChebyshevU{T}(degree(p)+1)
index(p::ChebyshevUPolynomial) = degree(p)+1


# indicator of the L2 norm ball with given radius

"""
  IndBallL2(r::Real=1.0)

Returns the function `g = ind{x : ||x|| ⩽ r}`, for a real parameter `r > 0`.
"""

immutable IndBallL2{R <: Real} <: IndicatorConvex
  r::R
  function IndBallL2(r::R)
    if r <= 0
      error("parameter r must be positive")
    else
      new(r)
    end
  end
end

IndBallL2{R <: Real}(r::R=1.0) = IndBallL2{R}(r)

function (f::IndBallL2){T <: RealOrComplex}(x::AbstractArray{T})
  if vecnorm(x) - f.r > 1e-14
    return +Inf
  end
  return 0.0
end

function prox!{T <: RealOrComplex}(f::IndBallL2, x::AbstractArray{T}, y::AbstractArray{T}, gamma::Real=1.0)
  scal = f.r/vecnorm(x)
  if scal > 1
    y[:] = x
    return 0.0
  end
  for k in eachindex(x)
    y[k] = scal*x[k]
  end
  return 0.0
end

fun_name(f::IndBallL2) = "indicator of an L2 norm ball"
fun_dom(f::IndBallL2) = "AbstractArray{Real}, AbstractArray{Complex}"
fun_expr(f::IndBallL2) = "x ↦ 0 if ||x|| ⩽ r, +∞ otherwise"
fun_params(f::IndBallL2) = "r = $(f.r)"

function prox_naive{T <: RealOrComplex}(f::IndBallL2, x::AbstractArray{T}, gamma::Real=1.0)
  normx = vecnorm(x)
  if normx > f.r
    y = (f.r/normx)*x
  else
    y = x
  end
  return y, 0.0
end


using GeometryTypes
using SurfaceTopology
using LinearAlgebra

function quadraticform(vects,vnormal)

    Lx = [0 0 0; 0 0 -1; 0 1 0]
    Ly = [0 0 1; 0 0 0; -1 0 0]
    Lz = [0 -1 0; 1 0 0; 0 0 0]

    d = [0,0,1] + vnormal
    d /= norm(d)

    Ln = d[1]*Lx + d[2]*Ly + d[3]*Lz
    R = exp(pi*Ln)

    vects = copy(vects)
    for vj in 1:length(vects)
        vects[vj] = R*vects[vj]
    end

    ### Construction of the system
    A = Array{Float64}(undef,3,3)
    B = Array{Float64}(undef,3)

    vects_norm2 = Array{Float64}(undef,length(vects))
    for vj in 1:length(vects)
       vects_norm2[vj] = norm(vects[vj])^2
    end

    A[1,1] = sum((v[1]^4 for v in vects) ./ vects_norm2)
    A[1,2] = sum((v[1]^3*v[2] for v in vects) ./ vects_norm2)
    A[1,3] = sum((v[1]^2*v[2]^2 for v in vects) ./ vects_norm2)
    A[2,1] = A[1,2]
    A[2,2] = A[1,3]
    A[2,3] = sum( (v[2]^3*v[1] for v in vects) ./vects_norm2)
    A[3,1] = A[1,3]
    A[3,2] = A[2,3]
    A[3,3] = sum((v[2]^4 for v in vects) ./vects_norm2)

    B[1] = sum((v[3]*v[1]^2 for v in vects) ./vects_norm2)
    B[2] = sum((v[1]*v[2]*v[3] for v in vects) ./vects_norm2)
    B[3] = sum((v[2]^2*v[3] for v in vects) ./vects_norm2)

    C,D,E = A\B
    return C,D,E
end

function meancurvature(points,topology)
    curvatures = Array{Float64}(undef,length(points))
    for v in 1:length(points)

        s = Point(0,0,0)
        for (v1,v2) in EdgeRing(v,topology)
            s += cross(points[v2],points[v1])
        end
        normal = s ./ norm(s)

        vring = collect(VertexRing(v,topology))
        vects = [points[vi] - points[v] for vi in vring]

        C,D,E = quadraticform(vects,normal)

        A = [C D/2;D/2 E]
        k1,k2 = eigvals(-A)
        H = (k1 + k2)/2

        ### Multiplier by is 2 an empirical fix
        curvatures[v] = 2*H
    end
    return curvatures
end

function normals(vertices,topology)
    n = Point{3,Float64}[]
    for v in 1:length(vertices)
        s = Point(0,0,0)
        for (v1,v2) in EdgeRing(v,topology)
            s += cross(vertices[v2],vertices[v1])
        end
        normal = s ./ norm(s)
        push!(n,normal)
    end
    return n
end

function vertexareas(points,topology)
    vareas = zeros(Float64,length(points))
    for face in Faces(topology)
        v1,v2,v3 = face
        area = norm(cross(points[v2]-points[v1],points[v3]-points[v1])) /2
        vareas[v1] += area/3
        vareas[v2] += area/3
        vareas[v3] += area/3
    end
    return vareas
end

function surfacevolume(points,topology)
    # Calculate face normal
    # Calculate projected area
    # Calculate ordinary volume
    # Calculate volume between projected and real area
    # (+) if normal is outwards

    normal0 = [0,0,1]
    s = 0
    for face in Faces(topology)
        y1 = points[face[1]]
        y2 = points[face[2]]
        y3 = points[face[3]]

        normaly = cross(y2-y1,y3-y1)
        normaly /= norm(normaly)

        area = norm(cross(y2-y1,y3-y1))/2
        areaproj = dot(normaly,normal0)*area
        volume = dot(y1 + y2 + y3,normal0)/3*areaproj

        s += volume
    end

    return s
end


"""
	random_density(n)

Returns a random ``N × N`` density matrix, distributed according to the
Hilbert-Schmidt measure.

NOTE: No guarantee about the quality of the algorithm
For more details see [V Al Osipov et al 2010 J. Phys. A: Math. Theor. 43 055302]
(https://doi.org/10.1088/1751-8113/43/5/055302)
"""
function random_density_matrix(n)
    A = randn(n,n) + 1im * randn(n,n)
    rho = A * A'
    return rho / tr(rho)
end

"""
	unitary_op(H, t)

Returns the unitary evolution operator ``U = \\exp[-i H t]``.
"""
function unitary_op(H, t)
   return exp(-1im * Matrix(H) * t)
end

"""
    commutator(A, B)

Returns the commutator [A,B] = A * B - B * A between two matrices
"""
function commutator(a, b)
	a * b - b * a
end

"""
	localized_state(N, i)

Returns the density operator for a state ``|i⟩`` in an ``N``-dimensional Hilbert
space
"""
function localized_state(n, i)
   sparse([i,n],[i,n],[1,0])
end


# # Chebyshev transform
# This demonstrates the Chebyshev transform and inverse transform,
# explaining precisely the normalization and points

using FastTransforms
n = 20

# First kind points $\to$ first kind polynomials
p_1 = chebyshevpoints(Float64, n, Val(1))
f = exp.(p_1)
f̌ = chebyshevtransform(f, Val(1))
f̃ = x -> [cos(k*acos(x)) for k=0:n-1]' * f̌
f̃(0.1) ≈ exp(0.1)

# First kind polynomials $\to$ first kind points
ichebyshevtransform(f̌, Val(1)) ≈ exp.(p_1)

# Second kind points $\to$ first kind polynomials
p_2 = chebyshevpoints(Float64, n, Val(2))
f = exp.(p_2)
f̌ = chebyshevtransform(f, Val(2))
f̃ = x -> [cos(k*acos(x)) for k=0:n-1]' * f̌
f̃(0.1) ≈ exp(0.1)

# First kind polynomials $\to$ second kind points
ichebyshevtransform(f̌, Val(2)) ≈ exp.(p_2)

# First kind points $\to$ second kind polynomials
p_1 = chebyshevpoints(Float64, n, Val(1))
f = exp.(p_1)
f̌ = chebyshevutransform(f, Val(1))
f̃ = x -> [sin((k+1)*acos(x))/sin(acos(x)) for k=0:n-1]' * f̌
f̃(0.1) ≈ exp(0.1)

# Second kind polynomials $\to$ first kind points
ichebyshevutransform(f̌, Val(1)) ≈ exp.(p_1)

# Second kind points $\to$ second kind polynomials
p_2 = chebyshevpoints(Float64, n, Val(2))[2:n-1]
f = exp.(p_2)
f̌ = chebyshevutransform(f, Val(2))
f̃ = x -> [sin((k+1)*acos(x))/sin(acos(x)) for k=0:n-3]' * f̌
f̃(0.1) ≈ exp(0.1)

# Second kind polynomials $\to$ second kind points
ichebyshevutransform(f̌, Val(2)) ≈ exp.(p_2)


# # Robust approximate fitting
# Section 6.4.2
# Boyd & Vandenberghe "Convex Optimization"
# Original by Lieven Vandenberghe
# Adapted for Convex by Joelle Skaf - 10/03/05
#
# Adapted for Convex.jl by Karanveer Mohan and David Zeng - 26/05/14
# Original cvx code and plots here:
# <http://web.cvxr.com/cvx/examples/cvxbook/Ch06_approx_fitting/html/fig6_15.html>
#
# Consider the least-squares problem:
#       minimize $\|(A + tB)x - b\|_2$
# where $t$ is an uncertain parameter in [-1,1]
# Three approximate solutions are found:
#
#   1. nominal optimal (i.e. letting t=0)
#   2. stochastic robust approximation:
#           minimize $\mathbb{E}\|(A+tB)x - b\|_2$
#      assuming $u$ is uniformly distributed on [-1,1].
#      (reduces to minimizing $\mathbb{E} \|(A+tB)x-b\|^2 = \|A*x-b\|^2  + x^TPx$
#        where $P = \mathbb{E}(t^2) B^TB = (1/3) B^TB$ )
#   3. worst-case robust approximation:
#           minimize $\mathrm{sup}_{-1\leq u\leq 1} \|(A+tB)x - b\|_2$
#      (reduces to minimizing $\max\{\|(A-B)x - b\|_2, \|(A+B)x - b\|_2\}$ ).
#
using Convex, LinearAlgebra, SCS

# Input Data
m = 20;
n = 10;
A = randn(m, n);
(U, S, V) = svd(A);
S = diagm(exp10.(range(-1, stop = 1, length = n)));
A = U[:, 1:n] * S * V';

B = randn(m, n);
B = B / norm(B);

b = randn(m, 1);
x = Variable(n)

# Case 1: Nominal optimal solution
p = minimize(norm(A * x - b, 2))
solve!(p, SCS.Optimizer; silent_solver = true)
x_nom = evaluate(x)

# Case 2: Stochastic robust approximation
P = 1 / 3 * B' * B;
p = minimize(square(pos(norm(A * x - b))) + quadform(x, Symmetric(P)))
solve!(p, SCS.Optimizer; silent_solver = true)
x_stoch = evaluate(x)

# Case 3: Worst-case robust approximation
p = minimize(max(norm((A - B) * x - b), norm((A + B) * x - b)))
solve!(p, SCS.Optimizer; silent_solver = true)
x_wc = evaluate(x)

# Plot residuals:
parvals = range(-2, stop = 2, length = 100);

errvals(x) = [norm((A + parvals[k] * B) * x - b) for k in eachindex(parvals)]
errvals_ls = errvals(x_nom)
errvals_stoch = errvals(x_stoch)
errvals_wc = errvals(x_wc)

using Plots
plot(parvals, errvals_ls, label = "Nominal problem")
plot!(parvals, errvals_stoch, label = "Stochastic Robust Approximation")
plot!(parvals, errvals_wc, label = "Worst-Case Robust Approximation")
plot!(
    title = "Residual r(u) vs a parameter u for three approximate solutions",
    xlabel = "u",
    ylabel = "r(u) = ||A(u)x-b||_2",
)


import Pkg; Pkg.activate(joinpath(@__DIR__, "../../../../../"))
using Distributions
using PyPlot

"""
Paper: Bayesian inference for finite mixtures of univariate and multivariate
       skew-normal and skew-t distributions, Biostatistics 2010.
skew (delta): a real number in (-1, 1)
"""
function rand_skewnormal(loc, scale, skew)
    z = rand(TruncatedNormal(0, 1, 0, Inf))
    return loc + scale * skew * z + scale * sqrt(1 - skew ^ 2) * randn()
end

function rand_skewnormal(loc, scale, skew, dims...)
    z = rand(TruncatedNormal(0, 1, 0, Inf), dims...)
    return loc .+ scale * skew * z + scale * sqrt(1 - skew ^ 2) * randn(dims...)
end

x = rand_skewnormal(1, .5, -.97, 100000)
plt.hist(x, bins=100);

mean(x .< 0)

mean(x), std(x)


module MDL

setprecision(BigFloat, 64)

"""
scalable binomial coefficient
"""
function choose(n::Int64, k::Int64)
    if 0 < k <= n
        p = BigInt(1)
        for t=0:(min(k,n-k)-1)
            p = floor((p * (n - t)) / (t + 1))
        end
        return p
    else
        return 0.
    end
end

"""
binary logarithm
"""
function log2(n::Real)
    return log(2, n)
end

"""
binary logarithm returning 0. at zero
"""
function log2_zero(n::Real)
    if n == 0
        return 0.
    else
        return log2(n)
    end
end

"""
binary logarithm of binomial coefficient
"""
function log2_choose(n::Int64, k::Int64)
    if n == 0 || k == 0
        return 0.
    else
        return log2_zero(choose(n,k))
    end
end

"""
iterated logarithm (log-star)
"""
function log2_star(n::Union{Int64,Float64,BigInt,BigFloat})
    if n <= 1
        return 0.
    else
        return 1. + log2_star(log2(n))
    end
end

"""
:param z: an int >= 1
:return: the encoded size of the int according to Rissanen 1983's universal code for integers
"""
function universal_integer(z::Int64)
    NORMALIZATION_CONSTANT = 2.865064
    if z <= 0
        return 0.
    end
    c = log2(NORMALIZATION_CONSTANT)
    logstar_z = log2_star(z)
    return logstar_z + c
end

end


"""
    treynor_ratio(asset_returns, benchmark_returns; multiplier=1.0, risk_free=0.0)

Calculates the Treynor ratio as the ratio of excess return divided by the CAPM beta. This ratio is similar to the Sharpe Ratio, but instead of dividing by the volatility, we devide by the CAPM beta as risk proxy.

# Formula

    TR = E[asset_returns - risk_free] / beta * multiplier

# Arguments
- `asset_returns`:      Vector of asset returns.
- `benchmark_returns`:  Vector of benchmark returns (e.g. market portfolio returns).
- `multiplier`:         Optional scalar multiplier, i.e. use `12` to annualize monthly returns, and use `252` to annualize daily returns.
- `risk_free`:          Optional vector or scalar value denoting the risk-free return (must have same frequency as the provided returns, e.g. daily).
"""
function treynor_ratio(asset_returns, benchmark_returns; multiplier=1.0, risk_free=0.0)
    α, β = capm(asset_returns, benchmark_returns; risk_free=risk_free)
    mean(asset_returns .- risk_free) / β * multiplier
end


using DifferentialEquations
using SimpleDiffEq
using Random
using Distributions
using GpABC
using Distances
using ApproxBayes
using Plots

function sir_ode!(du,u,p,t)
    (S,I,R,C) = u
    (β,c,γ) = p
    N = S+I+R
    infection = β*c*I/N*S
    recovery = γ*I
    @inbounds begin
        du[1] = -infection
        du[2] = infection - recovery
        du[3] = recovery
        du[4] = infection
    end
    nothing
end;

tmax = 40.0
δt = 1.0
tspan = (0.0,tmax)
obstimes = 1.0:δt:tmax;
u0 = [990.0,10.0,0.0,0.0]; # S,I.R,C
p = [0.05,10.0,0.25]; # β,c,γ

prob_ode = ODEProblem(sir_ode!,u0,tspan,p)
sol_ode = solve(prob_ode,saveat=δt)
out_ode = Array(sol_ode)
C = out_ode[4,:]
X = C[2:end] .- C[1:(end-1)];

Random.seed!(1234)
Y = rand.(Poisson.(X));

bar(obstimes,Y)
plot!(obstimes,X)

function simdata(x)
    (i0,β) = x
    I = i0*1000.0
    prob = remake(prob_ode,u0=[1000-I,I,0.0,0.0],p=[β,10.0,0.25])
    sol = solve(prob,Tsit5(),saveat=δt)
    out = Array(sol)
    C = out[4,:]
    X = C[2:end] .- C[1:(end-1)]
    transpose(X)
end;

priors = [Uniform(0.0,0.1),Uniform(0.0,0.1)];

Yt = transpose(float.(Y));

n_particles = 2000
threshold = 80.0
sim_rej_result = SimulatedABCRejection(
    Yt, # data
    simdata, # simulator
    priors, # priors
    threshold, # threshold distance
    n_particles; # particles required
    max_iter=convert(Int, 1e7),
    distance_function = Distances.euclidean,
    write_progress=false);

plot(sim_rej_result)

n_design_points = 500
emu_rej_result = EmulatedABCRejection(Yt,
    simdata,
    priors,
    threshold,
    n_particles,
    n_design_points;
    max_iter=convert(Int, 1e7),
    distance_function = Distances.euclidean,
    write_progress=false);

plot(emu_rej_result)

threshold_schedule = [110.0,100.0,90.0,80.0];

sim_smc_result = SimulatedABCSMC(Yt,
    simdata,
    priors,
    threshold_schedule,
    n_particles;
    max_iter=convert(Int, 1e7),
    distance_function = Distances.euclidean,
    write_progress=false);

population_colors=["#FF2F4E", "#D0001F", "#A20018", "#990017"]
plot(sim_smc_result, population_colors=population_colors)

emu_smc_result = EmulatedABCSMC(Yt,
    simdata,
    priors,
    threshold_schedule,
    n_particles,
    n_design_points;
    distance_metric = Distances.euclidean,
    batch_size=1000,
    write_progress=false,
    emulator_retraining = PreviousPopulationThresholdRetraining(n_design_points, 100, 10),
    emulated_particle_selection = MeanVarEmulatedParticleSelection());

plot(emu_smc_result, population_colors=population_colors)

function simdist(x, constants, y)
  s = transpose(simdata(x))
  Distances.euclidean(s, y), 1
end;

ab_rej_setup = ABCRejection(simdist, #simulation function
  2, # number of parameters
  threshold, #target ϵ
  Prior(priors); # Prior for each of the parameters
  maxiterations = 10^7, #Maximum number of iterations before the algorithm terminates
  nparticles = n_particles
  );

ab_rej = runabc(ab_rej_setup,
            Y,
            verbose = true,
            progress = true,
            parallel = true);

plot(ab_rej)

ab_smc_setup = ABCSMC(simdist, #simulation function
  2, # number of parameters
  threshold, #target ϵ
  Prior(priors), #Prior for each of the parameters
  maxiterations=convert(Int,1e7),
  nparticles=n_particles,
  α = 0.3,
  convergence = 0.05,
  kernel = uniformkernel
  );

ab_smc = runabc(ab_smc_setup,
            Y,
            verbose = true,
            progress = true,
            parallel = true);

plot(ab_smc)


"""
    angularspectrum(Uin, d, λ, lx, ly)

calculate the propagation light field based on the angular spectrum.

## Arguments

- `Uin::AbstractArray{<:Number,2}`: Complex array of input complex amplitude.
- `d::Real`: Distance to propagate in metres.
- `λ::Real`: Wavelength of light to propagate.
- `lx::Real`: The physical size of the input data along the x-axis.
- `ly::Real`: The physical size of the input data along the y-axis.

## Returns

- `::Array{<:Number,2}`: Complex amplitude data after propagation.
"""
function angularspectrum(Uin::AbstractArray{<:Number,2}, d::Real, λ::Real, lx::Real, ly::Real)
    (n, m) = size(Uin)
    ua = (-n/2:n/2-1)/lx*λ
    va = (-m/2:m/2-1)/ly*λ
    ifft(fft(Uin).*ifftshift([1.0-u^2-v^2>=0.0 ? exp(2im*pi*d/λ*sqrt(1.0-u^2-v^2)) : 0.0+0.0im for u in ua, v in va]))
end

"""
    angularspectrum(Uin::MonoLightField2D, d::Unitful.Length) -> MonoLightField2D

calculate the propagation light field for [`MonoLightField2D`](@ref) based on the angular spectrum.
"""
function angularspectrum(Uin::MonoLightField2D, d::Unitful.Length)
    Uout = angularspectrum(Uin.data, auval(d), auval(Uin.wavelength), auval(Uin.size[1]), auval(Uin.size[2]))
    MonoLightField2D(Uin, data=Uout)
end

"""
    angularspectrum(d)

represent an [`angularspectrum`](@ref) propagation functor with the specified distance.
"""
angularspectrum(d) = x->angularspectrum(x, d)


module InclusionsGenerator

export generate_inclusions

# sphere volume
vol(r) = 4/3*π*r^3

# random inclusoin center
new_center(lx, ly, lz) = (rand()*lx, rand()*ly, rand()*lz)

# 3D euclidean distance
distance(p1, p2) = sqrt((p1[1]-p2[1])^2 + (p1[2]-p2[2])^2 + (p1[3]-p2[3])^2)

# check if the new inclusion collides with any pre-existent inclusion
function sphere_collision(coords, new_coord, r; buffer = 0.2)
    for c in coords
        if distance(c, new_coord) ≤ 2*r*(1 + buffer)
            return true
        end
    end

    return false
end

function wall_collision(new_coord, r, lx, ly, lz; buffer = 0.2)
    x, y, z = new_coord
    buffer += 1
    if (x ≤ r*buffer) || (x ≥ (lx - r*buffer))
        return true

    elseif (y ≤ r*buffer) || (y ≥ (ly - r*buffer))
        return true

    elseif (z ≤ r*buffer) || (z ≥ (lz - r*buffer))
        return true
    end

    return false
end

function generate_inclusions(lx, ly, lz, ϕ, r;  buffer_ball = 0.2, buffer_wall = 0.2)
    # first sphere coordinate
    coords = [new_center(lx, ly, lz)]

    tot_vol = vol(r)
    vol_fraction = tot_vol/(lx*ly*lz)

    iter = 0
    while vol_fraction ≤ ϕ
        iter += 1
        # temptative coordinate
        new_coord = new_center(lx, ly, lz)
        collision1 = sphere_collision(coords, new_coord, r; buffer=buffer_ball)
        collision2 = wall_collision(new_coord, r, lx, ly, lz; buffer=buffer_wall)
        if !collision1 && !collision2
            push!(coords, new_coord)
            tot_vol += vol(r)
            vol_fraction = tot_vol/(lx*ly*lz)
            iter = 0
        end
        if iter > 500
            println("Forcing exit. Generator stalling")
            break
        end
    end

    println("Done with $(length(coords)) inclusions and $(vol_fraction*100) % volume fraction.")

    return coords, vol_fraction
end

end # module


using Nemo

# _Fast_ composed sums and composed products of polynomials,
# using the algorithm described in:
# "Fast computation of special resultants"
# by Bostan, Flajolet, Salvy, and Schost

# derivative of polynomial
derivative(c::Vector) = c[2:end] .* (1:length(c)-1)

function polyinv(coeffs::Vector, n)
	R, x = Nemo.PowerSeriesRing(Nemo.FlintQQ, n, "x")
	a = R(map(Nemo.FlintQQ, coeffs), length(coeffs), n, 0)
	ai = inv(a)
	return Nemo.fmpq[coeff(ai,i) for i=0:n-1]
end

# compute newton power series of polynomial given with coefficients coeff,
# in base field R,x.
# See fig.1 in reference
function to_newton(coeffs::Vector{BigInt},n,R,x)
	# first, make monic.
	coeffs = coeffs//coeffs[end]

	d = length(coeffs)-1
	a_cfs = reverse(derivative(coeffs))
	b_cfs = reverse(coeffs)

	# initialize power series polynomials
	a = R(map(Nemo.FlintQQ, a_cfs))
	b = R(map(Nemo.FlintQQ, b_cfs))
	b0 = R(polyinv(b_cfs, n))

	c  = truncate(a*b0, d)

	r = R()
	x_power = R(1)
	x_d = x^d

	l = round(Int64,floor(n/d))
	for j = 0 : l
		r += c*x_power
		x_power *= x_d
		c = -mullow(shift_right(b*c,d),b0,d)
	end
	return r
end

to_array(p) = Rational{BigInt}[Rational(coeff(p,i)) for i=0:Nemo.degree(p)]

# tr: traces i.e. newton series
# This algorithm is based on the Leverrier-Faddeev algorithm
# see: http://math.stackexchange.com/questions/405822/what-is-the-fastest-way-to-find-the-characteristic-polynomial-of-a-matrix
function from_newton(tr::Vector{T}) where {T<:Number}
	# special case
	if tr==[1]
		return [0,1]
	end
	n = length(tr)
	c = Array{T}(UndefInitializer(),n)
	c[end] = one(T)
	for k = 1 : n-1
		next_c = -sum(tr[2:(k+1)].*c[end-k+1:end])/k
		c[end-k] = next_c
	end
	return c
end

# Hadamard (element-wise) product of two polynomials
function hadm(p,q,R)
	n = max(Nemo.degree(p),Nemo.degree(q))
	R([Nemo.coeff(p,i)*Nemo.coeff(q,i) for i=0:n])
end

# composed product of two polynomials, given as coeffs p and q
function composed_product(p::Vector{BigInt},q::Vector{BigInt})
	# compute newton series
	n = (length(p)-1)*(length(q)-1)+1
	R, x = Nemo.PolynomialRing(Nemo.FlintQQ, "x")
	a = to_newton(p,n,R,x)
	b = to_newton(q,n,R,x)

	# multiply newton series and invert
	pq = from_newton(to_array(hadm(a,b,R)))

	# convert to integer and return
	return map(numerator, pq*lcm(map(denominator, pq)))
end

# composed sum of two polynomials, given as coeffs p and q
function composed_sum(p::Vector{BigInt},q::Vector{BigInt})
	# compute newton series
	n = (length(p)-1)*(length(q)-1)+1
	R, x = Nemo.PolynomialRing(Nemo.FlintQQ, "x")
	a = to_newton(p,n,R,x)
	b = to_newton(q,n,R,x)

	# exp series
	ee  = R([Nemo.FlintQQ(1//factorial(BigInt(i))) for i=0:(n-1)])
	eei = R([Nemo.FlintQQ(   factorial(BigInt(i))) for i=0:(n-1)])

	# multiply newton series and invert
	m = truncate(hadm(a,ee,R)*hadm(b,ee,R),n)
	pq = from_newton(to_array(hadm(m,eei,R)))

	# convert to integer and return
	return map(numerator, pq*lcm(map(denominator, pq)))
end


"""Julia program to implement Minimum XOR value.
Given an array of integers, we have to find the minimum XOR value, a pair in the array makes."""

function minimumXOR(arr, n)
    # Sort the array, so that elements with minimum bit difference are present consecutive.
    sort(arr)
    minXOR = typemax(Int64)
    val = 0

    # Compute the XOR values of the consecutive elements of the array.
    for i in 1:(n-1)
        val = arr[i] ⊻ arr[i+1]
        minXOR = min(minXOR, val)
    end

    return minXOR
end

print("How many numbers are present in the array? ")
n = readline()
n = parse(Int, n)
if (n <= 0)
    println("No numbers are present in the array!!!")
    exit()
end
arr = Int[]
print("Enter the numbers: ")
arr = [parse(Int, num) for num in split(readline())]
res = minimumXOR(arr, n)
print("The minimum XOR value present in the given array is $res")

"""
SAMPLE INPUT AND OUTPUT

SAMPLE I

How many numbers are present in the array? 8
Enter the numbers: 25 45 36 458 651 89 52 95
The minimum XOR value present in the given array is 9

SAMPLE II

How many numbers are present in the array? 5
Enter the numbers: 5 4 3 2 1
The minimum XOR value present in the given array is 1

"""


# ------------------------------------------------------------------- #
# CES Producer
# ------------------------------------------------------------------- #
immutable CESProducer <: AbstractProducer
    η::Float64
    ν::Float64
end

# call(p::CESProducer, k, l) = (p.η*k.^p.ν + (1-p.η)*l.^(p.ν)).^(1./p.ν)
Base.produce(p::CESProducer, k, l) = (p.η*k.^p.ν + (1-p.η)*l.^(p.ν)).^(1./p.ν)
f_k(p::CESProducer, k, l) = p.η * (produce(p, k, l) ./ k).^(1.-p.ν)

# ------------------------------------------------------------------- #
# Cobb-Douglas Producer
# ------------------------------------------------------------------- #
immutable CobbDouglassProducer <: AbstractProducer
    α_k::Float64
    α_l::Float64
end

is_crs(i::CobbDouglassProducer) = i.α_k == 1. - i.α_l


struct ShamirSecretSharing
    prime::BigInt
end

ShamirSecretSharing() = ShamirSecretSharing(BigInt(2)^127 - 1)

"""
n is the number of shares.
k is the threshold.

"""
function Base.split(scheme::ShamirSecretSharing, secret::BigInt, n::Int64, k::Int64)
    prime = scheme.prime
    polynomial = [secret]
    for i in 2:k
        push!(polynomial, rand(RandomDevice(), 0:(prime-1)))
    end

    points = []
    for i in 1:n
        push!(points, (i, eval_poly(polynomial, i, prime)))
    end
    return points
end

function Base.join(scheme::ShamirSecretSharing, points)
    prime = scheme.prime
    xs, ys = zip(points...)
    return lagrange_interpolate(0, [xs...], [ys...], prime)
end


# https://juliasmoothoptimizers.github.io/ADNLPModels.jl/stable/tutorial/

using ADNLPModels
using NLPModelsIpopt

f(x) = (x[1] - 1)^2 + 100*(x[2] - x[1]^2)^2

x0 = [-1.2; 1.0]
uvar = [10.0; 10.0]
lvar = [-10.0; -10.0]

c(x) = [x[1]^2 + x[2]^2; x[1]*x[2]]
ucon = [0.75; 0.3]
lcon = [0.0; 0.0]

nlp = ADNLPModel(f, x0, lvar, uvar, c, lcon, ucon)

#output = ipopt(nlp, print_level=0)
output = ipopt(nlp)

println(output.objective)
println(output.solution)


function Green(xr,xz,yr,yz; derivative=0)
    #=
    Fundamental solution (Green function) and derivatives of the Grad-Shafranov Delta^* operator
    -Delta_x^* G(x,y)=delta(x-y)
    y=(yr,yz) is the source point (a filament with unit current)
    x=(xr,xz) is the evaluation point
    Derivatives from Itagaki, Engineering Analysis with Boundary Elements, 33 (2009) 845-857

    derivative=0 only Green function value
    derivative=1 + gradient wrt to x and y
    derivative=2 + second order cross derivatives (see return outputs below)
    =#

    mu0=4*pi*10^-7
    #modulus at points x and y
    k_xy=sqrt(4*xr*yr/((xr+yr)^2+(xz-yz)^2))

    #Complete elliptic integrals of first and second kind
    K_xy,E_xy=Elliptic.ellipke(k_xy^2)

    #Green function
    G_xy = (mu0/pi*((xr*yr)^(1/2))/k_xy)*((1-0.5*k_xy^2)*K_xy - E_xy)

    #println("derivative=",derivative)

    if(derivative==0)
        return (G_xy=G_xy,)
        #return G_xy
    end

    if(derivative>0)
        # gradient with respect to y
        dyrG_xy = (mu0/(2*pi))*((yr   )/(sqrt( (xr+yr)^2 + (xz-yz)^2 )))*(K_xy - ((yr^2-xr^2+(xz-yz)^2 )/( (xr-yr)^2 + (xz-yz)^2))*E_xy)
        dyzG_xy = (mu0/(2*pi))*((yz-xz)/(sqrt( (xr+yr)^2 + (xz-yz)^2 )))*(K_xy - ((yr^2+xr^2+(xz-yz)^2 )/( (xr-yr)^2 + (xz-yz)^2))*E_xy)

        # gradient with respect to x
        dxrG_xy = (mu0/(2*pi))*((xr   )/(sqrt( (xr+yr)^2 + (xz-yz)^2 )))*(K_xy - ((xr^2-yr^2+(xz-yz)^2 )/( (xr-yr)^2 + (xz-yz)^2))*E_xy)
        dxzG_xy = (mu0/(2*pi))*((xz-yz)/(sqrt( (xr+yr)^2 + (xz-yz)^2 )))*(K_xy - ((xr^2+yr^2+(xz-yz)^2 )/( (xr-yr)^2 + (xz-yz)^2))*E_xy)

        if(derivative==1)
            #return G_xy, dyrG_xy, dyzG_xy, dxrG_xy, dxzG_xy
            return (G_xy=G_xy, dyrG_xy=dyrG_xy, dyzG_xy=dyzG_xy, dxrG_xy=dxrG_xy, dxzG_xy=dxzG_xy)
        end

    end

    if(derivative==2)

        # second order cross derivatives
        dxrdyrG_xy = -((xr+yr)/((xr+yr)^2+(xz-yz)^2))*dxrG_xy +
            (mu0/2*pi)*xr*(1/((xr+yr)^2+(xz-yz)^2)^(1/2))*(1/((xr-yr)^2+(xz-yz)^2))*
            ( ((xr^2-yr^2+(xz-yz)^2)/((xr+yr)^2+(xz-yz)^2))*((xr-yr)*K_xy+(xr+yr)*E_xy) -2*xr*E_xy + ((4*yr*(xz-yz)^2)/((xr-yr)^2+(xz-yz)^2))*E_xy )

        dxzdyrG_xy = -((xr+yr)/((xr+yr)^2+(xz-yz)^2))*dxzG_xy +
            (mu0/2*pi)*xr*(xz-yz)*(1/((xr-yr)^2+(xz-yz)^2)^(1/2))*(1/((xr-yr)^2+(xz-yz)^2))*
            ( ((xr^2-yr^2+(xz-yz)^2)/((xr+yr)^2+(xz-yz)^2))*(K_xy+E_xy) -2*E_xy - ((4*yr*(xr-yr))/((xr-yr)^2+(xz-yz)^2))*E_xy )

        dxrdyzG_xy = ((xz-yz)/((xr+yr)^2+(xz-yz)^2))*dxrG_xy +
            (mu0/pi)*yr*xr*(xz-yz)*
            (1/((xr+yr)^2+(xz-yz)^2)^(1/2))*(1/((xr-yr)^2+(xz-yz)^2))*
            ( (1/((xr+yr)^2+(xz-yz)^2))*((xr-yr)*K_xy+(xr+yr)*E_xy) +
            ((-2*(xr-yr))/((xr-yr)^2+(xz-yz)^2))*E_xy )

        dxzdyzG_xy = (((-1)/(xz-yz))+((xz-yz)/((xr+yr)^2+(xz-yz)^2)))*dxzG_xy +
            (mu0/pi)*yr*xr*(xz-yz)^2 *
            (1/((xr+yr)^2+(xz-yz)^2)^(1/2))*(1/((xr-yr)^2+(xz-yz)^2))*
            ( (1/((xr+yr)^2+(xz-yz)^2))*(K_xy+E_xy) +
            (-2/((xr-yr)^2+(xz-yz)^2))*E_xy )

        #return G_xy, dyrG_xy, dyzG_xy, dxrG_xy, dxzG_xy, dxrdyrG_xy, dxzdyrG_xy, dxrdyzG_xy, dxzdyzG_xy
        return (G_xy=G_xy, dyrG_xy=dyrG_xy, dyzG_xy=dyzG_xy, dxrG_xy=dxrG_xy, dxzG_xy=dxzG_xy, dxrdyrG_xy=dxrdyrG_xy, dxzdyrG_xy=dxzdyrG_xy, dxrdyzG_xy=dxrdyzG_xy, dxzdyzG_xy=dxzdyzG_xy)
    end

end


# This function implements Kalman Filter on a state-space model of the form:
# α_{t} = T * α_{t-1} + η_{t}
# y_{t} = Z * α_{t} + ϵ_{t}
# E[η_{t}η'_{t}] = Q, E[ϵ_{t}ϵ'_{t}] = H

function KalmanFilter(T::Matrix, Z::Matrix, Q::Matrix, H::Matrix, n::Integer, α0, y)

    # Initialize state variables
    αhat = repeat([zeros(size(α0))], n);
    αhat[1] = T*α0;

    yhat = repeat([zeros(size(y[1]))], n);
    v = copy(yhat);

    P0 = Q * Q;
    d = 10;
    tol = 10^(-8);
    maxiter = 1000;
    iter = 1;
    while d > tol && iter <= maxiter
        P1 = T * P0 * T' + Q;
        d = maximum(abs.(P0 - P1));
        P0 = P1;
        iter = iter + 1;
    end

    P = repeat([zeros(size(Q))], n);
    P[1] = T * P0 * T' + Q;

    F = repeat([zeros(size(H))], n)
    G = repeat([zeros(size(Z * Q))], n)
    K = repeat([zeros(size(T * G[1]' * ones(size(H))))], n)

    for i in 1:n-1
        yhat[i] = Z * αhat[i];
        v[i] = y[i] - yhat[i];

        F[i] = Z * P[i] * Z' + H;
        G[i] = Z * P[i];
        K[i] = T * G[i]' * inv(F[i]);

        αhat[i + 1] = T * αhat[i] + K[i] * v[i];
        P[i + 1] = T * (P[i] - G[i]' * inv(F[i]) * G[i]) * T' + Q;
    end

    yhat[n] = Z * αhat[n];
    v[n] = y[n] - yhat[n];

    F[n] = Z * P[n] * Z' + H;
    G[n] = Z * P[n];
    K[n] = T * G[n]' * inv(F[n]);

    return αhat, v, yhat, P, F, G, K
end


# A tolerance graph is a graph whose vertices are represented by pairs
# (J,t) where J is a closed interval and t is a real number. Vertices
# are adjacent if and only if length(J * JJ) >= min(t,tt).

export ToleranceGraph

"""
`ToleranceGraph(Jlist, tlist)` creates a tolerance graph where `Jlist`
is a list of intervals and `tlist` is a list of tolerances.
"""
function ToleranceGraph(
    Jlist::Vector{ClosedInterval{S}},
    tlist::Vector{T},
) where {S<:Real,T<:Real}
    n = length(Jlist)
    G = IntGraph(n)

    for i = 1:n-1
        J = Jlist[i]
        t = tlist[i]
        for j = i+1:n
            JJ = Jlist[j]
            tt = tlist[j]

            K = J * JJ
            if !isempty(K) && length(K) >= min(t, tt)
                add!(G, i, j)
            end
        end
    end
    return G
end

"""
`ToleranceGraph(f)` creates a tolerance graph where `f` is a `Dict`
mapping vertices to pairs `(J,t)` where `J` is a `ClosedInterval` and
`t` is a `Real` tolerance.
"""
function ToleranceGraph(f::Dict{X,Tuple{ClosedInterval{S},T}}) where {S<:Real,T<:Real,X}
    vtcs = collect(keys(f))
    G = SimpleGraph{X}()
    for v in vtcs
        add!(G, v)
    end

    n = length(vtcs)
    for i = 1:n-1
        v = vtcs[i]
        J, t = f[v]
        for j = i+1:n
            w = vtcs[j]
            JJ, tt = f[w]

            K = J * JJ
            if !isempty(K) && length(K) >= min(t, tt)
                add!(G, v, w)
            end
        end
    end
    return G
end


# calculate varianve of sensor data

module CalcVariance
    using DataFrames, CSV, Statistics

    function main()
        data_path = joinpath(split(@__FILE__, "src")[1], "data/sensor_data_200.txt")
        df_200_mm = CSV.read(data_path, DataFrame,
                             header=["date", "time", "ir", "lidar"],
                             delim=' ')

        # calculate from definition
        zs = df_200_mm.lidar # observation
        mean_def = sum(zs) / length(zs) # mean
        diff_square = [(z - mean_def)^2 for z in zs]

        sampling_var = sum(diff_square) / length(zs) # sampling variance
        unbiased_var = sum(diff_square) / (length(zs) - 1) # unbiased variance

        println("Sampling Variance = $sampling_var")
        println("Unbiased Variance = $unbiased_var")

        # calculate by Statistics
        stats_sampling_var = Statistics.var(zs, corrected=false)
        stats_unbiased_var = Statistics.var(zs, corrected=true)

        println("Sampling Variance by Statistics = $stats_sampling_var")
        println("Unbiased Variance by Statistics = $stats_unbiased_var")
    end
end


# Linear Isotropic Covariance Function

@doc """
# Description
Constructor for the isotropic linear kernel (covariance)

k(x,x') = xᵀx'/ℓ²
# Arguments:
* `ll::Float64`: Log of the length scale ℓ
""" ->
type LinIso <: Kernel
    ll::Float64      # Log of Length scale
    LinIso(ll::Float64) = new(ll)
end

function cov(lin::LinIso, x::Vector{Float64}, y::Vector{Float64})
    ell = exp(lin.ll)
    K = dot(x,y)/ell^2
    return K
end

function cov(lin::LinIso, X::Matrix{Float64}, data::EmptyData)
    ell = exp(lin.ll)
    return ((1/ell^2).*X') * X
end

get_params(lin::LinIso) = Float64[lin.ll]
get_param_names(lin::LinIso) = [:ll]
num_params(lin::LinIso) = 1

function set_params!(lin::LinIso, hyp::Vector{Float64})
    length(hyp) == 1 || throw(ArgumentError("Linear isotropic kernel only has one parameter"))
    lin.ll = hyp[1]
end

function grad_kern(lin::LinIso, x::Vector{Float64}, y::Vector{Float64})
    ell = exp(lin.ll)

    dK_ell = -2.0*dot(x,y)/ell^2
    dK_theta = [dK_ell]
    return dK_theta
end

function grad_stack!(stack::AbstractArray, lin::LinIso, X::Matrix{Float64}, data::EmptyData)
    ell = exp(lin.ll)
    stack[:,:,1] = ((-2.0/ell^2).*X') * X
end


function CholeskyL(A::Array)

  #A = the matrix of coefficients 'A' must be squared

    m,n =size(A)

    if m != n
        error("Matrix must be squared")
    end

    # System of equations
    L=zeros(m,m)

    L[1,1]=A[1,1]^0.5

    for k=2:m

    for i=1:k-1
       L[k,i]=(A[k,i]-(L[i,1:i-1]'*L[k,1:i-1]))/(L[i,i])

    end
          L[k,k]=(A[k,k]-(L[k,1:k-1]'*L[k,1:k-1]))^0.5
    end

    return L

end


module FeynmanKleinert

# Bring out the major leagues... https://www.youtube.com/watch?v=_E6DDktoPhg
using Optim

using QuadGK # For numerical integration

export K,Va2,Wtilde,verboseWtilde,W # export all the things.

"Integrand of (4) in Feynman and Kleinert"
K(xp,x,a2,V)=1/sqrt(2*π*a2)*exp(-(x-xp)^2/(2*a2))*V(xp)
"(4) in Feynman and Kleinert"
Va2(x,a2,V)=quadgk(xp->K(xp,x,a2,V), -Inf, +Inf) #, reltol=0, abstol=1e-5)

# (5) naively as written, in Feynman and Kleinert
#Wtilde(x0,a2,Ω,g,β)=(1/β) * log( sinh(β*Ω/2)/(β*Ω/2) ) - (Ω^2/2) * a2 + Va2(x0,a2,g)[1]
# Large β limit; by taking sinh(x)=e^x-e^-x , and then dropping e^-x
#Wtilde(x0,a2,Ω,g,β)=(1/β) * (β*Ω/2 - log(β*Ω/2)) - (Ω^2/2) * a2 + Va2(x0,a2,g)[1]

# Branch depending on β; for accuracy + lack of infinities
"""
    Wtilde(x0,a2,Ω,V,β)

(5) in Feynman and Kleinert. Switches depending on D=β*Ω/2 to use either an
exponential approximation for sinh(x) (D>7.5) or an explicit evaluation and
then take logarithms. Without doing this you get collapse to infinity for large
β.
"""
function Wtilde(x0,a2,Ω,V,β)
    D=β*Ω/2
    if D>7.5
        F=(1/β) * (D - log(2*D)) # Large limit;  by taking sinh(x)=e^x-e^-x , and then dropping e^-x
    else
        F=(1/β)*( log(sinh(D)/D) ) # explicit evaluation, naively as written in (5)
    end
    V0=Va2(x0,a2,V)[1]
    mid=(Ω^2/2) * a2
    return F-mid+V0
end

"""
   verboseWtilde(x0,a2,Ω,g,β)

Expanded version of Wtilde, reporting the decomposition - this was used to explain what was going wrong
"""
function verboseWtilde(x0,a2,Ω,g,β)
    D=β*Ω/2
    println("D=β*Ω/2 = $D")
    F= (1/β)*( log(sinh(D)) - log(D) )
    #F= (1/β)*( (log(2D)+D^2/6+D^4/180) - log(D) ) # generalised Puiseux series
    # Large D limit, through sinh(D) --> e^D [through away e^-D part]
    F = (1/β) * (D - log(2*D))
    V=Va2(x0,a2,g)[1]
    mid=(Ω^2/2) * a2
    println("Wtilde(x0=$x0,a2=$a2,Ω=$Ω,g=$g,β=$β) \n\t= F - mid + V")
    println("\t= $F - $mid + $V = ",F-mid+V)
    return F-mid+V
end

"""
    W(x0,g,β)

    Implementing optimisation / minimisation described in Feynman-Kleinert (6)
    Calculates optimal W for x0, with given g and Beta
    Use's (7) (a2 in terms of omega) to fix numerical noise from trying to do
    it directly...
"""
function W(x0,g,β)
    #println("W(x0=$x0,g=$g,β=$β)")

    a2(Ω)=1/(β*Ω^2)*( (β*Ω)/2 * coth(β*Ω/2)-1 ) #(7)
    myf(x)=Wtilde(x0,a2(x[1]),x[1],g,β)

    # Define bounds on solution
    lower=[0.0]
    upper=[5.0]
    initial=[1.0]

    res=optimize(OnceDifferentiable(myf, initial; autodiff = :forward), initial, lower, upper, Fminbox();
    optimizer=GradientDescent)

    minimum=Optim.minimum(res)
    Ω=Optim.minimizer(res)[1]
    mya2=a2(Ω)[1]
    #println("Optimised for W(x0=$x0,g=$g,β=$β). \n\tΩ=$Ω, a2=$mya2, minimum=$minimum.")
    #show(res)
    return minimum
end

end # module


# A comparison between ℓ_0 and ℓ_1 trend filtering

using Plots
using Convex
using Mosek
using EllZeroTrendFiltering

g = snp500_data()[1:300]

N = length(g)

## ℓ_1 trend filtering, see "ℓ_1 Trend Filtering", Kim et al. (2009)
H =  spdiagm(0=>ones(N-2), 1=>-2*ones(N-2), 2=>ones(N-2))

x = Variable(N)

λ = 2.0
problem = minimize(sumsquares(x - g) + λ * norm(H*x, 1))

solve!(problem, Mosek.Optimizer)
#solve!(problem, SCS.Optimizer)
println(problem.status)

x_l1 = evaluate(x)[:]
f_l1 = sum((x_l1 - g).^2)

## ℓ_0 trend filtering
ζ = 0.02
I_l0, Y_l0, f_l0 = fit_pwl_regularized(g, ζ)

plot(g, label="SP500", size=(1000,550));
plot!(x_l1, lw=2.5, label="l1 reg. lambda = $λ");
plot!(I_l0, Y_l0, lw=2.5, label="l0 reg. zeta = $ζ")

# plot(H*x.value);
# plot!(I_l0[2:end], 0.1*diff(Y_l0) ./ diff(I_l0), m=:o, lw=0)
#


# Generate random unitaries and density matrices for testing

export random_unitary, random_densitymatrix

""" `U = random_unitary(dim, real=true/false, complex=true/false)`

Generate a random unitary of dimension *dim*. Depending on *real* and *complex* the matrix has real and complex parts.

"""

function random_unitary(dim::Integer; real::Bool = true, complex::Bool = false)

  	@assert dim > 1 "Dimension must be at least 2."

  	generated = randn(dim, dim)

  	if complex || !real
    		generated += im * randn(dim, dim)
  	end

	# Compute the QR decomposition
  	(Q, R) = qr(generated)

	# Get the upper diagonal part of R
  	F = diagm(diag(R))

  	R = sign(diagm(diag(R)))
  	R[find((x)->  x == 0, R)] = 1
  	F = diagm(diag(R))

  	return Q*F
end

""" `rho = random_densitymatrix(dim)`

Generate a random density matrix in dimension *dim*.

"""

function random_densitymatrix(dim::Integer)

	@assert dim > 1 "Dimension must be at least 2."

        # Draw random numbers
	rList = rand(dim);

	# Normalize to form a distribution
	s = sum(rList);
	nList = rList/s;

	# Construct a matrix with the numbers on the diagonal
	M = diagm(nList);

	# Draw a random unitary
   	U = random_unitary(dim, real = true, complex = true)

	# Roate the diagonal matrix
	return(U*M*U');
end


# All matrices are DxN, where N is the number of positions and D is the dimensionality

# Here, P is the probe (to be rotated) and Q is the refereence
# https://en.wikipedia.org/wiki/Kabsch_algorithm
# This has been generalized to support weighted points:
# https://igl.ethz.ch/projects/ARAP/svd_rot.pdf

# assuming P and Q are already centered at the origin
# returns the rotation for alignment
function kabsch_centered(P,Q,w)
    @assert size(P) == size(Q)
    W = diagm(w/sum(w)) # here, the weights are assumed to sum to 1
    H = P*W*Q'
    D = Matrix{Float64}(I,size(H,1), size(H,2))
    U,Σ,V = svd(H)
    D[end] = sign(det(V*U'))
    return LinearMap(V * D * U')
end

# transform DxN matrices
function (tform::Translation)(A::AbstractMatrix)
    return hcat([tform(A[:,i]) for i=1:size(A,2)]...)
end

# P and Q are not necessarily centered
# returns the transformation for alignment
function kabsch(P, Q, w::AbstractVector=ones(size(P,2)))
    @assert !any(w.<0) && sum(w)>0
    wn = w/sum(w) # weights should sum to 1 for computing the centroid
    centerP, centerQ = center_translation(P,wn), center_translation(Q,wn)
    R = kabsch_centered(centerP(P), centerQ(Q), wn)
    return inv(centerQ) ∘ R ∘ centerP
end

function kabsch(P,Q,matches::AbstractVector{<:Tuple{Int,Int}},wp=ones(size(P,2)),wq=ones(size(Q,2)))
    matchedP, matchedQ = matched_points(P,Q,matches)
    w = [wp[i]*wq[j] for (i,j) in matches]
    return kabsch(matchedP, matchedQ, w)
end

# centroid of positions in A, weighted by weights in w (assumed to sum to 1)
centroid(A, w=fill(1/size(A,2), size(A,2))) = A*w

# translation moving centroid to origin
center_translation(A, w=fill(1/size(A,2), size(A,2))) = Translation(-centroid(A,w))


function rosenbrock(x, g=nothing)
    n = length(x)
    f = 0.0
    for i=1:(n-1)
        f += (1-x[i])^2 + 100*(x[i+1]-x[i]^2)^2
    end
    if g != nothing
        for i=1:n
            g[i] = ((i==n ? 0 : -2*(1-x[i]) - 400*(x[i+1]-x[i]^2)*x[i]) +
                    (i==1 ? 0 : 200*(x[i]-x[i-1]^2)))
        end
    end
    return f
end

function gtest(x=randn(10);eps=1e-5)
    g = similar(x)
    @show f0 = rosenbrock(x,g)
    for i=1:length(x)
        xi = x[i]
        x[i] = xi+eps; f1 = rosenbrock(x)
        x[i] = xi-eps; f2 = rosenbrock(x)
        x[i] = xi; df = (f1-f2)/(2*eps)
        @show (i,g[i],df)
    end
end

using Base.LinAlg

function eq10(f, x0; gtol=0.95, gamma=1e-8)
    g0, x1, rb, rr, r = [ zeros(x0) for i=1:5 ]
    f0 = f(x0, g0)
    n0 = vecnorm(g0)
    nf = cosa = 0
    while cosa < gtol
        f1 = f(axpy!(gamma, randn!(r), copy!(x1,x0)))
        nf += 1
        b = f1 - f0
        axpy!(b, r, rb)
        axpy!(1, r.*r, rr)
        g1 = rb ./ (gamma * rr)
        n1 = vecnorm(g1)
        cosa = dot(g0,g1)/(n0*n1)
    end
    return nf
end


# Bogumił Kamiński, 2021

# Codes for chapter 5

# Code for section 5.1.1

x = [1 2 3]
y = [1, 2, 3]
x * y

a = [1, 2, 3]
b = [4, 5, 6]
a * b

a .* b

map(*, a, b)
[a[i] * b[i] for i in eachindex(a, b)]

eachindex(a, b)

eachindex([1, 2, 3], [4, 5])

map(*, [1, 2, 3], [4, 5])

[1, 2, 3] .* [4, 5]

# Code for section 5.1.2

[1, 2, 3] .^ [2]

[1, 2, 3] .^ 2

[1, 2, 3, 4, 5, 6, 7, 8, 9, 10] .* [1 2 3 4 5 6 7 8 9 10]

["x", "y", "z"] .=> [sum minimum maximum]

abs.([1, -2, 3, -4])

abs([1, 2, 3])

string(1, 2, 3)

string.("x", 1:10)

f(i::Int) = string("got integer ", i)
f(s::String) = string("got string ", s)
f.([1, "1"])

# Code for section 5.1.3

in(1, [1, 2, 3])
in(4, [1, 2, 3])

1 in [1, 2, 3]
4 in [1, 2, 3]

in([1, 3, 5, 7, 9], [1, 2, 3, 4])

in([1, 3, 5, 7, 9], [1, 2, 3, 4, [1, 3, 5, 7, 9]])

in.([1, 3, 5, 7, 9], [1, 2, 3, 4])

in.([1, 3, 5, 7, 9], Ref([1, 2, 3, 4]))

isodd.([1, 2, 3, 4, 5, 6, 7, 8, 9, 10] .+ [1 2 3 4 5 6 7 8 9 10])

Matrix{Any}(isodd.([1, 2, 3, 4, 5, 6, 7, 8, 9, 10] .* [1 2 3 4 5 6 7 8 9 10]))

# Code for section 5.1.4

aq = [10.0   8.04  10.0  9.14  10.0   7.46   8.0   6.58
       8.0   6.95   8.0  8.14   8.0   6.77   8.0   5.76
      13.0   7.58  13.0  8.74  13.0  12.74   8.0   7.71
       9.0   8.81   9.0  8.77   9.0   7.11   8.0   8.84
      11.0   8.33  11.0  9.26  11.0   7.81   8.0   8.47
      14.0   9.96  14.0  8.1   14.0   8.84   8.0   7.04
       6.0   7.24   6.0  6.13   6.0   6.08   8.0   5.25
       4.0   4.26   4.0  3.1    4.0   5.39  19.0  12.50
      12.0  10.84  12.0  9.13  12.0   8.15   8.0   5.56
       7.0   4.82   7.0  7.26   7.0   6.42   8.0   7.91
       5.0   5.68   5.0  4.74   5.0   5.73   8.0   6.89]
using Statistics

mean.(eachcol(aq))

mean(eachcol(aq))

function R²(x, y)
    X = [ones(11) x]
    model = X \ y
    prediction = X * model
    error = y - prediction
    SS_res = sum(v -> v ^ 2, error)
    mean_y = mean(y)
    SS_tot = sum(v -> (v - mean_y) ^ 2, y)
    return 1 - SS_res / SS_tot
end

function R²(x, y)
    X = [ones(11) x]
    model = X \ y
    prediction = X * model
    SS_res = sum((y .- prediction) .^ 2)
    SS_tot = sum((y .- mean(y)) .^ 2)
    return 1 - SS_res / SS_tot
end

# Code for section 5.2

[]
Dict()

Float64[1, 2, 3]

Dict{UInt8, Float64}(0 => 0, 1 => 1)

UInt32(200)

Real[1, 1.0, 0x3]

v1 = Any[1, 2, 3]
eltype(v1)
v2 = Float64[1, 2, 3]
eltype(v2)
v3 = [1, 2, 3]
eltype(v3)
d1 = Dict()
eltype(d1)
d2 = Dict(1 => 2, 3 => 4)
eltype(d2)

p = 1 => 2
typeof(p)

# Code for section 5.2.1

[1, 2, 3] isa AbstractVector{Int}
[1, 2, 3] isa AbstractVector{Real}

AbstractVector{<:Real} == AbstractVector{T} where T<:Real

# Code for section 5.2.2

using Statistics
function ourcov(x::AbstractVector{<:Real},
                y::AbstractVector{<:Real})
    len = length(x)
    @assert len == length(y) > 0
    return sum((x .- mean(x)) .* (y .- mean(y))) / (len - 1)
end

ourcov(1:4, [1.0, 3.0, 2.0, 4.0])
cov(1:4, [1.0, 3.0, 2.0, 4.0])

ourcov(1:4, Any[1.0, 3.0, 2.0, 4.0])

x = Any[1, 2, 3]
identity.(x)
y = Any[1, 2.0]
identity.(y)

# Code for section 5.3

using Random
Random.seed!(1234);
cluster1 = randn(100, 5) .- 1
cluster2 = randn(100, 5) .+ 1

data5 = vcat(cluster1, cluster2)

using PyCall
manifold = pyimport("sklearn.manifold")

# Optional code to run if the pyimport("sklearn.manifold") fails
# There is no need to run it if the above operation worked
using Conda
Conda.add("scikit-learn")

tsne = manifold.TSNE(n_components=2, init="random",
                     learning_rate="auto", random_state=1234)
data2 = tsne.fit_transform(data5)

using Plots
scatter(data2[:, 1], data2[:, 2];
        color=[fill("black", 100); fill("gold", 100)],
        legend=false)


module BernsteinEllipses

export jouk, ijouk

"""
    rsmo(x)

Evaluate `√(x^2-1)` with branch cut along `[-1,1]`.

The function name is the abbreviation of "root of square minus one".
"""
rsmo(x) = rsmo(float(x))
rsmo(x::Union{T,Complex{T}}) where {T<:AbstractFloat} = ifelse(!signbit(real(x)),1,-1)*sqrt(x^2-1)

"""
    jouk(z)

Joukowsky map `(z+z^-1)/2`.
"""
jouk(z) = (z+inv(z))/2

"""
    ijouk(x; halfplane=Val(false), branch=Val(true)) -> z

Inverse Joukowsky map `z = x ± √(x^2-1)`. The branch to evaluate is seleced as
follows.

+-----------+---------------+----------------+
|           | `!halfplane`  |  `halfplane`   |
+-----------+---------------+----------------+
| ` branch` | `abs(z) >= 1` | `imag(z) >= 0` |
| `!branch` | `abs(z) <= 1` | `imag(z) <= 0` |
+-----------+---------------+----------------+

"""
ijouk(x; halfplane=Val(false), branch=Val(true)) = ijouk(x,halfplane,branch)
ijouk(x,::Val{false},::Val{true} ) = x + rsmo(x)
ijouk(x,::Val{false},::Val{false}) = x - rsmo(x)
ijouk(x,::Val{true} ,::Val{true} ) = x + im*sqrt(1-x^2)
ijouk(x,::Val{true} ,::Val{false}) = x - im*sqrt(1-x^2)

"""
   radius(x; kwargs...) = abs(ijouk(x; kwargs...))
"""
radius(x; kwargs...) = abs(ijouk(x; kwargs...))

"""
   semimajor(x; kwargs...)

Semi-major axis of the Bernstein ellipse through `x`.

See [`ijouk`](@ref) regardings `kwargs`.
"""
semimajor(x; kwargs...) = (w = abs(ijouk(x; kwargs...)); (w+inv(w))/2)

"""
   semiminor(x; kwargs...)

Semi-minor axis of the Bernstein ellipse through `x`.

See [`ijouk`](@ref) regardings `kwargs`.
"""
semiminor(x; kwargs...) = (w = abs(ijouk(x; kwargs...)); (w-inv(w))/2)

"""
    semiminor(x,xr)

Imaginary component of the point in the upper half-plane
where the line `{ x̃ | real(x̃) = xr }` intersects the Bernstein
ellipse through `x`.
"""
function semiminor(x::Number,xr::Real; kwargs...)
    a = semimajor(x; kwargs...)
    b = semiminor(x; kwargs...)
    return b*sqrt(1 - (xr/a)^2)
end

end # module


using StatsModelComparisons, StanSample
using CSV, Printf, StatsPlots

ProjDir = @__DIR__

df = CSV.read(joinpath(ProjDir, "roachdata.csv"), DataFrame)
df.roach1 = df.roach1 / 100

roaches1_stan = "
data {
  int<lower=0> N;
  int<lower=0> K;
  vector[N] exposure2;
  vector[N] roach1;
  vector[N] senior;
  vector[N] treatment;
  int y[N];
}
transformed data {
  vector[N] log_expo;
  log_expo = log(exposure2);
}
parameters {
  vector[K] beta;
}
transformed parameters {
   vector[N] eta;
   eta = log_expo + beta[1] + beta[2] * roach1 + beta[3] * treatment
                + beta[4] * senior;
}
model {
  y ~ poisson_log(eta);
}
generated quantities {
  vector[N] log_lik;
  for (i in 1:N)
    log_lik[i] = poisson_log_lpmf(y[i] | eta[i]);
}
";

data = (N = size(df, 1), K = 4, y = Int.(df.y), roach1=df.roach1,
  exposure2=df.exposure2, senior=df.senior, treatment=df.treatment)
n = size(df, 1)

tmpdir = mktempdir()
sm1 = SampleModel("roaches1", roaches1_stan; tmpdir)
rc1 = stan_sample(sm1; data)

if success(rc1)
  stan_summary(sm1, true)
  nt1 = read_samples(sm1)

  # Compute LOO and standard error
  log_lik1 = nt1.log_lik'
  loo1, loos1, pk1 = psisloo(log_lik1)
  elpd_loo1 = sum(loos1)
  se_elpd_loo1 = std(loos1) * sqrt(n)
  @printf(">> elpd_loo = %.1f, SE(elpd_loo) = %.1f\n", elpd_loo1, se_elpd_loo1)

  # Check the shape parameter k of the generalized Pareto distribution
  if all(pk1 .< 0.5)
      println("All Pareto k estimates OK (k < 0.5).")
  else
    pk_good = sum(pk1 .<= 0.5)
    pk_ok = length(pk1[pk1 .<= 0.7]) - pk_good
    pk_bad = length(pk1[pk1 .<= 1]) - pk_good - pk_ok
    println((good=pk_good, ok=pk_ok, bad=pk_bad, very_bad=sum(pk1 .> 1)))
  end
end

begin
  scatter(pk1, xlab="Datapoint", ylab="Pareto shape k",
    marker=2.5, lab="Pk points")
  hline!([0.5], lab="pk = 0.5");hline!([0.7], lab="pk = 0.7")
  hline!([1], lab="pk = 1.0")
  title!("PSIS diagnostic plot for poisson-log model.")
  savefig(joinpath(ProjDir, "diag_plot_1.png"))
end

#=
*Simple negative binomial regression example
*using the 2nd parametrization of the negative
*binomial distribution, see section 40.1-3 in the Stan
*reference guide
=#

roaches2_stan = "
data {
  int<lower=0> N;
  int<lower=0> K;
  vector[N] exposure2;
  vector[N] roach1;
  vector[N] senior;
  vector[N] treatment;
  int y[N];
}
transformed data {
  vector[N] log_expo;
  log_expo = log(exposure2);
}
parameters {
  real phi;
  vector[K] beta;
}
transformed parameters {
   vector[N] eta;
   eta = log_expo + beta[1] + beta[2] * roach1 + beta[3] * treatment
                + beta[4] * senior;
}
model {
  phi ~ normal(0, 10);
  beta[1] ~ cauchy(0,10);   //prior for the intercept following Gelman 2008
  for(i in 2:K)
   beta[i] ~ cauchy(0,2.5); //prior for the slopes following Gelman 2008
  y ~ neg_binomial_2_log(eta, phi);
}
generated quantities {
 vector[N] log_lik;
 for(i in 1:N){
  log_lik[i] <- neg_binomial_2_log_lpmf(y[i] | eta[i], phi);
 }
}
";

tmpdir=joinpath(ProjDir, "tmp")
sm2 = SampleModel("roaches2", roaches2_stan; tmpdir)
rc2 = stan_sample(sm2; data)

if success(rc2)
  read_summary(sm2, true)
  nt2 = read_samples(sm2)

  # Compute LOO and standard error
  log_lik2 = nt2.log_lik'
  loo2, loos2, pk2 = psisloo(log_lik2)
  elpd_loo2 = sum(loos2)
  se_elpd_loo2 = std(loos2) * sqrt(n)
  @printf(">> elpd_loo = %.1f, SE(elpd_loo) = %.1f\n", elpd_loo2, se_elpd_loo2)

  # Check the shape parameter k of the generalized Pareto distribution
  pk_qualify(pk2) |> display

end

begin
  scatter(pk2, xlab="Datapoint", ylab="Pareto shape k",
    marker=2.5, lab="Pk points")
  hline!([0.5], lab="pk = 0.5");hline!([0.7], lab="pk = 0.7")
  hline!([1], lab="pk = 1.0")
  title!("PSIS diagnostic plot for neg-binomial model.")
  savefig(joinpath(ProjDir, "diag_plot_2.png"))
end


"""
    isclockwise(polygon::Polygon)

The shoelace formula, also known as Gauss's area formula, is used to determine
whether polygon traversal is clockwise or counter-clockwise. If the sum over the
edges is greater than 0, the curve is clockwise. If the sum over the edges is
lesser than 0, the curve is counter-clockwise. When the sum over the edges
equals 0, the polygon traverses like a figure eight.

Note: The right-handed coordinate system is used.
"""
function isclockwise(polygon::Polygon)
    sum([(b[1] - a[1]) * (b[2] + a[2]) for (a, b) in edges(polygon)]) > 0
end


#############
# In this example, we explore integration of a harmonic function:
#
#   f(x,y) = (x^2-y^2+1)/[(x^2-y^2+1)^2+(2xy+1)^2],
#
# over the unit disk. In this case, we know from complex analysis that the
# integral of a holomorphic function is equal to π × f(0,0).
# We analyze the function on an N×M tensor product grid defined by:
#
#   rₙ = cos[(n+1/2)π/2N], for 0 ≤ n < N, and
#
#   θₘ = 2π m/M, for 0 ≤ m < M;
#
# we convert the function samples to Chebyshev×Fourier coefficients using
# `plan_disk_analysis`; and finally, we transform the Chebyshev×Fourier
# coefficients to disk harmonic coefficients using `plan_disk2cxf`.
#
# For the storage pattern of the arrays, please consult the documentation.
#############

using FastTransforms, LinearAlgebra

f = (x,y) -> (x^2-y^2+1)/((x^2-y^2+1)^2+(2x*y+1)^2)

N = 5
M = 4N-3

r = [sinpi((N-n-0.5)/(2N)) for n in 0:N-1]
θ = (0:M-1)*2/M # mod π.

# On the mapped tensor product grid, our function samples are:
F = [f(r*cospi(θ), r*sinpi(θ)) for r in r, θ in θ]

P = plan_disk2cxf(F)
PA = plan_disk_analysis(F)

# Its Zernike coefficients are:
U = P\(PA*F)

# The Zernike coefficients are useful for integration. The integral of f(x,y)
# over the disk should be π/2 by harmonicity. The coefficient of Z_0^0
# multiplied by √π is:
U[1, 1]*sqrt(π)

# Using an orthonormal basis, the integral of [f(x,y)]^2 over the disk is
# approximately the square of the 2-norm of the coefficients:
norm(U)^2


function chebyshev_extrema(n::S,domain = [1.0,-1.0]) where S <: Integer

  if n <= 0
    error("The number of nodes must be positive.")
  end

  nodes = zeros(n)

  if isodd(n)
    nodes[Int((n-1)/2)+1] = (domain[1]+domain[2])/2.0
  end

  for i = 1:div(n,2)
    x = -cos((i-1)*pi/(n-1))*(domain[1]-domain[2])/2.0
    nodes[i]       = (domain[1]+domain[2])/2.0  + x
    nodes[end-i+1] = (domain[1]+domain[2])/2.0  - x
  end

  return nodes

end


export resampleα, randomindex

"""
	resampleα(α::Float64, N::Int, K::Int)

Resample α given the approach by Escobar and West. (page 585)
"""
function resampleα(α::Float64, N::Int, K::Int; k = 1.0, Θ = 1.0, maxiter = 1)

    w = zeros(2)
    w[2] = N

    f0 = k + K
    f1 = Θ + K - 1

    newα = α

    for i in 1:maxiter

        η = rand(Beta(newα + 1, N))
        f3 = Θ - log(η)

        w[1] = f1 / f3

        z = randomindex(w)

        if z == 1
            newα = rand(Gamma(f0, 1/f3))
        else
            newα = rand(Gamma(f1, 1/f3))
        end
    end

    return newα
end

function resampleα(α::Float64, N::Array{Int}, K::Array{Int}; k = 1.0, Θ = 1.0, maxiter = 1)
    # implementation according to escobar - west page
    totalK = sum(K)
    num = length(N)

    η = zeros(num)

		newα = α
    for i in 1:maxiter

        for j in 1:num
            η[j] = rand(Beta(newα + 1, N[j]))
        end

        z = rand(num) .* (newα + N) .< N

        g_a = k + totalK - sum(z)
        g_b = Θ - sum(log(η))

        newα = rand(Gamma(g_a, 1/g_b))
    end

    newα
end

"""
	randomindex(p::Vector{Float64})

Randomly select an index propotional to its probability.
"""
function randomindex(p::Vector{Float64})

    max = sum(p)
    csum = 0.0
    thres = rand()

    @inbounds for i in 1:length(p)
        csum += p[i]
        if csum >= thres * max
            return convert(Int, i)
        end
    end

    return length(p)
end

function randnumtable(weights::Array{Float64}, table::Array{Int})

    numtable = zeros(Int, size(table))

    B = unique( table )
    w = log( weights )

    for i in 1:length(B)

        max = B[i]
        if max > 0

            m = 1:max

            stirnums = map( x -> abs(stirlings1(max, x)), m)
            stirnums /= maximum(stirnums)

            for (idx, j) in enumerate(findall(table .== max))

                clike = m .* w[idx]
                clike = cumsum(stirnums .* exp(clike - maximum(clike)))

                numtable[j] = 1+sum(rand() * clike[max] .> clike)
            end

        end

    end

    return numtable
end


# interpolate point at (x, y) from 4 nearby pixel values
function interp_bilinear{T <: Union(Real, Integer)}(dat::Array{T, 2},
                                                    x::Float64, y::Float64,
                                                    x1, x2, y1, y2)
    q11 = dat[y1, x1]
    q12 = dat[y2, x1]
    q21 = dat[y1, x2]
    q22 = dat[y2, x2]
    if x1 != x2
        r1 = (x2 - x) / (x2 - x1) * q11 + (x - x1) / (x2 - x1) * q21
        r2 = (x2 - x) / (x2 - x1) * q12 + (x - x1) / (x2 - x1) * q22
    else
        # special case of x1 == x2, no interpolation needed
        r1 = q11
        r2 = q12
    end
    if y1 != y2
        r = (y - y1) / (y2 - y1) * r1 + (y2 - y) / (y2 - y1) * r2
    else
        # special case of y1 == y2, no interpolation needed
        r = r1
    end
    if typeof(r) <: Real
        r = round(r)
    end
    r = convert(T, r)
    return r
end

function imresize{T <: Union(Real, Integer)}(dat::Array{T, 2},
                                             new_size::(Int, Int))
    new_dat = similar(dat, new_size)
    h, w = size(dat)
    new_h, new_w = new_size
    for new_j=1:new_w, new_i=1:new_h
        # coordinates in original image
        x = new_j * w / new_w
        y = new_i * h / new_h
        # coordinates of 4 points to interpolate from
        x1, x2 = max(1, floor(x)), min(w, ceil(x))
        y1, y2 = max(1, floor(y)), min(h, ceil(y))
        new_dat[new_i, new_j] = interp_bilinear(dat, x, y, x1, x2, y1, y2)
    end
    return new_dat
end

function imresize{T <: Union(Real, Integer)}(dat::Array{T, 2},
                                             new_size...)
    return imresize(dat, new_size)
end

## if isinteractive()
##     using Images
##     using ImageView
##     im = imread("../../data/CK/faces_aligned/S999_003_00000054.png")
##     dat = convert(Array, im)
## end

using Images, Cairo

function imresize_julia!(resized, original)
    scale1 = (size(original,1)-1)/(size(resized,1)-0.999f0)
    scale2 = (size(original,2)-1)/(size(resized,2)-0.999f0)
    for jr = 0:size(resized,2)-1
        jo = scale2*jr
        ijo = itrunc(jo)
        fjo = jo - oftype(jo, ijo)
        @inbounds for ir = 0:size(resized,1)-1
            io = scale1*ir
            iio = itrunc(io)
            fio = io - oftype(io, iio)
            tmp = (1-fio)*((1-fjo)*original[iio+1,ijo+1] +
                           fjo*original[iio+1,ijo+2])
            + fio*((1-fjo)*original[iio+2,ijo+1] +
                   fjo*original[iio+2,ijo+2])
            resized[ir+1,jr+1] = convertsafely(eltype(resized), tmp)
        end
    end
    resized
end
imresize_julia(original, new_size) =
    imresize_julia!(similar(original,
                            new_size), original)

convertsafely{T<:FloatingPoint}(::Type{T}, val) = convert(T, val)
convertsafely{T<:Integer}(::Type{T}, val::Integer) = convert(T, val)
convertsafely{T<:Integer}(::Type{T}, val::FloatingPoint) =
    itrunc(T,
           val+oftype(val, 0.5))

function imresize_cairo(dat::Array{Uint32, 2}, new_size::(Int, Int))
    cs = CairoImageSurface(dat, 0)
    new_dat = zeros(Uint32, new_size)
    new_cs = CairoImageSurface(new_dat, 0)
    pat = CairoPattern(cs)
    pattern_set_filter(pat, Cairo.FILTER_BILINEAR)
    c = CairoContext(new_cs)
    h, w = size(dat)
    new_h, new_w = new_size
    scale(c, new_h / h, new_w / w)
    set_source(c, pat)
    paint(c)
    return new_cs.data
end

## img = rand(0x00:0xff, 774, 512)
## new_size = (3096, 2048)
## imresize_cairo(convert(Array{Uint32}, img), new_size)
## println("Cairo:")
## @time imresize_cairo(convert(Array{Uint32}, img), new_size)
## imresize_julia(img, new_size)
## println("Julia:")
## @time imresize_julia(img, new_size)


# p20.m - 2nd-order wave eq. in 2D via FFT (compare p19.m)

# Grid and initial data:
N = 24; x = y = cheb(N)[2];
dt = 6/N^2;
xx = yy = x[end:-1:1];
plotgap = round(Int,(1/3)/dt); dt = (1/3)/plotgap;
vv = @. exp(-40*((x'-.4)^2 + y^2));
vvold = vv;  clf();

# Time-stepping by leap frog formula:
for n = 0:3*plotgap
    t = n*dt;
    if rem(n+.5,plotgap)<1     # plots at multiples of t=1/3
          i = n/plotgap+1;
          subplot(2,2,i,projection="3d");
          xxx = yyy = -1:1/16:1;
          s = Spline2D(xx,yy,reduce(flipdim,vv,1:2));
          vvv = evalgrid(s,xxx,yyy);
          surf(xxx,yyy,vvv); xlim(-1,1); ylim(-1,1); zlim(-0.15,1);
          view(-37.5,30); title("t = $(signif(t,5))");
    end
    uxx = zeros(N+1,N+1); uyy = zeros(N+1,N+1);
    ii = 2:N;
    for i = 2:N                # 2nd derivs wrt x in each row
          v = vv[i,:]; V = [v;flipdim(v[ii],1)];
          U = real(fft(V));
          W1 = real(ifft(1im*[0:N-1;0;1-N:-1].*U)); # diff wrt theta
          W2 = real(ifft(-[0:N;1-N:-1].^2.*U));     # diff^2 wrt theta
          uxx[i,ii] = W2[ii]./(1-x[ii].^2) - x[ii].*W1[ii]./(1-x[ii].^2).^(3/2);
    end
    for j = 2:N                # 2nd derivs wrt y in each column
          v = vv[:,j]; V = [v; flipdim(v[ii],1)];
          U = real(fft(V));
          W1 = real(ifft(1im*[0:N-1;0;1-N:-1].*U));# diff wrt theta
          W2 = real(ifft(-[0:N;1-N:-1].^2.*U));    # diff^2 wrt theta
          uyy[ii,j] = W2[ii]./(1-y[ii].^2) - y[ii].*W1[ii]./(1-y[ii].^2).^(3/2);
    end
    vvnew = 2*vv - vvold + dt^2*(uxx+uyy);
    vvold = vv; vv = vvnew;
end


module FunctionsCollection

    function function1(x)
        return 2 * exp.(-2 * (x .- 1).^2) .- exp.(-(x .- 1).^2)
    end

    function function2(x)
        return x.^(1/2)
    end

    function function3(x)
        return exp.(-x) .* sin.(2 * x)
    end

    function function4(x)
        return log.(log.(x))
    end

    function function5(x)
        return 6 * exp.(-2 * x) + 2 * sin.(x) - cos.(x)
    end

end # FunctionsCollection


abstract type InnerProduct end

dot_arg(f,p::InnerProduct) = evaluate(f,p)
dot_arg(f::AbstractArray,p::InnerProduct) = [f(x) for x in evalpoints(p), f in f]
dot_arg(f::AbstractBasis,p::InnerProduct) = Matrix(f,evalpoints(p))

LinearAlgebra.norm(f, p::InnerProduct) = norm(dot_arg(f,p), p)
LinearAlgebra.dot(f,g,p::InnerProduct) = dot(dot_arg.((f,g),Ref(p))...,p)

"""
    WeightedL2(x,w) <: InnerProduct

# Example
```
p = WeightedL2([0,1],[2,3])
norm(identity,p) -> sqrt(0^2*2 + 1^2*3)
```
"""
struct WeightedL2{X,W} <: InnerProduct
    points::X
    weights::W
end

evalpoints(p::WeightedL2) = p.points
evaluate(f,p::WeightedL2) = f.(evalpoints(p))
dot_arg(f::AbstractArray{<:Number},p::WeightedL2) = f

LinearAlgebra.norm(f::AbstractVecOrMat{<:Number}, p::WeightedL2) = sqrt(dot(f,f,p))

function LinearAlgebra.dot(
    f::AbstractVecOrMat{<:Number},
    g::AbstractVecOrMat{<:Number},
    p::WeightedL2
)
    w = p.weights
    # Workaround for https://github.com/JuliaLang/julia/issues/35424
    length(w) == 0 && return zero(eltype(f))*zero(eltype(w))*zero(eltype(g))
    return f'*Diagonal(w)*g
end

using FastGaussQuadrature

LegendreInnerProduct(n) = WeightedL2(gausslegendre(n)...)
JacobiInnerProduct(n,a,b) = WeightedL2(gaussjacobi(n,a,b)...)


using GR
using LaTeXStrings

𝒩(μ, σ) = 1 / (σ * √(2π)) * exp.(-0.5 * ((x .- μ) / σ) .^ 2)

x = LinRange(-5, 5, 500);
y = hcat(𝒩(0, √0.2), 𝒩(0, √1), 𝒩(0, √5), 𝒩(-2, √0.5));

plot(x, y, xlabel=L"\mathcal{X}", ylabel=L"\mathcal{N}(\mu,\,\sigma^{2})", title=L"\frac{1}{\sigma\sqrt{2\pi}} e^{-\frac{1}{2} \left({\frac{x-\mu}{\sigma}}\right)^2}", labels=(L"\mu=0, \sigma^2=0.2", L"\mu=0, \sigma^2=1", L"\mu=0, \sigma^2=5", L"\mu={-2}, \sigma^2=0.5"), xlim=(-5.2, 5.2), ylim=(-0.05, 1.05), linewidth=3)


"""
    uniform_bsplines_eval_basis( spline_degree, normalized_offset, bspl )

# UNIFORM B-SPLINE FUNCTIONS

## Evaluate all non vanishing uniform B-Splines in unit cell.

Returns an array with the values of the b-splines of the
requested degree, evaluated at a given cell offset. The cell size is
normalized between 0 and 1, thus the offset given must be a number
between 0 and 1.

Output:

```math
bspl(1:d+1)= B_d(-(d+1)/2+d+x),...,B_d(-(d+1)/2+x)
```

with ``d``=`spline_degree` and ``x``=`normalized_offset`
where ``B_d=B_{d-1}*B_0`` and ``B_0=1_[-1/2,1/2]`` and `*` is convolution
the following FORTRAN code can be used for comparison with
[deboor](http://pages.cs.wisc.edu/~deboor/)

```fortran
do i=-d,d+1
    t(i+d+1)=real(i,8)
end do
call bsplvb(t,d+1,1,normalized_offset,d+1,out)
```

We also have the property (from the symmetry of the B-spline)
```math
out[1:d+1]= B_d(-(d+1)/2+xx),...,B_d(-(d+1)/2+d+xx),...,
```
where ``xx=1-`` `normalized_offset`

"""
function uniform_bsplines_eval_basis(spline_degree::Int, normalized_offset::Float64)
    @assert spline_degree >= 0
    @assert normalized_offset >= 0.0
    @assert normalized_offset <= 1.0

    bspl = zeros(Float64, spline_degree + 1)

    bspl[1] = 1.0
    @inbounds for j in 1:spline_degree
        xx = -normalized_offset
        j_real = Float64(j)
        inv_j = 1.0 / j_real
        saved = 0.0
        for r in 0:(j - 1)
            xx = xx + 1.0
            temp = bspl[r + 1] * inv_j
            bspl[r + 1] = saved + xx * temp
            saved = (j_real - xx) * temp
        end
        bspl[j + 1] = saved
    end

    return bspl
end

function uniform_bsplines_eval_basis!(
    bspl::Vector{Float64}, spline_degree::Int, normalized_offset::Float64
)
    bspl[1] = 1.0
    @inbounds for j in 1:spline_degree
        xx = -normalized_offset
        j_real = Float64(j)
        inv_j = 1.0 / j_real
        saved = 0.0
        for r in 0:(j - 1)
            xx = xx + 1.0
            temp = bspl[r + 1] * inv_j
            bspl[r + 1] = saved + xx * temp
            saved = (j_real - xx) * temp
        end
        bspl[j + 1] = saved
    end
end

function uniform_bsplines_eval_basis!(
    bspl::Array{Float64,2},
    spline_degree::Int,
    normalized_offset_x::Float64,
    normalized_offset_y::Float64,
)
    bspl[1, 1] = 1.0
    bspl[1, 2] = 1.0
    @inbounds for j in 1:spline_degree
        xx = -normalized_offset_x
        yy = -normalized_offset_y
        j_real = Float64(j)
        inv_j = 1.0 / j_real
        saved_x = 0.0
        saved_y = 0.0
        for r in 0:(j - 1)
            xx = xx + 1
            yy = yy + 1
            temp_x = bspl[r + 1, 1] * inv_j
            temp_y = bspl[r + 1, 2] * inv_j
            bspl[r + 1, 1] = saved_x + xx * temp_x
            bspl[r + 1, 2] = saved_y + yy * temp_y
            saved_x = (j_real - xx) * temp_x
            saved_y = (j_real - yy) * temp_y
        end
        bspl[j + 1, 1] = saved_x
        bspl[j + 1, 2] = saved_y
    end
end

export eval_uniform_periodic_spline_curve

"""
    eval_uniform_periodic_spline_curve( degree, scoef )

Evaluate uniform periodic spline curve defined by coefficients scoef at
knots (which are the grid points)

"""
function eval_uniform_periodic_spline_curve(degree::Int, scoef::Vector{Float64})

    # get bspline values at knots
    bspl = uniform_bsplines_eval_basis(degree, 0.0)
    n = length(scoef)
    sval = similar(scoef)
    for i in 1:n
        val = 0.0
        for j in 1:degree
            imj = mod1(i - j + n, n)
            val = val + bspl[j] * scoef[imj]
        end
        sval[i] = val
    end

    return sval
end


@doc """
      Heun's Method (with iteration)
      Supported by:
      -PhD Kelvyn B. Sánchez kelvyn.baruc@gmail.com
 """ ->
function heun(f, tspan, y0; h = 1, AbsTol = 0.00001)

t = tspan[1]:h:tspan[end]
n = length(t)

      if t[n] < tspan[end]
            t = [t, tspan[end]]
            n += 1
      end

y1 = f(t[1],y0)
eq = length(y1) # Test the number of equations in f
y = zeros(Int(n),eq)
y[1,:] = y0
y1 = zeros(Int(n),eq)
y2 = zeros(Int(n),eq)
y3 = zeros(Int(n),eq)

	for i=1:n-1 # Solve for each x
            for j=1:eq # Solve for each equation in f
                  s1 = f(t[i],y[i,:])
                  y1[i,j] = y[i,j] + s1[j] * (t[i+1] - t[i])

                  s2 = f(t[i+1],y1[i,:])
                  y2[i,j] = y[i,j] + ((s1[j] + s2[j])/2) * (t[i+1] - t[i])

                  err = 100
                  while err >= AbsTol
                        s1 = f(t[i],y[i,:])
                        s2 = f(t[i+1],y2[i,:])
                        y3[i,j] = y[i,j] + ((s1[j] + s2[j]) / 2) * (t[i+1] - t[i])
                        err = abs((y3[i,j] - y2[i,j])/y3[i,j]) * 100
                        y2[i,j] = y3[i,j]
                        y[i+1,j] = y3[i,j]
                  end
            end
	end
	return t, y
end


using Distributions
using LightBayes

function simulate(r::Float64, n::Int)
    # Simulate Δb for H0 & H1
    #
    # generate two clusters of samples, respectively centered
    # at (r, 0) and (-r, 0), where each cluster contains n
    # samples.
    #
    # This function computes the values of Δb, respectively
    # for the following two hypotheses:
    #
    # - H0: all samples are generated by a single component
    # - H1: each cluster of samples are generated by a different
    #       component
    #
    # Prior model: θ ~ N(0, σ0^2) with σ0 = 10.0
    # Likelihood:  x ~ N(θ, σ^2) with σ = 1.0
    #

    # data
    n = 1000
    X1 = [r, 0.0] .+ rand(2, 1000)
    X2 = [-r, 0.0] .+ rand(2, 1000)

    X = hcat(X1, X2)
    @assert size(X) == (2, 2n)

    # prior
    β0 = zeros(2)
    σ0 = 10.0
    κ0 = inv(σ0^2)
    pri = MvNormalCanon(β0, κ0)

    # likelihood model
    σ = 1.0
    g = IsoGaussModel(σ)

    # H0: single component
    pp0 = posterior(pri, suffstats(g, X))
    Δb0 = logpar(pp0) - logpar(pri)

    # H1: two components:
    pp1 = posterior(pri, suffstats(g, X1))
    pp2 = posterior(pri, suffstats(g, X2))
    Δb1 = (logpar(pp1) - logpar(pri)) + (logpar(pp2) - logpar(pri))

    return (Δb0, Δb1)
end

# main script

n = 1000

for r in [0.0, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0]
    Δb0, Δb1 = simulate(r, n)
    @printf("With r = %5.1f:  H0.Δb = %8.2f,  H1.Δb = %8.2f\n", r, Δb0, Δb1)
end