# Special functions

# See:
#   Martin Maechler (2012) "Accurately Computing log(1 − exp(− |a|))"
#   http://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf

# log(1-exp(x))
# NOTE: different than Maechler (2012), no negation inside parantheses
log1mexp(x::Real) = x >= loghalf ? log(-expm1(x)) : log1p(-exp(x))
# log(1+exp(x))
log1pexp(x::Real) = x > zero(x) ? x + log1p(exp(-x)) : log1p(exp(x))
log1pexp(x::Float64) = x <= 18.0 ? log1p(exp(x)) : x <= 33.3 ? x + exp(-x) : x
log1pexp(x::Float32) = x <= 9f0 ? log1p(exp(x)) : x <= 16f0 ? x + exp(-x) : x
log1pexp(x::Integer) = log1pexp(float(x))
# log(exp(x)-1)
logexpm1(x::Real) = x + log1mexp(-x)
logexpm1(x::Float64) = x <= 18.0 ? log(expm1(x)) : x <= 33.3 ? x - exp(-x) : x
logexpm1(x::Float32) = x <= 9f0 ? log(expm1(x)) : x <= 16f0 ? x - exp(-x) : x
logexpm1(x::Integer) = logexpm1(float(x))

# log(1+x^2)
log1psq(x::FloatingPoint) = (ax = abs(x); ax < maxintfloat(x) ? log1p(ax*ax) : 2*log(ax))

φ(z::Real) = exp(-0.5*z*z)/√2π
logφ(z::Real) = -0.5*(z*z + log2π)

Φ(z::Real) = 0.5*erfc(-z/√2)
Φc(z::Real) = 0.5*erfc(z/√2)
logΦ(z::Real) = z < -1.0 ? log(0.5*erfcx(-z/√2)) - 0.5*z*z : log1p(-0.5*erfc(z/√2))
logΦc(z::Real) = z > 1.0 ? log(0.5*erfcx(z/√2)) - 0.5*z*z : log1p(-0.5*erfc(-z/√2))

import Base.Math.@horner

# Rational approximations for the inverse cdf, from:
#   Wichura, M.J. (1988) Algorithm AS 241: The Percentage Points of the Normal Distribution
#   Journal of the Royal Statistical Society. Series C (Applied Statistics), Vol. 37, No. 3, pp. 477-484
for (fn,arg) in ((:Φinv,:p),(:logΦinv,:logp))
    @eval begin
        function $fn($arg::Float32)
            if $(fn == :Φinv)
                q = p - 0.5f0
            else
                q = exp(logp) - 0.5f0
            end
            if abs(q) <= 0.425f0
                r = 0.180625f0 - q*q
                return q * @horner(r,
                                   3.38713_27179f0,
                                   5.04342_71938f1,
                                   1.59291_13202f2,
                                   5.91093_74720f1,
                                   ) /
                @horner(r,
                        1.0f0,
                        1.78951_69469f1,
                        7.87577_57664f1,
                        6.71875_63600f1)
            else
                if $(fn == :Φinv)
                    if p <= 0f0
                        return p == 0f0 ? -inf(Float32) : nan(Float32)
                    elseif p >= 1f0
                        return p == 1f0 ? inf(Float32) : nan(Float32)
                    end
                    r = sqrt(q < 0f0 ? -log(p) : -log1p(-p))
                else
                    if logp == -Inf
                        return -inf(Float32)
                    elseif logp >= 0f0
                        return logp == 0f0 ? inf(Float32) : nan(Float32)
                    end
                    r = sqrt(qf0 < 0 ? -logp : -log1mexp(logp))
                end
                if r < 5.0f0
                    r -= 1.6f0
                    z = @horner(r,
                                1.42343_72777f0,
                                2.75681_53900f0,
                                1.30672_84816f0,
                                1.70238_21103f-1) /
                    @horner(r,
                            1.0f0,
                            7.37001_64250f-1,
                            1.20211_32975f-1)
                else
                    r -= 5.0f0
                    z = @horner(r,
                                6.65790_51150f0,
                                3.08122_63860f0,
                                4.28682_94337f-1,
                                1.73372_03997f-2) /
                    @horner(r,
                            1.0f0,
                            2.41978_94225f-1,
                            1.22582_02635f-2)
                end
                return copysign(z,q)
            end
        end

        function $fn($arg::Real)
            if $(fn == :Φinv)
                q = p - 0.5
            else
                q = exp(logp) - 0.5
            end
            if abs(q) <= 0.425
                r = 0.180625 - q*q
                return q * @horner(r,
                                   3.38713_28727_96366_6080e0,
                                   1.33141_66789_17843_7745e2,
                                   1.97159_09503_06551_4427e3,
                                   1.37316_93765_50946_1125e4,
                                   4.59219_53931_54987_1457e4,
                                   6.72657_70927_00870_0853e4,
                                   3.34305_75583_58812_8105e4,
                                   2.50908_09287_30122_6727e3) /
                @horner(r,
                        1.0,
                        4.23133_30701_60091_1252e1,
                        6.87187_00749_20579_0830e2,
                        5.39419_60214_24751_1077e3,
                        2.12137_94301_58659_5867e4,
                        3.93078_95800_09271_0610e4,
                        2.87290_85735_72194_2674e4,
                        5.22649_52788_52854_5610e3)
            else
                if $(fn == :Φinv)
                    if p <= 0.0
                        return p == 0.0 ? -inf(Float64) : nan(Float64)
                    elseif p >= 1.0
                        return p == 1.0 ? inf(Float64) : nan(Float64)
                    end
                    r = sqrt(q < 0 ? -log(p) : -log1p(-p))
                else
                    if logp == -Inf
                        return -inf(Float64)
                    elseif logp >= 0.0
                        return logp == 0.0 ? inf(Float64) : nan(Float64)
                    end
                    r = sqrt(q < 0 ? -logp : -log1mexp(logp))
                end
                if r < 5.0
                    r -= 1.6
                    z = @horner(r,
                                1.42343_71107_49683_57734e0,
                                4.63033_78461_56545_29590e0,
                                5.76949_72214_60691_40550e0,
                                3.64784_83247_63204_60504e0,
                                1.27045_82524_52368_38258e0,
                                2.41780_72517_74506_11770e-1,
                                2.27238_44989_26918_45833e-2,
                                7.74545_01427_83414_07640e-4) /
                    @horner(r,
                            1.0,
                            2.05319_16266_37758_82187e0,
                            1.67638_48301_83803_84940e0,
                            6.89767_33498_51000_04550e-1,
                            1.48103_97642_74800_74590e-1,
                            1.51986_66563_61645_71966e-2,
                            5.47593_80849_95344_94600e-4,
                            1.05075_00716_44416_84324e-9)
                else
                    r -= 5.0
                    z = @horner(r,
                                6.65790_46435_01103_77720e0,
                                5.46378_49111_64114_36990e0,
                                1.78482_65399_17291_33580e0,
                                2.96560_57182_85048_91230e-1,
                                2.65321_89526_57612_30930e-2,
                                1.24266_09473_88078_43860e-3,
                                2.71155_55687_43487_57815e-5,
                                2.01033_43992_92288_13265e-7) /
                    @horner(r,
                            1.0,
                            5.99832_20655_58879_37690e-1,
                            1.36929_88092_27358_05310e-1,
                            1.48753_61290_85061_48525e-2,
                            7.86869_13114_56132_59100e-4,
                            1.84631_83175_10054_68180e-5,
                            1.42151_17583_16445_88870e-7,
                            2.04426_31033_89939_78564e-15)
                end
                return copysign(z,q)
            end
        end
    end
end

# Multidimensional gamma / partial gamma function
function lpgamma(p::Int, a::Float64)
    res::Float64 = p * (p - 1.0) / 4.0 * log(pi)
    for ii in 1:p
        res += lgamma(a + (1.0 - ii) / 2.0)
    end
    return res
end


export SLBFGS, apply
using Distributions, OnlineStats
Base.@kwdef mutable struct SLBFGS{F,Σ}
    m::Int = 20# Memory length, was ∈ [2, 54] in paper
    λ::Float64 = 1. # regularizer, was ∈ [0.04, 1] in paper
    ρ::Bool = false # Algorithm indicator
    κ::Float64 = 0.5 # Step-length scaling factor ∈ (0,1)
    ᾱ::Float64 = 1. # maximum allowed step-length ∈ (0,1)  TODO: can be function of k
    α::Float64 = ᾱ
    fold::Float64 = Inf
    k::Int = 1 # Internal iteration index
    σ²::Σ = OnlineStats.Variance(weight=McclainWeight(.1))
    xold::Vector{Float64}
    gold::Vector{Float64}
    f::F
    Y::Matrix{Float64}
    S::Matrix{Float64}
end

function SLBFGS(f,x0; m=20, kwargs...)
    d = length(x0)
    Y = zeros(d,m)
    S = zeros(d,m)
    SLBFGS(;m=m, xold=x0, gold=zeros(d), f=f, Y=Y, S=S, kwargs...)
end

function apply(opt::SLBFGS, g, x)
    ρ, κ, ᾱ = opt.ρ, opt.κ, opt.ᾱ
    update_YS!(opt, g, x)
    p = search_direction(opt, g)
    while true
        ξ = x + opt.α*p
        c,fξ = acceptance_indicator(opt, ξ)
        # Update vars
        if c
            x = ξ
            opt.fold = fξ
        end
        opt.α = c*(1/opt.k)^ρ * ᾱ + (1-c)κ*opt.α
        opt.k += 1
        c && break
    end
    x
end

function update_YS!(opt, g, x)
    opt.k == 1 && (opt.gold .= g)
    Y,S = opt.Y, opt.S
    ind = ((opt.k-1) % opt.m) + 1
    Y[:,ind] .= g .- opt.gold
    S[:,ind] .= x .- opt.xold
    opt.gold .= g
    opt.xold .= x
end

