module PolynomialZeros

using Polynomials
import PolynomialRoots

using PolynomialFactors
import Roots: bisection64

using AbstractAlgebra
const AA = AbstractAlgebra

using Printf

export poly_roots
export Over

# export agcd, multroot # qualify these

include("utils.jl")
include("over.jl")
include("special_cases.jl")

include("agcd/agcd.jl")        # AGCD.agcd
include("agcd/multroot.jl")    # MultRoot.multroot
include("amvw/AMVW.jl")        # AMVW.poly_roots
include("RealRoots/real_roots.jl")  # RealRoots.real_roots(p::Poly)

using .AGCD
using .MultRoot

"""

`poly_roots(f, domain)`: Find zeros of the polynomial `f` within the specified domain.

* `f` can be an instance of `Poly` type (from `Polynomials.jl`) or a callable object which can be so converted. Will throw a `DomainError` if the object `f` can not be converted to `Poly{T}`.

* `domain` is one of

    - `over.C` (the default) for solving over complex values
      (`Complex{Float64}`). Use `Over.CC{T}` to specfy a type `T<:
      AbstractFloat` other than `Float64`. The default method is from
      `PolynomialRoots`. Pass the argument `method=:roots` to use the
      `roots` function from `Polynomials.jl`. Pass the argument
      `method=:amvw` to use an algorithm by Aurentz, Mach, Vandebril,
      and Watkins. For a degree n polynomial over C, all n roots
      should be returned (including multiplicities).

    - `over.R` for solving over the real line (`Float64`). Use
      `Over.RR{T}` to specify a `T <: Integer` other than
      `Float64`. The algorithm assumes the polynomial is square free
      (none of its factors are squares over R). This is important for
      floating point coefficients. Pass the argument
      `square_free=false` to have an *approximate* gcd used to create
      a square-free version. Only unique real roots are returned (no
      multiplicities).

    - `over.Q` for solving over the rational numbers
      (`Rational{Int}`). Use `Over.RR{T}` to specify a `T <: Integer`
      other than `Int`. Only unique rational roots are returned (no
      multiplicities).

    - `over.Z` for solving over the integers (`Int`). Use `Over.ZZ{T}`
      to specify a `T` other than `Int`. Only unique integer roots are
      returned (no multiplicities).

    - `over.Zp{p}` for solving over the finite field `ZZ_p`, p a
      prime.  Only unique roots are returned (no multiplicities).

Returns an array of zeros, possibly empty. May throw error if polynomial type is inappropriate for specified domain. For an unspecified domain, the domain may reflect the element type of `f(0)`.

Examples:

```julia
using Polynomials, PolynomialZeros
x = variable(); p = x^5 - x - 1
poly_roots(p, Over.C)  # 5
poly_roots(p, Over.R)  # 1
poly_roots(p, Over.Q)  # empty

p = x^3 - 1
poly_roots(p, Over.C)  # 3
poly_roots(p, Over.R)  # 1
poly_roots(p, Over.Q)  # 1
poly_roots(p, Over.Z)  # 1
poly_roots(p, Over.Zp{7})  # 3
```

"""
function poly_roots(f; method=:PolynomialRoots)
    poly_roots(f, Over.C; method=method) # default
end

function poly_roots(f, ::Type{Over.C}; method=:PolynomialRoots)
    T = promote_type(Float64, e_type(f))
    poly_roots(f, Over.CC{T}, method=method)
end

function poly_roots(f, U::Type{Over.CC{T}}; method=:PolynomialRoots) where {T<:AbstractFloat}

    ps = poly_coeffs(T, f)
    fn = special_case(ps, U)

    if fn == identity
        if method == :PolynomialRoots && T <: Union{Float64, BigFloat}
            PolynomialRoots.roots(ps, polish=true)
        elseif method == :Roots && T <: Float64
            convert(Vector{Complex{T}}, Polynomials.roots(Poly(ps)))
        elseif (method == :AMVW || method == :amvw)
            AMVW.poly_roots(ps)
        else # default to AMVW
            AMVW.poly_roots(ps)
        end
    else
        fn(ps, U)
    end

end

function poly_roots(f, ::Type{Over.R};square_free=false)
    T = promote_type(Float64, e_type(f))
    poly_roots(f, Over.RR{T}, square_free=square_free)
end
function poly_roots(f, U::Type{Over.RR{T}}; square_free::Bool=false) where {T <: Real}

    ps = convert(Vector{T}, poly_coeffs(f))
    fn = special_case(ps, U)

    if fn == identity
        RealRoots.real_roots(as_poly(T, f), square_free=square_free)
    else
        fn(ps, U)
    end

end
# should I do an alias? Best to add if requestd, keeping name space light for now
## realroots(f; kwargs...) = poly_roots(f, Over.R; kwargs...) #real_roots?

_rational_T(::Type{Rational{T}}) where {T}= T
_rational_T(::Rational{T}) where {T} =T
function poly_roots(f, ::Type{Over.Q})
    T = promote_type(Int, e_type(f))

    if T <: Integer
        return poly_roots(f, Over.QQ{T})
    elseif T <: Rational
        S = promote_type(_rational_T(T), Int)
        p = as_poly(T, f)
        return poly_roots(f, Over.QQ{S})
    else
        throw(ArgumentError("Use Over.Q for polynomials with integer or rational coefficients"))
    end
end

# for v0.6 -> v0.7 transition on filter with Dicts

