module CoordinateGA

using Random
using LinearAlgebra: dot
using ..GAFramework
import ..GAFramework: fitness, crossover!, mutation!, selection, randcreature

export CoordinateCreature, FunctionModel, SumModel

# This is a CoordinateCreature, T is some coordinate type like Vector
# and objvalue is some objective value
struct CoordinateCreature{T}
    value :: T
    objvalue :: Float64
end

# Negative, since we are minimizing the objective value
fitness(x::CoordinateCreature{T}) where {T} = -x.objvalue

# We can use multiple models to optimize this creature.
# To demonstrate, below we use the FunctionModel model
# as well as the SumModel

# The following are some crossover functions for the CoordinateCreature
# Each model can choose its own crossover function
struct AverageCrossover end
function crossover!(::AverageCrossover, z::CoordinateCreature{T},
    x::CoordinateCreature{T}, y::CoordinateCreature{T},
    st::GAState, aux, rng::AbstractRNG) where {T}
    z.value .= 0.5 .* (x.value .+ y.value)
    CoordinateCreature(z.value, st.model)
end

struct SinglePointCrossover end
function crossover!(::SinglePointCrossover, z::CoordinateCreature{T},
    x::CoordinateCreature{T}, y::CoordinateCreature{T},
    st::GAState, aux, rng::AbstractRNG) where {T}
    N = length(x.value)
    i = rand(rng, 1:N)
    if rand(rng) < 0.5
        z.value[1:i] = x.value[1:i]
        z.value[i+1:end] = y.value[i+1:end]
    else
        z.value[1:i] = y.value[1:i]
        z.value[i+1:end] = x.value[i+1:end]
    end
    CoordinateCreature(z.value, st.model)
end

struct TwoPointCrossover end
function crossover!(::TwoPointCrossover, z::CoordinateCreature{T},
    x::CoordinateCreature{T}, y::CoordinateCreature{T},
    st::GAState, aux, rng::AbstractRNG) where {T}
    N = length(x.value)
    i,j = rand(rng, 1:N, 2)
    i,j = i > j ? (j,i) : (i,j)
    if rand(rng) < 0.5
        z.value[:] = x.value
        z.value[i+1:j] = y.value[i+1:j]
    else
        z.value[:] = y.value
        z.value[i+1:j] = x.value[i+1:j]
    end
    CoordinateCreature(z.value, st.model)
end

printfitness(curgen::Int, x::CoordinateCreature{T}) where {T} =
    println("curgen: $curgen value: $(x.value) obj. value: $(x.objvalue)")

# The following is the SumModel
# It finds a d-dimensional coordinate x such that dot(x, 1:length(x)) ≈ y
"""
    model = SumModel(5, 42.)
    state = GAState(model, ngen=500, npop=6_000, elite_fraction=0.1,
                       params=Dict(:mutation_rate=>0.9, :print_fitness_iter=>1))
    ga!(state)
"""
struct SumModel{T} <: GAModel
    d::Int
    target::T
end

function CoordinateCreature(value::Vector{T}, m::SumModel{T}) where {T}
    objval = Float64(abs(dot(value, 1:length(value)) - m.target))
    CoordinateCreature(value, objval)
end

function randcreature(m::SumModel{T}, aux, rng::AbstractRNG) where {T}
    value = rand(rng, m.d)
    CoordinateCreature(value, m)
end

function mutation!(x::CoordinateCreature,
    st::GAState{SumModel{T}}, aux, rng::AbstractRNG) where {T}
    gp = st.params
    if rand(rng) < get(gp, :mutation_rate, 0.1)
        x.value[rand(rng, 1:length(x.value))] += st.model.target * randn(rng,T)
        CoordinateCreature(x.value, st.model)
    else
        x
    end
end

function crossover!(z::CoordinateCreature,
    x::CoordinateCreature, y::CoordinateCreature,
    st::GAState{SumModel{T}}, aux, rng::AbstractRNG) where {T}
    crossover!(TwoPointCrossover(), z, x, y, st, aux, rng)
end

# The following is the FunctionModel
# It's a more complicated model than above
# It minimizes the objective value using a given objective function
"""
# E.g.
    model = FunctionModel(x -> abs(x[1] - 20.0), [-200.0], [200.0])
    state = GAState(model, ngen=500, npop=6_000, elite_fraction=0.1,
                       params=Dict(:mutation_rate=>0.9, :print_fitness_iter=>1))
    ga!(state)

    type T has to have properties
        y-x :: T
        0.25 .* (x+y) :: T
        randn(T) :: T
        z + 0.25*(y-x)*randn(T) :: T

"""
struct FunctionModel{F,T} <: GAModel
    f::F
    xmin::T
    xmax::T
    xspan::T # xmax-xmin
    clamp::Bool
end
function FunctionModel(f::F,xmin,xmax,clamp::Bool=true) where {F}
    xmin,xmax = promote(xmin,xmax)
    ET = eltype(xmin)
    N = length(xmin)
    xspan = xmax .- xmin
    # check that f(xmin), f(xmax) can be converted to Float64 without error
    z1 = Float64(f(xmin))
    z2 = Float64(f(xmax))
    # and that f(xspan), f(xmin), and f(xmax) has sane values maybe
    # z1!=Inf && z2!=Inf && !isnan(z1) && !isnan(z2) ||
    #    error("f(xmin) or f(xmax) objective function is either NaN or Inf")
    all(xspan .>= zero(ET)) || error("xmax[i] < xmin[i] for some i")
    FunctionModel{F,typeof(xspan)}(f,xmin,xmax,xspan,clamp)
end

CoordinateCreature(value::T, m::FunctionModel{F,T}) where {F,T} =
    CoordinateCreature(value, m.f(value))

function randcreature(m::FunctionModel{F,T}, aux, rng::AbstractRNG) where {F,T}
    ET = eltype(T)
    if T <: Vector
        xvalue = m.xmin .+ m.xspan .* rand(rng, ET, length(m.xspan))
    else
        xvalue = m.xmin .+ m.xspan .* rand(rng, T)
    end
    CoordinateCreature(xvalue, m)
end

function crossover!(z::CoordinateCreature{T},
    x::CoordinateCreature{T}, y::CoordinateCreature{T},
    st::GAState{FunctionModel{F,T}}, aux, rng::AbstractRNG) where {F,T}
    crossover!(TwoPointCrossover(), z, x, y, st, aux, rng)
end

# Mutate over all dimensions
function mutatenormal!(x::CoordinateCreature{T}, temp::Real,
    model::FunctionModel{F,T}, rng::AbstractRNG) where {F,T}
    x.value .+= temp .* model.xspan .* randn(rng,length(x.value))
    model.clamp && (x.value .= clamp.(x.value, model.xmin, model.xmax))
    CoordinateCreature(x.value, model)
end

# Mutate through a single dimension
function mutatenormaldim!(x::CoordinateCreature{T}, temp::Real, dim::Integer,
    model::FunctionModel{F,T}, rng::AbstractRNG) where {F,T}
    ET = eltype(T)
    x.value[dim] += temp * model.xspan[dim] * randn(rng,ET)
    model.clamp && (x.value[dim] = clamp(x.value[dim], model.xmin[dim], model.xmax[dim]))
    CoordinateCreature(x.value, model)
end

function mutation!(x::CoordinateCreature{T},
    st::GAState{FunctionModel{F,T}}, aux, rng::AbstractRNG) where {F,T}
    gp = st.params
    if rand(rng) < get(gp, :mutation_rate, 0.1)
        if rand(rng) < get(gp, :sa_rate, 0.0)
            sa(x, st.model, gp[:sa_k], gp[:sa_lambda],
                gp[:sa_maxiter], st.curgen, aux, rng)
        else
            N = length(x.value)
            mutatenormaldim!(x, 0.1, rand(rng, 1:N), st.model, rng)
        end
    else
        x
    end
end

# export sa,satemp,saprob,mutatenormal
function sa(x::CoordinateCreature{T}, model::FunctionModel{F,T},
    k::Real, lambda::Real, maxiter::Integer, curgen::Int,
    aux, rng::AbstractRNG) where {F,T}
    N = length(x.value)
    y = x
    numnoups = 0 # number of consecutive
    #curgen = 0
    iter = curgen*maxiter + 0
    while true
        temp = satemp(iter, k, lambda) + numnoups/maxiter
        #temp = satemp(iter, k, lambda) + lambda * log(1+numnoups)
        dim = rand(rng,1:N)
        yvdim_old = y.value[dim]
        yov_old = y.objvalue
        fitness_old = fitness(y)
        y = mutatenormaldim!(y, temp, dim, model, rng)
        diffe = fitness(y) - fitness_old
        if diffe >= 0
        elseif rand(rng) < saprob(diffe, temp)
        else
            y.value[dim] = yvdim_old
            y = CoordinateCreature(y.value, yov_old)
        end
        #println("temp: $temp newy: $(newy.objvalue) diffe: $diffe prob: $(saprob(diffe, temp))")
        #numnoups = ifelse(diffe > 0, 0, ifelse(diffe < 0, numnoups+1, numnoups))
        numnoups = ifelse(diffe > 0, 0, numnoups + 1)
        iter += 1
        iter > curgen*maxiter + maxiter && break
    end
    #println("pre: $(x.objvalue) post: $(y.objvalue)")
    y
end

satemp(iter::Integer, k::Real, lambda::Real) = k * exp(-lambda * iter)
# diff = newscore - oldscore
saprob(diff::Real, iter::Integer, k::Real, lambda::Real) =
    exp(diff / satemp(iter, k, lambda))
saprob(diff::Real, temp::Real) = exp(diff / temp)

function selection(pop::Vector, n::Integer,
    st::GAState{FunctionModel{F,T}}, rng::AbstractRNG) where {F,T}
    selection(TournamentSelection(2), pop, n, st, rng)
end

end # CoordinateGA


module LReg

"""This module wraps some functionality from [GLMNet.jl](https://github.com/simonster/GLMNet.jl)
for logistic regression."""
LReg

using GLMNet, Distributions

using StatsBase, StatsFuns

import StatsBase: predict
import Distributions: log1pexp

export LR, lreg, predict, linpred

"""The `LR` type contains the coefficent vector of a logistic regression fit, as well as indexes of
included columns in the design matrix.
"""
immutable LR{T<:AbstractFloat}
    β::Vector{T}
    idx::AbstractVector{Int}
    fitwithoffset::Bool
end

"""
Returns the linear component of predicted values from a logistic regression fit with a new X matrix

** Arguments **

* `lr` - an object of type `LR`
* `newx` - matrix of new covariates to predict on. Should be the same dimension as the design matrix used to fit `lr`

** Keyword Arguments **

* `offset` - offsets. The length should be the same as `size(newx, 1)` of `lr` was fit with an offset, or 0 otherwise.
"""
function linpred(lr::LR, newx; offset=Array(eltype(newx), 0))
    if size(newx, 2) == 1
        newx = reshape(newx, size(newx, 1), 1)
    end
    p = newx * lr.β
    if lr.fitwithoffset
        length(offset) == size(newx, 1) || throw(ArgumentError("fit with offset, 'offset' kw arg should have length size(newx, 1)"))
        for i in 1:length(p) #NumericExtensions
          @inbounds p[i] += offset[i]
        end
    else
        length(offset) == 0 || throw(ArgumentError("not fit with offset, 'offset' kw arg should have length 0"))
    end
    p
end

"""
Returns the predicted on the probability scale values from a logistic regression fit with a new X matrix

** Arguments **

* `lr` - an object of type `LR`
* `newx` - matrix of new covariates to predict on. Should be the same dimension as the design matrix used to fit `lr`

** Keyword Arguments

* `offset` - offsets. The length should be the same as `size(newx, 1)` of `lr` was fit with an offset, or 0 otherwise.
"""
predict(lr::LR, newx; offset=Array(eltype(newx), 0)) = logistic(linpred(lr, newx, offset=offset))

"""
Fits a logistic regression model

** Arguments **

* `x` - design matrix
* `y` - response vector. Should have the same length as `size(x, 1)`

** Keyword Arguments **

* `wts` - weight vector. Defaults to all 1s.
* `offset` - offset vector. Defaults to a vector of length 0 for no offset.
* `subset` - column indexes for `x` that should be included in the fit. Defaults to all columns.
* `convTol` - convergence criterion for relative change in deviance. Defaults to 1.0e-8.

** Details **

An intercept is not included by default. If you want one, make sure the first column of your design matrix
is all ones. Because of how `GLMNet` handles intercepts and how `LReg` interfaces with it,
this must be the first column.
If you do that, don't forget to add the column in the same place to `newx` when you call predict.

`subset` is useful if you want to include only some columns of the design matrix but you want to call `predict`
with a `newx` matrix with the same number of columns as `x`. The coefficients corresponding to columns of `x` which
are not used in the fit are set to zero. If you would like to call `predict` with a `newx` matrix that includes
only the columns that you fit on, you should subset `x` yourself before calling `lreg`.
"""
function lreg(x, y; wts=ones(y), offset=similar(y,0), subset=1:size(x,2), convTol=1.0e-8)
    if size(x, 2) == 1
        x = reshape(x, size(x, 1), 1)
    end
    subset=sort(subset)
    fitwithoffset = length(offset) > 0

    if subset == 1:size(x,2) || collect(1:size(x,2)) == subset
        subset = 1:size(x,2)
#         return LR(coef(fit(GeneralizedLinearModel, x, y, Binomial(); wts=wts, offset=offset, convTol=convTol)), subset, fitwithoffset)
        return LR(myfit(x, y, wts=wts, offset=offset, convTol=convTol), subset, fitwithoffset)
    else
        tempx = x[:, subset]
        β = zeros(eltype(x), size(x, 2))
#        β[subset] = coef(fit(GeneralizedLinearModel, tempx, y, Binomial(); wts=wts, offset=offset, convTol=convTol))
        β[subset] = myfit(tempx, y, wts=wts, offset=offset, convTol=convTol)
        return LR(β, subset, fitwithoffset)
    end
end

loss(y, xb) =
    y == one(y)? log1pexp(-xb) :
    y == zero(y)? log1pexp(xb) :
    y * log1pexp(-xb) + (one(y)-y) * log1pexp(xb)