function search_direction(opt, g)
    λ = opt.λ
    Y,S = opt.Y, opt.S
    # r4 = R1'\Y1'y
    # r5 = √(λ + y'y - r4'r4)
    # r6 = (y'Y2 - r4'R2)/r5
    # # R6TR6 = R4'R4 - r6'r6 + r3'r3 # TODO: one is really interested in R6 and this should be obtained by downdating rith r6'r6 and updating with r3'r3
    # R6 = downdate(R4, r6)
    # R6 = update(R6, r3)
    R = cholesky(Hermitian(Y'Y + λ*I))
    H̄ = I
    z = (1/λ)Y*(S'g) + H̄*g # QUESTION: they do not mention what H̄ they use
    w = R\(Y'z)
    p = -z + Y*w
    β = 2*max(0,p'g/(g'g))
    @. p -= β*g
end

function acceptance_indicator(opt, ξ)
    fξ = opt.f(ξ)
    ϵ = fξ - opt.fold
    ϵ < 0 && return (true, fξ)
    r = rand()
    # cdf(Normal(-ϵ,σ²)) # QUESTION: use online stats to calculate this variance?
    fit!(opt.σ²,fξ)
    r < Distributions.cdf(Normal(0,√(value(opt.σ²))), -ϵ), fξ # QUESTION: The paper is not clear about this acceptance probability, this is what I think they mean, but should really ask the authors to clarify
    # rand() < exp(ϵ/value(opt.σ²)), fξ
    # rand() < exp(ϵ), fξ
end

##
# using FluxOptTools, Optim, Zygote, Flux, Plots, Test, Statistics, Random
# m = Chain(Dense(1,3,tanh), Dense(3,3,tanh), Dense(3,3,tanh) , Dense(3,1))
# x = LinRange(-pi,pi,100)'
# y = sin.(x)
# sp = sortperm(x[:])
#
# loss() = mean(abs2, m(x) .- y)
#
# @show loss()
# # Zygote.refresh()
# pars = Flux.params(m)
#
# lossfun, gradfun, fg!, p0 = optfuns(loss, pars)
#
# opt = SLBFGS(lossfun,p0; m=3, ᾱ=1., ρ=false, λ=0.0001)
# function train(opt, p0, iters=20)
#     p = copy(p0)
#     g = zeros(veclength(pars))
#     trace = [loss()]
#     for i = 1:iters
#         g = gradfun(g,p)
#         p = apply(opt, g, p)
#         i % 20 == 0 && @show opt.fold
#         push!(trace, opt.fold)
#     end
#     trace
# end
#
# trace = train(opt,p0, 3000)
# plot(trace, yscale=:log10, xscale=:identity, size=(400,300))
# plot(x', [y' m(x)'])  |> display


module Vec

using StaticArrays
using LinearAlgebra
using Printf
import LinearAlgebra: ⋅, ×

export
    AbstractVec,
    VecE,  # an abstract euclidean-group vector
    VecSE, # an abstract special euclidean-group element

    VecE2, # two-element Float64 vector, {x, y}
    VecE3, # three-element Float64 vector, {x, y, z}
    VecSE2, # point in special euclidean group of order 2, {x, y, θ}

    polar, # construct a polar vector with (r,θ)

    proj,  # vector projection
           # proj(a::vec, b::vec, ::Type{Float64}) will do scalar projection of a onto b
           # proj(a::vec, b::vec, ::Type{Vec}) will do vector projection of a onto b

    lerp,  # linear interpolation between two vec's
    invlerp,
    lerp_angle,

    normsquared,         # see docstrings below
    scale_euclidean,
    clamp_euclidean,
    normalize_euclidian,

    dist,  # scalar distance between two vec's
    dist2, # squared scalar distance between two vec's

    rot,   # rotate the vector, always using the Right Hand Rule
    rot_normalized, # like rot, but assumes axis is normalized

    deltaangle, # signed delta angle
    angledist,  # distance between two angles

    inertial2body,
    body2inertial,

    orientation,
    are_collinear,
    get_intersection,

    Circ,
    AABB,
    OBB

abstract type AbstractVec{N, R} <: FieldVector{N, R} end
abstract type VecE{N, R} <: AbstractVec{N, R} end
abstract type VecSE{N, R} <: AbstractVec{N, R} end

lerp(a::Real, b::Real, t::Real) = a + (b-a)*t
invlerp(a::Real, b::Real, c::Real) = (c - a)/(b-a)

"The L2 norm squared."
normsquared(a::AbstractVec) = sum(x^2 for x in a)

"Scale the euclidean part of the vector by a factor b while leaving the orientation part unchanged."
scale_euclidean(a::VecE, b::Real) = b.*a

"Clamp each element of the euclidean part of the vector while leaving the orientation part unchanged."
clamp_euclidean(a::VecE, lo::Real, hi::Real) = clamp.(a, lo, hi)

"Normalize the euclidean part of the vector while leaving the orientation part unchanged."
normalize_euclidian(a::VecE, p::Real=2) = normalize(a, p)

include("common.jl")
include("vecE2.jl")
include("vecE3.jl")
include("vecSE2.jl")

include("geom/geom.jl")
include("coordinate_transforms.jl")
include("quat.jl")

Base.vec(v::Union{AbstractVec, Quat, RPY}) = convert(Vector{Float64}, v)

end # module


# Test various interconnections using MATLAB

using MATLAB, LinearAlgebra
include("StateSpace.jl");
include("Interconnections.jl");

# Test Parallel(...)
println("\nTesting Parallel(...)");
for i=1:10
    A1 = rand(3,3);
    B1 = rand(3,3);
    C1 = rand(3,3);
    D1 = rand(3,3);

    A2 = rand(3,3);
    B2 = rand(3,3);
    C2 = rand(3,3);
    D2 = rand(3,3);

    s1 = StateSpace(A1,B1,C1,D1);
    s2 = StateSpace(A2,B2,C2,D2);
    s3 = Parallel(s1,s2,Vector(1:2),Vector(1:2),Vector(1:3),Vector(1:3));

    # Get MATLAB's implementation.
    mat"s1 = ss($A1,$B1,$C1,$D1)";
    mat"s2 = ss($A2,$B2,$C2,$D2)";
    mat"s3 = parallel(s1,s2,[1:2],[1:2],[1:3],[1:3])";

    mat"Ap=s3.A; Bp=s3.B;Cp=s3.C;Dp=s3.D"
    @mget Ap Bp Cp Dp

    eA = norm(s3.A-Ap);
    eB = norm(s3.B-Bp)
    eC = norm(s3.C-Cp)
    eD = norm(s3.D-Dp)

    println("Norm of error: ", norm([eA,eB,eC,eD]));
end

# Test Feedback(...)
println("\n\nTesting Feedback(...)");
for i=1:10
    A1 = rand(3,3);
    B1 = rand(3,3);
    C1 = rand(3,3);
    D1 = rand(3,3);

    A2 = rand(3,3);
    B2 = rand(3,2);
    C2 = rand(2,3);
    D2 = rand(2,2);

    s1 = StateSpace(A1,B1,C1,D1);
    s2 = StateSpace(A2,B2,C2,D2);
    s3 = Feedback(s1,s2,Vector(2:3),Vector(2:3));

    # Get MATLAB's implementation.
    mat"s1 = ss($A1,$B1,$C1,$D1)";
    mat"s2 = ss($A2,$B2,$C2,$D2)";
    mat"s3 = feedback(s1,s2,[2:3],[2:3])";
    mat"Ap=s3.A; Bp=s3.B;Cp=s3.C;Dp=s3.D"
    @mget Ap Bp Cp Dp

    eA = norm(s3.A-Ap);
    eB = norm(s3.B-Bp)
    eC = norm(s3.C-Cp)
    eD = norm(s3.D-Dp)

    println("Norm of error: ", norm([eA,eB,eC,eD]));
end

# Test Series(...)
println("\n\nTesting Series(...)");
for i=1:10
    A1 = rand(3,3);
    B1 = rand(3,3);
    C1 = rand(3,3);
    D1 = rand(3,3);

    A2 = rand(3,3);
    B2 = rand(3,3);
    C2 = rand(2,3);
    D2 = rand(2,3);

    s1 = StateSpace(A1,B1,C1,D1);
    s2 = StateSpace(A2,B2,C2,D2);
    s3 = Series(s1,s2,Vector(2:3),Vector(1:2));
    s33 = Subsystem(s3,Vector(2:4),Vector(1:2)); # Matlab returns this.

    # Get MATLAB's implementation.
    mat"s1 = ss($A1,$B1,$C1,$D1)";
    mat"s2 = ss($A2,$B2,$C2,$D2)";
    mat"s3 = series(s1,s2,[2:3],[1:2])";
    mat"Ap=s3.A; Bp=s3.B;Cp=s3.C;Dp=s3.D"
    @mget Ap Bp Cp Dp

    # The ordering of states is different in Matlab.
    # We check sorted elements of matrices
    eA = sort(s33.A[:])-sort(Ap[:]);
    eB = sort(s33.B[:])-sort(Bp[:]);
    eC = sort(s33.C[:])-sort(Cp[:]);
    eD = sort(s33.D[:])-sort(Dp[:]);

    println("Norm of error: ",norm([eA;eB;eC;eD]))
end


import Dolo: Model, time_iteration, discretize
using LinearAlgebra, Plots

model = Model("examples/models/rbc_mc.yaml")

sol = time_iteration(model)

import Dolo: UnstructuredGrid, CartesianGrid, transition, n_nodes, n_inodes, inode, iweight, nodes, node, Point

dp = discretize(model.exogenous)

function ergodic_distribution(model, dr, exo_grid:: UnstructuredGrid, endo_grid:: CartesianGrid, dp)
    N_m = n_nodes(exo_grid)
    N_s = n_nodes(endo_grid)
    N = N_m*N_s
    Π = zeros(N_m, N_s, N_m, endo_grid.n...)
    s = nodes(endo_grid)
    p = SVector(model.calibration[:parameters]...)
    a = SVector(endo_grid.min...)
    b = SVector(endo_grid.max...)
    for i_m in 1:n_nodes(exo_grid)
        x = dr(i_m, s)
        m = node(exo_grid, i_m)
        for i_M in 1:n_inodes(dp, i_m)
            M = inode(Point, dp, i_m, i_M)
            w = iweight(dp, i_m, i_M)
            S = transition(model, m, s, x, M, p)
            S = [(S[n]-a)./(b-a) for n=1:length(S)]
            trembling_hand!(view(Π,i_m, :, i_M, :), S, w)
        end
    end
    Π0 = (reshape(Π,N,N))' - Matrix(I, N, N)
    Π0[end,:] .= 1
    B = zeros(N)
    B[end] = 1.0
    μ = Π0\B

    μ = reshape(μ, N_m, N_s)
    return Π, μ
end

function trembling_hand!(A, x, w)
    N,n0 = size(A)
    δ0 = 1.0./(n0-1.0)
    for n in 1:N
        print("hey")
        x0 = x[n][1]
        x0 = min.(max.(x0, 0.0),1.0)
        q0 = div.(x0, δ0)
        q0 = max.(0, q0)
        q0 = min.(q0, n0-2)
        λ0 = (x0./δ0-q0) # ∈[0,1[ by construction
        q0_ = round.(Int,q0) + 1
        println("n:", n, ":", q0_, ":", λ0, ":", w)
        A[n, q0_]   += (1-λ0)*w
        A[n, q0_+1] += λ0*w
    end
end

Π, μ = ergodic_distribution(model, dr, dr.grid_exo, dr.grid_endo, dp)

sum(reshape(Π, 40,40), dims=2)

μ = reshape(μ, 2, 20)

plot(μ[1,:])

#
# function trembling_hand(A, x, w)
#
#     N,n0,n1 = A.shape
#     δ0 = 1/(n0-1)
#     δ1 = 1/(n1-1)
#
#     for n in 1:N
#
#         x0 = x[n,0]
#         x0 = min(max(x0, 0),1)
#         q0 = np.floor_divide(x0, δ0)
#         q0 = max(0, q0)
#         q0 = min(q0, n0-2)
#
#         x1 = x[n,1]
#         x1 = min(max(x1, 0),1)
#         q1 = np.floor_divide(x1, δ1)
#         q1 = max(0, q1)
#         q1 = min(q1, n1-2)
#
#         λ0 = (x0-q0*δ0)/δ0 # ∈[0,1[ by construction
#         q0_ = int(q0)
#
#         λ1 = (x1-q1*δ1)/δ1 # ∈[0,1[ by construction
#         q1_ = int(q1)
#
#         A[n, q0_, q1_]   += (1-λ0)*(1-λ1)*w
#         A[n, q0_, q1_+1]   += (1-λ0)*(λ1)*w
#         A[n, q0_+1, q1_]   += (λ0)*(1-λ1)*w
#         A[n, q0_+1, q1_+1]   += (λ0)*(λ1)*w
#
#     end
# end
#

dr(1,s[1])

ss0

s[1]


using Revise
using NLSolvers
using LinearAlgebra
using SparseDiffTools
using SparseArrays
using IterativeSolvers
using ForwardDiff

# Start with a stupid example
n = 10
A = sprand(10, 10, 0.1)
A = A*A' + I
F(Fx, x) = mul!(Fx, A, x)

x0, = rand(10)
xp = copy(x0)
Fx = copy(xp)

function theta(x)
   if x[1] > 0
       return atan(x[2] / x[1]) / (2.0 * pi)
   else
       return (pi + atan(x[2] / x[1])) / (2.0 * pi)
   end
end

function F_powell!(Fx, x)
        if ( x[1]^2 + x[2]^2 == 0 )
            dtdx1 = 0;
            dtdx2 = 0;
        else
            dtdx1 = - x[2] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) );
            dtdx2 =   x[1] / ( 2 * pi * ( x[1]^2 + x[2]^2 ) );
        end
        Fx[1] = -2000.0*(x[3]-10.0*theta(x))*dtdx1 +
            200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[1]/sqrt( x[1]^2+x[2]^2 );
        Fx[2] = -2000.0*(x[3]-10.0*theta(x))*dtdx2 +
            200.0*(sqrt(x[1]^2+x[2]^2)-1)*x[2]/sqrt( x[1]^2+x[2]^2 );
        Fx[3] =  200.0*(x[3]-10.0*theta(x)) + 2.0*x[3];
    Fx
end

function F_jacobian_powell!(Fx, Jx, x)
    ForwardDiff.jacobian!(Jx, F_powell!, Fx, x)
    Fx, Jx
end
Fc, Jc = zeros(3), zeros(3,3)
F_jacobian_powell!(Fc, Jc, x0)

import NLSolvers: OnceDiffedJv
function OnceDiffedJv(F; seed, autodiff=false)
    JacOp = JacVec(F, seed; autodiff=false)
    OnceDiffedJv(F, JacOp)
end

jv = JacVec(F_powell!, rand(3); autodiff=false)
function jvop(x)
    jv.u .= x
    jv
end
prob_obj = NLSolvers.VectorObjective(F_powell!, nothing, F_jacobian_powell!, jvop)

prob = NEqProblem(prob_obj)
x0 = [-1.0, 0.0, 0.0]
res = solve(prob, copy(x0), LineSearch(Newton(), Backtracking()))
@test norm(res.info.best_residual) < 1e-15

using IterativeSolvers
function inexact_linsolve(x0, JvOp, Fx, ηₖ)
    krylov_iter = IterativeSolvers.gmres_iterable!(x0, JvOp, Fx; maxiter=50)
    res = copy(Fx)
    rhs = ηₖ*norm(Fx, 2)
    for item in krylov_iter
        res = krylov_iter.residual.current
        if res <= rhs
            break
        end
    end
    return x0, res
end
res = solve(prob, copy(x0), InexactNewton(inexact_linsolve, FixedForceTerm(0.001), 1e-4, 300), NEqOptions(maxiter=1000))
@test norm(res.info.best_residual) < 1e-10

function f_2by2!(F, x)
    F[1] = (x[1]+3)*(x[2]^3-7)+18
    F[2] = sin(x[2]*exp(x[1])-1)
    F
end

function g_2by2!(J, x)
    J[1, 1] = x[2]^3-7
    J[1, 2] = 3*x[2]^2*(x[1]+3)
    u = exp(x[1])*cos(x[2]*exp(x[1])-1)
    J[2, 1] = x[2]*u
    J[2, 2] = u
    J
end

FJacOp2 = OnceDiffedJv(f_2by2!; seed=rand(2))
solve(FJacOp2, rand(2), ResidualKrylov(FixedForceTerm(0.4), 1e-3, 300))

function f_sparse!(F, x)
    F[1] = (x[1]+3)*(x[2]^3-7)+18
    F[2] = sin(x[2]*exp(x[1])-1)
end

function g_sparse!(J, x)
    J[1, 1] = x[2]^3-7
    J[1, 2] = 3*x[2]^2*(x[1]+3)
    u = exp(x[1])*cos(x[2]*exp(x[1])-1)
    J[2, 1] = x[2]*u
    J[2, 2] = u
end

function helical_valley()
    tpi = 8*atan(1)
    c7 = 2.5e-1
    c8 = 5e-1

    function f!(fvec, x)
        if x[1] > 0
            temp1 = atan(x[2]/x[1])/tpi
        elseif x[1] < 0
            temp1 = atan(x[2]/x[1])/tpi + c8
        else
            temp1 = c7*sign(x[2])
        end
        temp2 = sqrt(x[1]^2+x[2]^2)
        fvec[1] = 10(x[3] - 10*temp1)
        fvec[2] = 10(temp2 - 1)
        fvec[3] = x[3]
    end
    function fj!(F, J, x)
        if x[1] > 0
            temp1 = atan(x[2]/x[1])/tpi
        elseif x[1] < 0
            temp1 = atan(x[2]/x[1])/tpi + c8
        else
            temp1 = c7*sign(x[2])
        end
        temp2 = sqrt(x[1]^2+x[2]^2)
        if !isa(F, Nothing)
            F[1] = 10*(x[3] - 10*temp1)
            F[2] = 10*(temp2 - 1)
            F[3] = x[3]
        end
        if !isa(J, Nothing)
            temp = x[1]^2 + x[2]^2
            temp1 = tpi*temp
            temp2 = sqrt(temp)
            J[1,1] = 100*x[2]/temp1
            J[1,2] = -100*x[1]/temp1
            J[1,3] = 10
            J[2,1] = 10*x[1]/temp2
            J[2,2] = 10*x[2]/temp2
            J[2,3] = 0
            J[3,1] = 0
            J[3,2] = 0
            J[3,3] = 1
            if !isa(F, nothing)
                return F, J
            else
                return J
            end
        end

    end
    f!, fj!, rand(3), [-1.; 0; 0]), [-1.; 0; 0], "Helical Valley")