function poly_roots(f, U::Type{Over.QQ{T}}) where {T <: Integer}
    p = as_poly(Rational{T}, f)
    fn = special_case(poly_coeffs(p), U)
    if fn == identity
        ps = poly_coeffs(p)
        l = lcm(denominator.(ps))
        qs = numerator.(l * ps)
        d = PolynomialFactors.poly_factor(qs)
        d = filter(kv -> AA.degree(kv[1]) == 1, d) ## is this broken
        rts = Rational{T}[]
        for k in keys(d)
            p0 = k(0); p1 = k(1) - p0
            push!(rts, -p0//p1)
        end
        rts
    else
        fn(p.a, U)
    end

end

## ----

function poly_roots(f, ::Type{Over.Z})
    T = eltype(as_poly(f)(0))
    poly_roots(f, Over.ZZ{T})
end

function poly_roots(f, U::Type{Over.ZZ{T}}) where {T <: Integer}

    p = as_poly(T, f)

    fn = special_case(poly_coeffs(p), U)
    if fn == identity
        d = PolynomialFactors.poly_factor(poly_coeffs(p))
        d = filter(kv -> AA.degree(kv[1]) == 1, d)
        rts = T[]
        for (k,v) in d
            x1, x0= k(1)-k(0), k(0)
            if iszero(rem(x0, x1))
                push!(rts, div(x0, x1))
            end
        end
        rts
    else
        fn(p.a, U)
    end
end

function poly_roots(f, U::Type{Over.Zp{q}}) where {q}
    # error if q is not prime?

    p = as_poly(e_type(f), f)
    paa = PolynomialFactors.as_poly(p.a)

    fn = special_case(poly_coeffs(p), U)
    if fn == identity
        fs = PolynomialFactors.factormod(paa,q)
        ls = filter(kv -> AA.degree(kv[1]) == 1, fs)
        rts = eltype(p.a)[]
        for (k,v) in ls
            x0 = k(0)
            x1 = k(1) - k(0)
            u = x0 * inv(x1)
            push!(rts, PolynomialFactors.value(u))
        end
        rts
    else
        fn(p.a, U)
    end

end

end # module


module AuxFunctions

using LinearAlgebra: mul!
using FFTW: fftshift
export peel, init!

function disc(centre::T,N::T,radius::P) where {T<:Signed,P<:AbstractFloat} # 1D method
    # This function constructs a disc in a vector:
    # 0 0 1 0 0 -> 0 1 0 1 0 - > 1 0 0 0 1
    D = Vector{P}(undef,N)
    @inbounds D = [sqrt((i-centre)^2)<=radius ? D[i] = 1.0 : D[i] = 0.0 for i = 1:N]
    return D
end

function disc(centre::Vector{T},N::T,radius::P) where {T<:Signed,P<:AbstractFloat} # 2D method
    # This function constructs a disc in a matrix, ex:
    # 0 0 1 0 0
    # 0 1 0 1 0
    # 0 0 1 0 0
    D = Matrix{P}(undef,N,N)
    @inbounds D = [sqrt((i-centre[1])^2+(j-centre[2])^2)<=radius ? D[i,j] = 1.0 : D[i,j] = 0.0 for i = 1:N, j = 1:N]
    return D
end

function peel(P,A::Vector{T},Ω) where {T<:AbstractFloat} # 1D method
    # This function constructs k_delay in 1D
    # Perform the Fourier Transform of the respective delay ring

    hN = Ω.N÷2+1 # Half of dim N
    B  = Vector{Complex{T}}(undef,Ω.rings1D*hN) # Each hN block corresponds to a delay ring
    Bi = Vector{Complex{T}}(undef,hN)

    # Ring loop. For each delay ring compute rfft
    @inbounds for i = 2:Ω.rings1D
        r1 = (i-1)*Ω.Δr
        r2 = i*Ω.Δr
        # fftshift is needed because the integral op A is a circulary convolution
        mul!(Bi,P,fftshift((disc(hN,Ω.N,r2) - disc(hN,Ω.N,r1)).*A))
        B[1+hN*(i-1):hN*i] .= Bi
    end

    mul!(Bi,P,fftshift(disc(hN,Ω.N,Ω.Δr).*A))
    B[1:hN] .= Bi
    return B
end

function peel(P,A::Matrix{T},Ω) where {T<:AbstractFloat} # 2D method
    # This function constructs k_delay in 2D
    # Perform the Fourier Transform of the respective delay ring

    hN = Ω.N÷2+1 # Half of dim N
    B  = Matrix{Complex{T}}(undef,hN*Ω.rings2D,Ω.N) # Each (hN X N) block corresponds to a delay ring
    Bi = Matrix{Complex{T}}(undef,hN,Ω.N)

    # Ring loop. For each delay ring compute rfft
    @inbounds for i = 2:Ω.rings2D
        r1 = (i-1)*Ω.Δr
        r2 = i*Ω.Δr
        # fftshift is needed because the integral op A is a circulary convolution
        mul!(Bi,P,fftshift((disc([hN,hN],Ω.N,r2) - disc([hN,hN],Ω.N,r1)).*A))
        B[1+hN*(i-1):hN*i,:] .= Bi
    end

    mul!(Bi,P,fftshift(disc([hN,hN],Ω.N,Ω.Δr).*A))
    B[1:hN,:] .= Bi
    return B
end

# Method init for 1D with V0 being a constant
function init!(v::Vector{<:Complex{<:T}},sv::Vector{<:Complex{<:T}},V0arr::Vector{T},P,S,V0::Real,Ω) where {T<:AbstractFloat}
    svu   = Vector{Complex{T}}(undef,Ω.N÷2+1)
    V0arr.= fill(V0,Ω.N) # Initial condition V(x,y,t=0)
    mul!(v,P,V0arr)      # Initial condition in the Fourier domain

    mul!(svu,P,S.(V0arr))
    sv .= repeat(svu,Ω.rings1D)
end

# Method init for 2D with V0 being a constant
function init!(v::Matrix{<:Complex{<:T}},sv::Matrix{<:Complex{<:T}},V0arr::Matrix{T},P,S,V0::Real,Ω) where {T<:AbstractFloat}
    svu   = Matrix{Complex{T}}(undef,Ω.N÷2+1,Ω.N)
    V0arr.= fill(V0,Ω.N,Ω.N) # Initial condition V(x,y,t=0)
    mul!(v,P,V0arr)          # Initial condition in the Fourier domain

    mul!(svu,P,S.(V0arr))
    sv .= repeat(svu,Ω.rings2D,1)
end

# Method init for 1D with V0 being a function
function init!(v::Vector{<:Complex{<:T}},sv::Vector{<:Complex{<:T}},V0arr::Vector{T},P,S,V0,Ω) where {T<:AbstractFloat}
    # V0 is a function
    svu   = Vector{Complex{T}}(undef,Ω.N÷2+1)
    V0arr.= [V0(i) for i in Ω.x] # Discretise V0 (Time domain)
    mul!(v,P,V0arr)              # Compute V0 in frequency domain

    mul!(svu,P,S.(V0arr))
    sv .= repeat(svu,Ω.rings1D)
end

# Method init for 2D with V0 being a function
function init!(v::Matrix{<:Complex{<:T}},sv::Matrix{<:Complex{<:T}},V0arr::Matrix{T},P,S,V0,Ω) where {T<:AbstractFloat}
    # V0 is a function
    svu   = Matrix{Complex{T}}(undef,Ω.N÷2+1,Ω.N)
    V0arr.= [V0(i,j) for j in Ω.y,i in Ω.x] # Discretise V0 (Time domain)
    mul!(v,P,V0arr)                         # Compute V0 in frequency domain

    mul!(svu,P,S.(V0arr))
    sv .= repeat(svu,Ω.rings2D,1)
end

end #end module


using TensorBoardLogger #import the TensorBoardLogger package
using Logging #import Logging package

logger = TBLogger("audiologs", tb_append) #create tensorboard logger

################log audio example: cos wave ################
samplerate = 44100
x = collect(1:samplerate*2)
x = cos.(440*pi*x/samplerate)
#using logger interface
with_logger(logger) do
    @info "cos/loggerinterface" x = TBAudio(x, samplerate)
end
#using explicit function interface
log_audio(logger, "cos/explicitinterface", x, samplerate, step = 0)

################log scalar example: sin+cos wave################
y = collect(1:samplerate*2)
y = sin.(440*pi*y/samplerate)
z = x+y
#using logger interface
with_logger(logger) do
    @info "sin+cos/loggerinterface" z = TBAudio(z, samplerate)
end
#using explicit function interface
log_audio(logger, "sin+cos/explicitinterface", z, samplerate, step = 0)


function prime_factors(n)
    pf = []
    temp = n
    i = 2
    while i <= temp
        if temp % i == 0
            push!(pf, i)
            temp ÷= i
        else
            i += 1
        end
    end
    pf
end


using VibronicToolkit: jackknife

using Random: MersenneTwister
using Statistics: cov, mean, std, var

# This uses the formulas from Peter Young's notes
# (https://arxiv.org/abs/1210.3781, https://arxiv.org/pdf/1210.3781v3.pdf).
function propagate_errors_nonlinear1(f, dfx, dfxx, xs)
    mu_x = mean(xs)
    biased_mean = f(mu_x)
    pre_bias = dfxx(mu_x)*var(xs)
    unbiased_mean = biased_mean - 0.5*pre_bias/length(xs)
    pre_err = dfx(mu_x)^2*var(xs)
    err = sqrt(pre_err/length(xs))
    unbiased_mean, err
end

function propagate_errors_nonlinear2(f, dfx, dfy, dfxx, dfxy, dfyy, xs, ys)
    mu_x, mu_y = mean(xs), mean(ys)
    biased_mean = f(mu_x,mu_y)
    pre_bias = dfxx(mu_x,mu_y)*var(xs) + 2dfxy(mu_x,mu_y)*cov(xs,ys) + dfyy(mu_x,mu_y)*var(ys)
    unbiased_mean = biased_mean - 0.5*pre_bias/length(xs)
    pre_err = dfx(mu_x,mu_y)^2*var(xs) + 2dfx(mu_x,mu_y)*dfy(mu_x,mu_y)*cov(xs,ys) + dfy(mu_x,mu_y)^2*var(ys)
    err = sqrt(pre_err/length(xs))
    unbiased_mean, err
end

compare_stats((mu1, err1), (mu2, err2)) = isapprox(mu1, mu2; rtol=1e-11) && isapprox(err1, err2; rtol=1e-6)

let rng = MersenneTwister(0x1234),
    xs = randn(rng, 10000),
    ys = randn(rng, 10000)

    # Linear function of a single variable.
    l1a(x) = 2x.-3.7
    @test compare_stats(jackknife(l1a, xs),
                        (l1a(mean(xs)),
                         2std(xs)/sqrt(length(xs))))
    l1b(y) = -1.8y.+9
    @test compare_stats(jackknife(l1b, ys),
                        (l1b(mean(ys)),
                         1.8std(ys)/sqrt(length(ys))))

    # Linear function of two variables.
    l2a(x, y) = l1a(x) + l1b(y)
    @test compare_stats(jackknife(l2a, xs, ys),
                        (l2a(mean(xs), mean(ys)),
                         sqrt((2std(xs))^2+(1.8std(ys))^2-2*2*1.8*cov(xs,ys))/sqrt(length(xs))))

    # Nonlinear function of one variable.
    n1a(y) = 1 ./ l1b(y)
    dn1ay(y) = 1.8 ./ l1b(y)^2
    dn1ayy(y) = 2*1.8^2 ./ l1b(y)^3
    @test compare_stats(jackknife(n1a, ys),
                        (propagate_errors_nonlinear1(n1a, dn1ay, dn1ayy, ys)))

    # Nonlinear function of two variables.
    n2a(x, y) = l1a(x) ./ l1b(y)
    dn2ax(x, y) = 2 ./ l1b(y)
    dn2ay(x, y) = 1.8*l1a(x) ./ l1b(y)^2
    dn2axx(x, y) = 0.0
    dn2axy(x, y) = 2*1.8 ./ l1b(y)^2
    dn2ayy(x, y) = 2*1.8^2*l1a(x) ./ l1b(y)^3
    @test compare_stats(jackknife(n2a, xs, ys),
                        propagate_errors_nonlinear2(n2a, dn2ax, dn2ay, dn2axx, dn2axy, dn2ayy, xs, ys))

    @test_throws DomainError jackknife(n2a)
    @test_throws DomainError jackknife(n2a, [1.0], [2.0, 3.0])
end


using Test
using TimeMachine

@testset "dtw 1" begin
    # dist = Distances.Euclidean()
    local cost = dtw([1, 2, 4, 1, 2], [1, 2, 4, 1, 2], (a, b) -> abs(a - b))
    @test 0.0 == cost

    cost = dtw([2, 6, 5, 3, 2], [1, 2, 4, 1, 2], (a, b) -> abs(a - b))

    @test 1 == cost
end

@testset "dtw 2" begin
    # dist = Distances.SqEuclidean()
    local a = [1, 1, 1, 2, 4, 6, 5, 5, 5, 4, 4, 3, 1, 1, 1]
    local b = [1, 1, 2, 4, 6, 6, 6, 5, 4, 4, 4, 3, 3, 3, 1]
    local cost = dtw(a, b, (a, b) -> (a - b)^2)

    @test 0.0 == cost

    a[end] += 2
    cost = dtw(a, b, (a, b) -> (a - b)^2)
    @test (2.0 / 11.0) == cost

    a = collect(1:10)
    b = a .+ 1
    cost = dtw(a, b, (a, b) -> (a - b)^2)
    @test (2.0 / 11.0) == cost
end


using RandomMatrixDistributions
using Test

n = 1000
p = 400
γ = p/n

# test white Wishart
for beta in [1, 4, 8]
    println("Testing white Wishart with β = ", beta)

    λs = randeigvals(SpikedWishart(beta, n, p))

    @test length(λs) == p
    @test eltype(λs) <: Real

    @test minimum(λs) > 0

    println("Testing Jacobi with β = ", beta)
    λs = randeigvals(Jacobi(beta, n, n, p))

    @test length(λs) == p
    @test eltype(λs) <: Real

    @test minimum(λs) > 0
end

println("Testing Spiked Wishart with β = 1")
s = 8
λs = randeigvals(SpikedWishart(1, n, p, [2, 4, s]))/n

@test length(λs) == p
@test eltype(λs) <: Real

# Check if the maximum spike is reasonable
# This tests can fail in principle, but only with negligible probability
μ = (1 + s) * (1 + γ/s)
σ = (1 + s) * sqrt(γ * 2 * (1 - γ/s^2))

@test abs(maximum(λs) - μ)/σ < 4


"""
    baryweights(x, y = nothing, csy = nothing) -> w

Compute the barycentric weights for the interpolation points `x`
and poles `[y; csy; conj.(csy)]`.
"""
function baryweights(x, y = nothing)
    n = length(x)

    wx = prodpot(x)
    wy = prodpot(x,y)
    s = LogNumber(true, mean(logabs.(wx)) - mean(logabs.(wy)))
    return float.(s.*wy./wx)
end

"""
    prodpot(x[,y::AbstractVector=x])

Evaluate the product potential `prod(x .- y)` at the scalar `x`,
or evaluate the potential pointwise if `x` is a vector. The result
is returned as a (vector of) `LogNumber`s to avoid over- or
underflow.

Pointwise evaluation is provided as a separate function rather than
through the dot syntax `prodpot.(x,Ref(y))` to allow this function
to be overloaded for point sets where `prodpot(x,y)` can be evaluated
more efficiently (e.g. Chebyshev points).
"""
prodpot(x) = default_prodpot(x)
prodpot(x,y) = default_prodpot(x,y)
prodpot(x,y::Nothing) = true

function default_prodpot(x)
    n = length(x)
    return (i -> prodpot(x[i], @view x[[1:i-1; i+1:n]])).(1:n)
end
default_prodpot(x,y) = default_prodpot.(x,Ref(y))
default_prodpot(x::Number,y) = mapreduce(yi->LogNumber(x - yi), *, y, init=LogNumber(one(x-one(eltype(y)))))

"""
    baryfactor(x̂,x,w) -> î,fac

If `all(x̂.!=x)`, returns `fac = sum(w./(x̂ .- x))` and an undefined `î`.
Otherwise, returns an `î` such that `x̂ == x[î]` and an undefined `fac`.
"""
function baryfactor(x̂,x,w)
    @assert length(x) == length(w)
    n = length(x)
    T = typeof(one(eltype(w))/(x̂-one(eltype(x))))
    fac = zero(T)
    for i = 1:n
        x[i] == x̂ && return i,fac
        fac += w[i]/(x̂-x[i])
    end
    return 0,fac
end

"""
    Barycentric(x, y = nothing, w = baryweights(x,y))

Barycentric basis for rational interpolation with interpolation points `x` and poles `y`.

If `isnothing(y)`, rational interpolation reduced to polynomial interpolation.
"""
struct Barycentric{X,Y,W} <: AbstractBasis
    points::X
    poles::Y
    weights::W
end
Barycentric(x, y = nothing) = Barycentric(x,y, baryweights(x,y))

Base.length(B::Barycentric) = length(B.points)

approxpoints(B::Barycentric) = B.points
approxtransform(B::Barycentric,f) = f

function iterate_basis(B::Barycentric, x̂, (i,î,fac) = (1,baryfactor(x̂,B.points,B.weights)...))
    i > length(B) && return nothing
    x = B.points
    w = B.weights
    î == 0 && return w[i]/((x̂ - x[i]) * fac), (i+1,î,fac)
    T = typeof(iterate_basis(B,x̂, (i,0,fac))[1])
    return (i == î ? one(T) : zero(T)), (i+1,î,fac)
end

evaluate_basis(B::Barycentric,i,x̂) = iterate_basis(B,x̂,(i,baryfactor(x̂,B.points,B.weights)...))[1]

# function evaltransform(B::Barycentric, x̂::Number, c::AbstractVector)
#     # TODO: implement this operation with single pass over data
# end


"""
    density(sp::AbstractSparseMatrix)

Compute the density of `sp`, the fraction of elements that
are non-zero. This assumes that there are no stored values equal to zero.
"""
function density(sp::AbstractSparseMatrix)
    return SparseArrays.nnz(sp) / length(sp)
end

"""
    summarystats(sp::SparseMatrixCSC)

Print some statistics for `sp`. The line
"number of non-zeros" counts structural non-zeros.
"""
function summarystats(sp::AbstractSparseMatrix)
    padding = 18
    for (description, statistic) in (
        ("size", size(sp)),
        ("num. elements", length(sp)),
        ("num. non-zeros", nnz(sp)),
        ("density", Printf.@sprintf("%.4e", density(sp))),
        ("num. stored zeros", count(iszero, SparseArrays.nonzeros(sp))))
        println(rpad(description, padding), ": ", statistic)
    end
    return nothing
end


"""
    map_to_canonical(x, xmin, xmax, basis::NodalBasis{Line})

Map the given coordinate `x` in the interval [`xmin`, `xmax`] to the corresponding
canonical coordinate `ξ` in the interval [-1, 1].
"""
function map_to_canonical(x, xmin, xmax, basis::NodalBasis{Line})
    ( 2x - (xmax + xmin) ) / (xmax - xmin)
end

"""
    map_to_canonical!(ξ, x, xmin, xmax, basis::NodalBasis{Line})

Map the given coordinate `x` in the interval [`xmin`, `xmax`] to the corresponding
canonical coordinate `ξ` in the interval [-1, 1], updating `ξ`.
"""
function map_to_canonical!(ξ, x, xmin, xmax, basis::NodalBasis{Line})
    @boundscheck begin
        @assert size(ξ) == size(x)
    end
    @inbounds for i in eachindex(ξ)
        ξ[i] = map_to_canonical(x[i], xmin, xmax, basis)
    end
    ξ
end

"""
    map_from_canonical(ξ, xmin, xmax, basis::NodalBasis{Line})

Map the given canonical coordinate `ξ` in the interval [-1, 1] to the corresponding
coordinate `x` in the interval [`xmin`, `xmax`].
"""
function map_from_canonical(ξ, xmin, xmax, basis::NodalBasis{Line})
    ( (xmax + xmin) + ξ * (xmax - xmin) ) / 2
end

"""
    map_from_canonical!(x, ξ, xmin, xmax, basis::NodalBasis{Line})

Map the given canonical coordinate `ξ` in the interval [-1, 1] to the corresponding
coordinate `x` in the interval [`xmin`, `xmax`], updating `x`.
"""
function map_from_canonical!(x, ξ, xmin, xmax, basis::NodalBasis{Line})
    @boundscheck begin
        @assert size(ξ) == size(x)
    end
    @inbounds for i in eachindex(x)
        x[i] = map_from_canonical(ξ[i], xmin, xmax, basis)
    end
    x
end


"""
    argmaxall()

Return the index of the maximum element in a collection. If there are multiple, all tied indexes are returned.
"""
function argmaxall(vals)
    maxvalue = maximum(vals)
    indexedvals = collect(zip(vals, eachindex(vals)))
    maxindexedvals = filter(v -> v[1] == maxvalue, indexedvals)
    last.(maxindexedvals)
end

"""
    invertdictionary()

Create a new dictionary using the values of an existing dictionary as keys and the values of the existing dictionary as
values.
"""
function invertdictionary(dictionary::IdDict{S, T}) where {S, T}
    newkeys = unique(values(dictionary))
    inverted = Dict{T, Vector{S}}()
    foreach(k -> inverted[k] = Vector{S}(), newkeys)
    foreach(kv -> push!(inverted[kv.second], kv.first), dictionary)
    inverted
end

"""
    inrange()

Evaluate if a value lies within the range provided, inclusive.
"""
within(value, lowerlimit, upperlimit) =
    value ≥ lowerlimit && value ≤ upperlimit

"""
    randomcoordinates2d()

Generate random coordinates within a radius.
"""
function randomcoordinates2d(count, radius)
    angles = rand(count) * τ
    radii = sqrt.(rand(count) .* (radius ^ 2)) # required for a uniform distribution around the circle
    x = radii .* cos.(angles)
    y = radii .* sin.(angles)
    z = zeros(count)
    return [x'; y'; z']
end

"""
    reindex()

Reindexes the array of molecules.
"""
reindex!(molecules) = foreach((m, i) -> m.index = i, molecules, eachindex(molecules))


# # Multivariate polynomials implementations

#md # [![](https://mybinder.org/badge_logo.svg)](@__BINDER_ROOT_URL__/generated/Extension/typed.ipynb)
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/Extension/typed.ipynb)
# **Contributed by**: Benoît Legat

# The SumOfSquares package is built on top of the [MultivariatePolynomials](https://github.com/JuliaAlgebra/MultivariatePolynomials.jl)
# abstract interface. [DynamicPolynomials](https://github.com/JuliaAlgebra/DynamicPolynomials.jl/)
# is an implementation of this abstract interface so it can be used with
# SumOfSquares. Moreover, any other implementation can be used as well. To
# illustrate, we solve Examples 3.38 of [BPT12] with
# [TypedPolynomials](https://github.com/JuliaAlgebra/TypedPolynomials.jl),
# another implementation of [MultivariatePolynomials](https://github.com/JuliaAlgebra/MultivariatePolynomials.jl).
#
# [BPT12] Blekherman, G. & Parrilo, P. A. & Thomas, R. R.
# *Semidefinite Optimization and Convex Algebraic Geometry*.
# Society for Industrial and Applied Mathematics, **2012**.

using Test #src
import TypedPolynomials
TypedPolynomials.@polyvar x y
using SumOfSquares
import CSDP
model = SOSModel(CSDP.Optimizer)
con_ref = @constraint(model, 2x^4 + 5y^4 - x^2*y^2 >= -2(x^3*y + x + 1))
optimize!(model)
solution_summary(model)

# We see that the problem is feasible. The Sum-of-Squares decomposition can be
# obtained as follows:

sos_decomposition(con_ref)

# Why is there several implementations ?
# Depending in the use-case, one implementation may be more appropriate than
# another one. [TypedPolynomials](https://github.com/JuliaAlgebra/TypedPolynomials.jl)
# is faster than [DynamicPolynomials](https://github.com/JuliaAlgebra/DynamicPolynomials.jl/)
# but it requires new compilation whenever the list of variables changes.
# This means that [TypedPolynomials](https://github.com/JuliaAlgebra/TypedPolynomials.jl)
# is not appropriate when the number of variables is dynamic or too large.
# However, for a small number of variables, it can be faster.
# When solving Sum-of-Squares programs, the time is mostly taken by the Semidefinite programming solver.
# The time taken by SumOfSquares/JuMP/MathOptInterface are usually negligible
# or it time is taken by manipulation of JuMP or MathOptInterface functions
# therefore using TypedPolynomials over DynamicPolynomials may not make much difference in most cases.

# One case for which using TypedPolynomials might be adequate is when
# using domain defined by equalities (possibly also with inequalities).
# Indeed, in that case, SumOfSquares computes the corresponding Gröbner basis which
# may take a non-negligible amount of time for large systems of equalities.

# To illustrate this, consider the computation of Gröbner basis for the
# following system from [CLO05, p. 17].
# The time taken by TypedPolynomials is below:
#
# [CLO05] Cox, A. David & Little, John & O'Shea, Donal
# *Using Algebraic Geometry*.
# Graduate Texts in Mathematics, **2005**.
# https://doi.org/10.1007/b138611

using BenchmarkTools
@btime let
    TypedPolynomials.@polyvar x y
    S = @set x^3 * y + x == 2x^2 * y^2 && 3x^4 == y
    SemialgebraicSets.computegröbnerbasis!(S.I)
end

# The time taken by DynamicPolynomials is as follows:

import DynamicPolynomials
@btime let
    DynamicPolynomials.@polyvar x y
    S = @set x^3 * y + x == 2x^2 * y^2 && 3x^4 == y
    SemialgebraicSets.computegröbnerbasis!(S.I)
end

# We see that TypedPolynomials is faster.
# The time is still negligible for this small system but for larger systems, choosing TypedPolynomials may be helpful.
# We can use this system in a Sum-of-Squares constraint as follows:

TypedPolynomials.@polyvar x y
S = @set x^3 * y + x == 2x^2 * y^2 && 3x^4 == y
poly = -6x - 4y^3 + 2x*y^2 + 6x^3 - 3y^4 + 13x^2 * y^2
model = Model(CSDP.Optimizer)
con_ref = @constraint(model, poly in SOSCone(), domain = S)
optimize!(model)
solution_summary(model)

# We obtain the following decomposition:

dec = sos_decomposition(con_ref, 1e-6)

# We can verify that it is correct as follows:

# TODO remove `polynomial` #src
@test SumOfSquares.MP.isapproxzero(rem(dec - poly, S.I), ztol = 1e-4) #src
rem(dec - poly, S.I)

# Note that the difference between `dec` and `poly` is larger
# than between the full gram matrix because `dec` is obtained by dropping
# the lowest eigenvalues with the threshold `1e-6`; see [`sos_decomposition`](@ref).

@test SumOfSquares.MP.isapproxzero(rem(gram_matrix(con_ref) - poly, S.I), ztol = 1e-6) #src
rem(gram_matrix(con_ref) - poly, S.I)


import Random: seed!
import LinearAlgebra: norm, dot

"""
    genprob(num_pts, u; sd=abs(rand(Int64)), K=5)

Generate a list of `num_pts + 1` functions to build function `min_{i =
1:num_pts+1} f_i`. The functions are defined inside the box `[0, u]` in
R^n. The function is created so that there are several (random)
quadratic functions with increasing minimum values. The `num_pts +
1`-th function is a plane that achieves the global solution at (0,
0). The values of the local minima are powers of `K`. To generate
always the same functions, one should provide the seed parameter `sd`.

Returns the list of functions and the bounds `l` and `u`.

"""
function genprob(num_pts, u; sd::Int64=abs(rand(Int64)), K=5)

    fmin_list = Vector{Function}(undef, num_pts + 1)

    n = length(u)

    l = zeros(Float64, n)

    sd = 38904589 + sd

    seed!(sd)

    K = 5.0

    for i = 1:num_pts

        fmin_list[i] = let

            local _n = n
            local k = K^i
            local d = 1000 * rand(n)
            local x0 = rand(n) .* u

            (x) -> k + 0.5 * sum((d[j] * (x[j] - x0[j])^2 for j = 1:_n))

        end

    end

    fmin_list[num_pts + 1] = let

        local c = rand(n)

        @. c = c / norm(c, 2)

        @views c[1] = (K^(num_pts + 2) - dot(c[2:end], u[2:end])) / u[1]

        (x) -> dot(c, x)

    end

    return fmin_list, l, u

end

"""
    gen_fmin(flist)

This function simply returns the `f_min` function, defined by `min
flist[i]`, which is usefull only for drawing purposes and intuition.

"""
function gen_fmin(flist)

    fmin = let

        fl = copy(flist)

        m = length(fl)

        (x) -> minimum((fl[i](x) for i = 1:m))

    end

    return fmin

end


using EngEconomics

# Given
# Option 1: 1,500,000 install; 200,000 every 10 years; i=0.05; lifespan = infinite
installation = 1500000
A = 200000
i = 0.05
ic = (1 + i)^(1 * 10)-1
N = 10
P = installation + A / ic


### A Pluto.jl notebook ###
# v0.14.7

using Markdown
using InteractiveUtils

# ╔═╡ d0c378ce-c5a1-11eb-0732-7bd1ff6303d2
md"""
# Naive Bayes

[Reference 1](https://en.wikipedia.org/wiki/Naive_Bayes_classifier)\
[Reference 2](https://scikit-learn.org/stable/modules/naive_bayes.html)
"""

# ╔═╡ c084dbc5-6f8f-47e2-9382-4df9234be5d2
md"""
Bayes' theorem:

$$P(y|x_1,...,x_n)=\frac{P(y)P(x_1,...,x_n|y)}{P(x_1,...,x_n)}$$

Classification:

$$\hat{y}=\text{arg}\max_yP(y)\prod^n_{i=1}P(x_i|y)$$

Gaussian Likelihood:

$$P(x_i|y)=\frac{1}{\sqrt{2\pi\sigma^2_{x_i,y}}}\exp\Bigg(-\frac{(x_i-\mu_{x_i,y})^2}{2\sigma^2_{x_i,y}}\Bigg)$$

where $\sigma_y$ and $\mu_y$ are estimated with maximum likelihood
"""

# ╔═╡ df15724c-422f-44d0-a921-1f0ef1c564fe
module tools include("../tools.jl") end

# ╔═╡ 92dca0e5-e95a-4060-bbd4-e464629473a9
JuTools = tools.JuTools

# ╔═╡ 7274c675-66c5-43a7-9c39-f7083de2619f
begin
X_data, Y_data = JuTools.data_generate_cluster_2d(pos1=(30.0, 80.0), pos2=(80.0, 30.0),
    radius1=5.0, radius2=10.0, random_scale=8.0, data_size=1000)
size(X_data), size(Y_data)
end

# ╔═╡ db6250c4-7e83-4a87-acf7-a2fc7051711f
mutable struct NaiveBayes
    n_features::Integer
    pY::Dict{Number,AbstractFloat}
    mean::Dict{Number,Array{AbstractFloat}}
    var::Dict{Number,Array{AbstractFloat}}
end

# ╔═╡ 2896ce68-e48e-4dc2-9317-18d6dfcd1365
function train(X_data::Array{T} where T<:Number, Y_data::Array{K} where K<:Number)::NaiveBayes
    @assert ndims(X_data) == ndims(Y_data) + 1
    @assert size(X_data)[1] == size(Y_data)[1]
    n_features = size(X_data)[2]
    # compute probability of Y
    pY = Dict{Number,AbstractFloat}()
    for Y_val in Y_data
        if haskey(pY, Y_val)
            pY[Y_val] += 1.0
        else
            pY[Y_val] = 1.0
        end
    end
    for Y_val in keys(pY)
        pY[Y_val] /= size(Y_data)[1]
    end
    # compute Xi mean and var for each Y
    mean = Dict{Number,Array{AbstractFloat}}()
    var = Dict{Number,Array{AbstractFloat}}()
    for Y_val in keys(pY)
        X_part = X_data[Y_data .== Y_val, :]
        mean[Y_val] = vec(sum(X_part, dims=1) ./ size(X_part)[1])
        var[Y_val] = vec(sum((X_part .- reshape(mean[Y_val], 1, :)).^2, dims=1) ./ size(X_part)[1])
    end
    return NaiveBayes(n_features, pY, mean, var)
end

# ╔═╡ 8935105e-72f1-46b6-8fe8-56f9afa6fa5f
function guassian(X_vec::Array, mean::Array, var::Array)::AbstractFloat
    @assert ndims(X_vec) == ndims(mean) == ndims(var) == 1
    @assert length(X_vec) == length(mean) == length(var)
    left = 1.0 ./ sqrt.((2.0 * pi) .* var)
    right = exp.(-(X_vec .- mean).^2 ./ (2.0 .* var))
    p = left .* right
    return prod(p)
end

# ╔═╡ 6eaa3999-30cd-4ab7-94cc-fab320f075d1
function predict(X_data::Array{T} where T<:Number, data::NaiveBayes)::Array
    if ndims(X_data) == 1
        X_data = reshape(X_data, 1, :)
    end
    @assert ndims(X_data) == 2
    @assert size(X_data)[2] == data.n_features
    prediction = Array{Number}(undef, size(X_data)[1])
    for i in 1:size(X_data)[1]
        X_vec = X_data[i, :]
        Y_dict = Dict{Number,AbstractFloat}()
        for Y_val in keys(data.pY)
            Y_dict[Y_val] = data.pY[Y_val] * guassian(X_vec, data.mean[Y_val], data.var[Y_val])
        end
        prediction[i] = sort(collect(Y_dict), by=m->m[2], rev=true)[1][1]
    end
    return prediction
end

# ╔═╡ b820e684-4cde-4353-b539-744bfa2a375b
trained = train(X_data, Y_data);

# ╔═╡ 4cdc095c-0d9e-44d0-8708-b3a3c880d153
prediction = predict(X_data, trained)

# ╔═╡ 8a60639b-6db4-4353-b69e-8e72c6961982
JuTools.compute_accuracy(prediction, Y_data)

# ╔═╡ Cell order:
# ╟─d0c378ce-c5a1-11eb-0732-7bd1ff6303d2
# ╟─c084dbc5-6f8f-47e2-9382-4df9234be5d2
# ╠═df15724c-422f-44d0-a921-1f0ef1c564fe
# ╠═92dca0e5-e95a-4060-bbd4-e464629473a9
# ╠═7274c675-66c5-43a7-9c39-f7083de2619f
# ╠═db6250c4-7e83-4a87-acf7-a2fc7051711f
# ╠═2896ce68-e48e-4dc2-9317-18d6dfcd1365
# ╠═8935105e-72f1-46b6-8fe8-56f9afa6fa5f
# ╠═6eaa3999-30cd-4ab7-94cc-fab320f075d1
# ╠═b820e684-4cde-4353-b539-744bfa2a375b
# ╠═4cdc095c-0d9e-44d0-8708-b3a3c880d153
# ╠═8a60639b-6db4-4353-b69e-8e72c6961982


# FeynmanAthermalAsymptoticComparison.jl - reproduce some 1950s numeric results

push!(LOAD_PATH,"../src/") # load module from local directory
using PolaronMobility

using Plots

# Here we collect all the data we need for plotting
αrange=0.2:0.1:10 # unstable below α=0.2
vs=[ feynmanvw(α)[1]  for α in αrange] # takes a while; computes all vw
ws=[ feynmanvw(α)[2]  for α in αrange]
# Urgh - can't seem to do this more elegantly; recalculating all the params just to strip the second component of the tuple

Es=map(F,vs,ws,αrange); # calculate polaron energy for these params

p=plot(xlim=αrange,ylim=(0,),
    xlabel="α parameter", ylabel="Reduced polaron units",
    title="Athermal (numeric, variational) v,w, and E for α" )

plot!(αrange,vs,label="v")
plot!(αrange,ws,label="w")
plot!(αrange,Es,label="E")
hline!([0.0],label="",color=:black)

# I like this plot!
# You can see how v=w=3 as α-->0.0
# w then tends down to 1.0 in a nice sigmoid form as α --> ~10+
# v increases monotonically

savefig("AthermalvwE.pdf")

# Feynman Stat Mech (1972), p. 240
smallαw(α)=3.0
P(α)=2/smallαw(α) * (sqrt(1+smallαw(α))-1) # Nb: error in Feynman Stat Mech: sqrt(1+W) , as in Feynman1955
smallαv(α)=3*(1 + 2*α*(1-P(α))/(3*smallαw(α)))

c=0.5772 # Euler Mascheroni constant

largeαw(α)=1.0
largeαv(α)=(4*α^2/9π) - (4*(log(2) + c/2)-1)

p=plot(xlim=αrange,ylim=(0,15), xlabel="α parameter", ylabel="v,w, Reduced polaron units", title="Asymtotic limits for v,w" )

plot!(αrange,α->smallαw(α),label="w, smallαw",style=:dash)
plot!(αrange,α->smallαv(α),label="v, smallαv",style=:solid)

plot!(αrange,α->largeαw(α),label="w, largeαw",style=:dash)
plot!(αrange,α->largeαv(α),label="v, largeαv",style=:solid)

plot!(αrange,vs,label="v, numeric",w=3,style=:solid)
plot!(αrange,ws,label="w, numeric",w=3,style=:dash)

savefig("AthermalvwAsymptotic.pdf")

smallαE(α)=-α-α^2/81
FeynmanStatMechlargeαE(α)=-α^2/(3π) - 3/2 *(2*log(2)+c)-3/4 # As in Feynman Stat Mech; with 2π corrected to 3π ...

Feynman1955largeαE(α)=-α^2/3π - 3*log(2) # as in Feynman I, 1955

p=plot(xlim=αrange,ylim=(-Inf,0), xlabel="α parameter", ylabel="Energy, Reduced polaron units", title="Asymtotic limits for E(α)" )

plot!(αrange,α->smallαE(α),label="smallαE")
plot!(αrange,α->FeynmanStatMechlargeαE(α),label="FeynmanStatMechlargeαE")
plot!(αrange,α->Feynman1955largeαE(α),label="Feynman1955largeαE")
plot!(αrange,Es,label="E (numeric integration)")

savefig("AthermalAsymptoticEnergy.pdf")


# Aodwt.jl
# 2019-02-23 Jeff Fessler, University of Michigan

using Plots
using LinearMaps
using Wavelets: dwt, idwt, wavelet, WT
#using MIRT: jim

"""
`A, levels, mfun = Aodwt(dims; level::Integer=3, wt=wavelet(WT.haar))`

create orthogonal discrete wavelet transform (ODWT) `LinearMap`

in
* `dims` tuple of dimensions

option
* `level` # of levels; default 3
* `wt` wavelet transform type (see `Wavelets` package); default Haar

out
* `A` a `LinearMap` object
* `scales` array of size `dims` showing the scale of each coefficient
which is useful when imposing scale-dependent regularization
* `mfun` convenience function for A*X when X is a Matrix or Array (not vector)

2019-02-23 Jeff Fessler, University of Michigan
"""
function Aodwt(dims; level::Integer=3, wt=wavelet(WT.haar))

	Afun = (level) -> LinearMap(
		x -> dwt(reshape(x, dims), wt, level)[:],
		y -> idwt(reshape(y, dims), wt, level)[:],
		prod(dims), prod(dims))

	A = Afun(level)

	mfun = (A, x) -> reshape(A * x[:], dims)

	scales = zeros(dims)
	for il=1:level
		Al = Afun(il)
		tmp = mfun(Al, ones(dims)) .== 0
		scales += il * (tmp .& (scales .== 0))
	end

	return A, scales, mfun
end

function Aodwt()
	@doc Aodwt
end

"""
`Aodwt_show(;M=32, N=64)`
show scales
"""
function Aodwt_show(; M=32, N=64)
	W, scales, mfun = Aodwt((M,N))
	jim(scales)
end

"""
`Aodwt(:test)`
self test

`Aodwt(:show)`
visualize
"""
function Aodwt(test::Symbol)
	if test == :show
		return Aodwt_show()
	end

	test != :test && throw(ArgumentError("test $test"))
	Aodwt()

	W,_,_ = Aodwt((8,16), level=2)
	@test Matrix(W)' == Matrix(W') # check adjoint

	Aodwt(:show)
	true