function myfit(x, y; wts=ones(y), offset=similar(y,0), convTol=1.0e-8)
    offsets = length(offset) == 0? nothing: offset
    try
        #annoyingly, glmnet won't do intercept only  models for some reason, even with intercept = true
        # and if there is a constant column in the design matrix, that will get a coefficient of zero
        if size(x, 2) == 1 && all(x .== 1)
            offsets==nothing || error("Fitting an intercept only model with an offset is not supported")
            #intercept only
            [logit(sum(y .* wts)/sum(wts))]
        elseif all(x[:, 1] .== 1)
            #first column is intercept
            est = glmnet(x[:, 2:end], [ones(length(y)) .- y y], Binomial(), weights=wts, offsets=offsets,
                         lambda=[0.0], intercept=true, tol=convTol)
            [est.a0; est.betas[:,1]]
        else
            #no intercept
            glmnet(x, [ones(length(y)) .- y y], Binomial(), weights=wts, offsets=offsets,
                   lambda=[0.0], intercept=false, tol=convTol).betas[:,1]
        end
    catch err
        if isdefined(Main, :LREG_DEBUG) &&
           Main.LREG_DEBUG
            fname = tempname()
            open(fname, "w") do f
                serialize(f, (x, y, wts, offset, convTol))
                info("Error in lreg. x, y, wts, offset and convTol written to $fname")
            end
        end
        rethrow(err)
    end
end

end


using CompScienceMeshes
using BEAST
using Test

m = meshrectangle(1.0, 1.0, 0.25)

X = raviartthomas(m, BEAST.Continuity{:none})
@test numfunctions(X) == 16*2*3
@test all(length.(X.fns) .== 1)

p = positions(X)

i = 12
c = X.fns[i][1].cellid
ch = chart(m, cells(m)[c])
ctr = cartesian(center(ch))
@test ctr ≈ p[i]

for (i,f) in enumerate(X.fns)
    @test f[1].coeff == 1.0
end


@doc raw"""
    CircleGroup <: GroupManifold{Circle{ℂ},MultiplicationOperation}

The circle group is the complex circle ([`Circle(ℂ)`](@ref)) equipped with
the group operation of complex multiplication ([`MultiplicationOperation`](@ref)).
"""
const CircleGroup = GroupManifold{ℂ,Circle{ℂ},MultiplicationOperation}

CircleGroup() = GroupManifold(Circle{ℂ}(), MultiplicationOperation())

Base.show(io::IO, ::CircleGroup) = print(io, "CircleGroup()")

invariant_metric_dispatch(::CircleGroup, ::ActionDirection) = Val(true)

default_metric_dispatch(::MetricManifold{ℂ,CircleGroup,EuclideanMetric}) = Val(true)

adjoint_action(::CircleGroup, p, X) = X

adjoint_action!(::CircleGroup, Y, p, X) = copyto!(Y, X)

function _compose(G::CircleGroup, p::AbstractVector, q::AbstractVector)
    return map(compose, repeated(G), p, q)
end

_compose!(G::CircleGroup, x, p, q) = copyto!(x, compose(G, p, q))

identity_element(G::CircleGroup) = 1.0
identity_element(::CircleGroup, p::Number) = one(p)
identity_element(::CircleGroup, p::AbstractArray) = map(i -> one(eltype(p)), p)

Base.inv(G::CircleGroup, p::AbstractVector) = map(inv, repeated(G), p)

function inverse_translate(
    ::CircleGroup,
    p::AbstractVector,
    q::AbstractVector,
    ::LeftAction,
)
    return map(/, q, p)
end
function inverse_translate(
    ::CircleGroup,
    p::AbstractVector,
    q::AbstractVector,
    ::RightAction,
)
    return map(/, q, p)
end

lie_bracket(::CircleGroup, X, Y) = zero(X)

lie_bracket!(::CircleGroup, Z, X, Y) = fill!(Z, 0)

translate_diff(::GT, p, q, X, ::ActionDirection) where {GT<:CircleGroup} = map(*, p, X)
function translate_diff(
    ::CircleGroup,
    ::Identity{MultiplicationOperation},
    q,
    X,
    ::ActionDirection,
)
    return X
end

function translate_diff!(G::CircleGroup, Y, p, q, X, conv::ActionDirection)
    return copyto!(Y, translate_diff(G, p, q, X, conv))
end

function exp_lie(::CircleGroup, X)
    return map(X) do imθ
        θ = imag(imθ)
        sinθ, cosθ = sincos(θ)
        return Complex(cosθ, sinθ)
    end
end

exp_lie!(G::CircleGroup, q, X) = (q .= exp_lie(G, X))

function log_lie(::CircleGroup, q)
    return map(q) do z
        cosθ, sinθ = reim(z)
        θ = atan(sinθ, cosθ)
        return θ * im
    end
end
log_lie(::CircleGroup, e::Identity{MultiplicationOperation}) = 0.0 * im

_log_lie!(G::CircleGroup, X, q) = (X .= log_lie(G, q))

function number_of_coordinates(G::CircleGroup, B::AbstractBasis)
    return number_of_coordinates(base_manifold(G), B)
end

@doc raw"""
    RealCircleGroup <: GroupManifold{Circle{ℝ},AdditionOperation}

The real circle group is the real circle ([`Circle(ℝ)`](@ref)) equipped with
the group operation of addition ([`AdditionOperation`](@ref)).
"""
const RealCircleGroup = GroupManifold{ℝ,Circle{ℝ},AdditionOperation}

RealCircleGroup() = GroupManifold(Circle{ℝ}(), AdditionOperation())

Base.show(io::IO, ::RealCircleGroup) = print(io, "RealCircleGroup()")

invariant_metric_dispatch(::RealCircleGroup, ::ActionDirection) = Val(true)

default_metric_dispatch(::MetricManifold{ℝ,RealCircleGroup,EuclideanMetric}) = Val(true)

_compose(::RealCircleGroup, p, q) = sym_rem(p + q)
function _compose(G::RealCircleGroup, p::AbstractVector, q::AbstractVector)
    return map(compose, repeated(G), p, q)
end

function _compose!(::RealCircleGroup, x, p, q)
    x .= sym_rem.(p .+ q)
    return x
end

identity_element(G::RealCircleGroup) = 0.0
identity_element(::RealCircleGroup, p::AbstractArray) = map(i -> zero(eltype(p)), p)

Base.inv(G::RealCircleGroup, p::AbstractVector) = map(inv, repeated(G), p)

function inverse_translate(
    ::RealCircleGroup,
    p::AbstractVector,
    q::AbstractVector,
    ::LeftAction,
)
    return map((x, y) -> sym_rem(x - y), q, p)
end
function inverse_translate(
    ::RealCircleGroup,
    p::AbstractVector,
    q::AbstractVector,
    ::RightAction,
)
    return map((x, y) -> sym_rem(x - y), q, p)
end

function exp_lie(::RealCircleGroup, X)
    return sym_rem(X)
end

exp_lie!(G::RealCircleGroup, q, X) = (q .= sym_rem(X))


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

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

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

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

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

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

function linspace2(s,e,n::Int)
    #this produces a linear spacing between
    #s and e, such that n total elements are
    #in the vector, but the last element
    #is e - spacing, rather than e, like in
    #the normal linspace function
    return [s:((e-s)/(n)):(e-((e-s)/(n)))]
end

function linspace2upsample(s,e,n::Int,m::Int)
    #this function upsamples what linspace2 would
    #return, such that the two vectors share the same
    #points at integer values of the spacing defined
    #by linspace2: (e-s)/m
    return [s:((e-s)/(n*m)):(e-((e-s)/(n*m)))]
end

function genZ(N::Int)
	return exp(im*linspace2(0.0,2*pi,N))
end

function genZArc(N::Int,Div::Int)
    return exp(im*linspace2(0.0,2*pi,N))
end

function genZUpSample(N::Int,Mult::Int)
    return exp(im*linspace2upsample(0.0,2*pi,N,Mult))
end

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

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


#=
Classical discrete time LQ optimal control and filtering problems. The
control problems we consider take the form

    \max \sum_{t = 0}^N \beta^t \{a_t y_t - h y^2_t / 2 - [d(L)y_t]^2 \}

subject to h > 0, 0 < \beta < 1 and

 * y_{-1},y_{-2}, \dots, y_{-m} (initial conditions)

 * d(L) = d_0 + d_1L + d_2L^2 + \dots + d_mL^m

The sequence {y_t} is scalar

Authors: Shunsuke Hori
=#

using Polynomials

struct LQFilter{TR<:Real, TI<:Integer, TF<:AbstractFloat}
    d::Vector{TF}
    h::TR
    y_m::Vector{TF}
    m::TI
    phi::Vector{TF}
    beta::TR
    phi_r::Union{Vector{TF},Void}
    k::Union{TI,Void}
end