end


@with_kw struct actM1{A,B,C} <: RTGmodel
    model::A = actM1_model()
    protein_lookup = get_protein_lookup(model)
    p::C = init_p(model; idx_hill_coefs=[])
    u::B = init_u(model, protein_lookup)
end

"""
Activation Model:
"""
function actM1_model()
    actM1 = @reaction_network begin
        # NLSact: Nuclear Localization Signal (NLS) activation
        (1), BmhMks → BmhMks
        (k2n + k1*(BmhMks)/(k1m+BmhMks)), Rtg13_a_c → Rtg13_i_c
        (k2, k2n), Rtg3_i_c ↔ Rtg3_a_c
        # Rtg1Binding: Rtg13 formation
        (k3, kn3), Rtg1_c + Rtg3_a_c ↔ Rtg13_a_c
        (k3, kn3), Rtg1_c + Rtg3_i_c ↔ Rtg13_i_c
        # Translocation
        (k4, kn4), Rtg13_a_c ↔ Rtg13_a_n
        (k4, kn4), Rtg3_a_c ↔ Rtg3_a_n
    end k1 k1m k2 k2n k3 kn3 k4 kn4
    return actM1
end


module MaterialModels

using Polynomials
using LinearAlgebra
using Dierckx

using ..Geometry
using ..CrossSection
using ..Mesh

export steel, ConstitutiveLaw

struct ConstitutiveLaw

    stress_strain::Dierckx.Spline1D
    E_tan_strain::Dierckx.Spline1D

end

function steel(σy, σy1, σu, σf, ϵy, ϵy1, ϵu, ϵf, n)

    # define elastic and yielding portion of the steel stress-strain curve

    seg_1 = norm([σy, ϵy])
    seg_2 = norm([(σy-σy1), (ϵy-ϵy1)])

    θ1 = rad2deg(atan(σy/ϵy))
    θ2 = rad2deg(atan((σy-σy1) / (ϵy-ϵy1)))

    ΔL = [seg_1, seg_2]
    θ = [θ1, θ2]  #degrees
    n_ep = [n[1], n[2]]          #hard code this for now
    radius = [0.001]  #hard code this for now
    n_radius = [400]     #hard code this for now
    closed_or_open = 1

   elastic_plastic = CrossSection.Feature(ΔL, θ, n_ep, radius, n_radius, closed_or_open)
   ϵ, σ = CrossSection.get_xy_coordinates(elastic_plastic)

    # #define plateau region
    # plateau_slope = (σy1 - σy) / (ϵy1 - ϵy)
    # ϵ_plateau = (ϵy + (ϵy1 - ϵy)/n[2]):(ϵy1 - ϵy)/n[2]:(ϵy1 - (ϵy1 - ϵy)/n[2])
    # σ_plateau = plateau_slope .* (ϵ_plateau .- ϵy) .+ σy

    # ϵ = [ϵ; ϵ_plateau]
    # σ = [σ; σ_plateau]

    #define plateau and hardening part of steel stress-strain curve
    p = fit([ϵy1, ϵu, ϵf],[σy1, σu, σf], 2)

    ϵ_range = ϵy1:(ϵf-ϵy1)./n[3]:ϵf

    σ_fit = zeros(length(ϵ_range))

    for i=1:length(ϵ_range)
        σ_fit[i] = p(ϵ_range[i])
    end

    ϵ = [ϵ; ϵ_range[2:end]]
    σ = [σ; σ_fit[2:end]]

    #calculate tangent elastic modulus
    num_points = length(ϵ)
    E_tan = zeros(Float64,  num_points)
    order = 1
    for i = 1:num_points

        if i < num_points - 2

            x = ϵ[i:i+2]
            x0 = x[1]

            stencil = Mesh.calculate_weights(order, x0, x)

            E_tan[i] = sum(stencil .* σ[i:i+2])

        elseif i > num_points - 2

            x = ϵ[i-2:i]
            x0 = x[3]

            stencil = Mesh.calculate_weights(order, x0, x)

            E_tan[i] = sum(stencil .* σ[i-2:i])

        end

        #tough to round the corner at yielding, revisit this with better model of proportional limit
        if E_tan[i] > E_tan[1]

            E_tan[i] = E_tan[i-1]

        end

    end

    stress_strain = Spline1D(ϵ, σ)
    E_tan_strain = Spline1D(ϵ, E_tan)

    steel_law = ConstitutiveLaw(stress_strain, E_tan_strain)

    return steel_law

end

end


using UncertaintyQuantification

MHexlikelihood(x) = x[1]^85 * (1-x[1])^(100-85)

prior(x) = 1

propdist = Normal(0, 0.2)
propdistsample() = [rand(propdist)]
propdistpdf(x) = pdf(propdist, x[1])

mcit=10000

rsample = mh(MHexlikelihood, prior, propdistpdf, propdistsample, [1], mcit, 100)


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

using Markdown
using InteractiveUtils

# ╔═╡ 9820cc3a-fe50-11ea-3fb6-991358edb7ff
using Pkg, DrWatson

# ╔═╡ 98210cc2-fe50-11ea-0fec-cfe81a89d0cb
begin
	#@quickactivate "StatisticalRethinkingStan"
	using StanSample, GLM
	using StatisticalRethinking
end

# ╔═╡ c5f141c2-fe4f-11ea-03d2-e5eb5d2349f6
md"## Clip-06-02-06s.jl"

# ╔═╡ 9831d8b8-fe50-11ea-38fd-c99a2b5fd0bc
md"### Snippet 6.1"

# ╔═╡ 98328740-fe50-11ea-3008-27d3dda98cd2
begin
	N = 100
	df = DataFrame(
		height = rand(Normal(10, 2), N),
		leg_prop = rand(Uniform(0.4, 0.5), N),
	)
	df.leg_left = df.leg_prop .* df.height + rand(Normal(0, 0.02), N)
	df.leg_right = df.leg_prop .* df.height + rand(Normal(0, 0.02), N)
end;

# ╔═╡ 9842396a-fe50-11ea-217d-2bba5fa44fb5
md"### Snippet 6.2"

# ╔═╡ 9842dbb8-fe50-11ea-158e-cbc2873c64bd
stan6_1 = "
data {
  int <lower=1> N;
  vector[N] H;
  vector[N] LL;
  vector[N] LR;
}
parameters {
  real a;
  real bL;
  real bR;
  real <lower=0> sigma;
}
model {
  vector[N] mu;
  mu = a + bL * LL + bR * LR;
  a ~ normal(10, 100);
  bL ~ normal(2, 10);
  bR ~ normal(2, 10);
  sigma ~ exponential(1);
  H ~ normal(mu, sigma);
}
";

# ╔═╡ 9851b4a8-fe50-11ea-083b-17e3c182a55a
begin
	m6_1s = SampleModel("m6.1s", stan6_1,
		method=StanSample.Sample(num_samples=1000))
	m_6_1_data = Dict(
	  :H => df[:, :height],
	  :LL => df[:, :leg_left],
	  :LR => df[:, :leg_right],
	  :N => size(df, 1)
	)
	rc6_1s = stan_sample(m6_1s, data=m_6_1_data)
	success(rc6_1s) && (part6_1s = read_samples(m6_1s, :particles))
end

# ╔═╡ eb921ad4-fe50-11ea-2450-7dfe19847755
if success(rc6_1s)
	(s0, p0) = plot_model_coef([m6_1s], [:a, :bL, :bR, :sigma];
		title="Multicollinearity between bL and bR")
	p0
end

# ╔═╡ b825edf6-00eb-11eb-039c-cfcb859dc43d
s0

# ╔═╡ 985265b0-fe50-11ea-0594-05ac311d2e87
if success(rc6_1s)
	post6_1s_df = read_samples(m6_1s, :dataframe)

	# Fit a linear regression

	m = lm(@formula(bL ~ bR), post6_1s_df)

	# estimated coefficients from the model

	coefs = coef(m)

	fig1 = plot(xlabel="bR", ylabel="bL", lab="bL ~ bR")
	plot!(post6_1s_df[:, :bR], post6_1s_df[:, :bL])
	fig2 = density(part6_1s.bR.particles + part6_1s.bL.particles, xlabel="sum of bL and bR",
		ylabel="Density", lab="bL + bR")
	plot(fig1, fig2, layout=(1, 2))
end

# ╔═╡ 9862bdb4-fe50-11ea-1402-bbac3257c25d
md"## End of clip-06-02-06s.jl"

# ╔═╡ Cell order:
# ╟─c5f141c2-fe4f-11ea-03d2-e5eb5d2349f6
# ╠═9820cc3a-fe50-11ea-3fb6-991358edb7ff
# ╠═98210cc2-fe50-11ea-0fec-cfe81a89d0cb
# ╟─9831d8b8-fe50-11ea-38fd-c99a2b5fd0bc
# ╠═98328740-fe50-11ea-3008-27d3dda98cd2
# ╟─9842396a-fe50-11ea-217d-2bba5fa44fb5
# ╠═9842dbb8-fe50-11ea-158e-cbc2873c64bd
# ╠═9851b4a8-fe50-11ea-083b-17e3c182a55a
# ╠═eb921ad4-fe50-11ea-2450-7dfe19847755
# ╠═b825edf6-00eb-11eb-039c-cfcb859dc43d
# ╠═985265b0-fe50-11ea-0594-05ac311d2e87
# ╟─9862bdb4-fe50-11ea-1402-bbac3257c25d


# Linear covariance function

include("lin_iso.jl")
include("lin_ard.jl")

@doc """
# Description
Constructors for linear kernel

# See also
LinIso, LinArd
""" ->
Lin(ll::Float64) = LinIso(ll)
Lin(ll::Vector{Float64}) = LinArd(ll)


using Squeezing
using Compat.Test

@testset "squeezing sum" begin
    @test @inferred @squeezing(sum(ones(3, 2), 1)) == [3.0, 3.0]
end

@testset "squeezing reducedim" begin
    A = reshape(Vector(1:16), (4,4))
    @test @inferred @squeezing(reducedim(max, A, 1)) == collect(4:4:16)
    @test @inferred @squeezing(reducedim(max, A, 2)) == collect(13:16)
end


path = "docs/src/assets/"