end


function rob(nums::Array{T,1}) where T
    table = Dict{T,UInt64}()
    function inner(nums, start, len, table)
        try
            return table[start]
        catch end

        if start == len
            table[start] = nums[start]
            return nums[start]
        elseif start == len + 1
            table[start] = 0
            return 0
        end

        a = max(
            maximum([inner(nums, x, len, table) for x in start+1:len]),
            nums[start] + inner(nums, start+2, len, table)
        )

        table[start] = a
        a
    end
    inner(nums, 1, length(nums), table)
end

using Test
@test rob([183, 219, 57, 193, 94, 233, 202, 154, 65, 240, 97, 234, 100, 249, 186, 66, 90, 238, 168, 128, 177, 235, 50, 81, 185, 165, 217, 207, 88, 80, 112, 78, 135, 62, 228, 247, 211]) == UInt64(3365)
@test rob([9, 297, 196, 336, 435, 2, 343, 159, 146, 359, 45, 470, 265, 131, 17, 271, 74, 242, 448, 402, 55, 423, 414, 240, 430, 135, 322, 468, 422, 351, 441, 463, 30, 399, 132, 439, 463, 260, 399, 32, 374, 383, 276, 166, 104, 315, 314, 445, 458, 422, 104, 251, 382, 230, 484, 127, 355, 332, 317, 362, 257, 493, 474, 401, 40, 93, 433, 464, 136, 342, 98, 159, 223, 110, 89, 47, 53, 179, 219, 314, 486, 301, 307, 453, 37, 366, 334, 355, 26, 484, 124, 408, 346, 133, 420, 280, 124, 210, 358, 140]) == UInt64(15632)