"""
Parameters
----------
d : Vector
        The order of the coefficients: [d_0, d_1, ..., d_m]
h : Real
        Parameter of the objective function (corresponding to the
        quadratic term)
y_m : Vector
        Initial conditions for y
r : Vector
        The order of the coefficients: [r_0, r_1, ..., r_k]
        (optional, if not defined -> deterministic problem)
beta : Real or nothing
        Discount factor (optional, default value is one)
h_eps :
"""
function LQFilter{TR<:Real}(d::Vector{TR},
                   h::TR,
                   y_m::Vector{TR};
                   r::Union{Vector{TR},Void}=nothing,
                   beta::Union{TR,Void}=nothing,
                   h_eps::Union{TR,Void}=nothing,
                   )

    m = length(d) - 1

    m == length(y_m) ||
        throw(ArgumentError("y_m and d must be of same length = $m"))

    #---------------------------------------------
    # Define the coefficients of phi up front
    #---------------------------------------------

    phi = Vector{TR}(2m + 1)
    for i in -m:m
        phi[m-i+1] = sum(diag(d*d', -i))
    end
    phi[m+1] = phi[m+1] + h

    #-----------------------------------------------------
    # If r is given calculate the vector phi_r
    #-----------------------------------------------------

    if r == nothing
        k=nothing
        phi_r = nothing
    else
        k = size(r,1) - 1
        phi_r = Vector{TR}(2k + 1)

        for i = -k:k
            phi_r[k-i+1] = sum(diag(r*r', -i))
        end

        if h_eps != nothing
            phi_r[k+1] = phi_r[k+1] + h_eps
        end
    end

    #-----------------------------------------------------
    # If beta is given, define the transformed variables
    #-----------------------------------------------------
    if beta == nothing
        beta = 1.0
    else
        d = beta.^(collect(0:m)/2) * d
        y_m = y_m * beta.^(- collect(1:m)/2)
    end

    return LQFilter(d,h,y_m,m,phi,beta,phi_r,k)
end

"""
This constructs the matrices W and W_m for a given number of periods N
"""
function construct_W_and_Wm(lqf::LQFilter, N::Integer)

    d, m = lqf.d, lqf.m

    W = zeros(N + 1, N + 1)
    W_m = zeros(N + 1, m)

    #---------------------------------------
    # Terminal conditions
    #---------------------------------------

    D_m1 = zeros(m + 1, m + 1)
    M = zeros(m + 1, m)

    # (1) Constuct the D_{m+1} matrix using the formula

    for j in 1:(m+1)
        for k in j:(m+1)
            D_m1[j, k] = dot(d[1:j,1], d[k-j+1:k,1])
        end
    end

    # Make the matrix symmetric
    D_m1 = D_m1 + D_m1' - diagm(diag(D_m1))

    # (2) Construct the M matrix using the entries of D_m1

    for j in 1:m
        for i in (j + 1):(m + 1)
            M[i, j] = D_m1[i-j, m+1]
        end
    end
    M

    #----------------------------------------------
    # Euler equations for t = 0, 1, ..., N-(m+1)
    #----------------------------------------------
    phi, h = lqf.phi, lqf.h

    W[1:(m + 1), 1:(m + 1)] = D_m1 + h * eye(m + 1)
    W[1:(m + 1), (m + 2):(2m + 1)] = M

    for (i, row) in enumerate((m + 2):(N + 1 - m))
        W[row, (i + 1):(2m + 1 + i)] = phi'
    end

    for i in 1:m
        W[N - m + i + 1 , end-(2m + 1 - i)+1:end] = phi[1:end-i]
    end

    for i in 1:m
        W_m[N - i + 2, 1:(m - i)+1] = phi[(m + 1 + i):end]
    end

    return W, W_m
end

"""
This function calculates z_0 and the 2m roots of the characteristic equation
associated with the Euler equation (1.7)
Note:
------
`poly(roots)` from Polynomial.jll defines a polynomial using its roots that can be
evaluated at any point by `polyval(Poly,x)`. If x_1, x_2, ... , x_m are the roots then
    polyval(poly(roots),x) = (x - x_1)(x - x_2)...(x - x_m)
"""
function roots_of_characteristic(lqf::LQFilter)
    m, phi = lqf.m, lqf.phi

    # Calculate the roots of the 2m-polynomial
    phi_poly=Poly(phi[end:-1:1])
    proots = roots(phi_poly)
    # sort the roots according to their length (in descending order)
    roots_sorted = sort(proots, by=abs)[end:-1:1]
    z_0 = sum(phi) / polyval(poly(proots), 1.0)
    z_1_to_m = roots_sorted[1:m]     # we need only those outside the unit circle
    lambdas = 1 ./ z_1_to_m
    return z_1_to_m, z_0, lambdas
end

"""
This function computes the coefficients {c_j, j = 0, 1, ..., m} for
    c(z) = sum_{j = 0}^{m} c_j z^j
Based on the expression (1.9). The order is
    c_coeffs = [c_0, c_1, ..., c_{m-1}, c_m]
"""
function coeffs_of_c(lqf::LQFilter)
    m = lqf.m
    z_1_to_m, z_0, lambdas = roots_of_characteristic(lqf)
    c_0 = (z_0 * prod(z_1_to_m) * (-1.0)^m)^(0.5)
    c_coeffs = coeffs(poly(z_1_to_m)) * z_0 / c_0
    return c_coeffs
end

"""
This function calculates {lambda_j, j=1,...,m} and {A_j, j=1,...,m}
of the expression (1.15)
"""
function solution(lqf::LQFilter)
    z_1_to_m, z_0, lambdas = roots_of_characteristic(lqf)
    c_0 = coeffs_of_c(lqf)[end]
    A = zeros(lqf.m)
    for j in 1:m
        denom = 1 - lambdas/lambdas[j]
        A[j] = c_0^(-2) / prod(denom[1:m .!= j])
    end
    return lambdas, A
end

"""
This function constructs the covariance matrix for x^N (see section 6)
for a given period N
"""
function construct_V(lqf::LQFilter; N::Integer=nothing)
    if N == nothing
        error("N must be provided!!")
    end
    if !(typeof(N) <: Integer)
        throw(ArgumentError("N must be Integer!"))
    end

    phi_r, k = lqf.phi_r, lqf.k
    V = zeros(N, N)
    for i in 1:N
        for j in 1:N
            if abs(i-j) <= k
                V[i, j] = phi_r[k + abs(i-j)+1]
            end
        end
    end
    return V
end

"""
Assuming that the u's are normal, this method draws a random path
for x^N
"""
function simulate_a(lqf::LQFilter, N::Integer)
    V = construct_V(N + 1)
    d = MVNSampler(zeros(N + 1), V)
    return rand(d)
end

"""
This function implements the prediction formula discussed is section 6 (1.59)
It takes a realization for a^N, and the period in which the prediciton is formed

Output:  E[abar | a_t, a_{t-1}, ..., a_1, a_0]
"""
function predict(lqf::LQFilter, a_hist::Vector, t::Integer)
    N = length(a_hist) - 1
    V = construct_V(N + 1)

    aux_matrix = zeros(N + 1, N + 1)
    aux_matrix[1:t+1 , 1:t+1 ] = eye(t + 1)
    L = chol(V)'
    Ea_hist = inv(L) * aux_matrix * L * a_hist

    return Ea_hist
end

"""
- if t is NOT given it takes a_hist (Vector or Array) as a deterministic a_t
- if t is given, it solves the combined control prediction problem (section 7)
　(by default, t == nothing -> deterministic)

for a given sequence of a_t (either determinstic or a particular realization),
it calculates the optimal y_t sequence using the method of the lecture

Note:
------
lufact normalizes L, U so that L has unit diagonal elements
To make things cosistent with the lecture, we need an auxiliary diagonal
matrix D which renormalizes L and U
"""

function optimal_y(lqf::LQFilter, a_hist::Vector, t = nothing)
    beta, y_m, m = lqf.beta, lqf.y_m, lqf.m

    N = length(a_hist) - 1
    W, W_m = construct_W_and_Wm(lqf, N)

    F = lufact(W, Val{true})

    L, U = F[:L], F[:U]
    D = diagm(1.0./diag(U))
    U = D * U
    L = L * diagm(1.0./diag(D))

    J = flipdim(eye(N + 1), 2)

    if t == nothing   # if the problem is deterministic
        a_hist = J * a_hist

        #--------------------------------------------
        # Transform the a sequence if beta is given
        #--------------------------------------------
        if beta != 1
            a_hist =  reshape(a_hist * (beta^(collect(N:0)/ 2)),N + 1, 1)
        end

        a_bar = a_hist - W_m * y_m           # a_bar from the lecutre
        Uy = \(L, a_bar)            # U @ y_bar = L^{-1}a_bar from the lecture
        y_bar = \(U, Uy)            # y_bar = U^{-1}L^{-1}a_bar
        # Reverse the order of y_bar with the matrix J
        J = flipdim(eye(N + m + 1), 2)
        y_hist = J * vcat(y_bar, y_m)     # y_hist : concatenated y_m and y_bar
        #--------------------------------------------
        # Transform the optimal sequence back if beta is given
        #--------------------------------------------
        if beta != 1
            y_hist = y_hist .* beta.^(- collect(-m:N)/2)
        end

    else           # if the problem is stochastic and we look at it
        Ea_hist = reshape(predict(a_hist, t), N + 1, 1)
        Ea_hist = J * Ea_hist

        a_bar = Ea_hist - W_m * y_m           # a_bar from the lecutre
        Uy = \(L, a_bar)            # U @ y_bar = L^{-1}a_bar from the lecture
        y_bar = \(U, Uy)            # y_bar = U^{-1}L^{-1}a_bar

        # Reverse the order of y_bar with the matrix J
        J = flipdim(eye(N + m + 1), 2)
        y_hist = J * vcat(y_bar, y_m)     # y_hist : concatenated y_m and y_bar
    end
    return y_hist, L, U, y_bar
end


using KernSmooth, RDatasets, Base.Test

expected = readcsv(joinpath(Pkg.dir("KernSmooth"), "test", "geyserfit.csv"), Float64, has_header=true)[1]

dat = dataset("MASS", "geyser")

x = array(dat[:Duration])
y = convert(Vector{Float64}, array(dat[:Waiting]))

resx, resy= locpoly(x, y, 0.25)

@test_approx_eq_eps(maximum(abs(resx - expected[:,1])), 0.0, sqrt(eps(Float64)))
@test_approx_eq_eps(maximum(abs(resy - expected[:,2])), 0.0, sqrt(eps(Float64)))

@test !any(isnan(resy))

#example from R function
@test_approx_eq dpill(x, y) 0.238350165926485


using LabelledArrays

export lorenz_lorenz_lorenz_transitive

@inline @inbounds function eom_lorenz_lorenz_lorenz_transitive(u, p, t)
    x₁, y₁, z₁, x₂, y₂, z₂, x₃, y₃, z₃ = (u...,)
    σ₁, σ₂, σ₃, ρ₁, ρ₂, ρ₃, β₁, β₂, β₃, c₁₂, c₂₃ = (p...,)

    # Subsystem 1
    dx₁ = σ₁*(y₁-x₁)
    dy₁ = ρ₁*x₁ - y₁ - x₁*z₁
    dz₁ = x₁*y₁ - β₁*z₁

    # Subsystem 2
    dx₂ = σ₂*(y₂-x₂) + c₁₂*(x₁ - x₂)
    dy₂ = ρ₂*x₂ - y₂ - x₂*z₂
    dz₂ = x₂*y₂ - β₂*z₂

    # Subsystem 3
    dx₃ = σ₃*(y₃-x₃) + c₂₃*(x₂ - x₃)
    dy₃ = ρ₃*x₃ - y₃ - x₃*z₃
    dz₃ = x₃*y₃ - β₃*z₃
    return SVector{9}(dx₁, dy₁, dz₁, dx₂, dy₂,dz₂, dx₃, dy₃, dz₃)
end

function lorenz_lorenz_lorenz_transitive(u₀, σ₁, σ₂, σ₃, ρ₁, ρ₂, ρ₃, β₁, β₂, β₃, c₁₂, c₂₃)
    p = @LArray [σ₁, σ₂, σ₃, ρ₁, ρ₂, ρ₃, β₁, β₂, β₃, c₁₂, c₂₃] (:σ₁, :σ₂, :σ₃, :ρ₁, :ρ₂, :ρ₃, :β₁, :β₂, :β₃, :c₁₂, :c₂₃)
    ContinuousDynamicalSystem(eom_lorenz_lorenz_lorenz_transitive, u₀, p)
end

"""
    lorenz_lorenz_lorenz_transitive(;u₀=rand(9),
                σ₁ = 10.0, σ₂ = 10.0, σ₃ = 10.0,
                ρ₁ = 28.0, ρ₂ = 28.0, ρ₃ = 28.0,
                β₁ = 8/3,  β₂ = 8/3,  β₃ = 8.3,
                c₁₂ = 1.0, c₂₃ = 1.0) → ContinuousDynamicalSystem

Initalise a dynamical system consisting of three coupled Lorenz attractors with
a transitive causality chain where X₁ → X₂ and X₂ → X₃. In total, the three
3D-subsystems create a 9-dimensional dynamical system.

The strength of the forcing X₁ → X₂ is controlled by the parameter `c₁`, and
the forcing from X₂ → X₃ by `c₂`. The remaining parameters are the usual
parameters for the Lorenz system, where the subscript `i` refers to the
subsystem Xᵢ.

## Equations of motion

The dynamics is generated by the following vector field

```math
\\begin{aligned}
\\dot{x_1} &= \\sigma_1(y_1 - x_1) \\\\
\\dot{y_1} &= \\rho_1 x_1 - y_1 - x_1 z_1 \\\\
\\dot{z_1} &= x_1 y_1 - \\beta_1 z_1 \\\\
\\dot{x_2} &=  \\sigma_2 (y_2 - x_2) + c_{12}(x_1 - x_2) \\\\
\\dot{y_2} &= \\rho_2 x_2 - y_2 - x_2 z_2 \\\\
\\dot{z_2} &= x_2 y_2 - \\beta_2 z_2 \\\\
\\dot{x_3} &= \\sigma_3 (y_3 - x_3) + c_{23} (x_2 - x_3) \\\\
\\dot{y_3} &= \\rho_3 x_3 - y_3 - x_3 z_3 \\\\
\\dot{z_3} &= x_3 y_3 - \\beta_3 z_3
\\end{aligned}
```

## Usage in literature

This system was studied by Papana et al. (2013) for coupling strengths
``c_{12} = 0, 1, 3, 5`` and ``c_{23} = 0, 1, 3, 5``.

## References

1. Papana et al., Simulation Study of Direct Causality Measures in Multivariate
    Time Series. Entropy 2013, 15(7), 2635-2661; doi:10.3390/e15072635
"""
lorenz_lorenz_lorenz_transitive(;u₀=rand(9),
            σ₁ = 10.0, σ₂ = 10.0, σ₃ = 10.0,
            ρ₁ = 28.0, ρ₂ = 28.0, ρ₃ = 28.0,
            β₁ = 8/3,  β₂ = 8/3,  β₃ = 8.3,
            c₁₂ = 1.0, c₂₃ = 1.0) =
    lorenz_lorenz_lorenz_transitive(u₀, σ₁, σ₂, σ₃, ρ₁, ρ₂, ρ₃, β₁, β₂, β₃, c₁₂, c₂₃)


module SkronTest

using Test
using DiffPrecTest: svec, skron
using SparseArrays, Random, LinearAlgebra

function Q(p::Int)
  _c = 1. / sqrt(2.)
  rp = div((p+1)*p,2)
  Q = spzeros( rp,  p*p )
  ind = 0
  for ci=1:p
    for ri=ci:p
      ind += 1
      if ri == ci
        Q[ind, (ci-1)*p + ri] = 1.
      else
        Q[ind, (ci-1)*p + ri] = _c
        Q[ind, (ri-1)*p + ci] = _c
      end
    end
  end
  Q
end

function Qt(p::Int)
  _c = 1. / sqrt(2.)
  rp = div((p+1)*p,2)
  Q = spzeros( p*p, rp )
  ind = 0
  for ci=1:p
    for ri=ci:p
      ind += 1
      if ri == ci
        Q[(ci-1)*p + ri, ind] = 1.
      else
        Q[(ci-1)*p + ri, ind] = _c
        Q[(ri-1)*p + ci, ind] = _c
      end
    end
  end
  Q
end

@testset "skron" begin
  p = 5
  x = randn(p)
  y = randn(p)

  num_row_reduced = div(p*(p+1),2)
  o = zeros( num_row_reduced )
  for i=1:num_row_reduced
      o[i] = svec(x, y, i)
  end

  @test o ≈ Q(p) * kron(x, y)

  o2 = zeros( num_row_reduced, num_row_reduced )
  for ri=1:num_row_reduced
      for ci=1:num_row_reduced
          o2[ri, ci] = skron(x, y, ri, ci)
      end
  end

  @test o2 ≈ Q(p) * ( kron(x*x', y*y') + kron(y*y', x*x') ) * Qt(p) / 2.

  @test skron(x, y, 1:4, 3:6) ≈ (Q(p) * ( kron(x*x', y*y') + kron(y*y', x*x') ) * Qt(p) / 2.)[1:4, 3:6]

  # Q_Sx_skron_Sy_Q
  X = randn(100, 5)
  Y = randn(100, 5)
  Sx = Symmetric( X'*X / 100 )
  Sy = Symmetric( Y'*Y / 100 )

  @test skron(Sx, Sy) ≈ Q(p) * ( kron(Sx, Sy) + kron(Sy, Sx) ) * Qt(p) / 2.

  I = 1:num_row_reduced
  @test skron(Sx, Sy, I) ≈ Q(p) * ( kron(Sx, Sy) + kron(Sy, Sx) ) * Qt(p) / 2.

  I = 3:10
  J = 4:7
  @test skron(Sx, Sy, I, J) ≈ (Q(p) * ( kron(Sx, Sy) + kron(Sy, Sx) ) * Qt(p) / 2.)[I, J]

  I = [2, 5, 7, 8]
  @test skron(Sx, Sy, I) ≈ (Q(p) * ( kron(Sx, Sy) + kron(Sy, Sx) ) * Qt(p) / 2.)[I, I]

  I = 10:15
  @test skron(Sx, Sy, I) ≈ (Q(p) * ( kron(Sx, Sy) + kron(Sy, Sx) ) * Qt(p) / 2.)[I, I]

  # Q Sx skron yy' Qt
  yy = view(Y, 1, :) * view(Y, 1, :)'
  @test skron(Sx, view(Y, 1, :), I, I) ≈ (Q(p) * ( kron(Sx, yy) + kron(yy, Sx) ) * Qt(p) / 2.)[I, I]
end

@testset "svec" begin
  p = 5
  X = randn(100, p)
  Y = randn(100, p)
  Sx = Symmetric( X'*X / 100 )
  Sy = Symmetric( Y'*Y / 100 )

  # Q_Sx_minus_Sy
  ox = svec(Sx)
  oy = svec(Sy)
  @test ox - oy ≈ Q(p) * vec( Sx - Sy )

  I = 10:15
  @test svec(Sx, I) ≈ (Q(p) * vec( Sx ))[I]

  I = [2, 5, 7, 8]
  @test svec(Sy, I) ≈ (Q(p) * vec( Sy ))[I]

  @test svec(view(X, 1, :), view(Y, 1, :), I) ≈ (Q(p) * kron(view(X, 1, :), view(Y, 1, :)))[I]
end

end


using ConvexBodyProximityQueries
using CurveProximityQueries
using Plots
import Random: seed!

pgfplots()

seed!(22)
poly = randpoly([-0.5, 0.4]; n=20, scale=0.5)

curve = Bernstein([[0.944223, 0.177264], [0.46193, 0.105887], [0.791169, 0.172386], [0.256546, 0.463735], [0.70966, 0.3173], [0.664713, 0.830956], [0.952688, 0.736242], [0.00763912, 0.97189], [0.270595, 0.827442], [0.16636, 0.996615], [0.381073, 0.543077], [0.225381, 0.754049], [0.570043, 0.638543], [0.254696, 0.988563]])

pts = closest_points(poly, curve)

plot(poly,
     fillalpha=1.0,
     fillcolor=Gray(0.7),
    )
plot!(curve,
      linewidth=2.0,
      linecolor=:black,
     )
plot!(pts,
      linestyle=:dash,
      markersize=8.0,
      linewidth=2.0,
     )
savefig("out/fig1a.pdf")


"""Gaussian process examples."""

using Distributions
using LinearAlgebra
using PyPlot
using Random

μ(X, m) = [m(x) for x in X]
Σ(X, k) = [k(x, xp) for x in X, xp in X]
K(X, Xp, k) = [k(x, xp) for x in X, xp in Xp]

# Mean functions.
mean_const(β) = (x) -> β
mean_lin(β) = (x) -> dot(x, β)
mean_zero() = (x) -> 0.0

# TODO(kykim): More mean functions.
# - Periodic mean: m(x)=a cos(2πx/p)+b sin(2πx/p).
# - Polynomical mean.

# Covariance functions.
cov_const(σ) = (x, xp) -> σ
cov_exp(l) = (x, xp) -> exp(-(x - xp) / l)
cov_gamma_exp(γ, l) = (x, xp) -> exp(-((x - xp) / l)^γ)
cov_sq_exp(l) = (x, xp) -> exp(-(x - xp)^2 / (2 * l^2))

mutable struct GaussianProcess
    m  # mean
    k  # covariance function
    X  # design points
    y  # objective values
    ν  # noise variance
end

function mvnrand(μ, Σ, inflation=1e-6)
    N = MvNormal(μ, Σ + inflation * I)
    return rand(N)
end
Base.rand(gp, X) = mvrand(μ(X, gp.m), Σ(X, gp.k))

function predict(gp, X_pred)
    X_pred = (isa(X_pred, Vector) ? X_pred : [X_pred])
    m, k, ν = gp.m, gp.k, gp.ν
    tmp = K(X_pred, gp.X, k) / (K(gp.X, gp.X, k) + ν * I)
    μ_p = μ(X_pred, m) + tmp * (gp.y - μ(gp.X, m))
    S = K(X_pred, X_pred, k) - tmp * K(gp.X, X_pred, k)
    ν_p = diag(S) .+ eps()
    return (μ_p, ν_p)
end

function plot_gp(gp, lb, ub, true_y::Function=nothing)
    xl = [x for x in lb:0.05:ub]

    # Plot observations.
    plot(gp.X, gp.y, "o", markersize=3, label="obs")

    # Plot the true y if available.
    if true_y != nothing
        plot(xl, [true_y(x) for x in xl], label="true y", lw=0.5)
    end

    # Plot predicted function mean with confidence interval.
    gp_pred = predict(gp, xl)
    gp_pred_mean, gp_pred_var = first(gp_pred), last(gp_pred)
    plot(xl, gp_pred_mean, label="pred mean", lw=0.5)
    gp_pred_c95 = 1.96 * gp_pred_var
    fill_between(xl, gp_pred_mean - gp_pred_c95, gp_pred_mean + gp_pred_c95,
                 alpha=0.1)
    xlabel("x"); ylabel("y"); legend();
    # savefig("myplot.png")
end

#################################
# Below are experiment functions.
#################################

# An example from stor-i.github.io/GaussianProcesses.jl/latest/Regression.
function simple_regression()
    Random.seed!(0)

    n = 10
    x = 2π * rand(n)
    y = sin.(x) + 0.05 * rand(n)

    gp = GaussianProcess(mean_zero(), cov_sq_exp(1.0), x, y, 1.0)
    true_y = (x) -> sin(x)
    figure(1, figsize=(9.0, 6.0))
    plot_gp(gp, -2.0, 8.0, true_y)
end

simple_regression()


function grad{T <: Number}(
    f::Function,
    x::T;
    method::Symbol = :finite,
    direction::Symbol = :central,
) # -> T
    if method == :finite
        return finitegrad(f, x, direction)
    elseif method == :ad
        return autograd(f, x)
    else
        throw(ArgumentError("method must be :finite or :ad"))
    end
end

function grad!{T <: Number}(
    gr::Vector{T},
    f::Function,
    x::Vector{T};
    method::Symbol = :finite,
    direction::Symbol = :central,
) # -> Vector{T}
    if method == :finite
        finitegrad!(gr, f, x, direction)
    elseif method == :ad
        # TODO: Raise a warning here that memory is being allocated?
        autograd!(gr, Array(Dual{T}, length(x)), f, x)
    else
        throw(ArgumentError("method must be :finite or :ad"))
    end
    return gr
end

function grad!{T <: Number}(
    gr::Vector{T},
    x_d::Vector{Dual{T}},
    f::Function,
    x::Vector{T};
    method::Symbol = :finite,
    direction::Symbol = :central,
) # -> Vector{T}
    if method == :finite
        finitegrad!(gr, f, x, direction)
    elseif method == :ad
        autograd!(gr, x_d, f, x)
    else
        throw(ArgumentError("method must be :finite or :ad"))
    end
    return gr
end

function grad{T <: Number}(
    f::Function,
    x::Vector{T};
    method::Symbol = :finite,
    direction::Symbol = :central,
) # -> Vector{T}
    n = length(x)
    gr = Array(T, n)
    if method == :ad
        x_d = Array(Dual{T}, n)
        grad!(gr, x_d, f, x, method = method, direction = direction)
    else
        grad!(gr, f, x, method = method, direction = direction)
    end
    return gr
end


#! /usr/bin/julia

# Rosetta Code, Linear congruential generator

function lcg_maker{T<:Integer}(r::T, a::T, c::T, m::T, sh::T)
    state = r
    function lcg_rand()
        state = mod(a*state + c, m)
        return state >> sh
    end
    return lcg_rand
end

snum = 10
seed = 0
bsd_rand = lcg_maker(seed, 1103515245, 12345, 2^31, 0)

print("The first ", snum, " results for a BSD rand() seeded with ")
println(seed, ":")

for i in 1:snum
    println(@sprintf "%14d" bsd_rand())
end

ms_rand = lcg_maker(seed, 214013, 2531011, 2^31, 16)

println()
print("The first ", snum, " results for a M\$ rand() seeded with ")
println(seed, ":")

for i in 1:snum
    println(@sprintf "%14d" ms_rand())
end


using Knet, KnetModules

# Model defition
type MLP <: KnetModule
    hiddens::Array{Linear, 1}
    output::Linear
    actfn::Activation
end

function MLP(ninput::Int, nhidden::Int, noutput::Int;
             nhiddens=1, actfn=ReLU())
    hiddens = []
    for i = 1:nhiddens
        push!(hiddens, Linear(nhidden, ninput))
        ninput = nhidden
    end
    output = Linear(noutput, nhidden)
    return MLP(hiddens, output, actfn)
end

function forward(ctx, m::MLP, x)
    for hidden in m.hiddens
        x = @mc m.actfn(@mc hidden(x))
    end
    return @mc m.output(x)
end

# Training epuch
function train(model, gradfn, data, optim)
    for (x, y) in data
        grads = gradfn(x, y)
        update!(model, grads, optim)
    end
end

include(Knet.dir("data","mnist.jl"))
xtrn, ytrn, xtst, ytst = mnist()
dtrn = minibatch(xtrn, ytrn, 100)
dtst = minibatch(xtst, ytst, 100)

model = MLP(784, 64, 10)
gradfn = grad(model, nll)
optim = optimizers(model, Adam)

println("epoch 0",
        "\ntrn accuracy: ", accuracy(model, dtrn),
        "\ntst accuracy: ", accuracy(model, dtst))
for epoch=1:10
    train(model, gradfn, dtrn, optim)
    println("\nepoch $epoch",
            "\ntrn accuracy: ", accuracy(model, dtrn),
            "\ntst accuracy: ", accuracy(model, dtst))
end

# Open this to play with serialization
const ser_demo = false
if ser_demo
    # Serialization
    const filename = "mnist_mlp.jld"
    info("Saving model")
    # save_module filename, model and optional addition state (here, it is optim)
    save_module(filename, model, Dict(["optim"=>optim]))

    reset_ctx!() # clear old model for memory efficiency
    model, state = load_module(filename)
    println("Accuracy of restored ", accuracy(model, dtst))

    optim = state["optim"]
    gradfn = grad(model, nll)
    train(model, gradfn, dtrn, optim)
    println("Test accuracy after one step more training ", accuracy(model, dtst))
end


#   Copyright (c) 2020-2022 Velexi Corporation
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

"""
reject.jl defines methods for the reject(x, y) function
"""

# --- Exports

export reject

# --- Method definitions

using LinearAlgebra: ⋅

"""
    reject(vectors::Matrix, B::AbstractBlade; normalize=false)::Matrix

Compute rejections of `vectors` from `B`. When `normalize` is true, the
rejection vectors are normalized.
"""
function reject end

# ------ Specializations involving a Blade instance

# vectors::Matrix, B::Blade
function reject(vectors::Matrix, B::Blade; normalize::Bool=false)
    # --- Check arguments

    if size(vectors, 1) != dim(B)
        throw(DimensionMismatch("`dim(vectors)` not equal to `dim(B)`"))
    end

    # --- Compute rejections using modified Gram-Schmidt algorithm

    # Initialize rejections
    rejections = Matrix{typeof(volume(B))}(vectors)

    # Remove basis(B) from rejections
    for idx_B in 1:grade(B)
        B_column = basis(B)[:, idx_B]
        for idx in 1:size(rejections, 2)
            rejections[:, idx] -=
                (rejections[:, idx] ⋅ B_column) * B_column
        end
    end

    # Normalize rejection vectors
    if normalize
        for idx in 1:size(rejections, 2)
            norm = LinearAlgebra.norm(rejections[:, idx])
            if norm > 0
                rejections[:, idx] /= norm
            end
        end
    end

    return rejections
end


"""
    Wishart(ν, S)
```julia
ν::Real           degrees of freedom (whole number or a real number greater than p - 1)
S::AbstractPDMat  p x p scale matrix
```
The [Wishart distribution](http://en.wikipedia.org/wiki/Wishart_distribution)
generalizes the gamma distribution to ``p\\times p`` real, positive semidefinite
matrices ``\\mathbf{H}``.

If ``\\nu>p-1``, then ``\\mathbf{H}\\sim \\textrm{W}_p(\\nu, \\mathbf{S})``
has rank ``p`` and its probability density function is

```math
f(\\mathbf{H};\\nu,\\mathbf{S}) = \\frac{1}{2^{\\nu p/2} \\left|\\mathbf{S}\\right|^{\\nu/2} \\Gamma_p\\left(\\frac {\\nu}{2}\\right ) }{\\left|\\mathbf{H}\\right|}^{(\\nu-p-1)/2} e^{-(1/2)\\operatorname{tr}(\\mathbf{S}^{-1}\\mathbf{H})}.
```

If ``\\nu\\leq p-1``, then ``\\mathbf{H}`` is rank ``\\nu`` and it has
a density with respect to a suitably chosen volume element on the space of
positive semidefinite matrices. See [here](https://doi.org/10.1214/aos/1176325375).

For integer ``\\nu``, a random matrix given by

```math
\\mathbf{H} = \\mathbf{X}\\mathbf{X}^{\\rm{T}},
\\quad\\mathbf{X} \\sim \\textrm{MN}_{p,\\nu}(\\mathbf{0}, \\mathbf{S}, \\mathbf{I}_{\\nu})
```

has ``\\mathbf{H}\\sim \\textrm{W}_p(\\nu, \\mathbf{S})``.
For non-integer ``\\nu``, Wishart matrices can be generated via the
[Bartlett decomposition](https://en.wikipedia.org/wiki/Wishart_distribution#Bartlett_decomposition).
"""
struct Wishart{T<:Real, ST<:AbstractPDMat, R<:Integer} <: ContinuousMatrixDistribution
    df::T          # degree of freedom
    S::ST          # the scale matrix
    logc0::T       # the logarithm of normalizing constant in pdf
    rank::R        # rank of a sample
    singular::Bool # singular of nonsingular wishart?
end

#  -----------------------------------------------------------------------------
#  Constructors
#  -----------------------------------------------------------------------------

function Wishart(df::T, S::AbstractPDMat{T}) where T<:Real
    df > 0 || throw(ArgumentError("df must be positive. got $(df)."))
    p = dim(S)
    singular = df <= p - 1
    if singular
        isinteger(df) || throw(
            ArgumentError("df of a singular Wishart distribution must be an integer (got $df)")
        )
    end
    rnk::Integer = ifelse(singular, df, p)
    logc0 = wishart_logc0(df, S, rnk)
    _df, _logc0 = promote(df, logc0)
    Wishart{typeof(_df), typeof(S), typeof(rnk)}(_df, S, _logc0, rnk, singular)
end

function Wishart(df::Real, S::AbstractPDMat)
    T = Base.promote_eltype(df, S)
    Wishart(T(df), convert(AbstractArray{T}, S))
end

Wishart(df::Real, S::Matrix) = Wishart(df, PDMat(S))
Wishart(df::Real, S::Cholesky) = Wishart(df, PDMat(S))

#  -----------------------------------------------------------------------------
#  REPL display
#  -----------------------------------------------------------------------------

show(io::IO, d::Wishart) = show_multline(io, d, [(:df, d.df), (:S, Matrix(d.S))])

#  -----------------------------------------------------------------------------
#  Conversion
#  -----------------------------------------------------------------------------

function convert(::Type{Wishart{T}}, d::Wishart) where T<:Real
    P = convert(AbstractArray{T}, d.S)
    Wishart{T, typeof(P), typeof(d.rank)}(T(d.df), P, T(d.logc0), d.rank, d.singular)
end
function convert(::Type{Wishart{T}}, df, S::AbstractPDMat, logc0, rnk, singular) where T<:Real
    P = convert(AbstractArray{T}, S)
    Wishart{T, typeof(P), typeof(rnk)}(T(df), P, T(logc0), rnk, singular)
end

#  -----------------------------------------------------------------------------
#  Properties
#  -----------------------------------------------------------------------------

insupport(::Type{Wishart}, X::AbstractMatrix) = ispossemdef(X)
function insupport(d::Wishart, X::AbstractMatrix)
    size(X) == size(d) || return false
    if d.singular
        return ispossemdef(X, rank(d))
    else
        return isposdef(X)
    end
end

dim(d::Wishart) = dim(d.S)
size(d::Wishart) = (p = dim(d); (p, p))
rank(d::Wishart) = d.rank
params(d::Wishart) = (d.df, d.S)
@inline partype(d::Wishart{T}) where {T<:Real} = T

mean(d::Wishart) = d.df * Matrix(d.S)

function mode(d::Wishart)
    r = d.df - dim(d) - 1.0
    r > 0.0 || throw(ArgumentError("mode is only defined when df > p + 1"))
    return Matrix(d.S) * r
end

function meanlogdet(d::Wishart)
    d.singular && return -Inf
    p = dim(d)
    df = d.df
    v = logdet(d.S) + p * logtwo
    for i = 1:p
        v += digamma(0.5 * (df - (i - 1)))
    end
    return v
end

function entropy(d::Wishart)
    d.singular && throw(ArgumentError("entropy not defined for singular Wishart."))
    p = dim(d)
    df = d.df
    -d.logc0 - 0.5 * (df - p - 1) * meanlogdet(d) + 0.5 * df * p
end

#  Gupta/Nagar (1999) Theorem 3.3.15.i
function cov(d::Wishart, i::Integer, j::Integer, k::Integer, l::Integer)
    S = Matrix(d.S)
    d.df * (S[i, k] * S[j, l] + S[i, l] * S[j, k])
end

function var(d::Wishart, i::Integer, j::Integer)
    S = Matrix(d.S)
    d.df * (S[i, i] * S[j, j] + S[i, j] ^ 2)
end

#  -----------------------------------------------------------------------------
#  Evaluation
#  -----------------------------------------------------------------------------

function wishart_logc0(df::Real, S::AbstractPDMat, rnk::Integer)
    p = dim(S)
    if df <= p - 1
        return singular_wishart_logc0(p, df, S, rnk)
    else
        return nonsingular_wishart_logc0(p, df, S)
    end
end

function logkernel(d::Wishart, X::AbstractMatrix)
    if d.singular
        return singular_wishart_logkernel(d, X)
    else
        return nonsingular_wishart_logkernel(d, X)
    end
end

#  Singular Wishart pdf: Theorem 6 in Uhlig (1994 AoS)
function singular_wishart_logc0(p::Integer, df::Real, S::AbstractPDMat, rnk::Integer)
    h_df = df / 2
    -h_df * (logdet(S) + p * typeof(df)(logtwo)) - logmvgamma(rnk, h_df) + (rnk*(rnk - p) / 2)*typeof(df)(logπ)
end

function singular_wishart_logkernel(d::Wishart, X::AbstractMatrix)
    df = d.df
    p = dim(d)
    r = rank(d)
    L = eigvals(Hermitian(X), (p - r + 1):p)
    0.5 * ((df - (p + 1)) * sum(log.(L)) - tr(d.S \ X))
end

#  Nonsingular Wishart pdf
function nonsingular_wishart_logc0(p::Integer, df::Real, S::AbstractPDMat)
    h_df = df / 2
    -h_df * (logdet(S) + p * typeof(df)(logtwo)) - logmvgamma(p, h_df)
end

function nonsingular_wishart_logkernel(d::Wishart, X::AbstractMatrix)
    df = d.df
    p = dim(d)
    Xcf = cholesky(X)
    0.5 * ((df - (p + 1)) * logdet(Xcf) - tr(d.S \ X))
end

#  -----------------------------------------------------------------------------
#  Sampling
#  -----------------------------------------------------------------------------

function _rand!(rng::AbstractRNG, d::Wishart, A::AbstractMatrix)
    if d.singular
        A .= zero(eltype(A))
        A[:, 1:rank(d)] = randn(rng, dim(d), rank(d))
    else
        _wishart_genA!(rng, dim(d), d.df, A)
    end
    unwhiten!(d.S, A)
    A .= A * A'
end

function _wishart_genA!(rng::AbstractRNG, p::Int, df::Real, A::AbstractMatrix)
    # Generate the matrix A in the Bartlett decomposition
    #
    #   A is a lower triangular matrix, with
    #
    #       A(i, j) ~ sqrt of Chisq(df - i + 1) when i == j
    #               ~ Normal()                  when i > j
    #
    A .= zero(eltype(A))
    for i = 1:p
        @inbounds A[i,i] = rand(rng, Chi(df - i + 1.0))
    end
    for j in 1:p-1, i in j+1:p
        @inbounds A[i,j] = randn(rng)
    end
end

#  -----------------------------------------------------------------------------
#  Test utils
#  -----------------------------------------------------------------------------

function _univariate(d::Wishart)
    check_univariate(d)
    df, S = params(d)
    α = df / 2
    β = 2Matrix(S)[1]
    return Gamma(α, β)
end

function _rand_params(::Type{Wishart}, elty, n::Int, p::Int)
    n == p || throw(ArgumentError("dims must be equal for Wishart"))
    ν = elty( n - 1 + abs(10randn()) )
    S = (X = 2rand(elty, n, n) .- 1; X * X')
    return ν, S
end


function smolyak_grid_full(node_type::Function,d::S,mu::S) where {S<:Integer}

  T = typeof(1.0)

  multi_index        = generate_multi_index(d,mu)
  unique_multi_index = sort(unique(multi_index))
  unique_node_number = m_i(unique_multi_index)

  # Create base nodes to be used in the sparse grid

  base_nodes   = Array{Array{T,1},1}(undef,length(unique_node_number))
  base_weights = Array{Array{T,1},1}(undef,length(unique_node_number))
  @inbounds for i = 1:length(unique_node_number)
    base_nodes[i], base_weights[i] = node_type(unique_node_number[i])
  end

  # Select the relevant polynomials from the multi index

  ii = (sum(multi_index,dims=2) .>= max(d,mu+1)).*(sum(multi_index,dims=2) .<= d+mu)
  multi_index_full = zeros(S,sum(ii),size(multi_index,2))
  j = 1
  @inbounds for i = 1:size(multi_index,1)
    if ii[i] == true
      multi_index_full[j,:] = multi_index[i,:]
      j += 1
    end
  end

  # Construct the sparse grid from the nodes

  nodes = Array{T,2}(undef,determine_grid_size_full(multi_index_full))
  l = 1
  for j = 1:size(multi_index_full,1)
    new_nodes = base_nodes[multi_index_full[j,1]]  # Here new_nodes is a 1d array
    for i = 2:d
      new_nodes = combine_nodes(new_nodes,base_nodes[multi_index_full[j,i]])  # Here new_nodes becomes a 2d array
    end
    m = size(new_nodes,1)
    nodes[l:l+m-1,:] = new_nodes
    l += m
  end

  if d == 1
    nodes = nodes[:]
  end

  return nodes, multi_index_full

end

function smolyak_grid_full(node_type::Function,d::S,mu::S,domain::Union{Array{T,1},Array{T,2}}) where {S<:Integer, T<:AbstractFloat}

  if size(domain,2) != d
    error("domain is inconsistent with the number of dimensions")
  end

  nodes, multi_index_full = smolyak_grid_full(node_type,d,mu)

  nodes = scale_nodes(nodes,domain)

  return nodes, multi_index_full

end

function determine_grid_size_full(mi)

  temp = similar(mi)

  @inbounds for i = 1:size(mi,1)
    @inbounds for j = 1:size(mi,2)
      if mi[i,j] == 1
        temp[i,j] = 1
      else
        temp[i,j] = 2^(mi[i,j]-1)+1
      end
    end
  end

  s = 0
  @inbounds for i = 1:size(mi,1)
    t = 1
    @inbounds for j = 1:size(mi,2)
      t *= temp[i,j]
    end
    s += t
  end

  return (s, size(mi,2))

end

function master_index(multi_index::Array{S,2}) where {S<:Integer}

  temp_ind   = similar(multi_index)
  master_ind = zeros(S,size(multi_index,1),2)

  @inbounds for i in eachindex(multi_index)
    if multi_index[i] == 1
      temp_ind[i] = 1
    else
      temp_ind[i] = m_i(multi_index[i])
    end
  end

  master_ind[1,1] = 1
  master_ind[1,2] = prod(temp_ind[1,:])

  @inbounds for i = 2:size(master_ind,1)
    master_ind[i,1] = master_ind[i-1,1] + master_ind[i-1,2]
    master_ind[i,2] = prod(temp_ind[i,:])
  end

  return master_ind

end

function cheb_poly(order::S,x::T) where {S<:Integer, T<:AbstractFloat}

  p  = 1.0
  p1 = 0.0
  p2 = 0.0

  for i = 2:order+1
    if i == 2
	  p1, p = p, x
    else
      p2, p1 = p1, p
  	  p  = 2*x*p1-p2
    end
  end

  return p

end

function prod_cjs(max_grid::Union{Array{T,1},Array{T,2}},min_grid::Union{Array{T,1},Array{T,2}},poly_grid::Array{T,2}) where {T<:AbstractFloat}

  cjs = ones(size(poly_grid))

  @inbounds for i = 1:size(poly_grid,1)
    @inbounds for j = 1:size(poly_grid,2)
      if poly_grid[i,j] == max_grid[j] || poly_grid[i,j] == min_grid[j]
        cjs[i,j] *= 2
      end
    end
  end

  return prod(cjs,dims=2)

end

function compute_scale_factor(multi_index::Array{S,1}) where {S<:Integer}

  scale_factor = 1.0

  @inbounds for i = 1:length(multi_index)
    if multi_index[i] > 1
      scale_factor *= 2.0/(m_i(multi_index[i])-1)
    end
  end

  return scale_factor

end

function smolyak_weights_full(y_f::Array{T,1},grid::Union{Array{T,1},Array{T,2}},multi_index::Array{S,2}) where {S<:Integer, T<:AbstractFloat}

  mi = sum(multi_index,dims=2)
  d  = size(multi_index,2)
  mu = maximum(mi)-d

  max_grid = maximum(grid,dims=1)
  min_grid = minimum(grid,dims=1)

  weights = Array{Array{T,1},1}(undef,size(multi_index,1))
  g_ind = master_index(multi_index)

  @inbounds for i = 1:size(g_ind,1) # This loops over the number of polynomials
    ws = zeros(g_ind[i,2])
    poly_grid = grid[g_ind[i,1]:g_ind[i,1]+g_ind[i,2]-1,:]
    poly_y    = y_f[g_ind[i,1]:g_ind[i,1]+g_ind[i,2]-1]
	  if size(grid,1) == 1 # This is to accommodate the mu = 0 case
	    cjs_prod = ones(size(poly_grid,1))
	  else
      cjs_prod = prod_cjs(max_grid,min_grid,poly_grid)
    end
    @inbounds for l = 1:g_ind[i,2] # This loops over the weights
      ll = CartesianIndices(Tuple(m_i(multi_index[i,:])))[l]
      @inbounds for j = 1:g_ind[i,2] # This loops over the nodes
        rhs_term = cheb_poly(ll[1]-1,poly_grid[j,1])*poly_y[j]
        @inbounds for k = 2:size(poly_grid,2) # This loops over the polynomial terms in the product
          rhs_term *= cheb_poly(ll[k]-1,poly_grid[j,k])
        end
        ws[l] += rhs_term/cjs_prod[j]
      end
      scale_factor = compute_scale_factor(multi_index[i,:])
      ws[l] = scale_factor*(1/cjs_prod[l])*ws[l]
    end
    weights[i] = (-1)^(d+mu-mi[i])*factorial(d-1)/(factorial(d+mu-mi[i])*factorial(-1-mu+mi[i]))*ws
  end

  return weights

end

function smolyak_weights_full(y_f::Array{T,1},grid::Union{Array{T,1},Array{T,2}},multi_index::Array{S,2},domain::Array{T,2}) where {S<:Integer, T<:AbstractFloat}

  d = size(multi_index,2)
  grid = copy(grid)
  for i = 1:d
    grid[:,i] = normalize_node(grid[:,i],domain[:,i])
  end

  weights = smolyak_weights_full(y_f,grid,multi_index)

  return weights

end

function smolyak_evaluate_full(weights::Array{Array{T,1},1},point::Array{T,1},multi_index::Array{S,2}) where {S<:Integer, T<:AbstractFloat}

  d = size(multi_index,2)

  evaluated_polynomials = zeros(size(multi_index,1))

  @inbounds for i = 1:size(multi_index,1) # This loops over the number of polynomials
  	@inbounds for l = 1:length(weights[i])
      ll = CartesianIndices(Tuple(m_i(multi_index[i:i,:])))[l]
      temp = weights[i][l]*cheb_poly(ll[1]-1,point[1])
      @inbounds for k = 2:d
  	    temp *= cheb_poly(ll[k]-1,point[k])
      end
      evaluated_polynomials[i] += temp
    end
    #evaluated_polynomials[i] *= (-1)^(d+mu-mi[i])*factorial(d-1)/(factorial(d+mu-mi[i])*factorial(-1-mu+mi[i]))
  end

  return sum(evaluated_polynomials)

end

function smolyak_evaluate_full(weights::Array{Array{T,1},1},point::Array{T,1},multi_index::Array{S,2},domain::Array{T,2}) where {S<:Integer, T<:AbstractFloat}

  d = size(multi_index,2)
  point = copy(point)
  for i = 1:d
    point[i] = normalize_node(point[i],domain[:,i])
  end

  estimate = smolyak_evaluate_full(weights,point,multi_index)

  return estimate

end

function deriv_cheb_poly(order::S,x::T) where {S<:Integer, T<:AbstractFloat}

  p0 = one(T)
  p1 = zero(T)
  p2 = zero(T)
  pd0 = zero(T)
  pd1 = zero(T)
  pd2 = zero(T)

  for i = 2:order+1
    if i == 2
	  p1, p0 = p0, x
	  pd1, pd0 = pd0, one(T)
    else
      p2, p1 = p1, p0
      pd2, pd1 = pd1, pd0
  	  p0  = 2*x*p1-p2
      pd0 = 2*p1+2*x*pd1-pd2
    end
  end

  return pd0

end

function smolyak_derivative_full(weights::Array{Array{T,1},1},point::Array{T,1},multi_index::Array{S,2},pos::S) where {S<:Integer, T<:AbstractFloat}

  mi = sum(multi_index,dims=2)
  d  = size(multi_index,2)
  mu = maximum(mi)-d

  evaluated_polynomials = zeros(size(multi_index,1))

  for i = 1:size(multi_index,1) # This loops over the number of polynomials
    for l = 1:length(weights[i])
      ll = CartesianIndices(Tuple(m_i(multi_index[i:i,:])))[l]
      if pos == 1
        temp = weights[i][l]*deriv_cheb_poly(ll[1]-1,point[1])
      else
        temp = weights[i][l]*cheb_poly(ll[1]-1,point[1])
      end
      for k = 2:d
        if k == pos
          temp *= deriv_cheb_poly(ll[k]-1,point[k])
        else
          temp *= cheb_poly(ll[k]-1,point[k])
        end
	    end
	    evaluated_polynomials[i] += temp
    end
	  #evaluated_polynomials[i] *= (-1)^(d+mu-mi[i])*factorial(d-1)/(factorial(d+mu-mi[i])*factorial(-1-mu+mi[i]))
  end

  evaluated_derivative = sum(evaluated_polynomials)

  return evaluated_derivative

end

function smolyak_derivative_full(weights::Array{Array{T,1},1},point::Array{T,1},multi_index::Array{S,2},domain::Array{T,2},pos::S) where {S<:Integer, T<:AbstractFloat}

  d = size(multi_index,2)
  point = copy(point)
  for i = 1:d
    point[i] = normalize_node(point[i],domain[:,i])
  end

  evaluated_derivative = smolyak_derivative_full(weights,point,multi_index,pos)

  return evaluated_derivatives

end

function smolyak_gradient_full(weights::Array{Array{T,1},1},point::Array{T,1},multi_index::Array{S,2}) where {S<:Integer, T<:AbstractFloat}

  d = size(multi_index,2)

  gradient = Array{T,2}(undef,1,d)
  for i = 1:d
    gradient[i] = smolyak_derivative_full(weights,point,multi_index,i)
  end

  return gradient

end

function smolyak_gradient_full(weights::Array{Array{T,1},1},point::Array{T,1},multi_index::Array{S,2},domain::Array{T,2}) where {S<:Integer, T<:AbstractFloat}

  d = size(multi_index,2)

  gradient = Array{T,2}(undef,1,d)
  for i = 1:d
    gradient[i] = smolyak_derivative_full(weights,point,multi_index,domain,i)
  end

  return gradient

end

###########################################################################

function smolyak_evaluate_full(weights::Array{T,1},multi_index_full::Array{S,2}) where {T<:AbstractFloat,S<:Integer}

  function goo(x::Array{T,1}) where {T <: AbstractFloat}

    return smolyak_evaluate_full(weights,x,multi_index_full)

  end

  return goo

end

function smolyak_evaluate_full(weights::Array{T,1},multi_index_full::Array{S,2},domain::Union{Array{T,1},Array{T,2}}) where {T<:AbstractFloat,S<:Integer}

  function goo(x::Array{T,1}) where {T <: AbstractFloat}

    return smolyak_evaluate_full(weights,x,multi_index_full,domain)

  end

  return goo

end

function smolyak_derivative_full(weights::Array{T,1},multi_index_full::Array{S,2}) where {T<:AbstractFloat,S<:Integer}

  function goo(x::Array{T,1}) where {T <: AbstractFloat}

    return smolyak_derivative_full(weights,x,multi_index_full)

  end

  return goo

end

function smolyak_derivative_full(weights::Array{T,1},multi_index_full::Array{S,2},domain::Union{Array{T,1},Array{T,2}}) where {T<:AbstractFloat,S<:Integer}

  function goo(x::Array{T,1}) where {T <: AbstractFloat}

    return smolyak_derivative_full(weights,x,multi_index_full,domain)

  end

  return goo

end

function smolyak_derivative_full(weights::Array{T,1},multi_index_full::Array{S,2},pos::Array{S,1}) where {T<:AbstractFloat,S<:Integer}

  function goo(x::Array{T,1}) where {T <: AbstractFloat}

    return smolyak_derivative_full(weights,x,multi_index_full,pos)

  end

  return goo

end

function smolyak_derivative_full(weights::Array{T,1},multi_index_full::Array{S,2},domain::Union{Array{T,1},Array{T,2}},pos::Array{S,1}) where {T<:AbstractFloat,S<:Integer}

  function goo(x::Array{T,1}) where {T <: AbstractFloat}

    return smolyak_derivative_full(weights,x,multi_index_full,domain,pos)

  end

  return goo

end


#=

Filename: Model.jl.
Author: Shunsuke Hori

Defining an instance "Model" storing parameters, functional forms, and grid.

=#

struct Model{TF <: AbstractFloat, TR <: Real, TI <: Integer}
    alpha::TR    # Productivity parameter
    beta::TF     # Discount factor
    gamma::TR    # risk aversion
    mu::TR       # First parameter in lognorm(mu, sigma)
    s::TR        # Second parameter in lognorm(mu, sigma)
    grid_min::TR # Smallest grid point
    grid_max::TR # Largest grid point
    grid_size::TI # Number of grid points
    u::Function       # utility function
    u_prime::Function # derivative of utility function
    f::Function       # production function
    f_prime::Function # derivative of production function
    grid::Vector{TR}   # grid
end

"""
construct Model instance using the information of parameters and functional form

arguments: see above

return: Model type instance
"""
function Model(; alpha::Real=0.65,   # Productivity parameter
                 beta::AbstractFloat=0.95,    # Discount factor
                 gamma::Real=1.0,    # risk aversion
                 mu::Real=0.0,       # First parameter in lognorm(mu, sigma)
                 s::Real=0.1,        # Second parameter in lognorm(mu, sigma)
                 grid_min::Real=1e-6,# Smallest grid point
                 grid_max::Real=4.0, # Largest grid point
                 grid_size::Integer=200, # Number of grid points
                 u::Function= c->(c^(1-gamma)-1)/(1-gamma), # utility function
                 u_prime::Function = c-> c^(-gamma),        # u'
                 f::Function = k-> k^alpha,                 # production function
                 f_prime::Function = k -> alpha*k^(alpha-1) # f'
                 )

    grid=collect(linspace(grid_min, grid_max, grid_size))

    if gamma == 1   # when gamma==1, log utility is assigned
        u_log(c) = log(c)
        m = Model(alpha, beta, gamma, mu, s, grid_min, grid_max,
                grid_size, u_log, u_prime, f, f_prime, grid)
    else
        m = Model(alpha, beta, gamma, mu, s, grid_min, grid_max,
                  grid_size, u, u_prime, f, f_prime, grid)
    end
    return m
end


mutable struct Fact
    uin::Vector{Float64}
    uout::Vector{Ptr{Float64}}
    neigs::Vector{Int}
    w::Vector{Float64}
    b::Int
end

Fact() = Fact(Float64[], Ptr{Float64}[], Int[], Float64[], 1)

deg(f::Fact) = length(f.uin)

mutable struct FGBMatching
    N::Int
    γ::Float64
    fnodes::Vector{Fact}
    adjlist::Vector{Vector{Int}}
end

show(io, ::FGBMatching) = show(io, "FGBMatching(...)")

function FGBMatching(net::Network; γ=Inf)
    @assert has_eprop(net, "w")
    @assert has_vprop(net, "b")

    N = nv(net)
    fnodes = [Fact() for i=1:N]
    wmap = eprop(net, "w")
    bmap = vprop(net, "b")

    # Prune Graph
    w = [Float64[] for i=1:N]
    adjlist = [Int[] for i=1:N]
    for i=1:N
        for e in edges(net, i)
            wij = wmap[e]
            if wij < γ
                push!(w[i], wij)
                push!(adjlist[i], dst(e))
            end
        end
    end

    for (i, f) in enumerate(fnodes)
        resize!(f.uin, length(adjlist[i]))
        resize!(f.uout, length(adjlist[i]))
        resize!(f.w, length(adjlist[i]))
    end

    for (i, f) in enumerate(fnodes)
        f.b = bmap[i]
        for (ki, j) in enumerate(adjlist[i])
            f.uin[ki] = 0
            kj = findfirst(==(i), adjlist[j])
            fnodes[j].uout[kj] = getref(f.uin, ki)
            f.w[ki] = w[i][ki]
        end
    end

    FGBMatching(N, γ, fnodes, adjlist)
end

function initrand!(g::FGBMatching)
    for f in g.fnodes
        for k=1:deg(f)
            f.uin[k] = randn()
        end
    end
end

"""
    topk(x, k; rev=false)

Return the indexes of the largest `k` elements in `x`.
If `rev=true`, return the smallest elements.
"""
function topk(x, k; rev=false)
    is = partialsortperm(x, 1:k; rev=!rev)
    return is # xs[is]
end

function update!(f::Fact)
    @unpack w, uin, uout, b = f
    h = w .- uin
    is = topk(h, b+1, rev=true)
    m1 = h[is[b]]
    m2 = h[is[b+1]]

    Δ = abs(uout[is[b]][] - m2)
    for i=1:deg(f)
        uout[i][] = m1
    end
    for i in is[1:b]
        uout[i][] = m2
    end
    return Δ
end

findmatch(f::Fact) = findmatch(f, f.b)

function findmatch(f::Fact, k)
    h = f.w .- f.uin
    is = topk(h, k, rev=true)
    return is
end

function oneBPiter!(g::FGBMatching)
    Δ = 0.
    for a in randperm(g.N)
        d = update!(g.fnodes[a])
        Δ = max(Δ, d)
    end
    return Δ
end

function converge!(g::FGBMatching; maxiters=100, ϵ=1e-8, verbose=true)
    Eold = 0.
    tstop = 0
    it = 0
    while it < maxiters
        it += 1
        Δ = oneBPiter!(g)
        E, matchmap, nviolations = energy(g)
        verbose && @printf("it=%d  E=%.5f  nviolations=%d \tΔ=%f \n", it, E, nviolations, Δ)

        if abs(Eold - E) < ϵ && nviolations == 0
            tstop += 1
            if tstop == 10
                verbose && println("Found ground state")
                break
            end
        else
            tstop = 0
        end

        Eold = E
    end
    return Eold, it
end

function energy(g::FGBMatching)
    @unpack fnodes, N, adjlist = g
    E = 0.
    matchmap = Vector{Vector{Int}}(undef, N)
    for i=1:N
        f = fnodes[i]
        ks = findmatch(f)
        E += sum(f.w[ks])
        matchmap[i] = adjlist[i][ks]
    end
    E /= 2
    nviolations = 0
    for i=1:N
        for j in matchmap[i]
            nviolations += i ∉ matchmap[j]
        end
    end
    return E, matchmap, nviolations
end

"""
    run_bp(net::Network; [γ, maxiters, ϵ, seed, verbose])

Computes the minimum weight perfect b-matching on graph `net`.
The input `net` has to contain the edge property `w` and
the vertex property `b`.

Can optionally impose a cutoff `γ` on the edge costs `w`
in order to sparsify the graph and improve performance.

The output object contains the following fields:

- `energy`: the optimal cost.
- `match`: where `match[i]` is the list of neighbors of vertex
           `i` in the optimal matching.
- `iters`: number of iterations performed.
- `ok`: whether the algorithm was successful or not.
- `num_violations`: the number of inconsistencies in the solution.
- `bpgraph`: a type storing the BP messages.
"""
function run_bp(net::Network;
                γ = Inf,
                maxiters = 10000,
                ϵ = 1e-4,
                seed = -1,
                verbose=true)
    seed > 0 && Random.seed!(seed)
    g = FGBMatching(net; γ)
    initrand!(g)
    E, iters = converge!(g; maxiters, ϵ, verbose)
    E, matchmap, nviolations = energy(g)

    res = (energy = E,
           match = matchmap,
           bpgraph = g,
           num_violations =  nviolations,
           iters = iters,
           ok = (nviolations == 0),)

    return res
end


# This file is part of the ValidatedNumerics.jl package; MIT licensed

# Newton method

# What is this guarded_mid for? Shouldn't it be checking if f(m)==0?
doc"""Returns the midpoint of the interval x, slightly shifted in case
it is zero"""
function guarded_mid{T}(x::Interval{T})
    m = mid(x)
    if m == zero(T)                     # midpoint exactly 0
        α = convert(T, 0.46875)      # close to 0.5, but exactly representable as a floating point
        m = α*x.lo + (one(T)-α)*x.hi   # displace to another point in the interval
    end

    @assert m ∈ x

    m
end

function N{T}(f::Function, x::Interval{T}, deriv::Interval{T})
    m = guarded_mid(x)
    m = Interval(m)

    m - (f(m) / deriv)
end

doc"""If a root is known to be inside an interval,
`newton_refine` iterates the interval Newton method until that root is found."""
function newton_refine{T}(f::Function, f_prime::Function, x::Interval{T};
                          tolerance=eps(T), debug=false)

    debug && (print("Entering newton_refine:"); @show x)

    while diam(x) > tolerance  # avoid problem with tiny floating-point numbers if 0 is a root
        deriv = f_prime(x)
        Nx = N(f, x, deriv)

        debug && @show(x, Nx)

        Nx = Nx ∩ x
        Nx == x && break

        x = Nx
    end

    debug && @show "Refined root", x

    return [Root(x, :unique)]
end

doc"""`newton` performs the interval Newton method on the given function `f`
with its optional derivative `f_prime` and initial interval `x`.
Optional keyword arguments give the `tolerance`, `maxlevel` at which to stop
subdividing, and a `debug` boolean argument that prints out diagnostic information."""

function newton{T}(f::Function, f_prime::Function, x::Interval{T}, level::Int=0;
                   tolerance=eps(T), debug=false, maxlevel=30)

    debug && (print("Entering newton:"); @show(level); @show(x))

    isempty(x) && return Root{T}[]  #[(x, :none)]

    z = zero(x.lo)
    tolerance = abs(tolerance)

    # Tolerance reached or maximum level of bisection
    (diam(x) < tolerance || level >= maxlevel) && return [Root(x, :unknown)]

    deriv = f_prime(x)
    debug && @show(deriv)

    if !(z in deriv)
        Nx = N(f, x, deriv)
        debug && @show(Nx, Nx ⊆ x, Nx ∩ x)

        isempty(Nx ∩ x) && return Root{T}[]

        if Nx ⊆ x
            debug && (print("Refining "); @show(x))
            return newton_refine(f, f_prime, Nx, tolerance=tolerance, debug=debug)
        end

        m = guarded_mid(x)

        debug && @show(x,m)

        # bisect:
        roots = vcat(
            newton(f, f_prime, Interval(x.lo, m), level+1,
                   tolerance=tolerance, debug=debug, maxlevel=maxlevel),
            newton(f, f_prime, Interval(m, x.hi), level+1,
                   tolerance=tolerance, debug=debug, maxlevel=maxlevel)
            # note the nextfloat here to prevent repetition
            )

        debug && @show(roots)

    else  # 0 in deriv; this does extended interval division by hand

        y1 = Interval(deriv.lo, -z)
        y2 = Interval(z, deriv.hi)

        if debug
            @show (y1, y2)
            @show N(f, x, y1)
            @show N(f, x, y2)
        end

        y1 = N(f, x, y1) ∩ x
        y2 = N(f, x, y2) ∩ x

        debug && @show(y1, y2)

        roots = vcat(
            newton(f, f_prime, y1, level+1,
                   tolerance=tolerance, debug=debug, maxlevel=maxlevel),
            newton(f, f_prime, y2, level+1,
                   tolerance=tolerance, debug=debug, maxlevel=maxlevel)
            )

        debug && @show(roots)

    end

    roots = clean_roots(f, roots)

    return roots

end

# use automatic differentiation if no derivative function given:
newton{T}(f::Function, x::Interval{T};  args...) =
    newton(f, x->D(f,x), x; args...)

# newton for vector of intervals:
newton{T}(f::Function, f_prime::Function, xx::Vector{Interval{T}}; args...) =
    vcat([newton(f, f_prime, @interval(x); args...) for x in xx]...)

newton{T}(f::Function,  xx::Vector{Interval{T}}, level; args...) =
    newton(f, x->D(f,x), xx, 0, args...)

newton{T}(f::Function,  xx::Vector{Root{T}}; args...) =
    newton(f, x->D(f,x), [x.interval for x in xx], args...)


### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ a9bbd3ea-fc34-11ea-282a-43f4fc159d39
using Pkg, DrWatson

# ╔═╡ a9e16b46-fc34-11ea-36de-ef647e4e4f6f
begin
	#@quickactivate "StatisticalRethinkingStan"
	using StanSample, StanQuap
	using StatisticalRethinking
end

# ╔═╡ feebbe4e-fc33-11ea-2ee7-d14f977c6497
md"## Clip-05-01-02s.jl"

# ╔═╡ a9e54766-fc34-11ea-1246-93a25f748d1d
md"### snippet 5.1"

# ╔═╡ 10471c94-fc55-11ea-0fec-656bd563d513
md"##### D (Divorce rate), A (MediumAgeMarriage) and M (Marriage rate) are all standardized."

# ╔═╡ a9f32456-fc34-11ea-3af8-c7e0de78a85d
begin
	df = CSV.read(sr_datadir("WaffleDivorce.csv"), DataFrame);
	scale!(df, [:Divorce, :MedianAgeMarriage, :Marriage])
end;

# ╔═╡ 4d0ca900-fc53-11ea-0f6c-e7d09d775387
md"##### The model m5.1s represents a regression of Divorce on MedianAgeMarriage and is defined as:"

# ╔═╡ 25c70d5a-fc54-11ea-3910-a9276dc7b696
md"
```
model {
  vector[N] mu;               // mu is a vector
  a ~ normal(0, 0.2);         // Priors
  bA ~ normal(0, 0.5);
  sigma ~ exponential(1);
  mu = a + bA * A;
  D ~ normal(mu , sigma);     // Likelihood
}
```
"

# ╔═╡ b8ae47cc-01af-11eb-3cbd-a749b6e3d327
md"##### The Stan language model."

# ╔═╡ b5431248-01af-11eb-03ac-ed6a17e17e39
stan5_1_priors = "
	data {
	 int N; // Sample size
	 vector[N] D; // Outcome
	 vector[N] A; // Predictor
	}

	parameters {
	 real a; // Intercept
	 real bA; // Slope (regression coefficients)
	 real < lower = 0 > sigma;    // Error SD
	}

	model {
	  vector[N] mu;               // mu is a vector
	  a ~ normal(0, 0.2);         // Priors
	  bA ~ normal(0, 0.5);
	  sigma ~ exponential(1);
	  mu = a + bA * A;
	  D ~ normal(mu , sigma);   // Likelihood
	}
";

# ╔═╡ db59a030-423e-11eb-3eb2-5d1f7d3dd0f5
begin
	m5_1s_priors = SampleModel("m5.1.priors", stan5_1_priors)
	rc5_1s_priors = stan_sample(m5_1s_priors;
		data = Dict("N" => 0, "D" => [], "A" => []))
	success(rc5_1s_priors) &&
		(part5_1s_priors = read_samples(m5_1s_priors, :particles))
end

# ╔═╡ 55ef3d46-423f-11eb-1d83-55276a40b702
if success(rc5_1s_priors)
	priors5_1s_df = read_samples(m5_1s_priors, :dataframe)
	xi = -3.0:0.1:3.0
	plot(xlab="Medium age marriage (scaled)", ylab="Divorce rate (scaled)",
		title="Showing 50 regression lines")
	for i in 1:50
		local yi = mean(priors5_1s_df[i, :a]) .+ priors5_1s_df[i, :bA] .* xi
		plot!(xi, yi, color=:lightgrey, leg=false)
	end
	scatter!(df[:, :MedianAgeMarriage_s], df[!, :Divorce_s], color=:darkblue)
end

# ╔═╡ fad134cc-423f-11eb-3398-b7d0e22aad9c
stan5_1 = "
	data {
	 int N; // Sample size
	 vector[N] D; // Outcome
	 vector[N] A; // Predictor
	}

	parameters {
	 real a; // Intercept
	 real bA; // Slope (regression coefficients)
	 real < lower = 0 > sigma;    // Error SD
	}

	model {
	  vector[N] mu;               // mu is a vector
	  a ~ normal(0, 0.2);         // Priors
	  bA ~ normal(0, 0.5);
	  sigma ~ exponential(1);
	  mu = a + bA * A;
	  D ~ normal(mu , sigma);   // Likelihood
	}
";

# ╔═╡ aa241f5e-fc34-11ea-21d5-1b03dfafd34f
md"##### Result rethinking:"

# ╔═╡ aa2a005e-fc34-11ea-09ab-290a8d6b6a9a
rethinking = "
	   mean   sd  5.5% 94.5%
a      0.00 0.10 -0.16  0.16
bA    -0.57 0.11 -0.74 -0.39
sigma  0.79 0.08  0.66  0.91
";

# ╔═╡ cbd5fe10-01af-11eb-055b-bd009d0dca55
begin
	data = (N = size(df, 1), D = df.Divorce_s, A = df.MedianAgeMarriage_s,
		M = df.Marriage_s)
	init = (a=1.0, bA=1.0, bM=1.0, sigma=10.0)
	q5_1s, m5_1s, om5_1s = stan_quap("m5.1s", stan5_1; data, init)
	if !isnothing(q5_1s)
		quap5_1s_df = sample(q5_1s)
		PRECIS(quap5_1s_df)
	end
end

# ╔═╡ 9547a2a0-01c4-11eb-169c-17d6dbace4d9
md"##### Compare below figure with the corresponding figure in clip-05-01-05s.jl."

# ╔═╡ 07cfcaec-01c4-11eb-38b7-31dca6cafefb
if !isnothing(m5_1s)
	post5_1s_df = read_samples(m5_1s, :dataframe)
	plot(xlab="Medium age marriage (scaled)", ylab="Divorce rate (scaled)",
		title="Showing 50 regression lines")
	for i in 1:50
		local yi = mean(post5_1s_df[i, :a]) .+ post5_1s_df[i, :bA] .* xi
		plot!(xi, yi, color=:lightgrey, leg=false)
	end
	scatter!(df[:, :MedianAgeMarriage_s], df[!, :Divorce_s], color=:darkblue)
end

# ╔═╡ aa064588-fc34-11ea-2153-59bc6c7f18a8
md"### snippet 5.2"

# ╔═╡ aa36edc8-fc34-11ea-3701-434917b6f7a3
if !isnothing(q5_1s)

	# Plot regression line D on A

	title1 = "Divorce rate vs. median age at marriage" *
		"\nshowing predicted and quantile range"
	fig1 = plotbounds(
		df, :MedianAgeMarriage, :Divorce,
		quap5_1s_df, [:a, :bA, :sigma];
		title=title1,
		colors=[:lightblue, :darkgrey]
	)
end

# ╔═╡ 683093dc-fc54-11ea-3be9-fdad0a8812f6
md"##### The model m5.2s represents a regression of Divorce on Marriage and is defined as:"

# ╔═╡ 77cb9670-fc54-11ea-38fe-3b8c48e2ce09
md"
```
model {
  vector[N] mu;               // mu is a vector
  a ~ normal(0, 0.2);         // Priors
  bM ~ normal(0, 0.5);
  sigma ~ exponential(1);
  mu = a + bM * M;
  D ~ normal(mu , sigma);     // Likelihood
}
```
"

# ╔═╡ 3c3888ec-fc55-11ea-3542-c7df3ecc74b4
md"##### Both D (Divorce rate) and A (Marriage rate) are standardized."

# ╔═╡ 864e52fc-fc39-11ea-2c30-43e6b888bf60
stan5_2 = "
data {
  int N;
  vector[N] D;
  vector[N] M;
}
parameters {
  real a;
  real bM;
  real<lower=0> sigma;
}
model {
  vector[N] mu = a + bM * M;
  a ~ normal( 0 , 0.2 );
  bM ~ normal( 0 , 0.5 );
  sigma ~ exponential( 1 );
  D ~ normal( mu , sigma );
}
";

# ╔═╡ 4a4b1284-80f7-11eb-01a4-452b4e2597e7
begin
	q5_2s, m5_2s, om = stan_quap("m5.2s", stan5_2; data)
	if !isnothing(q5_2s)
		quap5_2s_df = sample(q5_2s)
	end
	PRECIS(quap5_2s_df)
end

# ╔═╡ c206cdc6-fc55-11ea-2835-71c3067485f8
if !isnothing(m5_2s)
	post5_2s_df = read_samples(m5_2s, :dataframe)
	PRECIS(post5_2s_df)
end

# ╔═╡ 279006d8-fc56-11ea-2792-d1c2678a1e08
if !isnothing(q5_2s)

	# Plot regression line D on M

	title2 = "Divorce rate vs. marriage rate" * "\nshowing predicted and hpdi range"
	fig2 = plotbounds(
		df, :Marriage, :Divorce,
		quap5_2s_df, [:a, :bM, :sigma];
		title=title2,
		colors=[:lightblue, :darkgrey]
	)

end

# ╔═╡ 41bc0716-fc56-11ea-1cfe-3db82349a2d2
	plot(fig2, fig1, layout=(1,2), title="")

# ╔═╡ aa37b78a-fc34-11ea-11e5-7d1ef7bdf603
md"## End of clip-05-01-02s.jl"

# ╔═╡ Cell order:
# ╟─feebbe4e-fc33-11ea-2ee7-d14f977c6497
# ╠═a9bbd3ea-fc34-11ea-282a-43f4fc159d39
# ╠═a9e16b46-fc34-11ea-36de-ef647e4e4f6f
# ╟─a9e54766-fc34-11ea-1246-93a25f748d1d
# ╟─10471c94-fc55-11ea-0fec-656bd563d513
# ╠═a9f32456-fc34-11ea-3af8-c7e0de78a85d
# ╟─4d0ca900-fc53-11ea-0f6c-e7d09d775387
# ╟─25c70d5a-fc54-11ea-3910-a9276dc7b696
# ╟─b8ae47cc-01af-11eb-3cbd-a749b6e3d327
# ╠═b5431248-01af-11eb-03ac-ed6a17e17e39
# ╠═db59a030-423e-11eb-3eb2-5d1f7d3dd0f5
# ╠═55ef3d46-423f-11eb-1d83-55276a40b702
# ╠═fad134cc-423f-11eb-3398-b7d0e22aad9c
# ╟─aa241f5e-fc34-11ea-21d5-1b03dfafd34f
# ╠═aa2a005e-fc34-11ea-09ab-290a8d6b6a9a
# ╠═cbd5fe10-01af-11eb-055b-bd009d0dca55
# ╟─9547a2a0-01c4-11eb-169c-17d6dbace4d9
# ╠═07cfcaec-01c4-11eb-38b7-31dca6cafefb
# ╟─aa064588-fc34-11ea-2153-59bc6c7f18a8
# ╠═aa36edc8-fc34-11ea-3701-434917b6f7a3
# ╟─683093dc-fc54-11ea-3be9-fdad0a8812f6
# ╟─77cb9670-fc54-11ea-38fe-3b8c48e2ce09
# ╟─3c3888ec-fc55-11ea-3542-c7df3ecc74b4
# ╠═864e52fc-fc39-11ea-2c30-43e6b888bf60
# ╠═4a4b1284-80f7-11eb-01a4-452b4e2597e7
# ╠═c206cdc6-fc55-11ea-2835-71c3067485f8
# ╠═279006d8-fc56-11ea-2792-d1c2678a1e08
# ╠═41bc0716-fc56-11ea-1cfe-3db82349a2d2
# ╟─aa37b78a-fc34-11ea-11e5-7d1ef7bdf603


using Plots,  LaTeXStrings

function plot_fitness(step=0.1, imagename="fitness-function.png")
    x = 0:step:1
    fitness(i) = 2/(i+1) - 1

    plot(x, fitness.(x), color = :orange, line = 3, guidefontsize=14, legend = false,
         xlabel = L"t_{et}(i)", ylabel = L"f(i)", title = "Função fitness")

    savefig(imagename)
end

function plot_peaks_surface(step=0.1, imagename="peaks-function.png")
    xs = -3:step:3
    ys = -3:step:3
    x_grid = [x for x = xs for y = ys]
    y_grid = [y for x = xs for y = ys]
    peaks(x, y) = 3*(1-x)^2 * exp(-x^2 -(y+1)^2) -10(x/5 - x^3 - y^5)*exp(-x^2-y^2)-1/3*exp(-(x+1)^2-y^2)

    plot(x_grid, y_grid, -1 * peaks.(x_grid, y_grid), guidefontsize=14, label = false,
        xlabel = L"w_i", ylabel = L"w_j", title = L"E(w_i, w_j)", cbar=false,
        camera=(75,45), st = :surface, color = :bluesreds)

    savefig(imagename)
end


function relorientation(face, simplex)

    v = setdiff(simplex, face)
    length(v) == 1 || return 0

    # find the position of the missing vertex
    v = v[1]
    #i = Base.findfirst(simplex, v)
    i = something(findfirst(isequal(v), simplex),0)
    s = (-1)^(i-1)

    # remove that vertex from the simplex
    face2 = Array{Int}(undef,length(simplex)-1)
    for j in 1 : i-1
        face2[j] = simplex[j]
    end
    for j in i : length(simplex)-1
        face2[j] = simplex[j+1]
    end

    # get the permutation that maps face to face2
    #p = indexin(face, face2)
    #p = [ findfirst(face2,v) for v in face ]
    p = [something(findfirst(isequal(v),face2),0) for v in face]

    return s * levicivita(p) * i
end


####
#### Graph library
####

## This is a collection of some classes of graphs.
## Note that the empty graph is given by Graph(n), so no function
## for the empty graph is provided here.

"""
    cycle_graph(n::Integer)

Return the cycle graph of order `n`.
"""
function cycle_graph(n::Integer)
    g = Graph(n)
    for i in 1:n-1
        add_edge!(g, i, i + 1)
    end
    add_edge!(g, 1, n)
    return g
end

"""
    complete_graph(n::Integer)

Return the complete graph of order `n`.
"""
function complete_graph(n::Integer)
    g = Graph(n)
    for i in 1:n, j in 1:i-1
        add_edge!(g, i, j)
    end
    return g
end

"""
    star_graph(n::Integer)

The star graph of order `n`. This is a
central vertex with an edge to each of
`n-1` other vertices.
"""
function star_graph(n::Integer)
    g = Graph(n)
    for i in 2:n
        add_edge!(g, 1, i)
    end
    return g
end

"""
    star_cycle_graph(n::Integer)

Return the star graph of order `n`.
The star graph is a cycle graph of order `n-1` with
an additional vertex with and edge to all other vertices.
"""
function star_cycle_graph(n::Integer)
    g = cycle_graph(n - 1)
    add_vertex!(g)
    for i in 1:n-1
        add_edge!(g, i, n)
    end
    return g
end

## Note that Solomonoff and Rappaport introduced this graph and
## computed many important properties 10 years before Erd\"os and Renyi
## studied it.
## The Wikipedia article https://en.wikipedia.org/wiki/Random_graph mentions
## even earlier antecedents.
## Here, we use random_graph rather than a proper name.
"""
    random_graph(n::Integer, p::Real)

Return a `Graph` in which a fraction `p` of all possible edges
is present. This implementation is fast rather than exact.
For instance, it samples with replacement. This algorithm is
rather efficient for large graphs and small `p`.
"""
function random_graph(n::Integer, p::Real)
    p < 1 && p >= 0 || error("The value p=$p is not a probability less than 1")
    g = Graph(n)
    num_edges = round(Int, p * n * (n + 1) / 2)
    @show num_edges
    for i in 1:num_edges
        local v0, v1
        while true
            while true
                v0 = rand(1:n)
                v1 = rand(1:n)
                v0 != v1 && break
            end
            result = add_edge!(g, v0, v1)
            result != nothing && break # Retry if edge already added.
        end
    end
    return g
end


import Voronoi

#using Random
include("Draw.jl")

function randf(start, finish, n)
   v = rand(n)
   return map(x -> start + x*(finish-start), v)
end

function demoLloyd()
   WIDTH = 100.0
   HEIGHT = 100.0
   #n = 10

   #seed = 100
   #println("Random seed: ", seed)
   #Random.seed!(seed)

   command = nothing
   for n in 2:100
      points = convert(Array{Tuple{Real, Real}}, collect(zip(randf(1, WIDTH-1, n), randf(1, HEIGHT-1, n))))
      #n = 55
      #points = convert(Array{Tuple{Real, Real}}, [(90.0, 90.0),# (80.0, 50), (70.0, 55.0), (30.0, 85.0), (30, 45), (35.0, 44.0),
      #                                            #(1, 90),
      #                                            #(5, 90), (10, 90), (15, 90),
      #                                            (10, 5), (10,10), (10, 20), (80, 20), (30, 20),
      #                                            (1, 1.1), (2, 1.2), (3, 1.3),
      #                                           ])
      #points = convert(Array{Tuple{Real, Real}}, [ (100*0.2219056800797507, 100*0.7000658172436546), (100*0.49169960385155936, 100*0.6636482987025588), (100*0.7618642131758435, 100*0.7432736698473899), (100*0.8846773855975993, 100*0.6317667493965284) ])

      max = 20*n
      if max > 100
         max = 100
      end
      for i in 1:max
         # test for repeated points
         repeat = false
         for i in 1:length(points)
            for j in 1:length(points)
               if j == i
                  continue
               end

               if points[i][1] == points[j][1] && points[i][2] == points[j][2]
                  repeat = true
               end
            end
         end
         if repeat
            println("REPEATED POINTS!")
            continue
         end

         V = Voronoi.Fortune.compute(points, WIDTH, HEIGHT)

         Voronoi.Intersect.intersect(V, Voronoi.Intersect.Rectangle(WIDTH, HEIGHT))

         gradient = Voronoi.Optimization.∇f(V, points)
         println(Voronoi.Optimization.f(V), "\t\t\t", Voronoi.Optimization.norm2(gradient))

         Draw.init(WIDTH, HEIGHT)
         Draw.voronoiDiagram(V)

         points, areas = Voronoi.Diagram.centroidsAndAreas(V)

         if command != "a"
            println("Press Return for a step or enter \"a\" to animate until the end.")
            command = readline(stdin)
         else
            sleep(0.0000001) # make sure stuff is drawn
         end
         n += 1
      end
   end
end

demoLloyd()


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.1
#   kernelspec:
#     display_name: Julia 1.1.0
#     language: julia
#     name: julia-1.1
# ---

# # Streamfunction For Vertically Integrated Transport
#
# This notebook proceeds in several steps:
#
# - read `TrspX,TrspY` transport vector field from file and compute its convergence.
# - split `TrspX,TrspY` into rotational and divergent components (i.e., `Helmholtz Decomposition`).
# - compute streamfunction from the non-divergent component.
#
# _Note: this requires the `MeshArrays.jl v0.2.3` (upcoming release)_

# +
using MeshArrays, Plots

include(joinpath(dirname(pathof(MeshArrays)),"Plots.jl"))
include("prepare_transports.jl")

#1) get grid
!isdir("GRID_LLC90") ? run(`git clone https://github.com/gaelforget/GRID_LLC90`) : nothing
mygrid=GridSpec("LLC90"); GridVariables=GridLoad(mygrid);

#2) get transport
(TrspX, TrspY, TauX, TauY, SSH)=trsp_read(mygrid,"GRID_LLC90/")

#3) compute convergence
TrspCon=convergence(TrspX,TrspY);

#4) apply mask & avoid isolated Canyons / singular matrices
msk=1.0 .+ 0.0 * mask(view(GridVariables["hFacC"],:,1),NaN,0.0)
TrspCon=msk*TrspCon;
# -

# ### Helmholtz Decomposition

# +
#scalar potential
TrspPot=ScalarPotential(TrspCon)

#plot result
heatmap(TrspPot,clims=(-5e5,5e5))

#Divergent vector field
(TrspXdiv,TrspYdiv)=gradient(TrspPot,GridVariables)
TrspXdiv=TrspXdiv.*GridVariables["DXC"]
TrspYdiv=TrspYdiv.*GridVariables["DYC"];

# Check that divergences match:
tmpCon=convergence(TrspXdiv,TrspYdiv)
heatmap(tmpCon-TrspCon) #TrspCon-tmpCon should be << TrspCon

#Rotational vector field
TrspXrot = TrspX-TrspXdiv
TrspYrot = TrspY-TrspYdiv

#vector Potential
TrspPsi=VectorPotential(TrspX,TrspY,GridVariables)
heatmap(1e-6*msk*TrspPsi)
# -


# COP functions
function violationsSum(g::Vector, h::Vector; ε =0.0)
    sum_g = sum(max.(0.0, g))
    sum_h = 0.0

    for i = 1:length(h)
        if !isapprox(h[i], 0.0, atol=ε)
            sum_h += abs(h[i])
        end
    end

    return (sum_g/length(g)) + (sum_h / length(h))
end


export
    DiophantineEquationLinearXY,
    diophantine_equation_linear_xy

# Dx + Ey + F = 0
struct DiophantineEquationLinearXY{T<:Integer}
    cx::T  # D
    cy::T  # E
    c0::T  # F
end

diophantine_equation_linear_xy(;cx::T=0, cy::T=0, c0::T=0) where {T<:Integer} = DiophantineEquationLinearXY(cx, cy, c0)

Base.show(io::IO, eq::DiophantineEquationLinearXY) = print(io, "$(eq.cx)x + $(eq.cy)y + $(eq.c0) = 0")

evaluate(eq::DiophantineEquationLinearXY{T}, sol::DiophantineSolutionXY{T}) where {T<:Integer} = evaluate(eq, sol.x, sol.y)
evaluate(eq::DiophantineEquationLinearXY{T}, x::T, y::T) where {T<:Integer} = eq.cx * x + eq.cy * y + eq.c0

solve(eq::DiophantineEquationLinearXY{T}) where {T<:Integer} = begin
    cx, cy, c0 = eq.cx, eq.cy, eq.c0

    cx == cy == c0 == 0 && return diophantine_anyx_anyy(T)
    cx == cy == 0 && return diophantine_nonex_noney(T)

    cx == 0 && cy != 0 && return c0 % cy == 0 ?
        diophantine_anyx_oney(-(c0 ÷ cy)) :
        diophantine_nonex_noney(T)

    cx != 0 && cy == 0 && return c0 % cx == 0 ?
        diophantine_onex_anyy(-(c0 ÷ cx)) :
        diophantine_nonex_noney(T)

    g = gcd(cx, cy)
    c0 % g == 0 || return diophantine_nonex_noney(T)

    cx, cy, c0 = cx ÷ g, cy ÷ g, c0 ÷ g
    _, u, v = gcdx(cx, cy)

    diophantine_linearx_lineary(cy, -c0*u, -cx, -c0*v)
end


using Ratios
using Compat.Test

r = SimpleRatio(1,2)
@test convert(Float64, r) == 0.5
@test convert(Float32, r) == 0.5f0
@test convert(BigFloat, r) == BigFloat(1)/2

r2 = SimpleRatio(2,3)
@test r*r2 == SimpleRatio(2,6) == SimpleRatio(1,3)
@test r2*3 == 3*r2 == 2
@test r*false == false*r == 0
@test r/r2 == SimpleRatio(3,4)
@test r/2 == SimpleRatio(1,4)
@test 2/r == 4
@test 4 == 2/r
@test r+1 == 1+r == SimpleRatio(3,2)
@test r-1 == SimpleRatio(-1,2)
@test 1-r == r
@test r+r2 == SimpleRatio(7,6)
@test r-r2 == SimpleRatio(-1,6)
@test r^2 == SimpleRatio(1,4)
@test -r == SimpleRatio(-1,2)
@test 0.2*r ≈ 0.1
@test r == 0.5
@test 0.5 == r

@test_throws OverflowError -SimpleRatio(0x02,0x03)

@test r + SimpleRatio(0x02,0x03) == SimpleRatio(7,6)

@test SimpleRatio(11, 10) == 11//10
@test 1//3 + SimpleRatio(1, 5) == 8//15

@test isfinite(SimpleRatio(0,0)) == false
@test isfinite(SimpleRatio(1,0)) == false
@test isfinite(SimpleRatio(2,1)) == true


using Turing, TuringBenchmarks.TuringTools
using LinearAlgebra

include("lr_helper.jl")

X, Y = readlrdata()

@model lr_nuts(x, y, σ) = begin

    N,D = size(x)

    α ~ Normal(0, σ)
    β ~ MvNormal(zeros(D), ones(D)*σ)

    for n = 1:N
        y[n] ~ BinomialLogit(1, dot(x[n,:], β) + α)
    end
end

# Sampling parameter settings
n_adapts = 1_000

# Sampling
@tbenchmark chain = sample(lr_nuts(x, y, 100), NUTS(0.65), n_samples);


@doc raw"""
    TranslationGroup{T<:Tuple,𝔽} <: GroupManifold{Euclidean{T,𝔽},AdditionOperation}

Translation group $\mathrm{T}(n)$ represented by translation arrays.

# Constructor
    TranslationGroup(n₁,...,nᵢ; field = 𝔽)

Generate the translation group on
$𝔽^{n₁,…,nᵢ}$ = `Euclidean(n₁,...,nᵢ; field = 𝔽)`, which is isomorphic to the group itself.
"""
const TranslationGroup{T<:Tuple,𝔽} = GroupManifold{𝔽,Euclidean{T,𝔽},AdditionOperation}

function TranslationGroup(n::Int...; field::AbstractNumbers=ℝ)
    return TranslationGroup{Tuple{n...},field}(
        Euclidean(n...; field=field),
        AdditionOperation(),
    )
end

invariant_metric_dispatch(::TranslationGroup, ::ActionDirection) = Val(true)

function default_metric_dispatch(
    ::MetricManifold{𝔽,<:TranslationGroup,EuclideanMetric},
) where {𝔽}
    return Val(true)
end

function Base.show(io::IO, ::TranslationGroup{N,𝔽}) where {N,𝔽}
    return print(io, "TranslationGroup($(join(N.parameters, ", ")); field = $(𝔽))")
end