using ISA, Plots
𝐶₀ = AMFMtriplet(t->10t+100,t->25cos(t)+50,0.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(ψ₀; timeaxis=0.0:0.001:10.0)
png(path*"IS_exAMFM1.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->10t+100,t->25cos(t)+50,0.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(𝐶₀; timeaxis=0.0:0.001:10.0)
png(path*"IS_exAMFM2.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->10t+100,t->25cos(t)+50,0.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(𝐶₀; timeaxis=0.0:0.001:10.0, view="TF",
     left_margin=15Plots.mm, margin=5Plots.mm)
png(path*"IS_exAMFM3.png")

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

using ISA, Plots
𝐶₀ = AMFMtriplet(t->10sin(t),t->t^2+2t+50,1.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(ψ₀; timeaxis=0.0:0.001:10.0)
png(path*"IS_exAMFM4.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->10sin(t),t->t^2+2t+50,1.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(𝐶₀; timeaxis=0.0:0.001:10.0)
png(path*"IS_exAMFM5.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->10sin(t),t->t^2+2t+50,1.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(𝐶₀; timeaxis=0.0:0.001:10.0, view="TF",
     left_margin=15Plots.mm, margin=5Plots.mm)
png(path*"IS_exAMFM6.png")

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

using ISA, Plots
𝐶₀ = AMFMtriplet(t->5exp(0.1t)+15,t->t^2+t+25,0.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(ψ₀; timeaxis=0.0:0.001:10.0)
png(path*"IS_exAMFM7.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->5exp(0.1t)+15,t->t^2+t+25,0.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(𝐶₀; timeaxis=0.0:0.001:10.0)
png(path*"IS_exAMFM8.png")

using ISA, Plots
𝐶₀ = AMFMtriplet(t->5exp(0.1t)+15,t->t^2+t+25,0.0)
ψ₀ = AMFMcomp(𝐶₀)
plot(𝐶₀; timeaxis=0.0:0.001:10.0, view="TF",
     left_margin=15Plots.mm, margin=5Plots.mm)
png(path*"IS_exAMFM9.png")


n = rand(1:10)
@test matrixdepot("hilb", n) == matrixdepot("hilb", n, n)
@test_approx_eq matrixdepot("hilb" , Float16, n) matrixdepot("hilb", Float32, n)
x = Float64[1:n;]
y = x -1
@test_approx_eq matrixdepot("hilb", n) matrixdepot("cauchy", x, y)
println("'hilb' passed test...")


#=
  Sum of 2 dice.

  I asked how one should observe a sum of the two dice here:
  https://discourse.julialang.org/t/turing-observing-sum-of-two-random-dice/50486

=#

using Turing, StatsPlots, DataFrames
include("../jl_utils.jl")

@model function two_dice(ss)
    d1 ~ DiscreteUniform(1, 6)
    d2 ~ DiscreteUniform(1, 6)

    ss ~ Dirac(d1 + d2)

    return ss
end

# Note: Here we observe c=5 from outside
model = two_dice(4)

# chains = sample(model, MH(), 10000) # This don't work anymore! c is 1..5!
num_chains = 4 # 4
# This works: i.e. forces d1 and d2 to be 2
chains = sample(model, MH(), MCMCThreads(), 10000, num_chains)

# Both d1 and d2 have means 3.5
# chains = sample(model, IS(), MCMCThreads(), 1000, num_chains) # Nope

# Both d1 and d2 have means 3.5
# chains = sample(model, PG(20), 1000) # Nope

# # Both d1 and d2 have means 3.5
# chains = sample(model, SMC(1000), 1000) # Nope

# chains = sample(model, NUTS(), 1000) # Error
# chains = sample(model, HMC(0.1, 5), 1000) # Error
# chains = sample(model, HMCDA(0.15, 0.65), 1000) # ERROR

display(chains)
# display(plot(chains))
# display(histogram(chains))
# display(gelmandiag(chains))

show_var_dist_pct(chains, :d1)
show_var_dist_pct(chains, :d2)

println("\nhash d1+d2")
display(sort(make_hash([chains[:d1][i]+chains[:d2][i] for i in 1:length(chains)])))

gen = generated_quantities(model, chains)
println("After removing the 0s from gen")
# gen2 = [g for g in gen if g != 0]
show_var_dist_pct(gen, 120)


# This file is part of the TaylorSeries.jl Julia package, MIT license
#
# Luis Benet & David P. Sanders
# UNAM
#
# MIT Expat license
#

## Evaluating ##
"""
    evaluate(a, [dx])

Evaluate a `Taylor1` polynomial using Horner's rule (hand coded). If `dx` is
ommitted, its value is considered as zero. Note that the syntax `a(dx)` is
equivalent to `evaluate(a,dx)`, and `a()` is equivalent to `evaluate(a)`.
"""
function evaluate(a::Taylor1{T}, dx::T) where {T<:Number}
    @inbounds suma = a[end]
    @inbounds for k in a.order-1:-1:0
        suma = suma*dx + a[k]
    end
    suma
end
function evaluate(a::Taylor1{T}, dx::S) where {T<:Number,S<:Number}
    suma = a[end]*one(dx)
    @inbounds for k in a.order-1:-1:0
        suma = suma*dx + a[k]
    end
    suma
end
evaluate(a::Taylor1{T}) where {T<:Number} = a[0]

"""
    evaluate(x, δt)

Evaluates each element of `x::AbstractArray{Taylor1{T}}`,
representing the dependent variables of an ODE, at *time* δt. Note that the
syntax `x(δt)` is equivalent to `evaluate(x, δt)`, and `x()`
is equivalent to `evaluate(x)`.
"""
evaluate(x::AbstractArray{Taylor1{T}}, δt::S) where
    {T<:Number,S<:Number} = evaluate.(x, δt)
evaluate(a::AbstractArray{Taylor1{T}}) where {T<:Number} = evaluate.(a, zero(T))

"""
    evaluate!(x, δt, x0)

Evaluates each element of `x::AbstractArray{Taylor1{T}}`,
representing the Taylor expansion for the dependent variables
of an ODE at *time* `δt`. It updates the vector `x0` with the
computed values.
"""
function evaluate!(x::AbstractArray{Taylor1{T}}, δt::S,
        x0::AbstractArray{T}) where {T<:Number,S<:Number}

    @inbounds for i in eachindex(x, x0)
        x0[i] = evaluate( x[i], δt )
    end
    nothing
end

"""
    evaluate(a, x)

Substitute `x::Taylor1` as independent variable in a `a::Taylor1` polynomial.
Note that the syntax `a(x)` is equivalent to `evaluate(a, x)`.
"""
evaluate(a::Taylor1{T}, x::Taylor1{S}) where {T<:Number,S<:Number} =
    evaluate(promote(a,x)...)

function evaluate(a::Taylor1{T}, x::Taylor1{T}) where {T<:Number}
    if a.order != x.order
        a, x = fixorder(a, x)
    end
    @inbounds suma = a[end]*one(x)
    @inbounds for k = a.order-1:-1:0
        suma = suma*x + a[k]
    end
    suma
end

function evaluate(a::Taylor1{Taylor1{T}}, x::Taylor1{T}) where {T<:Number}
    @inbounds suma = a[end]*one(x)
    @inbounds for k = a.order-1:-1:0
        suma = suma*x + a[k]
    end
    suma
end
function evaluate(a::Taylor1{T}, x::Taylor1{Taylor1{T}}) where {T<:Number}
    @inbounds suma = a[end]*one(x)
    @inbounds for k = a.order-1:-1:0
        suma = suma*x + a[k]
    end
    suma
end

evaluate(p::Taylor1{T}, x::Array{S}) where {T<:Number,S<:Number} =
    evaluate.([p], x)

#function-like behavior for Taylor1
(p::Taylor1)(x) = evaluate(p, x)
(p::Taylor1)() = evaluate(p)

#function-like behavior for Vector{Taylor1}
if VERSION >= v"1.3"
    (p::AbstractArray{Taylor1{T}})(x) where {T<:Number} = evaluate.(p, x)
    (p::AbstractArray{Taylor1{T}})() where {T<:Number} = evaluate.(p)
else
    (p::Array{Taylor1{T}})(x) where {T<:Number} = evaluate.(p, x)
    (p::SubArray{Taylor1{T}})(x) where {T<:Number} = evaluate.(p, x)
    (p::Array{Taylor1{T}})() where {T<:Number} = evaluate.(p)
    (p::SubArray{Taylor1{T}})() where {T<:Number} = evaluate.(p)
end

## Evaluation of multivariable
function evaluate!(x::AbstractArray{TaylorN{T}}, δx::Array{T,1},
        x0::AbstractArray{T}) where {T<:Number}

    # @assert length(x) == length(x0)
    @inbounds for i in eachindex(x, x0)
        x0[i] = evaluate( x[i], δx )
    end
    nothing
end

function evaluate!(x::AbstractArray{TaylorN{T}}, δx::Array{Taylor1{T},1},
        x0::AbstractArray{Taylor1{T}}) where {T<:NumberNotSeriesN}

    # @assert length(x) == length(x0)
    @inbounds for i in eachindex(x, x0)
        x0[i] = evaluate( x[i], δx )
    end
    nothing
end

function evaluate!(x::AbstractArray{TaylorN{T}}, δx::Array{TaylorN{T},1},
        x0::AbstractArray{TaylorN{T}}; sorting::Bool=true) where {T<:NumberNotSeriesN}

    @inbounds for i in eachindex(x, x0)
        x0[i] = _evaluate( x[i], δx, Val(sorting) )
    end
    nothing
end

function evaluate!(x::AbstractArray{TaylorN{T}}, δt::T,
        x0::AbstractArray{TaylorN{T}}) where {T<:Number}

    # @assert length(x) == length(x0)
    @inbounds for i in eachindex(x, x0)
        x0[i] = evaluate( x[i], δt )
    end
    nothing
end

"""
    evaluate(a, [vals])

Evaluate a `HomogeneousPolynomial` polynomial at `vals`. If `vals` is ommitted,
it's evaluated at zero. Note that the syntax `a(vals)` is equivalent to
`evaluate(a, vals)`; and `a()` is equivalent to `evaluate(a)`.
"""
function evaluate(a::HomogeneousPolynomial{T}, vals::NTuple{N,S} ) where
        {T<:Number,S<:Number,N}

    @assert N == get_numvars()

    return _evaluate(a, vals)
end

function _evaluate(a::HomogeneousPolynomial{T}, vals::NTuple{N,S} ) where
        {T<:Number,S<:Number,N}

    ct = coeff_table[a.order+1]
    suma = zero(a[1])*vals[1]

    for (i,a_coeff) in enumerate(a.coeffs)
        iszero(a_coeff) && continue
        @inbounds tmp = prod( vals .^ ct[i] )
        suma += a_coeff * tmp
    end

    return suma
end

evaluate(a::HomogeneousPolynomial{T}, vals::Array{S,1} ) where
        {T<:Number,S<:NumberNotSeriesN} = _evaluate(a, (vals...,))

evaluate(a::HomogeneousPolynomial{T}, v, vals::Vararg{S,N}) where {T,S,N} = evaluate(a, (v, vals...,))

evaluate(a::HomogeneousPolynomial{T}, v) where {T<:Number} = evaluate(a, [v...])

function evaluate(a::HomogeneousPolynomial)
    a.order == 0 && return a[1]
    zero(a[1])
end

#function-like behavior for HomogeneousPolynomial
(p::HomogeneousPolynomial)(x) = evaluate(p, x)

(p::HomogeneousPolynomial)(x, v::Vararg{T, N}) where {T,N} = evaluate(p, (x, v...,))

(p::HomogeneousPolynomial)() = evaluate(p)

"""
    evaluate(a, [vals]; sorting::Bool=true)

Evaluate the `TaylorN` polynomial `a` at `vals`.
If `vals` is ommitted, it's evaluated at zero. The
keyword parameter `sorting` can be used to avoid
sorting (in increasing order by `abs2`) the
terms that are added.

Note that the syntax `a(vals)` is equivalent to
`evaluate(a, vals)`; and `a()` is equivalent to
`evaluate(a)`. No extension exists that incorporates
`sorting`.
"""
evaluate(a::TaylorN{T}, vals::NTuple; sorting::Bool=true) where {T<:Number} =
    _evaluate(a, vals, Val(sorting))

evaluate(a::TaylorN, vals; sorting::Bool=true) = _evaluate(a, (vals...,), Val(sorting))

evaluate(a::TaylorN, v, vals::Vararg{T, N}; sorting::Bool=true) where {T,N} =
    _evaluate(a, (v, vals...,), Val(sorting))

function _evaluate(a::TaylorN{T}, vals) where {T<:Number}
    @assert get_numvars() == length(vals)
    R = promote_type(T,typeof(vals[1]))
    a_length = length(a)
    suma = zeros(R, a_length)
    @inbounds for homPol in length(a):-1:1
        suma[homPol] = evaluate(a.coeffs[homPol], vals)
    end
    return suma
end

function _evaluate(a::TaylorN{T}, vals::NTuple, ::Val{true}) where {T<:Number}
    suma = _evaluate(a, vals)
    return sum( sort!(suma, by=abs2) )
end
function _evaluate(a::TaylorN{T}, vals::NTuple, ::Val{false}) where {T<:Number}
    suma = _evaluate(a, vals)
    return sum( suma )
end

function evaluate(a::TaylorN{T}, vals::NTuple{N,Taylor1{S}}) where
        {T<:Number,S<:NumberNotSeries,N}

    @assert N == get_numvars()

    R = promote_type(T,S)
    ord = maximum( get_order.(vals) )
    suma = Taylor1(zeros(R, ord))

    @inbounds for homPol in length(a):-1:1
        suma += evaluate(a.coeffs[homPol], vals)
    end

    return suma
end

evaluate(a::TaylorN{T}, vals::Array{Taylor1{S},1}) where
    {T<:Number,S<:NumberNotSeriesN} = evaluate(a, (vals...,))

function evaluate(a::TaylorN{Taylor1{T}}, vals::NTuple{N, Taylor1{T}}) where
        {T<:NumberNotSeries,N}

    @assert N == get_numvars()

    ord = maximum( get_order.(vals) )
    suma = Taylor1(zeros(T, ord))

    for homPol in length(a):-1:1
        suma += evaluate(a.coeffs[homPol], vals)
    end

    return suma
end

evaluate(a::TaylorN{Taylor1{T}}, vals::Array{Taylor1{T},1}) where
    {T<:NumberNotSeries} = evaluate(a, (vals...,))

function evaluate(a::TaylorN{T}, vals::NTuple{N, TaylorN{S}}) where
        {T<:Number,S<:NumberNotSeries,N}

    @assert length(vals) == get_numvars()

    suma = zero(constant_term(a))*vals[1]

    for homPol in length(a):-1:1
        suma += evaluate(a.coeffs[homPol], vals)
    end

    return suma
end

evaluate(a::TaylorN{T}, vals::Array{TaylorN{S},1}) where
    {T<:Number,S<:NumberNotSeries} = evaluate(a, (vals...,))

function evaluate(a::TaylorN{T}, s::Symbol, val::S) where
        {T<:Number, S<:NumberNotSeriesN}
    vars = get_variables(T)
    ind = lookupvar(s)
    vars[ind] = val
    evaluate(a, vars)
end

evaluate(a::TaylorN{T}, x::Pair{Symbol,S}) where {T<:Number,S<:NumberNotSeriesN} =
    evaluate(a, first(x), last(x))

evaluate(a::TaylorN{T}) where {T<:Number} = a[0][1]

#High-dim array evaluation
function evaluate(A::AbstractArray{TaylorN{T},N}, δx::Vector{S}) where {T<:Number,S<:Number,N}
    R = promote_type(T,S)
    return evaluate(convert(Array{TaylorN{R},N},A), convert(Vector{R},δx))
end
function evaluate(A::Array{TaylorN{T}}, δx::Vector{T}) where {T<:Number}
    Anew = Array{T}(undef, size(A)...)
    evaluate!(A, δx, Anew)
    return Anew
end
evaluate(A::AbstractArray{TaylorN{T}}) where {T<:Number} = evaluate.(A)

#function-like behavior for TaylorN
(p::TaylorN)(x) = evaluate(p, x)
(p::TaylorN)() = evaluate(p)
(p::TaylorN)(s::Symbol, x) = evaluate(p, s, x)
(p::TaylorN)(x::Pair) = evaluate(p, first(x), last(x))
(p::TaylorN)(x, v::Vararg{T, N}) where {T,N} = evaluate(p, (x, v...,))
(p::TaylorN)(b::Bool, x) = evaluate(p, x, sorting=b)
(p::TaylorN)(b::Bool, x, v::Vararg{T, N}) where {T,N} =
    evaluate(p, (x, v...,), sorting=b)

#function-like behavior for AbstractArray{TaylorN{T}}
if VERSION > v"1.1"
    (p::AbstractArray{TaylorN{T}})(x) where {T<:Number} = evaluate(p, x)
    (p::AbstractArray{TaylorN{T}})() where {T<:Number} = evaluate(p)
else
    (p::Array{TaylorN{T}})(x) where {T<:Number} = evaluate(p, x)
    (p::SubArray{TaylorN{T}})(x) where {T<:Number} = evaluate(p, x)
    (p::Array{TaylorN{T}})() where {T<:Number} = evaluate(p)
    (p::SubArray{TaylorN{T}})() where {T<:Number} = evaluate(p)
end


"""
Supertype for all prior distributions
"""
abstract type Prior end

"""
    logpdf(::Prior, θ)

Evaluate the log-probability density function at `θ` for a given prior.
"""
function Distributions.logpdf(::T, θ) where {T<:Prior}
    error("logpdf not implemented for prior $T.")
end

"""
Flat prior
"""
struct ImproperPrior <: Prior end
Distributions.logpdf(::ImproperPrior, θ) = 0.0

"""
Flat prior for positive coordinates
"""
struct ImproperPosPrior <: Prior end
Distributions.logpdf(::ImproperPosPrior, θ) = -sum(log.(θ))

"""
    struct StandardPrior{T} <: Prior
        dist::T
    end

A standard prior π(θ).
"""
struct StandardPrior{T} <: Prior
    dist::T
    StandardPrior(dist::T) where T = new{T}(dist)
end
Distributions.logpdf(prior::StandardPrior, θ) = logpdf(prior.dist, θ)

@doc raw"""
    struct ProductPrior{T,K} <: Prior
        dists::T
        idx::K
    end

Generic prior distribution over parameter vector θ written in a facorized form.
For instance if the prior may be written as
```math
π(θ) = π_1(θ_{1:k})π_2(θ_{(k+1):n}),
```
then `dists` would correspond to a list containing `π_1` and `π_2` and `idx`
would be a list containing `1:k` and `(k+1):n`.

    ProductPrior(dists, dims)

Base consructor specifying a list of prior dsitributions `dists` and a
corresponding number of dimensions to which each prior distribution refers to.
"""
struct ProductPrior{T,K} <: Prior
    dists::T
    idx::K

    function ProductPrior(dists, dims)
        dims_reformatted = Any[]
        last_idx = 1
        for dim in dims
            if dim == 1
                push!(dims_reformatted, dim)
                last_idx += 1
            else
                push!(dims_reformatted, last_idx:last_idx+dim-1)
                last_idx += dim
            end
        end
        idx = tuple(dims_reformatted...)
        dists = tuple(dists...)
        new{typeof(dists), typeof(idx)}(dists, idx)
    end
end

function Distributions.logpdf(prior::ProductPrior, θ)
    lp = 0.0
    for (dist, idx) in zip(prior.dists, prior.idx)
        lp += logpdf(dist, θ[idx])
    end
    lp
end


using Plots,JLD
function randomWalk(starting_position,L)
    mean_life_span=0
    TotData=[[0.0 for i in 1:L]]
    TotData[1][starting_position]=1
    death_p_sum=(TotData[1][1]+TotData[1][end])
    if death_p_sum>=0.9 && death_p_sum<=1
        return 0
    end
    step=2
    while true

        death_p_sum=0
        R=[0.0 for e in 1:L]
        for i in 2:(L-1)
            R[i-1:2:i+1].+=TotData[step-1][i]/2
        end
        R[1]+=TotData[step-1][1]
        R[end]+=TotData[step-1][end]
        push!(TotData, R)

        death_p_sum+=(R[1]+R[end])
        #println(R)
        # println(step)
        # println(R[1])
        # println(R[end])
        mean_life_span+=sum(R[2:19])
        if death_p_sum>=0.99999999
            break
        end
        step+=1
    end
    return mean_life_span
end

#time_=randomWalk(2,20)
MeanLifeSpan=[]
StartingPosList=[i for i in 1:20]
for StartingPos in StartingPosList
    println(StartingPos)
    push!(MeanLifeSpan,randomWalk(StartingPos,20))
end
save("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\chapter5\\4.4\\MeanLifeSpan_startpos_deterministic_data.jld",
 "data", MeanLifeSpan)
scatter(StartingPosList, MeanLifeSpan,legends=false, color=:black, alpha=0.5,dpi=400)
plot!(StartingPosList, MeanLifeSpan,guidefontsize=9, color=:pink2,
 xlabel="the Starting Position",ylabel="Mean Life-Span",
 title="Mean Life-Span Over the Starting Position Using the Deterministic Algorithm", titlefontsize=9)
 savefig("C:\\Users\\Narges\\Documents\\GitHub\\computational_physics\\chapter5\\4.4\\Figs\\MeanLifeSpan_startpos_deterministic.png")


#=
Helper function for score2pvalue. Find the lowest
score s such that P-value(M_ϵ,s) = P-value(M_ϵ,s-E).

Input:
    Q_m: a sorted dict
    E: the max column-wise error between matrix M and M_ϵ
Output:
=#
function find_s(Q_m::SortedDict{T,T}, E::T) where T <: Real
    keys_ = collect(keys(Q_m)); ℓ = length(keys_);
    if ℓ > 1
        for i = 1:(ℓ-1)
            @inbounds if keys_[i]+E < keys_[i+1]
                return keys_[i];
            end
        end
    end
    return nothing
end

"""
    score2pvalue(pwm, α, ϵ=1e-1, k=100, bg=[.25,.25,.25,.25])
Returns P-value(M,α) of a `pwm` with a given threshold `α`.

Input:
* `pwm`: a 4 x m matrix
* `α`: the score
* `ϵ`: initial granularity  (optional)
* `k`: Refinement parameter (optional)
* `bg`: multinomial background (optional)

Output:
* `pval`: p-value
"""
function score2pvalue(pwm::Matrix{T}, α::Real, ϵ=1e-2, k=10, bg=[.25,.25,.25,.25]) where T <: Real
    @assert size(pwm,1) == 4 "The input matrix must have 4 and only 4 rows"
    pwm_double64 = Double64.(pwm);
    α_double64 = Double64(α);
    bg_double64 = Double64.(bg);
    mpwm = min_score_range(pwm_double64);
    m = size(mpwm,2);
    β = best_score(mpwm)+1;
    pval = Double64(0);
    s = 0;
    i = 1;
    @inbounds while !(α_double64 ≈ s)
        # println(α_double64-s)
        ϵ = i == 1 ? ϵ : ϵ/k; i+=1;
        # isinf(ϵ) || isnan(ϵ) && (println(i); break);
        pwm_ϵ = round_pwm(mpwm, ϵ);
        # any(isnan.(pwm_ϵ)) && (println(i); break);
        E = calc_E(mpwm, pwm_ϵ);
        Q = create_Q(m);
        Q = score_distribution(pwm_ϵ, α-E, β, bg_double64);
        #=
        note:
        Sometimes the score range is small
        (i.e. β-(α-E) is very  small) and hence we won't be
        able to find such s. In such case, we return the
        calculated p-value.

        Note that this p-value is an underestimate. To get a
        more accurate result, set k to be a larger number,
        e.g. k=100.
        =#
        s = find_s(Q[m],E);
        if !isnothing(s)
            for (k,v) in Q[m]
                pval += k ≥ s ? v : 0;
            end
            β = s;
        else
            break;
        end
    end
    return pval
end


function utility_test()
	@testset "Lifetime utility" begin
		w = CollegeStratWorker.make_test_worker(3);
		hV = [1.8, 2.3];
		assetV = [-0.7, 0.9, 3.2];
		workStartAge = 3;

		# On a grid
		utilM = lifetime_utility_grid(w, workStartAge, assetV, hV);
		@test size(utilM) == (length(assetV), length(hV))
		# Increasing in h
		@test all(utilM[:,2] .> utilM[:,1])
		# Increasing in assets
		util2M = lifetime_utility_grid(w, workStartAge, assetV .+ 0.001, hV);
		@test all(util2M .> utilM)

		# Test against state by state
		util3M = similar(utilM);
		for (ih, h) in enumerate(hV)
            for (ia, a) in enumerate(assetV)
                # Must be qualified b/c exported by two packages
				util3M[ia, ih] = CollegeStratWorker.lifetime_utility(w, workStartAge, assetV[ia], hV[ih]);
			end
		end
		@test all(util3M .≈ utilM)

		# Vector
		h2V = fill(hV[1], length(assetV));
		util3V = lifetime_utility_vector(w, workStartAge, assetV, h2V);
		@test all(util3V .≈ utilM[:, 1])

		# MU wealth
		ltIncome = [1.2 2.3; 3.4 4.5];
		mu = CollegeStratWorker.mu_wealth(w, workStartAge, ltIncome);
		@test size(ltIncome) == size(mu)
		@test all(mu .> 0.0)
		# Test by perturbation
		dy = 1e-5;
		util1M = [csw.lifetime_utility(w, workStartAge, lty)  for lty in ltIncome];
		util2M =
			[csw.lifetime_utility(w, workStartAge, lty .+ dy)  for lty in ltIncome];
		muM = (util2M .- util1M) ./ dy;
		@test all(muM .> 0.0)
		@test isapprox(muM, mu, rtol = 1e-3)
	end
end

function lt_util_fct_test()
	@testset "Lifetime Utility Function" begin
		wk = CollegeStratWorker.make_test_worker(3);
		workStartAge = 3;
		nk = 30;
		nh = 20;
		kGridV = LinRange(-10.0, 25.0, nk);
		hGridV = LinRange(1.2, 4.7, nh);
		f = lifetime_utility_function(wk, workStartAge;
			kMin = first(kGridV), kMax = last(kGridV),
			hMin = first(hGridV), hMax = last(hGridV));

		for ik = 2 : 3 : (nk-1)
			for ih = 2 : 4 : (nh-1)
				k = kGridV[ik] + 0.3;
				h = hGridV[ih] - 0.1;
				u1 = f(k, h);
				u2 = csw.lifetime_utility(wk, workStartAge, k, h);
				@test isapprox(u1, u2, rtol = 1e-3);
			end
		end

        dev = interpolation_deviation(wk, workStartAge, f);
        @test dev > 0
        @test dev < 1e-3

        uGrid, kGrid, hGrid = lifetime_utility_function_grid(f);
        uGrid2 = lifetime_utility_grid(wk, workStartAge, kGrid, hGrid);
        dev2 = maximum(abs.(uGrid2 .- uGrid) ./ max.(1.0, abs.(uGrid)));
        @test dev2 < 1e-3
	end
end

@testset "Lifetime utility" begin
	utility_test();
	lt_util_fct_test();
end

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


#=
multigrid:
- Julia version: 0.6.2
- Author: jerrymei
- Date: 2018-05-22
=#

mutable struct gridWrapper{T}
    u::Array{T,2}
    newu::Array{T,2}
    dx::Array{T,2}
    dy::Array{T,2}
    rhs::Array{T,2}
    f::Array{T,2}
    fft::rFFTWrapper
end

function gridWrapper{T}(f::AbstractArray{T,2})
    fft=rFFTWrapper(f)
    u=zeros(f)
    newu=similar(u)
    dx=similar(u)
    dy=similar(u)
    rhs=similar(u)
    return gridWrapper(u,newu,dx,dy,rhs,f,fft)
end

function restriction!{T}(fine::AbstractArray{T,2}, coarse::AbstractArray{T,2})
    nx,ny=size(coarse)
    @inbounds for x in 1:nx
        mx=2 * x - 1
        lx=mx > 1?mx - 1:2 * nx
        ux=mx + 1
        for y in 1:ny
            my=2 * y - 1
            ly=my > 1?my - 1:2 * ny
            uy=my + 1
            coarse[x,y]=(fine[lx,ly] + 2. * fine[mx,ly] + fine[ux,ly])/16.+ (fine[lx,my] + 2. * fine[mx,my] + fine[ux,my])/8.+ (fine[lx,uy] + 2. * fine[mx,uy] + fine[ux,uy]) / 16.
        end
    end
end

function interpolation!{T}(coarse::AbstractArray{T,2}, fine::AbstractArray{T,2})
    nx,ny =size(coarse)
    @inbounds for x in 1:nx
        mx=2 * x - 1
        ux=x < nx?x + 1:1
        for y in 1:ny
            my=2 * y - 1
            uy=y < ny?y + 1:1
            fine[mx,my]=coarse[x,y]
            fine[mx,my + 1]=(coarse[x,y] + coarse[x,uy]) / 2.
            fine[mx + 1,my]=(coarse[x,y] + coarse[ux,y]) / 2.
            fine[mx + 1,my + 1]=(coarse[x,y] + coarse[x,uy] + coarse[ux,y] + coarse[ux,uy]) / 4.
        end
    end
end

function full_cycle{T}(iteration,grid1,grid2,grid3,grid4,eps::T,tol::T)
    iteration(grid4,eps)
    interpolation!(grid4.u,grid3.u)
    iteration(grid3,eps)
    restriction!(grid3.u,grid4.u)

    iteration(grid4,eps)
    interpolation!(grid4.u,grid3.u)
    iteration(grid3,eps)
    interpolation!(grid3.u,grid2.u)
    iteration(grid2,eps)
    restriction!(grid2.u,grid3.u)
    iteration(grid3,eps)
    restriction!(grid3.u,grid4.u)

    iteration(grid4,eps)
    interpolation!(grid4.u,grid3.u)
    iteration(grid3,eps)
    interpolation!(grid3.u,grid2.u)
    iteration(grid2,eps)
    interpolation!(grid2.u,grid1.u)
end


# ---
# title: 1641. Count Sorted Vowel Strings
# id: problem1641
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math, Dynamic Programming, Backtracking
# link: <https://leetcode.com/problems/count-sorted-vowel-strings/description/>
# hidden: true
# ---
#
# Given an integer `n`, return _the number of strings of length_`n` _that
# consist only of vowels (_`a` _,_`e` _,_`i` _,_`o` _,_`u` _) and are
# **lexicographically sorted**._
#
# A string `s` is **lexicographically sorted** if for all valid `i`, `s[i]` is
# the same as or comes before `s[i+1]` in the alphabet.
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 1
#     Output: 5
#     Explanation: The 5 sorted strings that consist of vowels only are ["a","e","i","o","u"].
#
#
# **Example 2:**
#
#
#
#     Input: n = 2
#     Output: 15
#     Explanation: The 15 sorted strings that consist of vowels only are
#     ["aa","ae","ai","ao","au","ee","ei","eo","eu","ii","io","iu","oo","ou","uu"].
#     Note that "ea" is not a valid string since 'e' comes after 'a' in the alphabet.
#
#
# **Example 3:**
#
#
#
#     Input: n = 33
#     Output: 66045
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 50`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


#
# Defaults
#
CIRCLE_DEFAULT_PLOT_POINTS = 720
WIREFRAME_DEFAULT = 32
SURFACE_RESOLUTION_DEFAULT = 32
#
# Plotting Recipe – Poincaré Ball
#
@recipe function f(
    M::Hyperbolic{2},
    pts::AbstractVector{P},
    vecs::Union{AbstractVector{T},Nothing}=nothing;
    circle_points=CIRCLE_DEFAULT_PLOT_POINTS,
    geodesic_interpolation=-1,
    hyperbolic_border_color=RGBA(0.0, 0.0, 0.0, 1.0),
) where {P<:PoincareBallPoint,T<:PoincareBallTVector}
    @series begin
        φr = range(0, stop=2 * π, length=circle_points)
        x = [cos(φ) for φ in φr]
        y = [sin(φ) for φ in φr]
        seriestype := :path
        seriescolor := hyperbolic_border_color
        label := ""
        x, y
    end
    show_axis --> false
    framestyle -> :none
    axis --> false
    xlims --> (-1.01, 1.01)
    ylims --> (-1.01, 1.01)
    grid --> false
    aspect_ratio --> :equal
    tickfontcolor --> RGBA(1.0, 1.0, 1.0, 1.0)
    x = []
    y = []
    if vecs === nothing
        if geodesic_interpolation < 0
            seriestype --> :scatter
            x = [p.value[1] for p in pts]
            y = [p.value[2] for p in pts]
        else
            lpts = empty(pts)
            for i in 1:(length(pts) - 1)
                # push interims points on geodesics between two points.
                push!(
                    lpts,
                    shortest_geodesic(
                        M,
                        pts[i],
                        pts[i + 1],
                        collect(range(0, 1, length=geodesic_interpolation + 2))[1:(end - 1)], # omit end point
                    )...,
                )
            end
            push!(lpts, last(pts)) # add last end point
            # split into x, y and plot as curve
            seriestype --> :path
            x = [p.value[1] for p in lpts]
            y = [p.value[2] for p in lpts]
        end
    else
        seriestype := :quiver
        quiver := ([v.value[1] for v in vecs], [v.value[2] for v in vecs])
        x = [p.value[1] for p in pts]
        y = [p.value[2] for p in pts]
    end
    return x, y
end
#
# Plotting Recipe – Poincaré Half plane
#
@recipe function f(
    M::Hyperbolic{2},
    pts::AbstractVector{P},
    vecs::Union{AbstractVector{T},Nothing}=nothing;
    geodesic_interpolation=-1,
) where {P<:PoincareHalfSpacePoint,T<:PoincareHalfSpaceTVector}
    aspect_ratio --> :equal
    framestyle --> :origin
    x = []
    y = []
    if vecs === nothing
        if geodesic_interpolation < 0
            seriestype --> :scatter
            x = [p.value[1] for p in pts]
            y = [p.value[2] for p in pts]
        else
            lpts = empty(pts)
            for i in 1:(length(pts) - 1)
                # push interims points on geodesics between two points.
                push!(
                    lpts,
                    shortest_geodesic(
                        M,
                        pts[i],
                        pts[i + 1],
                        collect(range(0, 1, length=geodesic_interpolation + 2))[1:(end - 1)], # omit end point
                    )...,
                )
            end
            push!(lpts, last(pts)) # add last end point
            # split into x, y, z and plot as curve
            seriestype --> :path
            x = [p.value[1] for p in lpts]
            y = [p.value[2] for p in lpts]
        end
    else
        quiver := ([v.value[1] for v in vecs], [v.value[2] for v in vecs])
        seriestype := :quiver
        x = [p.value[1] for p in pts]
        y = [p.value[2] for p in pts]
    end
    return x, y
end
#
# Plotting Recipe – Hyperboloid
#
@recipe function f(
    M::Hyperbolic{2},
    pts::Union{AbstractVector{P},Nothing}=nothing,
    vecs::Union{AbstractVector{T},Nothing}=nothing;
    geodesic_interpolation=-1,
    wireframe=true,
    wires=WIREFRAME_DEFAULT,
    wires_x=wires,
    wires_y=wires,
    wireframe_color=RGBA(0.0, 0.0, 0.0, 1.0),
    surface=false,
    surface_resolution=SURFACE_RESOLUTION_DEFAULT,
    surface_resolution_x=surface_resolution,
    surface_resolution_y=surface_resolution,
    surface_color=RGBA(0.9, 0.9, 0.9, 0.8),
) where {P,T}
    px = isnothing(pts) ? [-2.0, 2.0] : [p[1] for p in pts]
    py = isnothing(pts) ? [-2.0, 2.0] : [p[2] for p in pts]
    # part I: wire
    if wireframe
        x = range(min(px...), max(px...), length=wires_x)
        y = range(min(py...), max(py...), length=wires_y)
        z = sqrt.(1 .+ (x .^ 2)' .+ y .^ 2)
        @series begin
            seriestype := :wireframe
            seriescolor := wireframe_color
            x, y, z
        end
    end
    # part II: solid hyperboloid
    if surface
        x = range(min(px...), max(px...), length=surface_resolution_x)
        y = range(min(py...), max(py...), length=surface_resolution_y)
        z = sqrt.(1 .+ (x .^ 2)' .+ y .^ 2)
        @series begin
            seriestype := :surface
            seriescolor := surface_color
            x, y, z
        end
    end
    show_axis --> false
    framestyle -> :none
    axis --> false
    xlims --> (min(px...), max(px...))
    ylims --> (min(py...), max(py...))
    grid --> false
    colorbar --> false
    tickfontcolor --> RGBA(1.0, 1.0, 1.0, 1.0)
    #
    # just pts given -> plot points
    x = []
    y = []
    z = []
    if pts !== nothing
        if vecs === nothing
            if geodesic_interpolation < 0
                seriestype --> :scatter
                x = [p[1] for p in pts]
                y = [p[2] for p in pts]
                z = [p[3] for p in pts]
            else
                lpts = empty(pts)
                for i in 1:(length(pts) - 1)
                    # push interims points on geodesics between two points.
                    push!(
                        lpts,
                        shortest_geodesic(
                            M,
                            pts[i],
                            pts[i + 1],
                            collect(range(0, 1, length=geodesic_interpolation + 2))[1:(end - 1)], # omit end point
                        )...,
                    )
                end
                push!(lpts, last(pts)) # add last end point
                # split into x, y and plot as curve
                seriestype --> :path
                x = [p[1] for p in lpts]
                y = [p[2] for p in lpts]
                z = [p[3] for p in lpts]
            end
        else
            seriestype := :quiver
            quiver := ([v[1] for v in vecs], [v[2] for v in vecs], [v[3] for v in vecs])
            x = [p[1] for p in pts]
            y = [p[2] for p in pts]
            z = [p[3] for p in pts]
        end
        return x, y, z
    else
        return nothing
    end
end
#
# Plotting Recipe – Sphere
#
@recipe function f(
    M::Sphere{2,ℝ},
    pts::Union{AbstractVector{P},Nothing}=nothing,
    vecs::Union{AbstractVector{T},Nothing}=nothing;
    geodesic_interpolation=-1,
    wireframe=true,
    wires=WIREFRAME_DEFAULT,
    wires_lat=wires,
    wires_lon=wires,
    wireframe_color=RGBA(0.0, 0.0, 0.0, 1.0),
    surface=false,
    surface_resolution=SURFACE_RESOLUTION_DEFAULT,
    surface_resolution_lat=surface_resolution,
    surface_resolution_lon=surface_resolution,
    surface_color=RGBA(0.9, 0.9, 0.9, 0.8),
) where {P,T}
    # part I: wire
    if wireframe
        u = range(0, 2π, length=wires_lon + 1)
        v = range(0, π, length=wires_lat + 1)
        x = cos.(u) * sin.(v)'
        y = sin.(u) * sin.(v)'
        z = repeat(cos.(v)', outer=[wires_lon + 1, 1])
        @series begin
            seriestype := :wireframe
            seriescolor := wireframe_color
            return x, y, z
        end
    end
    # part II: solid sphere
    if surface
        u = range(0, 2π, length=surface_resolution_lon + 1)
        v = range(0, π, length=surface_resolution_lat + 1)
        x = cos.(u) * sin.(v)'
        y = sin.(u) * sin.(v)'
        z = repeat(cos.(v)', outer=[wires_lon + 1, 1])
        @series begin
            seriestype := :surface
            seriescolor := surface_color
            return x, y, z
        end
    end
    show_axis --> false
    framestyle -> :none
    axis --> false
    xlims --> (-1.05, 1.05)
    ylims --> (-1.05, 1.05)
    zlims --> (-1.05, 1.05)
    grid --> false
    colorbar --> false
    tickfontcolor --> RGBA(1.0, 1.0, 1.0, 1.0)
    x = []
    y = []
    z = []
    if pts !== nothing
        if vecs === nothing #plot points/geodesics
            if geodesic_interpolation < 0
                seriestype --> :scatter
                x = [p[1] for p in pts]
                y = [p[2] for p in pts]
                z = [p[3] for p in pts]
            else
                lpts = empty(pts)
                for i in 1:(length(pts) - 1)
                    # push interims points on geodesics between two points.
                    push!(
                        lpts,
                        shortest_geodesic(
                            M,
                            pts[i],
                            pts[i + 1],
                            collect(range(0, 1, length=geodesic_interpolation + 2))[1:(end - 1)], # omit end point
                        )...,
                    )
                end
                push!(lpts, last(pts)) # add last end point
                # split into x, y and plot as curve
                seriestype --> :path
                x = [p[1] for p in lpts]
                y = [p[2] for p in lpts]
                z = [p[3] for p in lpts]
            end
        else
            seriestype := :quiver
            quiver := ([v[1] for v in vecs], [v[2] for v in vecs], [v[3] for v in vecs])
            x = [p[1] for p in pts]
            y = [p[2] for p in pts]
            z = [p[3] for p in pts]
        end
    end
    return x, y, z
end


using Test, Random, EndowmentsLH

function quantile_test()
    rng = MersenneTwister(45);
    @testset "Quantiles" begin
        n = 50;
        ed = EndowmentsLH.make_test_endowment_draws(n);
        # pctV = LinRange(0.1, 0.9, n);
        for dName in names(ed)
            drawV = get_draws(ed, dName);
            if eltype(drawV) <: Real
                qV = endow_quantiles(ed, dName, [0.3, 0.7]);
                @test eltype(qV) == eltype(drawV)
                @test size(qV) == (2,)
                @test all(diff(qV) .> 0.0)

                q = endow_quantiles(ed, dName, 0.7);
                @test isapprox(qV[2], q)
                nLess = sum(get_draws(ed, dName) .<= q);
                @test abs(nLess / n - 0.7) < 1.5 / n
            end
        end
    end
end

@testset "Quantiles" begin
    quantile_test();
end

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


export ConvertInterval, DGP_FEOL_2

function ConvertInterval(ystar,thresholds)
    # This function translates a vector of latent variables to a vector of interval-censored outcomes.

    y = similar(ystar)
    for i in 1:length(ystar)
        y[i] = sum(ystar[i].>thresholds) + 1
    end
    return y
end
#df[:y] = ConvertInterval(df[:ystar],thresholds)

function DGP_FEOL_2(n, beta0k,thresholds1,thresholds2)

    T = 2 #hard-code T=2?
    K = 2 # also hard-coded
    nT = n*T

    beta0 = beta0k + zeros(K,1)
    D_X = Normal(0,1)
    D_u = Logistic()

    # Generate data frame.
    df = DataFrame(
        i = repeat(1:n, inner = T),
        t = repeat(1:T, outer = n),
        X1 = rand(D_X,n*T),
        X2 = rand(D_X,n*T),
        u = rand(D_u,n*T)
    );

    df_a = by(df, [:i], y-> 0.2*(mean(y[:X1]) + mean(y[:X2])))
    df_a[:a] = df_a[:x1]
    delete!(df_a,:x1)
    df = join(df,df_a,on = :i)
    df[:ystar] = 0 + df[:X1]*beta0[1] + df[:X2]*beta0[2] - df[:u] + df[:a]

    # This only works because the data is perfectly shaped and T=2.
    df[:y] = 0
    df[filter(isodd,1:nT),:y] = ConvertInterval(df[filter(isodd,1:nT),:ystar],thresholds1)
    df[filter(iseven,1:nT),:y] = ConvertInterval(df[filter(iseven,1:nT),:ystar],thresholds2)

    return df
end


@testset "Small graphs" begin
    @testset "by string" begin
        g = smallgraph("diamondgraph")
        @test nv(g) == 4 && ne(g) == 5

        g = smallgraph("diamond")
        @test nv(g) == 4 && ne(g) == 5
    end

    smallgraphs = [
        (:diamond        , 4, 5, false),
        (:bull           , 5, 5, false),
        (:chvatal        , 12, 24, false),
        (:cubical        , 8, 12, false),
        (:desargues      , 20, 30, false),
        (:dodecahedral   , 20, 30, false),
        (:frucht         , 12, 18, false),
        (:heawood        , 14, 21, false),
        (:house          , 5,6, false),
        (:housex         , 5, 8, false),
        (:icosahedral    , 12, 30, false),
        (:krackhardtkite , 10, 18, false),
        (:moebiuskantor  , 16, 24, false),
        (:octahedral     , 6, 12, false),
        (:pappus         , 18, 27, false),
        (:petersen       , 10, 15, false),
        (:sedgewickmaze  , 8, 10, false),
        (:tetrahedral    , 4, 6, false),
        (:truncatedcube  , 24, 36, false),
        (:truncatedtetrahedron , 12, 18, false),
        (:truncatedtetrahedron_dir , 12, 18, true),
        (:tutte          , 46, 69, false)
    ]

    @testset "$(s[1])" for s in smallgraphs
        g = smallgraph(s[1])
        (nvg, neg, dir) = s[2:4]
        @test nv(g) == nvg
        @test ne(g) == neg
        @test is_directed(g) == dir
    end

    @testset "karate" begin
        g = smallgraph(:karate)
        degree_sequence = sort([16, 9, 10, 6, 3, 4, 4, 4, 5, 2, 3, 1, 2, 5, 2,  2,  2,
                                 2, 2,  3, 2, 2, 2, 5, 3, 3, 2, 4, 3, 4, 4, 6, 12, 17])
        @test nv(g) == 34 && ne(g) == 78 && sort(degree(g)) == degree_sequence
    end
    @testset "nonexistent graph" begin
        @test_throws ArgumentError g = smallgraph(:nonexistent)
    end
end


function test_mat2mat_base(f::Function, M, N)
    srand(98765)
    srcmat = repmat(rand(M), 1, N)
    dstlen = length(f(srcmat[:,1]))
    dstmat = Array{Float64}(dstlen, N)
    for i = 1:N
        dstmat[:,i] = f(srcmat[:,i])
    end
    dstmat2 = f(srcmat)
    @test dstmat2 ≈ dstmat
end

function test_mat2mat()
    src = zeros(20)
    srcmat = zeros(20, 2)

    # all extended functions should accept function calls like:
    # 1. f(x) where x is a input vector
    # 2. f(x) where x is a input matrix
    for f in SPTK.vec2vec
        @show f
        @test applicable(eval(f), src)
        @test applicable(eval(f), srcmat)
    end

    # test_mat2mat_base(mcep, 1024, 2)
    # test_mat2mat_base(mfcc, 512, 2)
    test_mat2mat_base(lpc2c, 20, 2)
    test_mat2mat_base(mc2b, 20, 2)
    test_mat2mat_base(gc2gc, 20, 2)
    test_mat2mat_base(mgc2mgc, 20, 2)
end

println("-- test mat2mat conversion applicable (col-wise)")
test_mat2mat()


using LinearAlgebra
using SparseArrays

using FastGaussQuadrature

"Construct the Element-To-Element (EToE) and
Element-To-Face (EtoF) maps
"
function Connect1D(EToV)
    K = size(EToV)[1] # Number of cells
    Nfaces = 2

    TotalFaces = Nfaces*K  # Total number of faces
    Nv = K+1  # Number of vertices

    # Local face to local vertex connections
    vn = [1,2]

    FToVrows = 1:TotalFaces
    FToVcols = zeros(TotalFaces)
    FToVvals = ones(Int, TotalFaces)

    row = 1
    for k=1:K
        for face=1:Nfaces
            FToVcols[row] = EToV[k,vn[face]]
            row += 1
        end
    end
    FToV = sparse(FToVrows, FToVcols, FToVvals)

    # Global face to global face sparse array
    FToF = FToV * FToV' - I

    faces1, faces2, _ = findnz(FToF) # Find non-zero elements

    # Convert face global number to element and face numbers
    element1 = floor.(Int, (faces1 .- 1)/Nfaces) .+ 1
    face1 = mod.( faces1 .- 1, Nfaces ) .+ 1
    element2 = floor.(Int, (faces2 .- 1)/Nfaces) .+ 1
    face2 = mod.( faces2 .- 1, Nfaces ) .+ 1

    # Rearrange into Nelements x Nfaces sized arrays

    # This initialises all faces to connect to themselves
    # Faces which connect to other cells are then overwritten
    EToE = (1:K) * ones(Int, 1, Nfaces)
    EToF = ones(Int, K,1)*(1:Nfaces)'

    # this index code doesn't work as expected
    #ind = sub2ind([K,Nfaces], element1, face1)
    #EToE[ind] = element2
    #EToF[ind] = face2

    for i = 1:length(element1)
        EToE[element1[i], face1[i]] = element2[i]
        EToF[element1[i], face1[i]] = face2[i]
    end

    return EToE, EToF
end

"
K    Number of cells
Nfp  Number of nodes per face
Np   Number of nodes per cell
"
function BuildMaps1D(K::Int, Nfaces::Int, Nfp::Int, Np::Int, Fmask, EToE, EToF, x)

    nodeids = reshape(1:K*Np, Np, K)

    vmapM   = zeros(Int, Nfp, Nfaces, K)
    vmapP   = zeros(Int, Nfp, Nfaces, K)

    for k1=1:K
        for f1=1:Nfaces
            # find index of face nodes with respect to volume node ordering
            vmapM[:,f1,k1] .= nodeids[Fmask[f1], k1]
        end
    end

    for k1=1:K
        for f1=1:Nfaces
            # find neighbor
            k2 = EToE[k1,f1]
            f2 = EToF[k1,f1]

            # find volume node numbers of left and right nodes
            vidM = vmapM[:,f1,k1]
            vidP = vmapM[:,f2,k2]

            x1  = x[vidM]
            x2  = x[vidP]

            # Compute distance matrix
            D = sum(abs2, x1 -x2)
            if D < 1e-5
                vmapP[:,f1,k1] = vidP
            end
        end
    end

    # Find where vmapP == vmapM
    mapB = findall(iszero, vmapP - vmapM)
    vmapB = vmapM[mapB]

    # Create specific left (inflow) and right (outflow) maps
    mapI = 1
    mapO = K*Nfaces
    vmapI = 1
    vmapO = K*Np;

    return vmapP, vmapM, mapB, vmapB, mapI, mapO, vmapI, vmapO
end

# Compute LIFT
function Lift1D(Np, Nfaces, Nfp, V)
    print(Np,", ", Nfaces, ", ", Nfp)
    Emat = zeros(Np, Nfaces*Nfp)

    # Define Emat
    Emat[1,1] = 1.0
    Emat[Np,2] = 1.0

    # inv(mass matrix)*\s_n (L_i,L_j)_{edge_n}
    V*(V'*Emat);
end

struct Quadrature1D
    # locations in reference cell [-1,1]
    locations::Array{Float64, 1}
    weights::Array{Float64, 1}
end

"Compute LGL quadrature points on reference cell [-1,1]"
function gaussLobatto1D(npoints)
    r, w = gausslobatto(npoints)
    Quadrature1D(r, w)
end

function equidistant1D(npoints)
    r = range(-1, stop=1, length=npoints)
    w = fill(2.0 / npoints, npoints)
    Quadrature1D(r, w)
end

"Describes the basis functions and quadrature nodes on a reference cell.
The matrices derived on this cell can be applied to all cells
by including the geometrical scaling factors."
struct ReferenceCell
    basis
    quadrature::Quadrature1D

    V    # Vandermonde matrix on reference cell
    invV # Inverse Vandermonde matrix
    Dr   # Differentiation matrix

    function ReferenceCell(basis, quadrature)
        V = zeros(length(quadrature.locations), basis.order)
        Dr = zeros(length(quadrature.locations), basis.order)

        # Calculate basis function at each quadrature location
        # (Vandermonde matrix)
        #   V[i,j] = P_j(x_i)
        #
        # DVr is the differentiation matrix
        #
        for j=1:basis.order
            V[:,j] = basisFunction(basis, quadrature.locations, j)
            Dr[:,j] = basisGradient(basis, quadrature.locations, j)
        end
        Dr /= V

        new(basis, quadrature,
            V, inv(V), Dr)
    end
end

"Describes a 1D grid"
mutable struct Grid1D

    mesh::Mesh1D # Defines node locations and cells

    npoints::Int # Number of points per cell, including edges (Np)
    npoints_per_face::Int # Number of points on a face
    nfaces::Int # Number of faces per cell

    reference::ReferenceCell

    LIFT # Surface integral terms

    J    # Jacobian [point, cell]
    rx   # Inverse Jacobian

    x::Array{Float64, 2} # Physical coordinates of the grid points [point,cell]

    Fscale

    nx  # Outward pointing normal at element faces

    vmapP # Map of external points for each cell face
    vmapM # Map of internal points for each cell face

    mapB
    vmapB

    mapI
    mapO
    vmapI
    vmapO

    function Grid1D(mesh::Mesh1D, basis, quadrature::Quadrature1D)

        # Create an instance. Not yet initialised
        # This is to avoid a long constructor parameter list
        g = new()

        g.mesh = mesh

        g.npoints = length(quadrature.locations)
        g.nfaces = 2 # Number of faces per cell
        g.npoints_per_face = 1  # Number of points per face

        g.reference = ReferenceCell(basis, quadrature)

        N = length(quadrature.locations)-1

        g.LIFT = Lift1D(g.npoints, g.nfaces, g.npoints_per_face, g.reference.V)

        # Coordinates of all the vertices
        va = mesh.element_vertices[:,1]'
        vb = mesh.element_vertices[:,2]'
        g.x = ones(N+1,1)*mesh.vertex_coordinates[va] .+ 0.5*(quadrature.locations  .+ 1)*(mesh.vertex_coordinates[vb] - g.mesh.vertex_coordinates[va])

        # Geometric factors
        g.J = g.reference.Dr * g.x
        g.rx = 1.0 ./ g.J

        fmask1 = 1
        fmask2 = N+1
        Fmask = [fmask1;fmask2]
        Fx = g.x[Fmask,:]  # Locations of cell edges for each cell

        g.Fscale = 1.0 ./ g.J[Fmask,:]

        # Compute outward pointing normals
        g.nx = zeros(g.nfaces, mesh.nelements)
        g.nx[1,:] .= -1
        g.nx[2,:] .= 1

        # Build global connectivity arrays
        EToE, EToF = Connect1D(mesh.element_vertices)

        g.vmapP, g.vmapM, g.mapB, g.vmapB, g.mapI, g.mapO, g.vmapI, g.vmapO = BuildMaps1D(mesh.nelements, 2, 1, g.npoints, Fmask, EToE, EToF, g.x)

        return g
    end
end


# https://github.com/nalimilan/FreqTables.jl #
##############################################

import Pkg
Pkg.instantiate() # downloads all dependencies for the current project
Pkg.add("FreqTables")

using DataFrames, CSV, FreqTables, DelimitedFiles

df = CSV.read("../externals/Core.Math.Data/data/Pejcic_318.csv", copycols = true)

println(df.MJESTO)

af = freqtable(df.MJESTO)
println(af)

rf = prop(af)
println(rf)

af_2 = freqtable(df.SPOL, df.RAZRED)
println(af_2)


using SobolSearches, Test

@testset "simple consistency checks" begin
    y = [0.2, 0.7]
    f(x) = sum(abs2, x .- y)
    s = sobol_search(f, [-2, -1.0], [3, 1.0], 10, 10000)
    @test s[1].value == f(s[1].position)
    @test s[1].position ≈ y atol = 0.02
    @test IndexStyle(typeof(s)) ≡ IndexLinear()
    @test size(s) == (10, )
    @test length(s) == 10
end


using Test

using MathOptInterface
const MOI = MathOptInterface
const MOIT = MathOptInterface.Test
const MOIU = MathOptInterface.Utilities
const MOIB = MathOptInterface.Bridges

include("../utilities.jl")

mock = MOIU.MockOptimizer(MOIU.Model{Float64}())
config = MOIT.TestConfig()

@testset "SOCR" begin
    bridged_mock = MOIB.Constraint.SOCR{Float64}(mock)

    MOIT.basic_constraint_tests(
        bridged_mock, config,
        include = [(F, S)
                   for F in [MOI.VectorOfVariables, MOI.VectorAffineFunction{Float64},
                             MOI.VectorQuadraticFunction{Float64}]
                   for S in [MOI.SecondOrderCone]])

    mock.optimize! = (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [1.0, 1/√2, 1/√2],
                          (MOI.VectorAffineFunction{Float64}, MOI.RotatedSecondOrderCone) => [[1 - 1/√2, 1 + 1/√2, -1]],
                          (MOI.VectorAffineFunction{Float64}, MOI.Zeros)           => [[-√2]])
    MOIT.soc1vtest(bridged_mock, config)
    mock.optimize! = (mock::MOIU.MockOptimizer) -> MOIU.mock_optimize!(mock, [1.0, 1/√2, 1/√2],
                          (MOI.VectorAffineFunction{Float64}, MOI.RotatedSecondOrderCone) => [[1 - 1/√2, 1 + 1/√2, -1]],
                          (MOI.VectorAffineFunction{Float64}, MOI.Zeros)           => [[-√2]])
    MOIT.soc1ftest(bridged_mock, config)
    ci = first(MOI.get(
        bridged_mock,
        MOI.ListOfConstraintIndices{MOI.VectorAffineFunction{Float64},
                                    MOI.SecondOrderCone}()))
    test_delete_bridge(bridged_mock, ci, 3,
                       ((MOI.VectorAffineFunction{Float64},
                         MOI.RotatedSecondOrderCone, 0),))
end


# Try to use Wang-Landau to estimate "density of states" of "energy"=path length

using BilliardModels

function find_free_path(billiard_table, x, v)
    l = Vector2D(0, 0)
    p = ParticleOnLattice(x, v, l)
    num_collisions = 1
    xs, ls, free_paths = billiard_dynamics_on_lattice(p, billiard_table, num_collisions)

    free_paths[1]
end

unif_rand(delta) = delta*(rand()-0.5)

function propose(x, y, theta, radius, delta)
    x_new = x + unif_rand(delta)
    y_new = y + unif_rand(delta)

     # check outside disc:
    if x*x + y*y < radius^2
        x_new, y_new = x, y   # return to where it was
    else
        # periodize:
        (x_new > 0.5) && (x_new -= 1)
        (x_new < -0.5) && (x_new += 1)
        (y > 0.5) && (y -= 1)
        (y < -0.5) && (y += 1)
    end

    theta += unif_rand(delta)  # in any case change theta

    x, y, theta
end

acceptance_prob(x, y, theta) = 1.0  # standard random walk

function explore_phase_space(radius=0.1, delta=0.1, num_steps=1000)

    billiard_table = Sinai_billiard(radius, true, true)  # periodic in x and y

    xx, vv = initial_condition(billiard_table, -.5, .5, -.5, .5)

    theta = atan2(vv[2], vv[1])  # this is how I generated the velocity to start with (from an angle) -- a bit ridiculous

    xs = [xx[1]]
    ys = [xx[2]]
    thetas = [theta]
    free_paths = [find_free_path(billiard_table, xx, vv)]

    num_free_path_bins = 1000
    free_path_distribution = zeros(num_free_path_bins)

    for i in 1:num_steps
        x, y = xx[1], xx[2]

        x_new, y_new, theta_new = propose(x, y, theta, radius, delta)

        if rand() < acceptance_prob(x, y, theta)
            x, y, theta = x_new, y_new, theta_new
        end

        xx_new = Vector2D(x, y)
        vv = Vector2D(cos(theta), sin(theta))

        free_path = find_free_path(billiard_table, xx_new, vv)

        push!(xs, x)
        push!(ys, y)
        push!(thetas, theta)
        push!(free_paths, free_path)

        free_path_bin = ceil(free_path)
        (free_path_bin > num_free_path_bins) && (free_path_bin = num_free_path_bins)

        free_path_distribution[free_path_bin] += 1

        xx = xx_new

    end

    xs, ys, thetas, free_paths, free_path_distribution
end


@testset "Line" begin
    p1 = SVector(1.0, 1.0, 1.0); p2 = SVector(2.0, 2.0, 0.0);
    l1 = Line(p1, p2)

    @test l1[0] ≈ p1
    @test l1[end/2] ≈ SVector(1.5, 1.5, 0.5)
    @test l1[end] ≈ p2

    @test all(map(isapprox, l1[[0, length(l1)]], [p1, p2]))

    p3 = SVector(1.0, 1.0, 1.0); p4 = SVector(1.0, 1.0, 2.0);
    l2 = Line(p3, p4)

    l3 = Line(p1, p2) # New instance of l1
    @test isequal(l1, l3)
    @test !isequal(l1, l2)
    @test hash(l1) == hash(l3)

    @test distance(l2, SVector(1.0,1.0,1.0)) == 0
    @test distance(SVector(1,1,1.5), l2) == 0
    @test distance(SVector(3,2,3), l2) == sqrt(6)
    @test distance(l2, SVector(3,2,0)) == sqrt(6)

    @test distance(l2, SVector(1.0,1.0)) == 0
    @test distance(SVector(1,1), l2) == 0
    @test distance(SVector(3,2), l2) == sqrt(5)
    @test distance(l2, SVector(3,2)) == sqrt(5)

    l1 = Line(SVector(0.0, 0.0), SVector(1.0, 0.0))
    @test convert2d(l1) == l1
    @test convert3d(l1) == Line(SVector(0.0, 0.0, 0.0), SVector(1.0, 0.0, 0.0))
    @test convert3d(l1, 1.0) == Line(SVector(0.0, 0.0, 1.0), SVector(1.0, 0.0, 1.0))
    @test convert3d(Line(SVector(0.0, 0.0, 1.0), SVector(1.0, 0.0, 1.0)), 2.0) == Line(SVector(0.0, 0.0, 2.0), SVector(1.0, 0.0, 2.0))

    # boundingbox
    @test boundingbox(Line(SVector(1,2,3), SVector(4,5,6))) == BoundingBox(1,2,3,4,5,6)

    @testset "2D Line distances" begin
        # Intersecting lines
        @test distance(Line(SVector(0.0, 0.0), SVector(1.0, 1.0)), Line(SVector(0.0, 1.0), SVector(1.0, 0.0))) === 0.0
        @test distance(Line(SVector(0.3f0, 0.3f0), SVector(0.6f0, 0.6f0)), Line(SVector(0.3f0, 0.6f0), SVector(0.6f0, 0.3f0))) === 0f0

        # Intersecting at vertices
        @test distance(Line(SVector(0.0, 0.0), SVector(1.0, 1.0)), Line(SVector(1.0, 1.0), SVector(1.0, 0.0))) === 0.0
        @test closest_points(Line(SVector(0.0, 0.0), SVector(1.0, 1.0)), Line(SVector(1.0, 1.0), SVector(1.0, 0.0)))[1] === SVector(1.0, 1.0)
        @test closest_points(Line(SVector(0.0, 0.0), SVector(1.0, 1.0)), Line(SVector(1.0, 1.0), SVector(1.0, 0.0)))[2] === SVector(1.0, 1.0)

        @test distance(Line(SVector(0.3f0, 0.3f0), SVector(0.6f0, 0.6f0)), Line(SVector(0.6f0, 0.6f0), SVector(0.6f0, 0.9f0))) === 0f0

        # Parallel lines
        @test distance(Line(SVector(0.0, 0.0), SVector(1.0, 0.0)), Line(SVector(0.0, 0.0), SVector(1.0, 0.0))) === 0.0
        @test distance(Line(SVector(0.3f0, 0.3f0), SVector(0.6f0, 0.6f0)), Line(SVector(0.3f0, 0.3f0), SVector(0.6f0, 0.6f0))) === 0f0
        @test distance(Line(SVector(0.3f0, 0.3f0), SVector(0.6f0, 0.6f0)), Line(SVector(0.4f0, 0.4f0), SVector(0.5f0, 0.5f0))) === 0f0
        @test distance(Line(SVector(0.4f0, 0.4f0), SVector(0.8f0, 0.8f0)), Line(SVector(0.3f0, 0.3f0), SVector(0.6f0, 0.6f0))) === 0f0

        @test distance(Line(SVector(0.0, 0.5), SVector(1.0, 0.5)), Line(SVector(0.0, 0.5), SVector(1.0, 0.5))) === 0.0
        @test distance(Line(SVector(0.0, 0.5), SVector(1.0, 0.5)), Line(SVector(1.0, 0.5), SVector(2.0, 0.5))) === 0.0
        @test distance(Line(SVector(0.0, 0.5), SVector(1.0, 0.5)), Line(SVector(2.0, 0.5), SVector(3.0, 0.5))) === 1.0
        @test distance(Line(SVector(0.0, 0.0), SVector(1.0, 0.0)), Line(SVector(2.0, 0.5), SVector(3.0, 0.5))) ≈ sqrt(1.25)

        # Non-intersecting lines
        @test distance(Line(SVector(0.0, 0.0), SVector(1.0, 0.0)), Line(SVector(1.0, 1.0), SVector(1.0, 2.0))) === 1.0
        @test closest_points(Line(SVector(0.0, 0.0), SVector(1.0, 0.0)), Line(SVector(1.0, 1.0), SVector(1.0, 2.0)))[1] === SVector(1.0, 0.0)
        @test closest_points(Line(SVector(0.0, 0.0), SVector(1.0, 0.0)), Line(SVector(1.0, 1.0), SVector(1.0, 2.0)))[2] === SVector(1.0, 1.0)

    end

    @testset "3D Line distances" begin
        l1 = Line(SVector(0.0, 0.0, 0.0), SVector(1.0, 1.0, 0.0))
        l2 = Line(SVector(0.0, 1.0, 0.5), SVector(1.0, 0.0, 0.5))
        @test distance(l1, l2) ≈ 0.5
        @test closest_points(l1, l2)[1] ≈ SVector(0.5, 0.5, 0.0)
        @test closest_points(l1, l2)[2] ≈ SVector(0.5, 0.5, 0.5)
    end

    # TODO test transformations
end