# use this to test time
# @time rob([9, 297, 196, 336, 435, 2, 343, 159, 146, 359, 45, 470, 265, 131, 17, 271, 74, 242, 448, 402, 55, 423, 414, 240, 430, 135, 322, 468, 422, 351, 441, 463, 30, 399, 132, 439, 463, 260, 399, 32, 374, 383, 276, 166, 104, 315, 314, 445, 458, 422, 104, 251, 382, 230, 484, 127, 355, 332, 317, 362, 257, 493, 474, 401, 40, 93, 433, 464, 136, 342, 98, 159, 223, 110, 89, 47, 53, 179, 219, 314, 486, 301, 307, 453, 37, 366, 334, 355, 26, 484, 124, 408, 346, 133, 420, 280, 124, 210, 358, 140, 9, 297, 196, 336, 435, 2, 343, 159, 146, 359, 45, 470, 265, 131, 17, 271, 74, 242, 448, 402, 55, 423, 414, 240, 430, 135, 322, 468, 422, 351, 441, 463, 30, 399, 132, 439, 463, 260, 399, 32, 374, 383, 276, 166, 104, 315, 314, 445, 458, 422, 104, 251, 382, 230, 484, 127, 355, 332, 317, 362, 257, 493, 474, 401, 40, 93, 433, 464, 136, 342, 98, 159, 223, 110, 89, 47, 53, 179, 219, 314, 486, 301, 307, 453, 37, 366, 334, 355, 26, 484, 124, 408, 346, 133, 420, 280, 124, 210, 358, 140, 9, 297, 196, 336, 435, 2, 343, 159, 146, 359, 45, 470, 265, 131, 17, 271, 74, 242, 448, 402, 55, 423, 414, 240, 430, 135, 322, 468, 422, 351, 441, 463, 30, 399, 132, 439, 463, 260, 399, 32, 374, 383, 276, 166, 104, 315, 314, 445, 458, 422, 104, 251, 382, 230, 484, 127, 355, 332, 317, 362, 257, 493, 474, 401, 40, 93, 433, 464, 136, 342, 98, 159, 223, 110, 89, 47, 53, 179, 219, 314, 486, 301, 307, 453, 37, 366, 334, 355, 26, 484, 124, 408, 346, 133, 420, 280, 124, 210, 358, 140])


using HermiteNormalForm
using LinearAlgebra
using Test

@testset "HermiteNormalForm.jl" begin
    for m in 1:20, n in 1:20
        A = rand(-big(3):3, m, n)
        F = hnf(A)
        if F !== nothing
            H, U = F
            @test ishnf(H)
            @test abs(det(U.//1)) == 1
            @test LinearAlgebra.issuccess(lu(Rational.(A)', check=false))
        else
            @test !LinearAlgebra.issuccess(lu(Rational.(A)', check=false))
        end
    end
end


using LinearAlgebra

import Base: getindex, size
import LinearAlgebra: reflectorApply!

struct QR2{T,S<:AbstractMatrix{T},V<:AbstractVector{T}} <: Factorization{T}
    factors::S
    τ::V
    QR2{T,S,V}(factors::AbstractMatrix{T}, τ::AbstractVector{T}) where {T,S<:AbstractMatrix,V<:AbstractVector} =
        new(factors, τ)
end
QR2(factors::AbstractMatrix{T}, τ::Vector{T}) where {T} = QR2{T,typeof(factors)}(factors, τ)

size(F::QR2, i::Integer...) = size(F.factors, i...)

# Similar to the definition in base but applies the reflector from the right
@inline function reflectorApply!(A::StridedMatrix, x::AbstractVector, τ::Number) # apply conjugate transpose reflector from right.
    m, n = size(A)
    if length(x) != n
        throw(DimensionMismatch("reflector must have same length as second dimension of matrix"))
    end
    @inbounds begin
        for i in 1:m
            Aiv = A[i, 1]
            for j in 2:n
                Aiv += A[i, j]*x[j]
            end
            Aiv = Aiv*τ
            A[i, 1] -= Aiv
            for j in 2:n
                A[i, j] -= Aiv*x[j]'
            end
        end
    end
    return A
end

# FixMe! Consider how to represent Q

# immutable Q{T,S<:QR2} <: AbstractMatrix{T}
#     data::S
# end

# size(A::Q) = size(A.data)
# size(A::Q, i::Integer) = size(A.data, i)

# getindex{T}(A::QR2{T}, ::Type{Tuple{:Q}}) = Q{T,typeof(A)}(A)

function getindex(A::QR2{T}, ::Type{Tuple{:R}}) where T
    m, n = size(A)
    if m >= n
        UpperTriangular(view(A.factors, 1:n, 1:n))
    else
        throw(ArgumentError("R matrix is trapezoid and cannot be extracted with indexing"))
    end
end

function getindex(A::QR2{T}, ::Type{Tuple{:QBlocked}}) where T
    m, n = size(A)
    mmn = min(m,n)
    F = A.factors
    τ = A.τ
    Tmat = zeros(T, mmn, mmn)
    for j = 1:mmn
        for i = 1:j - 1
            Tmat[i,j] = τ[i]*(F[j,i] + dot(view(F, j + 1:m, i), view(F, j + 1:m, j)))
        end
    end
    LinearAlgebra.inv!(LinearAlgebra.UnitUpperTriangular(Tmat))
    for j = 1:min(m,n)
        Tmat[j,j] = τ[j]
        for i = 1:j - 1
            Tmat[i,j] = Tmat[i,j]*τ[j]
        end
    end
    HouseholderBlock{T,typeof(F),Matrix{T}}(F, UpperTriangular(Tmat))
end

# qrUnblocked!(A::StridedMatrix) = LinearAlgebra.qrfactUnblocked!(A)
function qrUnblocked!(A::StridedMatrix{T},
                      τ::StridedVector{T} = fill(zero(T), min(size(A)...))) where {T}

    m, n = size(A)

    # Make a vector view of the first column
    a = view(A, :, 1)
    # Construct the elementary reflector of the first column in-place
    τ1 = LinearAlgebra.reflector!(a)
    # Store reflector loading in τ
    τ[1] = τ1

    # Apply the reflector to the trailing columns if any
    if n > 1
        reflectorApply!(a, τ1, view(A, :, 2:n))
    end

    # Call the function recursively on the submatrix excluding first row and column or return
    if m > 1 && n > 1
        qrUnblocked!(view(A, 2:m, 2:n), view(τ, 2:min(m, n)))
    end

    return QR2{T,typeof(A),typeof(τ)}(A, τ)
end

function qrBlocked!(A::StridedMatrix{T},
                    blocksize::Integer = 12,
                    τ::StridedVector{T} = fill(zero(T), min(size(A)...)),
                    work = Matrix{T}(undef, blocksize, size(A, 2))) where T

    m, n = size(A)

    # Make a vector view of the first column block
    A1 = view(A, :, 1:min(n, blocksize))
    # Construct QR factorization of first column block in-place
    F = qrUnblocked!(A1, view(τ, 1:min(m, n, blocksize)))

    # Apply the QR factorization to the trailing columns (if any) and recurse
    if n > blocksize
        # Make a view of the trailing columns
        A2 = view(A, :, blocksize + 1:n)
        # Apply Q' to the trailing columns
        lmul!(F[Tuple{:QBlocked}]', A2, view(work, 1:min(m, blocksize), 1:(n - blocksize)))
    end

    # Compute QR factorization of trailing block
    if m > blocksize && n > blocksize
        qrBlocked!(view(A, (blocksize + 1):m, (blocksize + 1):n),
                   blocksize,
                   view(τ, (blocksize + 1):min(m, n)),
                   work)
    end

    return QR2{T,typeof(A),typeof(τ)}(A, τ)
end


"""
    prediction( sys::StateSpace, N::Integer )

Form the prediction matrix ``Γ`` that computes the states of the linear system `sys` over the horizon of length `N`
by using all the inputs across the entire horizon.

# Extended help

This matrix consists of the blocks ``A^i B`` where ``i`` is the diagonal number, forming the matrix
```math
⎡ B   •   •  •  ⋯⎤
  ⎢ AB  B   •  •  ⋯⎥
  ⎢ A²B AB  B  •  ⋯⎥
  ⎢ A³B A²B AB B  ⋯⎥
  ⎣ ⋮   ⋮   ⋮  ⋮ ⋱ ⎦
```

It is used to compute the component of the state vector for a horizon that is caused by the system inputs
during the horizon using the equation ``x = Γu``.

The full state vector for a horizon can be computed as ``x = Γu + Φx₀``, where ``Φ`` is the initial condition
propagation matrix from [`initialpropagation`](@ref).
"""
function prediction( sys::StateSpace, N::Integer )
    if( N < 1 )
        throw( DomainError( N, "Horizon length must be greater than 1" ) )
    end

    nu = ninputs( sys )
    nx = nstates( sys )
    mt = eltype( sys.A )

    mainDiag  = kron( I(N), sys.B )

    otherDiag = kron( I(N), to_matrix( mt, I, nx ) )

    D = zeros( mt, N, N )
    D[diagind( D, -1 )] .= 1

    otherDiag = otherDiag + kron( D, -sys.A )

    Γ = otherDiag\mainDiag

    return BlockArray{ mt }( Γ, [nx for i = 1:N], [nu for i = 1:N] )
end

"""
    initialpropagation( sys::StateSpace, N::Integer )

Form the matrix ``Φ`` that propagates the initial condition ``x₀`` in the linear system `sys` across the horizon `N`.

# Extended help

This matrix consists of a column of the blocks ``A^i`` where ``i`` is the row number, forming the matrix
```math
⎡ A  ⎤
  ⎢ A² ⎥
  ⎢ A³ ⎥
  ⎢ A⁴ ⎥
  ⎣ ⋮  ⎦
```

It is used to compute the component of the state vector for a horizon that is caused by the intial state of the system using
the equation ``x = Φx₀``.

The full state vector for a horizon can be computed as ``x = Γu + Φx₀``, where ``Γ`` is the prediction matrix
from [`prediction`](@ref).
"""
function initialpropagation( sys::StateSpace, N::Integer )
    if( N < 1 )
        throw( DomainError( N, "Horizon length must be greater than 1" ) )
    end

    blockGen = (sys.A^i for i=1:N)

    # The mortar function requires the passed array to have ndims=2, so reshape to
    # have an array with 2 dimensions even though there is only one column
    blocks = reshape( collect( blockGen ),  (N, 1) )

    return mortar( blocks )
end

function inputprediction( sys::StateSpace, K::AbstractArray, N::Integer )
    if( N < 1 )
        throw( DomainError( N, "Horizon length must be greater than 1" ) )
    end

    nu = ninputs( sys )
    nx = nstates( sys )
    mt = eltype( sys.A )

    # Create the prediction matrix for the controlled system
    Γ = prediction( sys, N )
    K̄ = blockkron( 1.0*I(N), K )

    zerorow = zeros( Float64, (nu, nu*N) )

    Γ = -K̄*Γ
    Γ = vcat( zerorow, Γ )

    eye = blockkron( 1.0*I(N), Matrix( 1.0*I(nu) ) )
    eye = vcat( eye, zerorow )

    Γ = Γ + eye
    Γ = Γ[1:N*nu, 1:N*nu]

    return BlockArray{ mt }( Γ, [nu for i = 1:N], [nu for i = 1:N] )
end

function inputinitialpropagation( sys::StateSpace, K::AbstractArray, N::Integer )
    if( N < 1 )
        throw( DomainError( N, "Horizon length must be greater than 1" ) )
    end

    blockGen = (-K*sys.A^i for i=0:(N-1))

    # The mortar function requires the passed array to have ndims=2, so reshape to
    # have an array with 2 dimensions even though there is only one column
    blocks = reshape( collect( blockGen ),  (N, 1) )

    return mortar( blocks )
end

"""
    hessian( clqr::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )

Form the complete Hessian for the constrained time-invariant LQR problem `clqr` using the fully
condensed form.
"""
function hessian( clqr::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )
    sys = getsystem( clqr )
    N   = clqr.N
    nx  = nstates( sys )
    nu  = ninputs( sys )

    Γ = prediction( sys, N )

    Q̅ = blockkron( 1.0*I(N), clqr.Qₖ )
    S̅ = blockkron( 1.0*I(N), clqr.S  )
    R̅ = blockkron( 1.0*I(N), clqr.R  )

    Q̅[Block( N, N )] = clqr.P
    S̅[Block( N, N )] = zeros( eltype( clqr.S ), size( clqr.S ) )

    # Pass the result through the Hermitian type to cleanup numerical errors that can occur
    # making it non-symmetric (they are on the order of 10e-17)
    return Hermitian( Γ'*Q̅*Γ + Γ'*S̅ + S̅'*Γ + R̅ )
end

"""
    linearcoefficients( clqr::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )

Form the vector of coefficients for the linear term of the fully condensed quadratic QP for the
CLQR problem defined by `clqr`.
"""
function linearcoefficients( clqr::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )
    sys = getsystem( clqr )
    N   = clqr.N
    nx  = nstates( sys )
    nu  = ninputs( sys )

    Γ = prediction( sys, N )
    Φ = initialpropagation( sys, N )

    Q̅ = blockkron( 1.0*I(N), clqr.Qₖ )
    S̅ = blockkron( 1.0*I(N), clqr.S  )

    Q̅[Block( N, N )] = clqr.P
    S̅[Block( N, N )] = zeros( eltype( clqr.S ), size( clqr.S ) )

    return Γ'*Q̅*Φ + S̅'*Φ
end

function inequalityconstraints( clqr::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )
    # For the initial computations we want the non-pre-stabilized system
    sys = getsystem( clqr, prestabilized = false )

    # Find the number of constraints and the system size
    N   = clqr.N
    nC  = size( clqr.E, 1 );    # Number of constraints
    n, m  = size( sys.B );

    # Get the condensed initial propagation matrix
    Φ = initialpropagation( sys, N )

    # Get the prediction matrix for the system
    Γ = prediction( sys, N )

    # Create the component matrices
    Ē = blockkron( 1.0*I(N), clqr.E )
    F̄ = blockkron( 1.0*I(N), clqr.F )

    # Compute the coefficient matrix of the linear inequality constraints
    G = Ē*Γ + F̄;

    # Compute the initial state matrix for the RHS of the constraints
    L = -Ē*Φ;

    if isprestabilized( clqr )
        # Get the pre-stabilized matrices
        sysₖ = getsystem( clqr )
        Γₖ   = inputprediction( sysₖ, clqr.K, N )
        Φₖ   = inputinitialpropagation( sysₖ, clqr.K, N )

        # Modify the initial state coefficients to be in the new input space
        L = L - G*Φₖ

        # Modify the coefficient matrix to be in the new input space
        G = G * Γₖ
    end

    # Compute the constant vector for the RHS
    g = blockkron( ones( Float64, (N) ), clqr.g )

    return G, L, g
end

function equalityconstraints( ::ConstrainedTimeInvariantLQR, ::Type{PrimalFullCondensing} )
    # The fully condensed form does not have any equality constraints from the problem
    return nothing
end


"""
    AnalyticalNonlinearShrinkage

Analytical nonlinear shrinkage estimator. See docs for
`analytical_nonlinear_shrinkage` for details.
"""
struct AnalyticalNonlinearShrinkage{TEigen<:Union{Eigen,Nothing}} <: CovarianceEstimator
    decomp::TEigen
    corrected::Bool
    function AnalyticalNonlinearShrinkage(;corrected=false)
        new{Nothing}(nothing, corrected)
    end
    function AnalyticalNonlinearShrinkage(decomp::Eigen; corrected=false)
        new{Eigen}(decomp, corrected)
    end
end

const SQRT5  = sqrt(5.0)
const INVPI  = 1.0 / π
const EPAN_1 = 3.0/(4.0 * SQRT5)
const EPAN_2 = EPAN_1 * INVPI
const EPAN_3 = 0.3 * INVPI

"""
    epanechnikov(x)

Return the Epanechnikov kernel evaluated at `x`.
"""
epanechnikov(x::Real) = EPAN_1 * max(0.0, 1.0 - x^2/5.0)

"""
    epnanechnikov_HT(x)

Return the Hilbert Transform of the Epanechnikov kernel evaluated at `x`
if `|x|≂̸√5`.
"""
function epanechnikov_HT1(x::Real)
    -EPAN_3 * x + EPAN_2 * (1.0 - x^2/5.0) * log(abs((SQRT5 - x)/(SQRT5 + x)))
end

"""
    epnanechnikov_HT2(x)
Return the Hilbert Transform of the Epanechnikov kernel evaluated at `x`
if `|x|=√5`.
"""
epanechnikov_HT2(x::Real) = -EPAN_3*x

"""
    analytical_nonlinear_shrinkage(S, n, p; decomp)

Internal implementation of the analytical nonlinear shrinkage. The
implementation is inspired from the Matlab code given in section C of
Olivier Ledoit and Michael Wolf's paper "Analytical Nonlinear
Shrinkage of Large-Dimensional Covariance Matrices". (Nov 2018)
http://www.econ.uzh.ch/static/wp/econwp264.pdf
"""
function analytical_nonlinear_shrinkage(S::AbstractMatrix{<:Real},
                                        n::Int, p::Int, est_mean::Bool;
                                        decomp::Union{Nothing,Eigen}=nothing)

    η = ifelse(p < n, n, n - Int(est_mean)) # effective sample size
    # sample eigenvalues sorted in ascending order and eigenvectors
    F    = isa(decomp, Nothing) ? eigen(S) : decomp # O(p^3)
    perm = sortperm(F.values)
    sample_perm = @view perm[max(1, (p - η) + 1):p]
    λ    = @view F.values[sample_perm]
    U    = F.vectors[:, perm]

    # dominant cost forming of S or eigen(S) --> O(max{np^2, p^3})

    # compute analytical nonlinear shrinkage kernel formula
    L = repeat(λ, outer=(1, min(p, η)))

    # Equation (4.9)
    h = η^(-1//3)
    H = h * L'
    x = (L .- L') ./ H

    # additional useful definitions
    γ  = p/η
    πλ = π * λ

    # Equation (4.7) in http://www.econ.uzh.ch/static/wp/econwp264.pdf
    f̃ = vec(mean(epanechnikov.(x) ./ H, dims=2))

    # Equation (4.8)
    Hf̃_tmp = epanechnikov_HT1.(x)
    # if by any chance there are x such that |x| ≈ √5 ...
    mask = (@. abs(x) ≈ SQRT5)
    any(mask) && (Hf̃_tmp[mask] = epanechnikov_HT2.(@view x[mask]))
    Hf̃ = vec(mean(Hf̃_tmp ./ H, dims=2))

    # dominant cost up to here: elementwise ops on x --> O(max{p^2, η^2})

    if p <= η
        # Equation (4.3)
        πγλ = γ * πλ
        denom = @. (πγλ * f̃)^2 + (1.0 - γ - πγλ * Hf̃)^2
        d̃ = λ ./ denom
    else
        # Equation (C.8)
        hs5  = h * SQRT5
        Hf̃0  = (0.3/h^2 + 0.75/hs5 * (1.0 - 0.2/h^2) * log((1+hs5)/(1-hs5)))
        Hf̃0 *= mean(1.0 ./ πλ)
        # Equation (C.5)
        d̃0 = INVPI / ((γ - 1.0) * Hf̃0)
        # Eq. (C.4)
        d̃1 = @. 1.0 / (π * πλ * (f̃^2 + Hf̃^2))
        d̃  = [fill(d̃0, (p - η, 1)); d̃1]
    end

    # Equation (4.4)
    return Symmetric(U*(d̃ .* U'))
end

"""
    cov(ans::AnalyticalNonlinearShrinkage, X; dims=1, mean=nothing)

Nonlinear covariance estimator derived from the sample covariance estimator `S`
and its eigenvalue decomposition (which can be given through `decomp`).
See Ledoit and Wolf's paper http://www.econ.uzh.ch/static/wp/econwp264.pdf
The keyword `mean` can be `nothing` (centering via estimated mean),
zero (no centering) or a provided vector. In the first case, a rank-1
modification is applied and therefore the effective sample size is decreased
by one (see `analytical_nonlinear_shrinkage`). In the latter two case the mean
cannot have been estimated on the data (otherwise the effective sample size
will be 1 larger than it should be resulting in numerical instabilities).
If you are unsure, use either `nothing` or provide an explicit (non-estimated)
vector (possibly a zero vector) and avoid the use of `mean=0`.

* Time complexity (including formation of `S`)
    - (p<n): O(np^2 + n^2) with moderate constant
    - (p>n): O(p^3) with low constant (dominated by eigendecomposition of S)
"""
function cov(ans::AnalyticalNonlinearShrinkage, X::AbstractMatrix{<:Real};
             dims::Int=1, mean=nothing)

    @assert dims ∈ [1, 2] "Argument dims can only be 1 or 2 (given: $dims)"

    szx    = size(X)
    (n, p) = ifelse(dims==1, szx, reverse(szx))

    # explained in the paper there must be at least 12 samples
    (n < 12) && throw(ArgumentError("The number of samples `n` must be at " *
                                    "least 12 (given: $n)."))

    S  = cov(SimpleCovariance(corrected=ans.corrected), X; dims=dims, mean=mean)
    return analytical_nonlinear_shrinkage(S, n, p, mean === nothing;
                decomp=ans.decomp)
end


using .RationalFunctionFields: load_generators

logger = Logging.SimpleLogger(stderr, Logging.Debug)
Logging.global_logger(logger)

@testset "Loading from file tests" begin

    set = load_generators("../data/Simple.txt")

    R = parent(first(set))
    x, y, z = AA.gens(AA.base_ring(R))

    @test set[1] == (x^2 - z^2) // (x + y)
    @test set[2] == (x*y) // (x + y)
    @test set[3] == (1) // z
    @test set[4] == z^3 // z

end


@timed_testset "turing_model" begin
    DATA_DIR = joinpath("..", "data")
    kidiq = CSV.read(joinpath(DATA_DIR, "kidiq.csv"), DataFrame)
    wells = CSV.read(joinpath(DATA_DIR, "wells.csv"), DataFrame)
    roaches = CSV.read(joinpath(DATA_DIR, "roaches.csv"), DataFrame)
    cheese = CSV.read(joinpath(DATA_DIR, "cheese.csv"), DataFrame)
    @timed_testset "Gaussian Model" begin
        f = @formula(kid_score ~ mom_iq * mom_hs)
        @testset "standardize=false" begin
            m = turing_model(f, kidiq)
            chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
            @test summarystats(chn)[:α, :mean] ≈ 29.30 atol = 2.0
            @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 0.533 atol = 0.2
            @test quantile(chn)[Symbol("β[2]"), Symbol("50.0%")] ≈ 0.593 atol = 0.2
        end

        @testset "standardize=true" begin
            m = turing_model(f, kidiq; standardize=true)
            chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
            @test summarystats(chn)[:α, :mean] ≈ 0.000 atol = 0.2
            @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 0.648 atol = 0.2
            @test quantile(chn)[Symbol("β[2]"), Symbol("50.0%")] ≈ 0.849 atol = 0.2
        end

        @testset "custom_priors" begin
            priors = CustomPrior(Normal(), Normal(28, 5), nothing)
            m = turing_model(f, kidiq; priors)
            chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
            @test summarystats(chn)[:α, :mean] ≈ 28.758 atol = 2.0
            @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 0.539 atol = 0.2
            @test quantile(chn)[Symbol("β[2]"), Symbol("50.0%")] ≈ 0.3863 atol = 0.2
        end
        @testset "explicit calling Normal" begin
            m = turing_model(f, kidiq; model=Normal)
            chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
            @test summarystats(chn)[:α, :mean] ≈ 29.30 atol = 2.0
            @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 0.533 atol = 0.2
            @test quantile(chn)[Symbol("β[2]"), Symbol("50.0%")] ≈ 0.593 atol = 0.2
        end
    end
    @timed_testset "TDist Model" begin
        f = @formula(kid_score ~ mom_iq * mom_hs)
        @testset "standardize=false" begin
            m = turing_model(f, kidiq; model=TDist)
            chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
            @test summarystats(chn)[:α, :mean] ≈ 40.380 atol = 2.0
            @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 0.478 atol = 0.2
            @test quantile(chn)[Symbol("β[2]"), Symbol("50.0%")] ≈ 0.736 atol = 0.2
            @test quantile(chn)[:ν, Symbol("50.0%")] ≈ 1.039 atol = 0.5
        end

        @testset "custom_priors" begin
            priors = CustomPrior(Normal(), Normal(28, 5), Exponential(2))
            m = turing_model(f, kidiq; model=TDist, priors)
            chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
            @test summarystats(chn)[:α, :mean] ≈ 35.506 atol = 2.0
            @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 0.522 atol = 0.2
            @test quantile(chn)[Symbol("β[2]"), Symbol("50.0%")] ≈ 0.628 atol = 0.2
            @test quantile(chn)[:ν, Symbol("50.0%")] ≈ 1.178 atol = 0.5
        end
    end
    @timed_testset "Bernoulli Model" begin
        f = @formula(switch ~ arsenic + dist + assoc + educ)
        @testset "standardize=false" begin
            m = turing_model(f, wells; model=Bernoulli)
            chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
            @test summarystats(chn)[:α, :mean] ≈ -0.153 atol = 0.2
            @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 0.467 atol = 0.2
            @test quantile(chn)[Symbol("β[2]"), Symbol("50.0%")] ≈ -0.009 atol = 0.2
        end

        @testset "custom_priors" begin
            priors = CustomPrior(Normal(), Normal(), nothing)
            m = turing_model(f, wells; model=Bernoulli, priors)
            chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
            @test summarystats(chn)[:α, :mean] ≈ -0.155 atol = 0.2
            @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 0.468 atol = 0.2
            @test quantile(chn)[Symbol("β[2]"), Symbol("50.0%")] ≈ -0.009 atol = 0.2
        end
    end
    @timed_testset "Poisson Model" begin
        f = @formula(y ~ roach1 + treatment + senior + exposure2)
        @testset "standardize=false" begin
            m = turing_model(f, roaches; model=Poisson)
            chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
            @test summarystats(chn)[:α, :mean] ≈ 2.969 atol = 0.5
            @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 0.006 atol = 0.2
            @test quantile(chn)[Symbol("β[2]"), Symbol("50.0%")] ≈ -0.5145 atol = 0.2
        end

        @testset "custom_priors" begin
            priors = CustomPrior(Normal(2, 5), Normal(), nothing)
            m = turing_model(f, roaches; model=Poisson, priors)
            chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
            @test summarystats(chn)[:α, :mean] ≈ 2.963 atol = 0.5
            @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 0.006 atol = 0.2
            @test quantile(chn)[Symbol("β[2]"), Symbol("50.0%")] ≈ -0.5145 atol = 0.2
        end
    end
    @timed_testset "NegativeBinomial Model" begin
        f = @formula(y ~ roach1 + treatment + senior + exposure2)
        @testset "standardize=false" begin
            m = turing_model(f, roaches; model=NegativeBinomial)
            chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
            @test summarystats(chn)[:α, :mean] ≈ 2.448 atol = 0.5
            @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 0.013 atol = 0.2
            @test quantile(chn)[Symbol("β[2]"), Symbol("50.0%")] ≈ -0.734 atol = 0.2
            @test quantile(chn)[:ϕ⁻, Symbol("50.0%")] ≈ 1.401 atol = 0.2
        end

        @testset "custom_priors" begin
            priors = CustomPrior(Normal(), Normal(2, 5), Exponential(0.5))
            m = turing_model(f, roaches; model=NegativeBinomial, priors)
            chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
            @test summarystats(chn)[:α, :mean] ≈ 2.422 atol = 0.5
            @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 0.013 atol = 0.2
            @test quantile(chn)[Symbol("β[2]"), Symbol("50.0%")] ≈ -0.732 atol = 0.2
            @test quantile(chn)[:ϕ⁻, Symbol("50.0%")] ≈ 3.56 atol = 0.2
        end
    end
    @timed_testset "Hierarchical Model" begin
        f = @formula(y ~ (1 | cheese) + background)
        m = turing_model(f, cheese)
        chn = sample(seed!(123), m, NUTS(), MCMCThreads(), 2_000, 2)
        @test summarystats(chn)[:α, :mean] ≈ 68.33 atol = 2.0
        @test summarystats(chn)[Symbol("β[1]"), :mean] ≈ 6.928 atol = 0.2
        @test summarystats(chn)[Symbol("zⱼ[1]"), :mean] ≈ 0.306 atol = 0.2
        @test quantile(chn)[Symbol("zⱼ[2]"), Symbol("50.0%")] ≈ -1.422 atol = 0.5
    end
    @testset "Unsupported Model Likelihoods" begin
        @test_throws ArgumentError turing_model(@formula(y ~ x), nt_str; model=Binomial)
    end
    @testset "NegativeBinomial2" begin
        @test T.NegativeBinomial2(0, 1) == T.NegativeBinomial(1, 1)
        @test T.NegativeBinomial2(2, 8) == T.NegativeBinomial(8, 0.8)
        @test_throws DomainError T.NegativeBinomial2(0, -1)
    end
end


using Flux, Flux.Data.MNIST, Statistics, BSON, Random, StatsPlots; pyplot()
using Flux: onehotbatch, onecold, crossentropy, @epochs

epochs = 30
eta = 1e-3
batchSize = 200
trainRange, validateRange = 1:1000, 1001:5000

function minibatch(x, y, idxs)
    xBatch = Array{Float32}(undef, size(x[1])..., 1, length(idxs))
    for i in 1:length(idxs)
        xBatch[:, :, :, i] = Float32.(x[idxs[i]])
    end
    return (xBatch, onehotbatch(y[idxs], 0:9))
end

trainLabels = MNIST.labels()[trainRange]
trainImgs = MNIST.images()[trainRange]
mbIdxs = Iterators.partition(1:length(trainImgs), batchSize)
trainSet = [minibatch(trainImgs, trainLabels, i) for i in mbIdxs]

validateLabels = MNIST.labels()[validateRange]
validateImgs = MNIST.images()[validateRange]
validateSet = minibatch(validateImgs, validateLabels, 1:length(validateImgs))

accuracy(x, y, model) = mean(onecold(model(x)) .== onecold(y))
loss(x, y, model) = crossentropy(model(x), y)

function evalAccuracy(dropP)
    model= Chain(Conv((5, 5), 1=>8, relu), MaxPool((2,2)),
                    Conv((3, 3), 8=>16, relu), MaxPool((2,2)),
                    flatten,
                    Dense(400, 200,relu), Dropout(dropP),
                    Dense(200, 200,relu), Dropout(dropP),
                    Dense(200, 200,relu), Dropout(dropP),
                    Dense(200, 10,relu), Dropout(dropP),
                    softmax)
    opt = ADAM(eta);
    @epochs epochs Flux.train!((x,y)->loss(x,y,model),params(model),trainSet,opt)
    accuracy(validateSet..., model)
end

pToTest = [0.0, 0.25, 0.5, 0.75]
n = 10
results = [[evalAccuracy(p) for _ in 1:n] for p in pToTest]
bestAcc, bestI = findmax(median.(results))
println("The best dropout probability is $(pToTest[bestI]).")
println("It achieves $(bestAcc) accuracy on average.")

boxplot(results,label="",
       xticks=([1:1:4;],string.(pToTest)),
       xlabel="Dropout Probability", ylabel = "Accuracy", legend = false)


using LinearAlgebra
include("../various_mle_things/incremental.jl")

"""
Computes remainder and minimizer for lsq problem min_x |hat(A)x-hat(b)|_2^2 where
ith row of A and b are discarded, given QR factorization of A,
remainder r and minimizer xstar of the full problem
"""
function lsq_loocv(A, AQR, r, xstar, i)
    n = size(A, 1)
    p = size(A, 2)
    @assert n>=p #more rows than columns
    @assert i>=1 && n>=i #i must be valid row
    Q = AQR.Q[1:n, 1:p] #get reduced QR factor
    R = AQR.R
    At = A'
    r = reshape(r, length(r), 1) #reshape into col vector
    xstar = reshape(xstar, length(xstar), 1) #reshape into col vector
    remainder = repeat(r, 1, 1) #initial value (r) is important to algorithm
    minimizer = repeat(xstar, 1, 1) #initial value (xstar) is important to algorithm
    rmv = @views remainder
    minv = @views minimizer
    Atv = @views At
    leverage_score = norm(Q[i, :])^2 #norm of ith row of Q squared
    gamma = r[i] / (1-leverage_score)
    minv[:, 1] .-= R\(R'\ Atv[:, i]) * gamma
    rmv[:, 1] .+= Q * (Q'[:, i]) * gamma
    return remainder, minimizer #col vectors
end

"""
LOOCV for least squares problem min_x |Ax-b|_2^2. Computes all n remainder and minimizer vectors. Primarily used for testing.
"""
function lsq_loocv(A, AQR, r, xstar)
    n = size(A, 1)
    p = size(A, 2)
    @assert n>=p #more rows than columns
    Q = AQR.Q[1:n, 1:p] #get reduced QR factor
    R = AQR.R
    At = A'
    r = reshape(r, length(r), 1) #reshape into col vector
    xstar = reshape(xstar, length(xstar), 1) #reshape into col vector
    remainders = repeat(r, 1, n) #length n b/c we have n subproblems
    minimizers = repeat(xstar, 1, n)
    rmv = @views remainders
    minv = @views minimizers
    Atv = @views At
    for i = 1:n
        leverage_score = norm(Q[i, :])^2 #norm of ith row of Q squared
        gamma = r[i] / (1-leverage_score)
        minv[:, i] .-= R\(R'\ Atv[:, i]) * gamma
        rmv[:, i] .+= Q * (Q'[:, i]) * gamma
    end
    return remainders, minimizers
end

function unitvec(n, i)
    v = zeros(n, 1)
    v[i] = 1
    return v
end

"""
loocv for linear system Kc = y, generalized to case where c is a matrix
    INPUTS: R: upper triangular Cholesky factor of K
            c: solution to Kc = y
            i: row up for deletion
"""
function lin_sys_loocv(c, R, i)
    n = size(R)[1]
    #println(n)
    if size(c, 2)==1
        c = reshape(c, length(c), 1)
    end
    @assert i <= size(c, 1)
    ret = Array{Float64, 2}(undef, size(c, 1) - 1, size(c, 2))
    for j = 1:size(ret, 2)
        #println(unitvec(n, i) )
        #println(R')
        ve = R'\ unitvec(n, i)
        ri = c[i, j] / norm(ve)^2
        ci = c[:, j] - ri * (R\ve)
        ret[:, j] = ci[[1:i-1;i+1:end]] #dimension of return vector is n-1, ith entry will be zero
    end
    return ret
end

"""
Use incremental cholesky in computation
"""
function lin_sys_loocv_IC(c::Array{T} where T, IC::IncrementalCholesky{Float64}, i::Int64)
    R = get_chol(IC).U
    L = get_chol(IC).L
    n = size(IC)[1]
    if size(c, 2)==1
        c = reshape(c, length(c), 1)
    end
    @assert i <= size(c, 1)
    ret = Array{Float64, 2}(undef, size(c, 1) - 1, size(c, 2))
    for j = 1:size(ret, 2)
        #println(unitvec(n, i) )
        #println(R')
        #ve = R'\ unitvec(n, i)
        ve = L\unitvec(n, i)
        ri = c[i, j] / norm(ve)^2
        ci = c[:, j] - ri * (R\ve)
        ret[:, j] = ci[[1:i-1;i+1:end]] #dimension of return vector is n-1, ith entry will be zero
    end
    return ret
end


@doc """
  lift(a::padic) -> fmpz

  Returns the positive canonical representative in Z. a needs
  to be integral.
""" ->
function lift(a::padic)
  b = fmpz()
  R = parent(a)
  ccall((:padic_get_fmpz, :libflint), Void, (Ptr{fmpz}, Ptr{padic}, Ptr{FlintPadicField}), &b, &a, &R)
  return b
end


using Cairo
using LSystems
using DataStructures

c = CairoRGBSurface(1024, 1024);
cr = CairoContext(c);

save(cr);

set_source_rgb(cr, 1.0, 1.0, 1.0);
rectangle(cr, 0.0, 0.0, 1024.0, 1024.0);
fill(cr);

restore(cr);
save(cr);

x, y, a, d = 50., 512., 0, 15
set_line_width(cr, 0.6);
move_to(cr, x, y);

sierpinski_start = "F+XF+F+XF"
sierpinski_trans = Dict([('X', "XF-F+F-XF+F+XF-F+F-X")])

function determine_new_position(x, y, a, d)
    x -= sind(a) * d
    y -= cosd(a) * d
    return x, y
end

s = Stack(Vector{Float64}) # [x, y, a]

level = 4
drawing_instructions = @task lindenmayer(sierpinski_start, sierpinski_trans, level)

for v in drawing_instructions
    if v == 'F'
        x, y = determine_new_position(x, y, a, d)
        line_to(cr, x, y)
    elseif v == '+'
        a -= 90
    elseif v == '-'
        a += 90
    elseif v == '['
        push!(s, [x, y, a])
    elseif v == ']'
        x, y, a = pop!(s)
        move_to(cr, x, y)
    end
end

stroke(cr);
write_to_png(c, "sierpinski_$level.png");


using LinearAlgebra
using CompScienceMeshes
using SauterSchwab3D

const pI = point(1,5,3)
const pII = point(2,5,3)
const pIII = point(7,1,0)
const pIV = point(5,1,-3)
const pV = point(0,0,0)

const P = simplex(pI,pII,pIII)
const Q = simplex(pI,pIV,pV)

const sing = SauterSchwab3D.singularity_detection(P,Q)

Accuracy2 = 30
cv_ref = CommonVertex4D(sing,SauterSchwab3D._legendre(Accuracy2,0.0,1.0))

function integrand(x,y)
      return ((x-pI)'*(y-pV))*exp(-im*1*norm(x-y))/(4pi*norm(x-y))
end

function INTEGRAND(u,v)
   n1 = neighborhood(P,u)
   n2 = neighborhood(Q,v)
   x = cartesian(n1)
   y = cartesian(n2)
   output = integrand(x,y)*jacobian(n1)*jacobian(n2)
   return(output)
end

print("Ref: ")
ref = sauterschwab_parameterized(INTEGRAND, cv_ref)
println(ref)
println()

res_tp =[]
res_sp = []
res_gm = []
n1 = []
n2 = []
n3 = []
#Grauss tensor product
for i in 2:1:25
   Accuracy = i
   cv = CommonVertex4D(sing,SauterSchwab3D._legendre(Accuracy,0.0,1.0))

   int_tp = sauterschwab_parameterized(INTEGRAND, cv)

   num_pts = length(cv.qps)^4
   push!(n1,num_pts)

   push!(res_tp,int_tp)
end

#ShunnHam simplex product
for i in 2:1:7
   Accuracy = i
   cv_s = CommonVertex4D_S(sing,SauterSchwab3D._shunnham2D(Accuracy))

   int_sp = sauterschwab_parameterized(INTEGRAND, cv_s)

   num_pts = length(cv_s.qps)^2
   push!(n2,num_pts)
   push!(res_sp,int_sp)
end

#GrundmannMoeller simplex product
for i in 2:1:12
   Accuracy = i
   cv_gm = CommonVertex4D_S(sing,SauterSchwab3D._grundmannMoeller2D(2*Accuracy-1))

   int_gm = sauterschwab_parameterized(INTEGRAND, cv_gm)

   num_pts = length(cv_gm.qps)^2
   push!(n3,num_pts)
   push!(res_gm,int_gm)
end

err_tp = norm.(res_tp.-ref)/norm(ref)
err_sp = norm.(res_sp.-ref)/norm(ref)
err_gm = norm.(res_gm.-ref)/norm(ref)

println(err_tp)
println(err_sp)
println(err_gm)

using Plots
plot( yaxis=:log, xaxis=:log, fontfamily="Times")
plot!(n1,err_tp, label="Gauss Tensor-Product",markershape=:circle)
plot!(n2,err_sp, label="Simplex Tensor-Product",markershape=:rect)
#plot!(n3,err_gm, label="Simplex-Product GM",markershape=:x)
plot!(xlims=(5,10^4),ylims=(1e-6,1))
plot!(xlabel="Quad. points/Func. evals.", ylabel="Rel. Error.", title="Common Vertex 4D",legend=:bottomleft)

using BenchmarkTools

ref = -1.0761994366251377 + 1.0886007687110666im

cv = CommonVertex4D(sing,SauterSchwab3D._legendre(6,0.0,1.0))

int_tp = sauterschwab_parameterized(INTEGRAND, cv)

num_pts = length(cv.qps)^4
println(num_pts)
err_tp = norm.(int_tp.-ref)/norm(ref)
println(err_tp)
@time for i in 1:100 sauterschwab_parameterized(INTEGRAND, cv) end

cv_s = CommonVertex4D_S(sing,SauterSchwab3D._shunnham2D(6))

int_sp = sauterschwab_parameterized(INTEGRAND, cv_s)

num_pts = length(cv_s.qps)^2
println(num_pts)
err_sp = norm.(int_sp.-ref)/norm(ref)
println(err_sp)
 @time for i in 1:100 sauterschwab_parameterized(INTEGRAND, cv_s) end

#=

cv_gm = CommonVertex4D_S(SauterSchwab3D._grundmannMoeller2D(2*6-1))

int_gm = sauterschwab_parameterized(INTEGRAND, cv_gm)

num_pts = length(cv_gm.qps)^2
println(num_pts)
err_gm = norm.(int_gm.-ref)/norm(ref)
println(err_gm)
=#


# using PyPlot
using MatrixColorings

#####
# Example evaluating a numerical a Jacobian.  Uses the Brusselator in IVPtestsuite.jl
####

const N = 10^4 # active gridpoints
dof = 2N # degrees of freedom (boundary points are not free)
dx = 1/(N+1)
x = dx:dx:1-dx
T = Float64

const alpha = 1/50
const gamma = alpha/dx^2

# BC
const one_ = one(T)
const three = 3*one_
ubc1(t) = one_
ubc2(t) = one_
vbc1(t) = three
vbc2(t) = three

# immutable Pair
#     i::Int
#     j::Int
# end
# inds(i) = Pair(2i-1, 2i)
inds(i) = (2i-1, 2i)
function fn!(t,y, dydt)
    # the ode function dydt = f(t,y)
    #
    # Note the u and v are staggered:
    # u = y[1:2:end], v=y[2:2:end]

    # setup for first step
    i = 1
    iu, iv = inds(i)

    ui0 = ubc1(t); vi0 = vbc1(t) # BC
    ui1 = y[iu];   vi1 = y[iv]
    ui2 = y[iu+2]; vi2 = y[iv+2]
    while i<=N
        dydt[iu] = 1 + ui1^2*vi1 - 4*ui1 + gamma*(ui0 - 2*ui1 + ui2)
        dydt[iv] = 3*ui1 - ui1^2*vi1     + gamma*(vi0 - 2*vi1 + vi2)
        # setup for next step:
        i += 1
        iu, iv = inds(i)
        ui0 = ui1; vi0 = vi1
        ui1 = ui2; vi1 = vi2
        if i<N
            ui2 = y[iu+2]
            vi2 = y[iv+2]
        else # on last step do BC
            ui2 = ubc2(t)
            vi2 = vbc2(t)
        end
    end
    return nothing
end
fn!(;T_::Type=T, dof_=dof) = zeros(T_,dof_) # returns array for storage for dydt

# Analytic Jacobian: it is a banded matrix with upper and lower
# bandwidth 2, c.f. W&H p.148, but Julia does not support banded
# matrices yet.  Thus use a sparse matrix.
function jac!(t,y,dfdy)
    # The Jacobian of fn
    ii = 1 # direct index into dfdy.nzval
    for i=1:N # iterate over double-columns
        iu, iv = inds(i)
        ui = y[iu]
        vi = y[iv]
        # iu: column belonging to d/du(i)
        if iu-2>=1
            # du'(i-1)/du(i)
            dfdy.nzval[ii] = gamma
            ii +=1
        end
        if iv-2>=1
            # dv'(i-1)/du(i)==0
            dfdy.nzval[ii] = 0
            ii +=1
        end
        # du'(i)/du(i)
        dfdy.nzval[ii]   = 2*ui*vi - 4 -2*gamma
        ii += 1
        # dv'(i)/du(i)
        dfdy.nzval[ii] = 3 - 2*ui*vi
        ii +=1
        if  iu+2<=2N
            # du'(i+1)/du(i)
            dfdy.nzval[ii] = gamma
            ii +=1
        end

        # iv: column belonging to d/dv(i)
        if iv-2>=1
            # dv'(i-1)/dv(i)
            dfdy.nzval[ii] = gamma
            ii +=1
        end
        # du'(i)/dv(i)
        dfdy.nzval[ii] = ui^2
        ii +=1
        # dv'(i)/dv(i)
        dfdy.nzval[ii]   = -ui^2 -2*gamma
        ii +=1
        if  iu+2<=2N
            # du'(i+1)/dv(i)==0
            dfdy.nzval[ii] = 0
            ii +=1
        end
        if iv+2<=2N
            # dv'(i+1)/dv(i)
            dfdy.nzval[ii] = gamma
            ii +=1
        end
    end
    return nothing
end
function jac!(;T_::Type=T, dof_=dof)
    # returns an initialized sparse matrix
    B = (ones(T_,dof_-2), ones(T_,dof_-1), ones(T_,dof_), ones(T_,dof_-1), ones(T_,dof_-2))
    spdiagm(B, (-2,-1,0,1,2))
end

####
# Numerical Jacobians
####

# Non-coloring method:
######################
# Copied & adapted from DASSL:
# generate a function that computes approximate Jacobian using forward
# finite differences
if VERSION < v"0.4.0-dev"
    # TODO: remove these two for 0.4
    rowvals(S::SparseMatrixCSC) = S.rowval
    nzrange(S::SparseMatrixCSC, col::Integer) = S.colptr[col]:(S.colptr[col+1]-1)
end
function numjac!(t, y, jac)
    # updates jac in-place

    edelta  = eps(4.0)*1e7 # trial and error
    dof = length(y)
    f1 = similar(y)
    tmpy = similar(y) # d+dy
    f0 = similar(y)
    fn!(t, y, f0)
    for i=1:length(y)
        copy!(tmpy, y)
        tmpy[i] += edelta
        fn!(t, tmpy, f1)
        for j=1:dof
            tmpy[j] = (f1[j]-f0[j])/edelta
        end
        for nz in nzrange(jac,i)
            nonzeros(jac)[nz] = tmpy[rowvals(jac)[nz]]
        end
        # which, for large matrices, is much faster than:
        # jac[:,i] = (f1-f0)/edelta
    end
    return nothing
end

# Coloring method:
function numjac_c!(t, y, jac, S)
    # updates jac in-place
    nc = size(S,2) # number of colors

    edelta  = eps(4.0)*1e7 # trial and error
    dof = length(y)
    f0 = similar(y)
    fn!(t, y, f0)
    ddy = deepcopy(y) # d+dy
    f1 = similar(y)   # f(d+dy)
    for c in 1:nc
        add_delta!(ddy, S, c, edelta)  # adds a delta to all indices of color c
        fn!(t, ddy, f1)
        for j=1:dof
            f1[j] = (f1[j]-f0[j])/edelta
        end
        recover_jac!(jac, f1, S, c) # recovers the columns of jac which correspond c
        remove_delta!(ddy, S, c, edelta) # remove delta again
    end
    return nothing
end

function numjac_c!(t, y, jac)
    # updates jac in-place, uses its sparsity structure for coloring
    S = color_jac(jac) # seed matrix
    numjac_c!(t, y, jac, S)
    return nothing
end

####
# Compare
####

t = 0.0 # not dependent on t
# initialize
dydt = fn!()
dfdy_ana = jac!()
dfdy_num = jac!()
dfdy_num_c = jac!()

y = fn!()
#IC
u0(x) = 1 + 0.5*sin(2π*x)
v0(x) = 3 * ones(length(x))
y[1:2:2N] = u0(x)
y[2:2:2N] = v0(x)
# warm ups:
@time 1
jac!(t, y, dfdy_ana)
numjac!(t, y, dfdy_num)

#ys = (y, rand(T, dof)*3+0.5)
for y in (y,)
    # calculate, time and compare
    fn!(t, y, dydt)
    println("@time for analytic:")
    @time jac!(t, y, dfdy_ana)
    println("@time for numeric with no coloring:")
    @time numjac!(t, y, dfdy_num)
    println("@time for numeric with coloring:")
    @time numjac_c!(t, y, dfdy_num_c)
    S = color_jac(dfdy_num_c) # seed matrix
    println("@time for numeric with coloring with provided seed matrix:")
    @time numjac_c!(t, y, dfdy_num_c,S)

    @show maximum(dfdy_ana-dfdy_num)
    @show maximum(dfdy_ana-dfdy_num_c)
end


module _XK2En
export XK2En
using Statistics: std, mean
    """
        XK2, Ci = XK2En(Sig)

    Returns the cross-Kolmogorov entropy estimates (`XK2`) and the correlation
    integrals (`Ci`) for m = [1,2] estimated between the data sequences
    contained in `Sig` using the default parameters:
    embedding dimension = 2, time delay = 1, distance threshold (r) = 0.2*SD(Sig),
    logarithm = natural

        XK2, Ci = XK2En(Sig::AbstractArray{T,2} where T<:Real; m::Int=2, tau::Int=1, r::Real=0.2*std(Sig,corrected=false), Logx::Real=exp(1))

    Returns the cross-Kolmogorov entropy estimates (`XK2`) estimated between
    the data sequences contained in `Sig` using the specified 'keyword' arguments:

    # Arguments:
    `m`     - Embedding Dimension, a positive integer [default: 2]  \n
    `tau`   - Time Delay, a positive integer         [default: 1]  \n
    `r`     - Radius Distance Threshold, a positive scalar  [default: 0.2*SD(`Sig`)]  \n
    `Logx`  - Logarithm base, a positive scalar      [default: natural]  \n

    # See also `XSampEn`, `XFuzzEn`, `XApEn`, `K2En`, `XMSEn`, `XDistEn`

    # References:
        [1]  Matthew W. Flood,
             "XK2En - EntropyHub Project"
             (2021) https://github.com/MattWillFlood/EntropyHub

    """
    function XK2En(Sig::AbstractArray{T,2} where T<:Real; m::Int=2,
        tau::Int=1, r::Real=0.2*std(Sig,corrected=false), Logx::Real=exp(1))

    size(Sig,2) > size(Sig,1) ? Sig = transpose(Sig) : nothing

    N = size(Sig,1)
    (N>10 && size(Sig,2)==2) ? nothing :  error("Sig:   must be a 2-columns matrix")
    (m > 0) ? nothing : error("m:     must be an integer > 0")
    (tau>0) ? nothing : error("tau:   must be an integer > 0")
    (Logx>0) ? nothing : error("Logx: must be a positive number > 0")
    (r>0) ? nothing : error("r:  must be 2 element tuple of positive values")

    m   += 1
    Zm1 = zeros(N,m)
    Zm2 = zeros(N,m)
    Ci = zeros(m)
    for n = 1:m
        N2 = N-(n-1)*tau
        Zm1[1:N2,n] = Sig[(n-1)*tau + 1:N,1]
        Zm2[1:N2,n] = Sig[(n-1)*tau + 1:N,2]
        Norm = zeros(N2,N2)
        for k = 1:N2
            Temp = repeat(transpose(Zm1[k,1:n]),outer=N2) .- Zm2[1:N2,1:n]
            Norm[k,:] = sqrt.(sum(Temp.*Temp,dims=2))
        end
        Ci[n] = mean(Norm .< r)
    end

    XK2 = log.(Logx, Ci[1:m-1]./Ci[2:m])/tau
    XK2[isinf.(XK2)] .= NaN

    return XK2, Ci

    end
end

"""
Copyright 2021 Matthew W. Flood, EntropyHub

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.

For Terms of Use see https://github.com/MattWillFlood/EntropyHub
"""


#Fit N-mixture model

using Random, Distributions, StatsFuns, Optim, NLSolversBase, ForwardDiff
using LinearAlgebra

using Unmarked

Random.seed!(123);

#Simulate data
N = 4000;
J = 5;
z = rand(Poisson(5),N);
r = 0.1;

y = Array{Int64,2}(undef, N, J);

for i = 1:N
  for j = 1:J
    p = 1 - (1-r)^z[i];
    y[i,j] = rand(Binomial(1, p));
  end
end

um = UmData(y)

@time mod = rn(@formula(λ~1), @formula(p~1), um, K=50)

#Display
truth = [log(5), logit(r)];
pnames = ["β[1]", "β[2]"];
cnames = ["param" "truth" "estimate" "se"];
[cnames; pnames truth coef(mod) stderror(mod)]


using LinearAlgebra

"""
Truncated SVD to find the matrix of rank k that is the closest to A

# Arguments
- `A::Array{Float64,2}`
- `k::Array{Float64,1}`

"""
function approx_matrix_svd(A,k)
    (m,n) = size(A)
    F = svd(A)
    A_SVD = F.U[1:m,1:k] * Diagonal(F.S)[1:k,1:k] * F.V'[1:k,1:n]
    return A_SVD
end


@diffusion_process FavettoSamson{T,K} begin
    :dimensions
    process --> 2
    wiener --> 2

    :parameters
    (α, β, λ, μ, σ1, σ2) --> T
    dose --> K

    :additional
    diagdiff --> true
    constdiff --> true
end

DiffusionDefinition.b(t, x, P::FavettoSamson) = @SVector [
    P.α * P.dose(t) - (P.λ +P.β)*x[1] + P.μ*x[2],
    P.λ*x[1] - P.μ*x[2]
]

DiffusionDefinition.σ(t, x, P::FavettoSamson) = SDiagonal(P.σ1, P.σ2)


#############################################################################################
# Calculating the Required Structural Thickness of Pipe Using Beam Stress Theory
#############################################################################################
# Input parameters
P = 850 # pressure
Do = 4.5 # outside diameter
S = 17900 # code allowable stress
E = .85 # weld joint efficieny
W = 1.0
Y = 0.4
tnom = .237
D = Do - (2*(tnom)) # inside diameter
psteel = .284 # density of the steel = 0.284 lb/in3
tjacket = 0.016 # jacket thicknesses
pjacket = 0.0975 # density of aluminum jacket ln/in3
tinsul = 3.0 # insulation thickness
pinsul = 0.004919 # density of insulation - note set at calcium silicate lb/in3
pproduct = 0.0361108 # product density (Water)  lb/in3
My = 0.1071 # Moment in the Y direction (up or down), in-lb
L = 14 # span in feet
l = L * 12 # span in inches
y = Do / 2 # distance from the center axis to the outer stress element, (1/2 Diameter), in

# Step 1 - Determine tpress , the thickness needed to contain the design pressure at the design temperature
# Pressure design thickness
t = (P * Do) / (2*((S * E * W) + (P * Y)))

# Step 2 -  Determine the weight of each component that contributes to the overall weight of the piping system.
# This includes the weight of the pipe, the fluid inside, the insulation, and the weight of any external items.

# Determine weight of pipe = determine the cross sectional area of the new pipe and multiply by the density of steel
wpipe = (π* psteel * (Do^2 - D^2) / 4)

# Determine weight of the insulation
winusl = (π* pinsul *((Do+tinsul*2)^2 - Do^2) / 4) #
wjacket = π* pjacket * (((Do+(tinsul*2))+(tjacket*2))^2 - (Do+tinsul*2)^2) / 4

# Determine weight of the fluid
Ai = π* D^2 / 4 # area
wprod = pproduct * Ai # lb/inches

#####
# Consideration for the weight from other external influences should be considered and added.   Items such as ice, snow, or even animals may be appropriate,
# depending on location and conditions.  If the additional weight is from a point load condition, refer to AISC Beam Deflection Table for the appropriate moment calculation.
######

wother = 0.0 # lb/inch

wmax = (wpipe + wprod + winusl + wjacket + wother) # maximum weight lb/in
W = wmax * l # total weight over the full span
R = (wmax * l) / 2 # reactions, R each support (2x)

# Step 3 -  Determine the Moments that will be used to calculate the bending stress in the pipe.  The Moments are loads created due to the weight of the overall piping system
My = My * (wmax) * (L * 12)^2 # AISC equation
#My =(wmax) * ((L * 12)^2) # AISC equation
#My = wmax * (L*12)^2 / 24
# wL^2/8 # for simply supported beam with a uniformly distributed load - Max shear force at either end of the beam (supports) - maximum bending moment - middle of beam
#My = wmax * (L*12)^2 / 12
#My = (R * (l/2)) - (wmax/2)

# Determine structural wall thickness
# Beam Flexural Stress Method
# The required minimum structural thickness of the pipe is derived from the Moment of Inertia needed to limit the bending stress to the maximum allowable stress levels
# Step 4
σmax = S # code allowable stress limit

# If the effect of wind loading is ignored, then σy = σmax because forces are only acting in the downward direction, -Y.
# σy = [ My * y ] / Ipipe
Ipipe_solve = (My * y)  / σmax # in4 - algebra to find Ipipe I = My/(S/y)

# Step 5 - Moment of inertia
# Ipipe = (π / 64) * (Do^4 - x^4) # use algebra to solve for x
# find x
# divide both sides by (π / 64)
solve = Ipipe_solve / (π / 64) # (Do^4 - x^4) = (28.77 / (π / 64))
# subtract Do^4 from both sides
solve = solve - Do^4 # -x^4 = solve - Do^4
solve = solve *-1
Di = (solve)^(1/4)
Ipipe = (π / 64) * (Do^4 - Di^4)

x = (((Ipipe_solve / (π / 64) - Do^4) * -1))^(1/4)

# Stress in y (up and down position)
σy = (My * y) / Ipipe

# maximum deflection
#δmax = 0.0065 * wmax * (L *12)^4 / (29000000 * Ipipe) # AISC equation
# defelction = (5 * wmax * 288^4) / (384 * 29000000 * Ipipe)

# Step 6 - Calcualte the minimum structural thickness
tstruct = (Do - Di) / 2

####################################### specific pipe case

#############################################################################################
# Calculating the Required Structural Thickness of Pipe Using Beam Stress Theory
#############################################################################################
# Input parameters
P = 285 # pressure
Do = 2.375 # outside diameter
S = 20000 # code allowable stress
E = 1 # weld joint efficieny
W = 1.0
Y = 0.4
tnom = 0.218
D = Do - (2*(tnom)) # inside diameter
psteel = .284 # density of the steel = 0.284 lb/in3
tinsul = 0.0 # insulation thickness
pinsul = 0.0049 # density of insulation - note set at calcium silicate lb/in3
pproduct = 8.5 # product density (Water)  lb/gal
My = 0.1071 # AISC equation
L = 24 # span in feet
l = L * 12 # span in inches
y = Do / 2 # distance from the center axis to the outer stress element, (1/2 Diameter), in

# Step 1 - Determine tpress , the thickness needed to contain the design pressure at the design temperature
# Pressure design thickness
t = (P * Do) / (2*((S * E * W) + (P * Y)))

# Step 2 -  Determine the weight of each component that contributes to the overall weight of the piping system.
# This includes the weight of the pipe, the fluid inside, the insulation, and the weight of any external items.

# Determine weight of pipe = determine the cross sectional area of the new pipe and multiply by the density of steel
wpipe = ((Do-(tnom / 2)) * π * tnom * psteel)

# Determine weight of the insulation
winsul = (Do + tinsul) * π * tinsul * pinsul

# Determine weight of the fluid
r = D/2
h = D
# Calculate volme of the fluid (cubic gal)
area = acos((r-h)/r) * r^2 - (r-h) * sqrt(2*r*h - h^2)
volume_cubic_inch = area * 12
volume_cubic_gal = volume_cubic_inch / 231

wprod = (volume_cubic_gal * pproduct) / 12 # lb/inches

#####
# Consideration for the weight from other external influences should be considered and added.   Items such as ice, snow, or even animals may be appropriate,
# depending on location and conditions.  If the additional weight is from a point load condition, refer to AISC Beam Deflection Table for the appropriate moment calculation.
######

wother = 0.0 # lb/inch

wmax = (wpipe + wprod + winsul + wother) # maximum weight lb/in

# Step 3 -  Determine the Moments that will be used to calculate the bending stress in the pipe.  The Moments are loads created due to the weight of the overall piping system
My = My * (wmax) * ((L * 12)^2) # AISC equation
#My = wmax * (L*12)^2 / 8 # wL^2/8 # for simply supported beam with a uniformly distributed load - Max shear force at either end of the beam (supports) - maximum bending moment - middle of beam

# Determine structural wall thickness
# Beam Flexural Stress Method
# The required minimum structural thickness of the pipe is derived from the Moment of Inertia needed to limit the bending stress to the maximum allowable stress levels
# Step 4
σmax = S # code allowable stress limit

# If the effect of wind loading is ignored, then σy = σmax because forces are only acting in the downward direction, -Y.
# σy = [ My * y ] / Ipipe
Ipipe_solve = (My * y)  / σmax # in4 - algebra to find Ipipe

# Step 5 - Moment of inertia
# Ipipe = (π / 64) * (Do^4 - x^4) # use algebra to solve for x
# find x
# divide both sides by (π / 64)
solve = Ipipe_solve / (π / 64) # (Do^4 - x^4) = (28.77 / (π / 64))
# subtract Do^4 from both sides
solve = solve - Do^4 # -x^4 = solve - Do^4
solve = solve *-1
Di = (solve)^(1/4)
Ipipe = (π / 64) * (Do^4 - Di^4)

# Stress in y (up and down position)
σy = (My * y) / Ipipe

# maximum deflection at code allowable S
δmax = 0.0065 * (wmax) * (L *12)^4 / (29000000 * Ipipe) # AISC equation
# defelction = (5 * wmax * 288^4) / (384 * 29000000 * Ipipe)

# Step 6 - Calcualte the minimum structural thickness
tstruct = (Do - Di) / 2
    print("Structural tmin - bend stress theory = ", tstruct)

###########
    #### iteration
    spans = collect(1:1:60)
    out = zeros(size(spans,1))

    for i = 1:size(spans,1)
        # Input parameters
        P = 285 # pressure
        Do = 2.375 # outside diameter
        S = 20000 # code allowable stress
        E = 1 # weld joint efficieny
        W = 1.0
        Y = 0.4
        tnom = 0.218
        D = Do - (2*(tnom)) # inside diameter
        psteel = .284 # density of the steel = 0.284 lb/in3
        tinsul = 0.0 # insulation thickness
        pinsul = 0.0049 # density of insulation - note set at calcium silicate lb/in3
        pproduct = 8.5 # product density (Water)  lb/gal
        My = 0.1071
        span = spans[i] # in feet
        y = Do / 2 # distance from the center axis to the outer stress element, (1/2 Diameter), in

        # Step 1 - Determine tpress , the thickness needed to contain the design pressure at the design temperature
    # Pressure design thickness
    t = (P * Do) / (2*((S * E * W) + (P * Y)))

    # Step 2 -  Determine the weight of each component that contributes to the overall weight of the piping system.
    # This includes the weight of the pipe, the fluid inside, the insulation, and the weight of any external items.

    # Determine weight of pipe = determine the cross sectional area of the new pipe and multiply by the density of steel
    wpipe = ((Do - (tnom / 2)) * π * tnom * psteel)

    # Determine weight of the insulation
    winsul = (Do + tinsul) * π * tinsul * pinsul

    # Determine weight of the fluid
    r = D/2
    h = D
    # Calculate volme of the fluid (cubic gal)
    area = acos((r-h)/r) * r^2 - (r-h) * sqrt(2*r*h - h^2)
    volume_cubic_inch = area * 12
    volume_cubic_gal = volume_cubic_inch / 231

    wprod = (volume_cubic_gal * pproduct) / 12 # lb/inches

    #####
    # Consideration for the weight from other external influences should be considered and added.   Items such as ice, snow, or even animals may be appropriate,
    # depending on location and conditions.  If the additional weight is from a point load condition, refer to AISC Beam Deflection Table for the appropriate moment calculation.
    ######

    wother = 0.0 # lb/inch

    # Step 3 -  Determine the Moments that will be used to calculate the bending stress in the pipe.  The Moments are loads created due to the weight of the overall piping system
    My = My * (wpipe + wprod + winsul + wother) * (span * 12)^2

    # Determine structural wall thickness
    # Beam Flexural Stress Method
    # The required minimum structural thickness of the pipe is derived from the Moment of Inertia needed to limit the bending stress to the maximum allowable stress levels
    # Step 4
    σmax = S # code allowable stress limit

    # If the effect of wind loading is ignored, then σy = σmax because forces are only acting in the downward direction, -Y.
    # σy = [ My * y ] / Ipipe
    Ipipe_solve = (My * y)  / σmax # in4 - algebra to find Ipipe

    # Step 5 - Moment of inertia
    # Ipipe = (π / 64) * (Do^4 - x^4) # use algebra to solve for x
    # find x
    # divide both sides by (π / 64)
    solve = Ipipe_solve / (π / 64) # (Do^4 - x^4) = (28.77 / (π / 64))
    # subtract Do^4 from both sides
    solve = solve - Do^4 # -x^4 = solve - Do^4
    solve = solve *-1
    Di = (solve)^(1/4)
    Ipipe = (π / 64) * (Do^4 - Di^4)

    # Step 6 - Calcualte the minimum structural thickness
    tstruct = (Do - Di) / 2

    out[i] = tstruct
end

for i = 1:size(out,1)
    if out[i] == 0.0
        out[i] = NaN
    end
end

using DataFrames
using CSV
using Gadfly
df = DataFrame(hcat(spans,out))
CSV.write("C:/Users/Andrew.Bannerman/Desktop/MARS/struct_out2.csv", df;delim=',')

plot(x=df.x1,y=df.x2, Geom.line,Theme(default_color=colorant"blue"),
Theme(background_color="white"),
Guide.title("Structural Required Thickness for Varying pipe spans - 2 inch sch 80"),
Guide.xlabel("Span"),
Guide.ylabel("tstructural"),
Guide.yticks(ticks = collect(0:.05:.75)),
Guide.xticks(ticks = collect(0:5:60)))


using EconPDEs, Distributions

Base.@kwdef struct BansalYaronModel
    # consumption process parameters
    μbar::Float64 = 0.018
    vbar::Float64 = 0.00073
    κμ::Float64 = 0.252
    νμ::Float64 = 0.528
    κv::Float64 = 0.156
    νv::Float64 = 0.00354

    # utility parameters
    ρ::Float64 = 0.024
    γ::Float64 = 7.5
    ψ::Float64 = 1.5
end

function (m::BansalYaronModel)(state::NamedTuple, y::NamedTuple)
    (; μbar, vbar, κμ, νμ, κv, νv, ρ, γ, ψ) = m
    (; μ, v) = state
    (; p, pμ_up, pμ_down, pv_up, pv_down, pμμ, pμv, pvv) = y

    # drift and volatility of c, μ, σ, p
    μc = μ
    σc = sqrt(v)
    μμ = κμ * (μbar - μ)
    σμ = νμ * sqrt(v)
    μv = κv * (vbar - v)
    σv = νv * sqrt(v)
    pμ = (μμ >= 0) ? pμ_up : pμ_down
    pv = (μv >= 0) ? pv_up : pv_down
    σp_Zμ = pμ / p * σμ
    σp_Zv = pv / p * σv
    σp2 = σp_Zμ^2 + σp_Zv^2
    μp = pμ / p * μμ + pv / p * μv + 0.5 * pμμ / p * σμ^2 + 0.5 * pvv / p * σv^2

    # Market price of risk κ
    κ_Zc = γ * σc
    κ_Zμ = - (1 - γ * ψ) / (ψ - 1) * σp_Zμ
    κ_Zv = - (1 - γ * ψ) / (ψ - 1) * σp_Zv
    κ2 = κ_Zc^2 + κ_Zμ^2 + κ_Zv^2

    # Risk free rate r
    r = ρ + μc / ψ - (1 + 1 / ψ) / 2 * γ * σc^2 - (γ * ψ - 1) / (2 * (ψ - 1)) * σp2

    # Market Pricing
    pt = - p * (1 / p + μc + μp - r - κ_Zc * σc - κ_Zμ * σp_Zμ - κ_Zv * σp_Zv)
    #pt = p * (1 / p - ρ + (1 - 1 / ψ) * (μc - 0.5 * γ * σc^2) + μp + 0.5 * (1 / ψ - γ) / (1 - 1 / ψ) * σp2)

    return (; pt)
end

# Bansal Yaron (2004)
m = BansalYaronModel()
μn = 30
νn = 30
μdistribution = Normal(m.μbar,  sqrt(m.νμ^2 * m.vbar / (2 * m.κμ)))
μs = range(quantile(μdistribution, 0.025), quantile(μdistribution, 0.975), length = μn)
νdistribution = Gamma(2 * m.κv * m.vbar / m.νv^2, m.νv^2 / (2 * m.κv))
vs = range(quantile(νdistribution, 0.025), quantile(νdistribution, 0.975), length = νn)
stategrid = OrderedDict(:μ => μs, :v => vs)
yend = OrderedDict(:p => ones(μn, νn))
result = pdesolve(m, stategrid, yend)
@assert result.residual_norm <= 1e-5

# Bansal Yaron (2009)
## m = BansalYaronModel(μbar = 0.018, vbar = 0.00062, κμ = 0.3, νμ = 0.456, κv = 0.012, νv = 0.00472, ρ = 0.0132, γ = 10, ψ = 1.5)
## stategrid = initialize_stategrid(m)
## yend = OrderedDict(:p => ones(length(stategrid[:μ]), length(stategrid[:v])))
## y, result, distance = pdesolve(m, stategrid, yend)