function eye(tpe::Type,n::S) where {S <: Integer}

    if tpe <: Number
        return Matrix{tpe}(I,n,n)
    end

end

eye(n::Integer) = eye(Float64,n::Integer)

function tracem(x::Array{T,2}) where {T <: Real}

  # Computes the matrix-trace as defined by Gomme and Klein (2011)

  # We require the number of rows to be greater than the number of columns, so
  # that m is greater than one.

  trans = false

  if size(x,1) < size(x,2)
    x = Matrix(x')
    trans = true
  end

  n = size(x,2)
  m = Int(size(x,1)/n)

  y = zeros(m,1)  # We want this to be a 2-d array for subsequent matrix multiplication

  @inbounds for i = 1:m
    @views y[i,1] = tr(x[(n*(i-1)+1):i*n,1:n])
  end

  if trans == true
    y = y'
  end

  return Matrix(y)

end

function kron_prod_times_vector(a::AbstractArray{T,2},b::AbstractArray{T,2},v::AbstractArray{T,1}) where {T <: Real}

    # This function efficiently computes kron(a,b)*v where a and b are matrices
    # and v is a vector.

    (n1,n2) = size(a)
    (n3,n4) = size(b)
    n5      = length(v)

    if n5 != n2*n4
        error("a, b, and v do not have the correct sizes")
    end

    p = vec(b*reshape(v,n4,n2)*a')

    return p

end

function vector_times_kron_prod(v::AbstractArray{T,1},a::AbstractArray{T,2},b::AbstractArray{T,2}) where {T <: Real}

    product = Matrix(kron_prod_times_vector(a',b',v')')

    return product

end

function kron_prod_times_matrix(a::AbstractArray{T,2},b::AbstractArray{T,2},v::AbstractArray{T,2}) where {T <: Real}

    # This function efficiently computes kron(a,b)*v where a, b, and v are
    # conformable matrices.

    (n1,n2) = size(a)
    (n3,n4) = size(b)
    (n5,n6) = size(v)

    if n5 != n2*n4
        error("a, b, and v do not have the correct sizes")
    end

    p = zeros(n1*n3,n6)
    for i = 1:n6
        @views p[:,i] = kron_prod_times_vector(a,b,v[:,i])
    end

    return p

end

function matrix_times_kron_prod(v::AbstractArray{T,2},a::AbstractArray{T,2},b::AbstractArray{T,2}) where {T <: Real}

    product = Matrix(kron_prod_times_matrix(a',b',v')')

    return product

end

function kron_prod_times_vector(a::AbstractArray{T,2},b::AbstractArray{T,2},c::AbstractArray{T,2},v::AbstractArray{T,1}) where {T <: Real}

    # This function efficiently computes kron(a,kron(b,c))*v where a, b, and c
    # are matrices and v is a vector.  Included for completeness, but not
    # actually used.

    (n1,n2) = size(a)
    (n3,n4) = size(b)
    (n5,n6) = size(c)
    n7      = length(v)

    if n7 != n2*n4*n6
        error("a, b, c, and v do not have the correct sizes")
    end

    v_tilda = reshape(v,n4*n6,n2)*a'
    p = vec(kron_prod_times_matrix(b,c,v_tilda))

    return p

end

function vector_times_kron_prod(v::AbstractArray{T,1},a::AbstractArray{T,2},b::AbstractArray{T,2},c::AbstractArray{T,2}) where {T <: Real}

    product = Matrix(kron_prod_times_vector(a',b',c',v')')

    return product

end

function kron_prod_times_matrix(a::AbstractArray{T,2},b::AbstractArray{T,2},c::AbstractArray{T,2},v::AbstractArray{T,2}) where {T <: Real}

    # This function efficiently computes kron(a,kron(b,c))*v where a, b, c and
    # v are comformable matrices.

    (n1,n2) = size(a)
    (n3,n4) = size(b)
    (n5,n6) = size(c)
    (n7,n8) = size(v)

    if n7 != n2*n4*n6
        error("a, b, c, and v do not have the correct sizes")
    end

    p = zeros(n1*n3*n5,n8)
    for i = 1:n8
        @views v_tilda = reshape(v[:,i],n4*n6,n2)*a'
        p[:,i] .= vec(kron_prod_times_matrix(b,c,v_tilda))
    end

    return p

end

function matrix_times_kron_prod(v::AbstractArray{T,2},a::AbstractArray{T,2},b::AbstractArray{T,2},c::AbstractArray{T,2}) where {T <: Real}

    product = Matrix(kron_prod_times_matrix(a',b',c',v')')

    return product

end

function dsylvester(a::AbstractArray{T,2}, b::AbstractArray{T,2}, c::Union{AbstractArray{T,1},AbstractArray{T,2}}) where {T <: Real}

    #= Uses the Hessenberg-Schur method to find the bounded solution of the
       discrete Sylvester equation:

             X + A*X*B = C

       Based on Golub, Nash, and Van Loan (1979).

    =#

    n = size(a, 1)
    m = size(b, 1)
    x = zeros(size(c))

    (s,u) = schur(Matrix(b'))
    (v,t) = hessenberg(a)

    c = v'*c*u

    j = m
    while j > 0
        j1 = j
        if j == 1
            block_size = 1
        elseif isequal(s[j,j-1],0.0) == false
            block_size = 2
            j -= 1
        else;
            block_size = 1
        end
        @views ajj = kron(s[j:j1,j:j1],t) + I
        @views rhs = vec(c[:,j:j1])
        if j1 < m
            @views rhs2 = t*(x[:,(j+1):m]*s[j:j1,(j+1):m]')
            rhs -= vec(rhs2)
        end
        w = ajj\rhs
        @views x[:,j] = w[1:n]
        if block_size == 2
            @views x[:,j1] = w[(n+1):2*n]
        end
        j -= 1
    end

    x = v*x*u'

    return x

end

function trm(x::AbstractArray{T,2}) where {T <: Real}

    # Computes the matrix trace as defined by Binning (2013).  Used for
    # computing the second-order terms, hss, gss.

    (n1,n2) = size(x)
    k = Int(round(sqrt(n2)))

    y = zeros(n1)
    for i = 1:k
        @views y += x[:,i+(i-1)*k]
    end

    return y

end

function trm2(x::AbstractArray{T,2}) where {T <: Real}

    # Computes the matrix trace as defined by Binning (2013).  Used for
    # computing the third-order terms, hssx, gssx.

    (n1,n2) = size(x)
    k = Int(round(n2^(1//3)))

    y = zeros(n1,k)
    for j = 1:k
        for i = 1:k
            @views y[:,j] += x[:,(j-1)+i+(i-1)*k^2]
        end
    end

    return y

end

function create_omega3(n::S) where {S <: Integer}

    # Creates the combination matrix for a third-order perturbation as defined in Levintal (2017).

    # This function is a simplified version of the create_OMEGA function originally written in
    # Matlab by Oren Levintal for his paper "Fifth Order Perturbation Solution to DSGE Models"
    # published in the Journal of Economic Dynamics and Control, 2017.  Permission to translate
    # this function into Julia and release it within the SolveDSGE module was granted by Oren
    # Levintal on February 5, 2020.

    ind = [1:n^3;]
    M = reshape(ind,1,n,n,n)
    Ix = eye(S,n^3)
    omega3 = (reshape(Ix[:,PermutedDimsArray(M,[1,4,2,3])],n^3,n^3)
           + reshape(Ix[:,PermutedDimsArray(M,[1,2,4,3])],n^3,n^3)
           + reshape(Ix[:,PermutedDimsArray(M,[1,2,3,4])],n^3,n^3))

    return omega3

end

function kron_prod_times_vector(A::Union{Array{Array{T,2},1},Array{Array{Complex{T},2},1}},x::Union{Array{T,1},Array{Complex{T},1}},n::Array{S,1},p::S) where {T <: Real, S <: Integer}

    # Computes y = (A[p] * A[p-1] * ... * A[1] )*x

    N = prod(n[1:p])
    z = copy(x)
    for i = 1:p
        z = (A[i]*reshape(z,n[i],Int(N/n[i])))'
    end
    y = reshape(Matrix(z),N)

    return y

end

function kron_prod_times_matrix(A::Union{Array{Array{T,2},1},Array{Array{Complex{T},2},1}},x::Union{Array{T,2},Array{Complex{T},2}},n::Array{S,1},p::S) where {T <: Real, S <: Integer}

    # Computes y = (A[p] * A[p-1] * ... * A[1] )*x

    N = prod(n[1:p])
    y = Array{Complex{T}}(undef,N,size(x,2))
    for j = 1:size(x,2)
        @views z = x[:,j]
        for i = 1:p
            z = (A[i]*reshape(z,n[i],Int(N/n[i])))'
        end
        y[:,j] .= reshape(Matrix(z),N)
    end

    return y

end

function KPShiftSolve(TT::Union{Array{Array{T,2},1},Array{Array{Complex{T},2},1}},n::Array{S,1},c::Union{Array{T,1},Array{Complex{T},1}},lambda::T,alpha::Union{T,Complex{T}}) where {T <: Real, S <: Integer}

    p = length(n)
    N = prod(n)

    c = copy(c)

    TT[p] = alpha*TT[p]
    if p == 1
        y = (TT[1] + lambda*Matrix{T}(I,n[1],n[1]))\c
    else
        y = Array{Complex{T}}(undef,N)
        mp = Int(N/n[p])
        for i = n[p]:-1:1
            idx = ((i-1)*mp+1):(i*mp)
            y[idx] = KPShiftSolve(TT[1:(p-1)],n[1:(p-1)],c[idx],lambda,TT[p][i,i])
            z = kron_prod_times_vector(TT[1:p],y[idx],n,p-1)
            for j = 1:(i-1)
                jdx = ((j-1)*mp+1):(j*mp)
                c[jdx] = c[jdx] - TT[p][j,i]*z
            end
        end
    end

    return y

end

function martin_van_loan(a::Array{T,2},b::Array{T,2},c::Array{T,2},d::Array{T,2},k::S) where {T <: Real, S <: Integer}

    #= Uses a recursive Schur method to find the bounded solution of the Sylvester
       equation:

       AX + B*X*(Kron^(k)C) = D

       Based on Martin and Van Loan (2006).  This is a simplified implementation
       of their algorithm, but it captures most of the gains over Golub, Nash,
       and van Loan (1979).

    =#

    a = copy(a)
    b = a\copy(b)
    c = copy(c)
    d = a\copy(d)

    (v,s) = hessenberg(b)       # v*s*v' = b
    (t,q) = schur(complex(c'))  # q*t*q' = c'

    v = Matrix(v)

    p = k + 2
    TT = Array{typeof(t)}(undef,p)
    TT[1] = s
    for i = 2:p
        TT[i] = conj(t)
    end

    Q     = fill(q,k+1)
    inv_Q = fill(Matrix(q'),k+1)

    n = fill(size(c,1),p)
    n[1] = size(b,1)
    N = prod(n)

    lambda = 1.0

    e = vec(v'*kron_prod_times_matrix(inv_Q,Matrix(d'),n[2:end],k+1)')

    y = reshape(KPShiftSolve(TT,n,e,lambda,1.0),size(d))
    x = real(Matrix(v*kron_prod_times_matrix(Q,Matrix(y'),n[2:end],k+1)'))

    return x

end

function dlyap(a::Array{T,2}, b::Array{T,2}) where {T <: Real}

    n = size(a,1)
    x = zeros(n,n)
    j = n

    (s,u) = schur(a)
    b = u'b*u

    while j > 0
        j1 = j
        if j == 1
            block = 1
        elseif !isequal(s[j,j-1],0.0)
            block = 2
            j -= 1
        else
            block = 1
        end
        @views lhs = kron(s[j:j1,j:j1],s) - I # I = eye(block*n)
        @views rhs = vec(b[:,j:j1])
        if j1 < n
            @views rhs2 = s*(x[:,(j1+1):n]*s[j:j1,(j1+1):n]')
            rhs += vec(rhs2)
        end
        w = -lhs\rhs
        @views x[:,j] = w[1:n]
        if block == 2
            @views x[:,j1] = w[(n+1):block*n]
        end
        j -= 1
    end

    x = u*x*u'

    return x

end

function ind2sub(i::S,dims::Tuple{S,Vararg{S}}) where {S <: Integer}

    if i < 1 || i > prod(dims)
        error("index is out of bounds.")
    end

    subs = Tuple(CartesianIndices(dims)[i])

    return subs

end

function compute_variances(soln::FirstOrderSolutionStoch)

    hx = soln.hx
    k = soln.k
    gx = soln.gx
    sigma = soln.sigma

    var_states = dlyap(hx,k*sigma*k')
    var_jumps  = gx*var_states*gx'

  return var_states,var_jumps

end

function compute_chebyshev_integrals(eps_nodes,eps_weights,nodes,order,rho,sigma)

  # A simplified implementation of the integration described in Judd et al (2017, section 5.1);
  # it uses the mean integral across the nodes.
  # Using integrals[i] = sum(exp.(sqrt(2)*sigma*(i-1)*eps_nodes).*eps_weights)*pi^(-1/2)
  # leads to something much simplier and equally accurate, but is less recognizably appropriate
  # for the case where the shocks are AR(1) processes and ordinary polynomials are not being used.

  terms_num  = Array{Float64}(undef,length(eps_nodes))
  integrals  = Array{Float64}(undef,order+1,length(nodes))
  integrals2 = Array{Float64}(undef,order+1)
  for i = 1:(order+1)
    integrals2[i] = sum(exp.(sqrt(2)*sigma*(i-1)*eps_nodes).*eps_weights)*pi^(-1/2)
    for j = 1:length(nodes)
      terms_num     .= rho*nodes[j] .+ sqrt(2)*sigma*eps_nodes
      terms_den      = rho*nodes[j]
      terms_num     .= chebyshev_polynomial(i,terms_num)[:,i]
      terms_den      = chebyshev_polynomial(i,terms_den)[i]
      integrals[i,j] = sum((terms_num/terms_den).*eps_weights)*pi^(-1/2)
     end
  end

  nodetoosmall = abs.(nodes) .< sqrt(eps())
  if sum(nodetoosmall) > 0
    if length(nodes) == 1
      integrals[:,1] .= integrals2
    else
      for i = 1:length(nodes)
        if nodetoosmall[i] == 1
          if i == 1
            integrals[:,i] .= integrals[:,i+1]
        elseif i == length(nodes)
            integrals[:,i] .= integrals[:,i-1]
          else
            integrals[:,i] .= (integrals[:,i-1]+integrals[:,i+1])/2
          end
        end
      end
    end
  end

  return reshape(sum(integrals,dims=2)/length(nodes),order+1)

#  integrals = Array{Float64}(undef,order+1)
#  for i = 1:(order+1)
#    integrals[i] = sum(exp.(sqrt(2)*sigma*(i-1)*eps_nodes).*eps_weights)*pi^(-1/2)
#  end
#  return integrals

end

function compute_smolyak_integrals(eps_nodes,eps_weights,nx,order,grid,RHO,sigma)

  integrals = ones(nx,order+1)
  for j = 1:size(sigma,2)
      nodes = unique(grid[:,j])
      integrals[j,:] .= compute_chebyshev_integrals(eps_nodes,eps_weights,nodes,order,RHO[j,j],sigma[j,j])
  end

#  for j = 1:size(sigma,2)
#      for i = 1:(order+1)
#          integrals[j,i] = sum(exp.(sqrt(2)*sigma[j,j]*(i-1)*eps_nodes).*eps_weights)*pi^(-1/2)
#      end
#  end

  return integrals

end

function weight_scale_factors(eps_nodes,eps_weights,multi_index,nx,grid,RHO,sigma)

  unique_multi_index = sort(unique(multi_index))
  unique_orders = SmolyakApprox.m_i(unique_multi_index).-1

  # Here we construct the base integrals

  base_integrals = Array{Array{Float64,2}}(undef,length(unique_orders))
  for i = 1:length(unique_orders)
    base_integrals[i] = compute_smolyak_integrals(eps_nodes,eps_weights,nx,unique_orders[i],grid,RHO,sigma)
  end

  # Compute the unique polynomial terms from the base polynomials

  unique_base_integrals = Array{Array{Float64,2}}(undef,length(unique_orders))
  for i = length(unique_orders):-1:2
    unique_base_integrals[i] = base_integrals[i][:,size(base_integrals[i-1],2)+1:end]
  end
  unique_base_integrals[1] = base_integrals[1]

  # Construct the first row of the interplation matrix

  new_integrals = unique_base_integrals[multi_index[1,1]][1,:]
  for i = 2:size(multi_index,2)
    new_integrals = kron(new_integrals,unique_base_integrals[multi_index[1,i]][i,:])
  end

  weight_scale_factor = copy(new_integrals)

  # Iterate over nodes, doing the above three steps at each iteration

  for j = 2:size(multi_index,1)

    new_integrals = unique_base_integrals[multi_index[j,1]][1,:]
    for i = 2:size(multi_index,2)
      new_integrals = kron(new_integrals,unique_base_integrals[multi_index[j,i]][i,:])
    end
    weight_scale_factor = [weight_scale_factor; new_integrals]

  end

  return weight_scale_factor

end

function scale_weights(weights,weight_scale_factor)

    scaled_weights = weights.*weight_scale_factor

    return scaled_weights

end

function compute_piecewise_linear_integrals(eps_nodes,eps_weights,sigma)

  integral = 1.0
  #integral = sum(exp.(sqrt(2)*sigma*eps_nodes).*eps_weights)*pi^(-1/2)

  return integral

end


"""
    eom_linear3d_nonlinearcoupling(x, p, n) -> Function

Equations of motion for a 3d linear system with nonlinear coupling [1].
The difference equations are

```math
\\begin{aligned}
x_1(t+1) &= a_1 x_1 (1-x_1(t))^2  e^{-x_2(t)^2} + 0.4 \\xi_{1}(t) \\
x_2(t+1) &= a_1 x_2 (1-x_2(t))^2  e^{-x_2(t)^2} + 0.4 \\xi_{2}(t) + b x_1 x_2 \\
x_3(t+1) &= a_3 x_3 (1-x_3(t))^2  e^{-x_3(t)^2} + 0.4 \\xi_{3}(t) + c x_{2}(t) \\
         &+ d x_{1}(t)^2.
\\end{aligned}
```

Here, ``\\xi_{1,2,3}(t)`` are independent normally distributed noise processes,
representing dynamical noise in the system, with zero mean and standard
deviations ``\\sigma_1``, ``\\sigma_2``, ``\\sigma_3``, respectively.

# References
1. Gourévitch, B., Le Bouquin-Jeannès, R., & Faucon, G. (2006). Linear and
nonlinear causality between signals: methods, examples and neurophysiological
applications. Biological Cybernetics, 95(4), 349–369.
"""
function eom_linear3d_nonlinearcoupling(x, p, n)
    x₁, x₂, x₃ = (x...,)
    a₁, a₂, a₃, b, c, d, σ₁, σ₂, σ₃ = (p...,)
    ξ₁ = rand(Normal(0, σ₁))
    ξ₂ = rand(Normal(0, σ₂))
    ξ₃ = rand(Normal(0, σ₃))

    dx₁ = a₁*x₁*(1-x₁)^2 * exp(-x₁^2) + 0.4*ξ₁
    dx₂ = a₂*x₂*(1-x₂)^2 * exp(-x₂^2) + 0.4*ξ₂ + b*x₁*x₂
    dx₃ = a₃*x₃*(1-x₃)^2 * exp(-x₃^2) + 0.4*ξ₃ + c*x₂ + d*x₁^2

    return SVector{3}(dx₁, dx₂, dx₃)
end

function linear3d_nonlinearcoupling(uᵢ, a₁, a₂, a₃, b, c, d, σ₁, σ₂, σ₃)
    p = [a₁, a₂, a₃, b, c, d, σ₁, σ₂, σ₃]
    s = DiscreteDynamicalSystem(eom_linear3d_nonlinearcoupling, uᵢ, p)
    return s
end

"""
    linear3d_nonlinearcoupling(;uᵢ = rand(3), σ₁ = 1.0, σ₂ = 1.0, σ₃ = 1.0,
        a₁ = 3.4, a₂ = 3.4, a₃ = 3.4, b = 0.5, c = 0.3, d = 0.5) -> DiscreteDynamicalSystem

A 3d linear system with nonlinear coupling [1]. The difference equations are

```math
\begin{aligned}
x_1(t+1) &= a_1 x_1 (1-x_1(t))^2  e^{-x_2(t)^2} + 0.4 \\xi_{1}(t) \\
x_2(t+1) &= a_1 x_2 (1-x_2(t))^2  e^{-x_2(t)^2} + 0.4 \\xi_{2}(t) + b x_1 x_2 \\
x_3(t+1) &= a_3 x_3 (1-x_3(t))^2  e^{-x_3(t)^2} + 0.4 \\xi_{3}(t) + c x_{2}(t) \\
         &+ d x_{1}(t)^2.
\end{aligned}
```

Here, ``\\xi_{1,2,3}(t)`` are independent normally distributed noise processes,
representing dynamical noise in the system, with zero mean and standard
deviations ``\\sigma_1``, ``\\sigma_2``, ``\\sigma_3``, respectively.

# References
1. Gourévitch, B., Le Bouquin-Jeannès, R., & Faucon, G. (2006). Linear and
nonlinear causality between signals: methods, examples and neurophysiological
applications. Biological Cybernetics, 95(4), 349–369.
"""
linear3d_nonlinearcoupling(;uᵢ = rand(3), σ₁ = 1.0, σ₂ = 1.0, σ₃ = 1.0,
        a₁ = 3.4, a₂ = 3.4, a₃ = 3.4, b = 0.5, c = 0.3, d = 0.5) =
    linear3d_nonlinearcoupling(uᵢ, a₁, a₂, a₃, b, c, d, σ₁, σ₂, σ₃)


"""
    trianglecircumcenter(pt1::Point, pt2::Point, pt3::Point)

Return the circumcenter of the triangle defined by `pt1`,
`pt2`, and `pt3`. The circumcenter is the center of a circle
that passes through the vertices of the triangle.

"""
function trianglecircumcenter(pt1::Point, pt2::Point, pt3::Point)
    d = 2(pt1.x * (pt2.y - pt3.y) + pt2.x * (pt3.y - pt1.y) + pt3.x * (pt1.y - pt2.y))
    x = ((pt1.x^2 + pt1.y^2) * (pt2.y - pt3.y) + (pt2.x^2 + pt2.y^2) * (pt3.y - pt1.y) + (pt3.x^2 + pt3.y * pt3.y) * (pt1.y - pt2.y)) / d
    y = ((pt1.x^2 + pt1.y^2) * (pt3.x - pt2.x) + (pt2.x^2 + pt2.y^2) * (pt1.x - pt3.x) + (pt3.x^2 + pt3.y * pt3.y) * (pt2.x - pt1.x)) / d
    return Point(x, y)
end

# """
#     trianglecircumcenter(pt1::Point, pt2::Point, pt3::Point)
# """
# function trianglecircumcenter(pt1::Point, pt2::Point, pt3::Point)
#     ap1, ap2 = perpendicular(pt1, pt2)
#     bp1, bp2 = perpendicular(pt2, pt3)
#     flag, ip = intersectionlines(ap1, ap2, bp1, bp2)
#     if flag
#         return ip
#     else
#         throw(error("trianglecircumcenter(): triangle is not valid"))
#     end
# end

"""
    triangleincenter(pt1::Point, pt2::Point, pt3::Point)

Return the incenter of the triangle defined by `pt1`, `pt2`,
and `pt3`. The incenter is the center of a circle inscribed
inside the triangle.
"""
function triangleincenter(pt1::Point, pt2::Point, pt3::Point)
    a = distance(pt2, pt3)
    b = distance(pt1, pt3)
    c = distance(pt1, pt2)
    x = ((a * pt1.x) + (b * pt2.x) + (c * pt3.x)) / (a + b + c)
    y = ((a * pt1.y) + (b * pt2.y) + (c * pt3.y)) / (a + b + c)
    return Point(x, y)
end

"""
    trianglecenter(pt1::Point, pt2::Point, pt3::Point)

Return the centroid of the triangle defined by `pt1`, `pt2`, and `pt3`.
"""
function trianglecenter(pt1::Point, pt2::Point, pt3::Point)
    return 1/3 * (pt1 + pt2  + pt3)
end

"""
    triangleorthocenter(pt1::Point, pt2::Point, pt3::Point)

Return the orthocenter of the triangle defined by `pt1`, `pt2`, and `pt3`.
"""
function triangleorthocenter(pt1::Point, pt2::Point, pt3::Point)
    return pt1 + pt2 + pt3 - 2trianglecircumcenter(pt1, pt2, pt3)
end


using Random, Statistics, Distributions, Plots; pyplot()
Random.seed!(0)

mu, sig = 50, 5
dist = Normal(mu, sig)
alpha = 0.01
nMax = 40

observations = rand(dist,1)
piLarray, piUarray = [], []

for _ in 2:nMax
    xNew = rand(dist)
    push!(observations,xNew)

    xbar, sd = mean(observations), std(observations)
    n = length(observations)
    tVal = quantile(TDist(n-1),1-alpha/2)
    delta = tVal * sd * sqrt(1+1/n)
    piL, piU = xbar - delta, xbar + delta

    push!(piLarray,piL); push!(piUarray,piU)
end

scatter(1:nMax, observations,
	c=:blue, msw=0, label="Observations")
plot!(2:nMax, piUarray,
	c=:red, shape=:xcross, msw=0, label="Prediction Interval")
plot!(2:nMax, piLarray,
	c=:red, shape=:xcross, msw=0, label="",
	ylims=(0,100), xlabel="Number of observations", ylabel="Value")


"""
    second-order cone
"""
function second_order_cone_projection(z)
    n = length(z)

    z0 = z[1]
    z1 = view(z, 2:n)

    if norm(z1) <= z0
        return z, true
    elseif norm(z1) <= -z0
        return zero(z), false
    else
        a = 0.5 * (1.0 + z0 / norm(z1))
        z_proj = zero(z)
        z_proj[1] = a * norm(z1)
        z_proj[2:end] = a * z1
        return z_proj, false
    end
end

function second_order_cone_product(z, s)
    n = length(z)
    SVector{n}([z' * s; z[1] * view(s, 2:n) + s[1] * view(z, 2:n)])
end

# check that second-order cone constraints are satisfied
function second_order_cone_check(x, idx_soc)
    for idx in idx_soc
        !second_order_cone_projection(view(x, idx))[2] && return false
    end
    return true
end

# check that inequality (non-negative orthant) constraints are satisfied
function inequality_check(x, idx_ineq)
    for i in idx_ineq
        if x[i] <= 0.0
            return false
        end
    end
    return true
end

function cone_check(x, idx_ineq, idx_soc)
    !inequality_check(x, idx_ineq) && return true
    !second_order_cone_check(x, idx_soc) && return true
end


import Remez
using SpecialFunctions: erfcx
using PyPlot

T = Float64

function g(y::T) where T <: AbstractFloat
    two = parse(T, "2")
    rootpi = sqrt(convert(T, π))
    if -one(T) ≤ y ≤ -one(T) + eps(T)
        return ( 1 / two - 1 / rootpi )
    elseif one(T) - eps(T) ≤ y ≤ one(T)
        return ( -1 + 1/rootpi ) / 2
    else
        x = (1+y) / (1-y)
        return ( 2*erfcx(x) / (1-y) - 1 ) / (1+y)
    end
end

function erfcx_approx(x)
    y = (x-1) / (x+1)
    return (1-y)*(1+(1+y)*g_approx(y)) / 2
end

n = 15
iterations = 6
pt, ddp, E = Remez.minimax(T, g, n, iterations)

g_approx(y) = Remez.Newton_poly(ddp[0:n], pt[0:n], y)

figure(1)
y = range(-one(T), one(T), length=201)

plot(y, g.(y))
grid(true)
xlabel(L"$y$")
ylabel(L"$g(y)$")

figure(2)
plot(y, g.(y)-g_approx.(y))
grid(true)
xlabel(L"$y$")
ylabel(L"error in $g$")

figure(3)
x = range(0, 10, length=201)
plot(x, erfcx_approx.(x) - erfcx.(x))
grid(true)
xlabel(L"$x$")
ylabel("error in erfcx")


using FdeSolver
using Plots

## inputs
tSpan = [0, 10]     # [intial time, final time]
x0 = [1, 1]             # intial value
β = 2            # order of the derivative
par = [16.0, 4.0] # [spring constant for a mass on a spring, inertial mass]
h = 0.01

## Equation
function F(t, n, β, x, par)

      K = par[1]
      m = par[2]

      - K ./ m .* x[n]

end

## Numerical solution
t, Yapp = FDEsolver(F, tSpan, x0, β, nothing, par, h = h)

#plot
plot(t, Yapp, linewidth = 5, title = "Simple harmonic motion (order=2)",
     xaxis = "Time (t)", yaxis = "x(t)", label = "Approximation")
a = x0[1] .* map(cos, sqrt(par[1] / par[2]) .* t) .+ x0[2] ./ sqrt(par[1] / par[2]) .* map(sin, sqrt(par[1] / par[2]) .* t)
plot!(t, a, lw = 3, ls = :dash, label = "Exact solution")


using SMM, DataFrames, CSV, HTTP, Random, LinearAlgebra, Statistics, Optim, DataFramesMeta, GLM

function PS7()
	#question 1
	url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2021/master/ProblemSets/PS1-julia-intro/nlsw88.csv"
	df = CSV.read(HTTP.get(url).body, DataFrames)
	X = [ones(size(df,1),1) df.age df.race.==1 df.collgrad.==1]
	y = log.(df.wage);

	function ols_gmm(beta, X, y)
		return (y.-X*beta)'*(y.-X*beta)
	end

	function ols_gmm_with_sigma(beta, X, y)
		g = vcat(y .- X*beta[1:end-1],( (size(y,1)-1)/(size(y,1)-size(X,2)) )*var(y .- X*beta[1:end-1]) .- beta[end]^2)
		return g'*g
	end

	beta_gmm = optimize(b -> ols_gmm(b, X, y), rand(size(X,2)), LBFGS(), Optim.Options(g_tol=1e-6, iterations=100_000))
	println(beta_gmm.minimizer)

	beta_gmm_sigma = optimize(b -> ols_gmm_with_sigma(b, X, y), rand(size(X,2)+1), LBFGS(), Optim.Options(g_tol=1e-6, iterations=100_000))
	println(beta_gmm_sigma.minimizer)

	#Question 2
	df = dropmissing(df, :occupation)
	df[df.occupation.>7,:occupation] .= 7
	X = [ones(size(df,1),1) df.age df.race.==1 df.collgrad.==1]
	y = df.occupation;

	#from PS2
	function mlogit_mle(alpha, X, y)

			K = size(X,2)
			J = length(unique(y))
			N = length(y)
			bigY = zeros(N,J)
			for j=1:J
				bigY[:,j] = y.==j
			end
			A = [reshape(alpha,K,J-1) zeros(K)]

			P = exp.(X*A) ./ sum.(eachrow(exp.(X*A)))

			return -sum( bigY.*log.(P) )
	end

	alpha_hat= optimize(a -> mlogit_mle(a, X, y), rand(6*size(X,2)), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))

	print(alpha_hat.minimizer)

	function mlogit_gmm(α, X, y)
			K = size(X,2)
			J = length(unique(y))
			N = length(y)
			bigY = zeros(N,J)
			for j=1:J
				bigY[:,j] = y.==j
			end
			bigα = [reshape(α,K,J-1) zeros(K)]

			P = exp.(X*bigα) ./ sum.(eachrow(exp.(X*bigα)))

			g = bigY[:] .- P[:]

			J = g'*I*g
			return J
	end

	α_hat_gmm = optimize(a -> mlogit_gmm(a, X, y), alpha_hat.minimizer, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=10_000, show_trace=true, show_every=100))
	print(α_hat_gmm.minimizer)

	td = TwiceDifferentiable(b -> mlogit_gmm(b, X, y), alpha_hat.minimizer; autodiff = :forward)
	α_hat_gmm_random=optimize(td, rand(6*size(X,2)), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=10_000, show_trace=true, show_every=100))

	print(α_hat_gmm_random.minimizer)

	#Question 3

	#Question4
	MA = SMM.parallelNormal()
	dc = SMM.history(MA.chains[1])
	dc = dc[dc[:accepted].==true, :]
	println(describe(dc))

	#Question 5
end

PS7()


module problem114

description = """
Counting block combinations I
A row measuring seven units in length has red blocks with a minimum length of three units placed on it,
 such that any two red blocks (which are allowed to be different lengths) are separated by at least one black square.
There are exactly seventeen ways of doing this.
How many ways can a row measuring fifty units in length be filled?
"""

function block_combinations(n :: Integer, memo = zeros(Int, n))
  # trivial cases
  if n < 1 return 1 end

  # Count the total with each block size, and block position,
  # recursively counting the ways to fill the remaining space at the end of the row.
  if memo[n] == 0
    total = 1
    for blocksize = 3:n, blockpos = 1:(n-blocksize+1)
      total += block_combinations(n - blockpos - blocksize, memo)
    end
    memo[n] = total
  end
  return memo[n]
end

using Base.Test

@test block_combinations(3) == 2 #empty or 1 block
@test block_combinations(4) == 4 #empty or length 3 at pos 1, or length 3 at pos 2, or length 4 at pos 1
@test block_combinations(7) == 17

end


using LinearAlgebra, LowRankApprox, Formatting, ImageFiltering, Images, Plots, ImageMagick, Colors, TestImages, ImageView, ImageTransformations

I=load("rgb.jpg");

G=Gray.(I);
# size(G)

A=imresize(G, ratio=1/3);
# rank(A), size(A), eltype(A)

(p,q)=size(A);
B=Array{Float64}(A);

using DSP, Wavelets
B1=dwt(B, wavelet(WT.haar), 1);
B2=dwt(B, wavelet(WT.haar), 2);
B3=dwt(B, wavelet(WT.haar), 3);
B4=dwt(B, wavelet(WT.haar), 4);

B1a=zeros(Float64, p, q);
B2a=zeros(Float64, p, q);
B3a=zeros(Float64, p, q);
B4a=zeros(Float64, p, q);

B1a[1:Int(p/2),1:Int(q/2)]=B1[1:Int(p/2),1:Int(q/2)];
B2a[1:Int(p/4),1:Int(q/4)]=B2[1:Int(p/4),1:Int(q/4)];
B3a[1:Int(p/8),1:Int(q/8)]=B3[1:Int(p/8),1:Int(q/8)];
B4a[1:Int(p/16),1:Int(q/16)]=B4[1:Int(p/16),1:Int(q/16)];
B1d=B1-B1a;
B2d=B2-B2a;
B3d=B3-B3a;
B4d=B4-B4a;

C1a=idwt(B1a, wavelet(WT.haar), 1);
C2a=idwt(B2a, wavelet(WT.haar), 2);
C3a=idwt(B3a, wavelet(WT.haar), 3);
C4a=idwt(B4a, wavelet(WT.haar), 4);
C1d=idwt(B1d, wavelet(WT.haar), 1);
C2d=idwt(B2d, wavelet(WT.haar), 2);
C3d=idwt(B3d, wavelet(WT.haar), 3);
C4d=idwt(B4d, wavelet(WT.haar), 4);

A1a=Gray.(C1a);
A1d=Gray.(C1d);
A2a=Gray.(C2a);
A2d=Gray.(C2d);
A3a=Gray.(C3a);
A3d=Gray.(C3d);
A4a=Gray.(C4a);
A4d=Gray.(C4d);

plot(Gray.(A),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 0: Original Grayscale Image")
savefig("grayscale_haar_level0.png")

plot(Gray.(B1),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 1 (haar)")
savefig("grayscale_haar_level1.png")

plot(Gray.(C1a),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 1 Approximation (haar)")
savefig("grayscale_haar_level1_a.png")

plot(Gray.(C2a),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 2 Approximation (haar)")
savefig("grayscale_haar_level2_a.png")

plot(Gray.(C3a),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 3 Approximation (haar)")
savefig("grayscale_haar_level3_a.png")

plot(Gray.(C4a),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 4 Approximation (haar)")
savefig("grayscale_haar_level4_a.png")

plot(Gray.(C4d),
    xaxis=false,
    xticks=false,
    yaxis=false,
    yticks=false,
    grid=false,
    title="Level 4 Detail (haar)")
savefig("grayscale_haar_level4_d.png")


"""
    MVCalc(μstar,μ,Σ)

Calculate the std and weights of a portfolio (with mean return μstar) on MVF of risky assets.

# Remark
- Only (λ,δ) and thus (w,stdRp) depend on μstar. We could therefore speed up the computations a bit
by doing the loop over different μstar values inside the function (and thus not recalculate Σ_1,a,b,c).
"""
function MVCalc(μstar,μ,Σ)
    n    = length(μ)
    Σ_1  = inv(Σ)
    a    = μ'Σ_1*μ
    b    = μ'Σ_1*ones(n)
    c    = ones(n)'Σ_1*ones(n)
    λ    = (c*μstar - b)/(a*c-b^2)
    δ    = (a-b*μstar)/(a*c-b^2)
    w    = Σ_1 *(μ*λ.+δ)
    StdRp = sqrt(w'Σ*w)
    return StdRp,w
end

"""
    MVCalcRf(μstar,μ,Σ,Rf)

Calculate the std and portfolio weights of a portfolio (with a given mean, μstar) on MVF of (risky assets,riskfree)
"""
function MVCalcRf(μstar,μ,Σ,Rf)
    μe    = μ .- Rf
    Σ_1   = inv(Σ)
    w     = (μstar-Rf)/(μe'Σ_1*μe) * Σ_1*μe
    StdRp = sqrt(w'Σ*w)
    return StdRp,w                    #std and portfolio weights
end

"""
    MVTangencyP(μ,Σ,Rf)

Calculate the tangency portfolio
"""
function MVTangencyP(μ,Σ,Rf)           #calculates the tangency portfolio
    n    = length(μ)
    μe   = μ .- Rf                    #expected excess returns
    Σ_1  = inv(Σ)
    w    = Σ_1 *μe/(ones(n)'Σ_1*μe)
    muT  = w'μ + (1-sum(w))*Rf
    StdT = sqrt(w'Σ*w)
    return w,muT,StdT                  #portolio weights, mean and std
end


function asinh(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && return x
    signbit(x) && return - asinh(-x)
    result = abs(x)
    result = result + sqrt(square(result) + 1.0)
    result = log(result)
    return copysign(result, x)
end

function acosh(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && !signbit(x) && return x
    (x < 1.0 || isinf(x)) && throw(DomainError("$x"))
    result = x + sqrt(square(x) - 1.0)
    result = log(result)
    return result
end

function atanh(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    (abs(x) > 1.0 || isinf(x)) && throw(DomainError("$x"))
    twox = DoubleFloat{T}(x.hi+x.hi, x.lo+x.lo)
    result = 1.0 + twox / (1.0 - x)
    result = log(result)
    result = DoubleFloat{T}(result.hi*0.5, result.lo*0.5)
    return result
end

function acsch(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && return zero(DoubleFloat{T})
    iszero(x) && return signbit(x) ? -DoubleFloat{T}(Inf) : DoubleFloat{T}(Inf)
    invx = inv(x)
    invx2 = inv(x*x)
    result = log(sqrt(1 + invx2) + invx)
    return result
end

function asech(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    (x < 0.0 || x > 1.0 || isinf(x)) && throw(DomainError("$x"))
    iszero(x) && return inf(DoubleFloat{T})
    isone(x) && return zero(DoubleFloat{T})
    invx = inv(x)
    result = fma(sqrt(invx + 1), sqrt(invx - 1), invx)
    return log(result)
end

function acoth(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && return copysign(zero(DoubleFloat{T}), x)
    isone(x) && return inf(DoubleFloat{T})
    isone(-x) && return -inf(DoubleFloat{T})
    (-1.0 < x < 1.0) && throw(DomainError("$x"))
    invx = inv(x)
    result = (log(1+invx) - log(1-invx))/2
    return result
end


## Julia program for Hierarchical Modeling of means and variances
## author: weiya <szcfweiya@gmail.com>
## date: 27 August, 2018

using Distributions
using SpecialFunctions
using StatsBase
using DelimitedFiles

function higibbs(Y, T, mu0 = 50.0, gamma20 = 25.0, nu0 = 1.0, sigma20 = 100.0, eta0 = 1.0, tau20 = 100.0, a = 1.0, b = 1/100.0, alpha = 1.0, NUMAX = 40)
    m = size(unique(Y[:,1]), 1)
    # starting value
    ybar = ones(m)
    sv = ones(m)
    n = ones(m)
    for j = 1:m
        yj = Y[ [Y[i,1] == j for i = 1:end], 2]
        ybar[j] = mean(yj)
        sv[j] = var(yj)
        n[j] = size(yj, 1)
    end
    theta = ybar
    sigma2 = copy(sv)
#    sigma20 = 1 / mean(sigma2)
#    nu0 = 2 * mean(sigma2)^2 / var(sigma2)
    mu = mean(theta)
    tau2 = var(theta)

    THETA = ones(T, m)
    SIGMA2 = ones(T, m)
    # mu tau2 sigma20 nu0
    MTSN = ones(T, 4)

    for t = 1:T
        # sample mu
        varmu = 1 / (m / tau2 + 1 / gamma20)
        meanmu = varmu * (m * mean(theta) / tau2 + mu0 / gamma20)
        rnorm = Normal(meanmu, sqrt(varmu))
        mu = rand(rnorm, 1)[1]

        # sample 1/tau2
        shapetau = (eta0 + m) / 2
        ratetau = ( eta0 * tau20 + sum((theta .- mu).^2) ) / 2
        rgamma = Gamma(shapetau, 1/ratetau)
        tau2 = 1 / rand(rgamma, 1)[1]

        # sample theta
        for j = 1:m
            vartheta = 1 / (n[j] / sigma2[j] + 1 / tau2)
            meantheta = vartheta * ( n[j] * mean(Y[ [Y[i,1] == j for i = 1:end], 2]) / sigma2[j] + mu / tau2)
            rnorm = Normal(meantheta, sqrt(vartheta))
            theta[j] = rand(rnorm, 1)[1]
        end
        THETA[t, :] .= theta

        # sample sigma2
        for j = 1:m
            shapesig = (nu0 + n[j])/2
            yj = Y[ [Y[i,1] == j for i = 1:end], 2]
            ratesig = ( nu0*sigma20 + sum( (yj .- theta[j]).^2 ) )/2
            rgamma = Gamma(shapesig, 1/ratesig)
            sigma2[j] = 1 / rand(rgamma, 1)[1]
        end
        SIGMA2[t, :] .= sigma2

        # sample sigma20
        shapesig = a + 0.5 * m * nu0
        ratesig = b + 0.5 * nu0 * sum(1 ./ sigma2)
        rgamma = Gamma(shapesig, 1/ratesig)
        sigma20 = rand(rgamma, 1)[1]

        # sample nu0
        x = 1:NUMAX
        lpnu0 = ones(NUMAX)
        lpnu0 .= m * ( .5 * x .* log.(sigma20 * x / 2) .- lgamma.(x/2) ) .+ (x / 2 .+ 1) * sum(log.(1 ./ sigma2)) .- x .* (alpha + .5 * sigma20 * sum(1 ./ sigma2))
        #println(lpnu0)
        nu0 = sample(x, pweights(exp.(lpnu0 .- maximum(lpnu0))))
#        println(pweights(exp.(lpnu0 .- maximum(lpnu0))))

        # store results
        MTSN[t, :] .= [mu, tau2, sigma20, nu0]
    end
    return THETA, SIGMA2, MTSN, sv, n
end

# run
Y = readdlm("math-score-Y.csv")
THETA, SIGMA2, MTSN, sv, n = higibbs(Y, 5000)

using PyPlot
# histogram
plt[:subplot](221)
plt[:hist](MTSN[:,2])
ylabel(L"$\mu$")
plt[:subplot](222)
plt[:hist](MTSN[:,2])
ylabel(L"$\tau^2$")
plt[:subplot](223)
plt[:hist](MTSN[:,3])
ylabel(L"$\nu_0$")
plt[:subplot](224)
plt[:hist](MTSN[:,4])
ylabel(L"$\sigma_0^2$")
plt[:tight_layout]()
show()

# shrinkage
f, (ax1, ax2) = plt[:subplots](1, 2)
ax1[:scatter](sv, SIGMA2[end,:])
ax1[:plot](sv, sv)
ax1[:set_xlabel](L"$s^2$")
ax1[:set_ylabel](L"$\hat \sigma^2$")
ax2[:scatter](n, sv-SIGMA2[end,:])
ax2[:plot](n, zeros(size(n, 1)))
ax2[:set_xlabel]("sample size")
ax2[:set_ylabel](L"$s^2-\hat \sigma^2$")
plt[:tight_layout]()
show()


using ForwardDiff
function cournot_newton(eta, c1, c2, guess, tolerance) # note that guess is a 1*2 array
    profit_1(guess) = (sum(guess))^(-1/eta)*guess[1] - 0.5*c1*guess[1]^2
    profit_2(guess) = (sum(guess))^(-1/eta)*guess[2] - 0.5*c2*guess[2]^2

    diff = 1e3
    while diff > tolerance #here diff is the distance matrix
        println("Now the guess for q1 is $(guess[1]), the guess for q2 is $(guess[2]).")
        f1 = guess -> ForwardDiff.gradient(profit_1,guess) # calculate FOC w.r.t q1
        f2 = guess -> ForwardDiff.gradient(profit_2,guess) # calculate FOC w.r.t q2
        f1_prime = guess -> ForwardDiff.hessian(profit_1,guess) # calculate SOC w.r.t q1
        f2_prime = guess -> ForwardDiff.hessian(profit_2,guess) # calculate SOC w.r.t q2
        new_guess_q1 = guess[1] - f1(guess)[1] / f1_prime(guess)[1,1]
        new_guess_q2 = guess[2] - f2(guess)[2] / f2_prime(guess)[2,2]
        diff = sqrt((guess[1]-new_guess_q1)^2 + (guess[2]-new_guess_q2)^2) # eclidean distance from old guess
        guess = [new_guess_q1 new_guess_q2] # form new guess
    end
    println("The optimal quantity for firm 1 is $(guess[1]), for firm 2 is $(guess[2])")
end


#=

A Pythagorean triplet is a set of three natural numbers, a < b < c, for which,
a2 + b2 = c2

For example, 3^2 + 4^2 = 9 + 16 = 25 = 5^2.

There exists exactly one Pythagorean triplet for which a + b + c = 1000.
Find the product abc.

=#

p = 12

function pythagorean_triplet(perimeter)
    for b in 2:Int(ceil((perimeter-1)/2))
        for a in 1:b-1
            if 2*perimeter*(a+b) == perimeter^2 + 2*a*b
                #println(a, ", ", b, ", ", perimeter-a-b)
                return a, b, perimeter-a-b
            end
            #println(a, ", ", b, ", ", perimeter-a-b)
        end
    end
end

res = pythagorean_triplet(1000)
prod(res)


function mandelbrot(a)
    z = 0
    for i=1:50
        z = z^2 + a
    end
    return z
end

for y=1.0:-0.05:-1.0
    for x=-2.0:0.0315:0.5
        abs(mandelbrot(complex(x, y))) < 2 ? print("*") : print(" ")
    end
    println()
end

# Taken from: https://rosettacode.org/wiki/Mandelbrot_set#Julia


using DrWatson
@quickactivate "StatReth"

# %% 3.1
Pr_Positive_Vampire = 0.95
Pr_Positive_Mortal = 0.01
Pr_Vampire = 0.001
Pr_Positive = Pr_Positive_Vampire * Pr_Vampire + Pr_Positive_Mortal * (1 - Pr_Vampire)
Pr_Vampire_Positive = Pr_Positive_Vampire * Pr_Vampire / Pr_Positive


export NumberTheory

module NumberTheory

export Z_star, Eulerφ, continued_fraction, mod_inverse, rand_primeto, factor_a_power_b
export is_order, order_from_float, find_order

"""
    Z_star(N::Int) -> Vector

returns the Z* group elements of `N`, i.e. {x | gcd(x, N) == 1}
"""
Z_star(N::Int) = filter(i->gcd(i, N)==1, 0:N-1)
Eulerφ(N) = length(Z_star(N))

"""
    continued_fraction(ϕ, niter::Int) -> Rational

obtain `s` and `r` from `ϕ` that satisfies `|s/r - ϕ| ≦ 1/2r²`
"""
continued_fraction(ϕ, niter::Int) = niter==0 ? floor(Int, ϕ) : floor(Int, ϕ) + 1//continued_fraction(1/mod(ϕ, 1), niter-1)
continued_fraction(ϕ::Rational, niter::Int) = niter==0 || ϕ.den==1 ? floor(Int, ϕ) : floor(Int, ϕ) + 1//continued_fraction(1/mod(ϕ, 1), niter-1)

"""
    mod_inverse(x::Int, N::Int) -> Int

Return `y` that `(x*y)%N == 1`, notice the `(x*y)%N` operations in Z* forms a group and this is the definition of inverse.
"""
function mod_inverse(x::Int, N::Int)
    for i=1:N
        (x*i)%N == 1 && return i
    end
    throw(ArgumentError("Can not find the inverse, $x is probably not in Z*($N)!"))
end

"""
    is_order(r, x, N) -> Bool

Returns true if `r` is the order of `x`, i.e. `r` satisfies `x^r % N == 1`.
"""
is_order(r, x, N) = powermod(x, r, N) == 1

"""
    find_order(x::Int, N::Int) -> Int

Find the order of `x` by brute force search.
"""
function find_order(x::Int, N::Int)
    findfirst(r->is_order(r, x, N), 1:N)
end

"""
    rand_primeto(N::Int) -> Int

Returns a random number `2 ≦ x < N` that is prime to `N`.
"""
function rand_primeto(N::Int)
    while true
        x = rand(2:N-1)
        d = gcd(x, N)
        if d == 1
            return x
        end
    end
end

"""
    order_from_float(ϕ, x, L) -> Int

Estimate the order of `x` to `L`, `r`, from a floating point number `ϕ ∼ s/r` using the continued fraction method.
"""
function order_from_float(ϕ, x, L)
    k = 1
    rnum = continued_fraction(ϕ, k)
    while rnum.den < L
        r = rnum.den
        if is_order(r, x, L)
            return r
        end
        k += 1
        rnum = continued_fraction(ϕ, k)
    end
    return nothing
end

"""
    factor_a_power_b(N::Int) -> (Int, Int) or nothing

Factorize `N` into the power form `a^b`.
"""
function factor_a_power_b(N::Int)
    y = log2(N)
    for b = 2:ceil(Int, y)
        x = 2^(y/b)
        u1 = floor(Int, x)
        u1^b == N && return (u1, b)
        (u1+1)^b == N && return (u1+1, b)
    end
end

end


using LinearAlgebra, Optim

export tracenorm, tracedist, frobdist, isunitary, ishermitian_tol,
        isspecial, offdiag, jacobi_simultaneous_diag

ϵ = sqrt(eps(Float64))

"""
Compute the trace norm of a matrix ``A``
"""
function tracenorm(A::Matrix{T}) where T<:Number
    #Compute the trace norm of a matrix A
    if ishermitian(A)
        return sum(abs.(eigvals(A)))
    else
        return sum(svd(A).S)
    end
end

"""
Compute the trace distance between matrices ``A``, ``B``.
"""
function tracedist(A::Matrix{T}, B::Matrix{T}) where T<:Number
    #Compute trace distance between matrices A, B
    return 0.5*tracenorm(A-B)
end

"""
Compute the Frobenius distance between matrices ``A``, ``B``.
"""
function frobdist(A::Matrix{T}, B::Matrix{T}) where T<:Number
    return norm(A-B, p=2)
end

"""
Check if a matrix is unitary up to a tolerance ``\epsilon``.
"""
function isunitary(A::Matrix{T}; tol=ϵ) where T <: Number
    return tracenorm(A'A - I) < tol
end

"""
Check if a matrix is Hermitian up to a tolerance ``\epsilon``.
"""
function ishermitian_tol(A::Matrix{T}; tol=ϵ) where T <: Number
    return tracedist(A, Matrix(A')) < tol
end

"""
Check if a matrix is special up to a tolerance ``\epsilon``.
"""
function isspecial(A::Matrix{T}; tol=ϵ) where T <: Number
    return isapprox(det(A), 1.0, rtol=tol)
end

"""
Sum of the off-diagonal elements of a matrix.
"""
function offdiag(A::Matrix{T}) where T <: Number
    return sum(tril(A,-1).^2) + sum(triu(A,1).^2)
end

"""
Function which implements the Jacobi algorithm for simultaneous (possibly approxiamate)
diagonalization of two real, symmetric matrices A, B.

Algorithm is from:

V. Kuleshov, A. Chaganty, P. Liang. "Simultaneous diagonalization: the asymmetric,
low-rank, and noisy settings" arXiv: 1051.06318 (2015)

J.-F. Cardoso, A. Soulomiac. "Jacobi angles for simultaneous diagonalization"
SIAM J. Matrix Anal. Appl. 17(1), 161-164. (1996)

"""
function jacobi_simultaneous_diag(A::Matrix{T}, B::Matrix{T};
        tol=sqrt(eps(T)), maxiter=1000) where T<:Real

    @assert ishermitian_tol(A) && ishermitian_tol(B) "A, B must be symmetric!"

    (m,mm) = size(A)
    done = false

    U = Float64.(Matrix(I, (m,mm)))
    Λa = A
    Λb = B

    #Cost function for minimization of "off-diagonal-ness"
    #TODO: I was not able to succesfully implement the exact solution for the minimizer from the
    #Cardoso paper...

    cost(X,Y) = offdiag(X)+offdiag(Y)

    start_cost = cost(A,B)

    function minims_lsq(X, Y, i, j)
        function f(x)
            g = LinearAlgebra.Givens(i,j, cos(x), sin(x))
            return cost(g'*X*g, g'*Y*g)
        end
        xopt = Optim.minimizer(optimize(f, -π/4, π/4)) #are these the correct limits?
        return (cos(xopt), sin(xopt))
    end

    target = tol*(norm(A,2)+norm(B,2))

    n = 1
    while ((cost(Λa,Λb) > tol) && (n < maxiter))
        for j=1:m-1
            for k=j+1:m
                (c, s) = minims_lsq(Λa, Λb, j, k)
                Γ = LinearAlgebra.Givens(j, k, c, s)
                U = U*Γ
                Λa = Γ'*Λa*Γ
                Λb = Γ'*Λb*Γ
            end
        end
    end

    final_cost = cost(Λa,Λb)

    success = true

    if ((n >= maxiter) || (final_cost > tol))
        @warn "Was unable to find a simultaneous diagonalization: start
                cost function: $(start_cost), final: $(final_cost)."
        success = false
    end

    return (U, Diagonal(Λa), Diagonal(Λb), success)
end


using SymPy

function ry_mat(ang)
    return [SymPy.cos(ang/2) -SymPy.sin(ang/2); SymPy.sin(ang/2) SymPy.cos(ang/2)]
end

function rz_mat(ang)
    return [SymPy.exp(-1im*ang/2) 0; 0 SymPy.exp(1im*ang/2)]
end

function gen_mat(a,b,c,d)
    return SymPy.exp(1im*d).*ry_mat(c)*rz_mat(b)*ry_mat(a)
end

function _params_zyz(mat::Matrix{<:Symbol})
    """Return the euler angles and phase for the ZYZ basis."""
    # We rescale the input matrix to be special unitary (det(U) = 1)
    # This ensures that the quaternion representation is real
    coeff = sympy.sqrt(sympy.det(mat))
    phase = -sympy.atan2(real(coeff),imag(coeff))
    su_mat = coeff * mat  # U in SU(2)
    # OpenQASM SU(2) parameterization:
    # U[0, 0] = exp(-i(phi+lambda)/2) * cos(theta/2)
    # U[0, 1] = -exp(-i(phi-lambda)/2) * sin(theta/2)
    # U[1, 0] = exp(i(phi-lambda)/2) * sin(theta/2)
    # U[1, 1] = exp(i(phi+lambda)/2) * cos(theta/2)
    theta = 2 * sympy.atan2(abs(su_mat[2, 1]), abs(su_mat[1, 1]))
    phiplambda = 2 * sympy.atan2(real(su_mat[2, 2]), imag(su_mat[2, 2]))
    phimlambda = 2 * sympy.atan2(real(su_mat[2, 1]), imag(su_mat[2, 1]))
    phi = (phiplambda + phimlambda) / 2
    lam = (phiplambda - phimlambda) / 2
    return theta, phi, lam, phase
end

function _params_zyz(mat::Matrix{<:Number})
    """Return the euler angles and phase for the ZYZ basis."""
    # We rescale the input matrix to be special unitary (det(U) = 1)
    # This ensures that the quaternion representation is real
    coeff = sqrt(det(mat) +0im)
    phase = -atan(real(coeff), imag(coeff))
    su_mat = coeff * mat  # U in SU(2)
    # OpenQASM SU(2) parameterization:
    # U[0, 0] = exp(-i(phi+lambda)/2) * cos(theta/2)
    # U[0, 1] = -exp(-i(phi-lambda)/2) * sin(theta/2)
    # U[1, 0] = exp(i(phi-lambda)/2) * sin(theta/2)
    # U[1, 1] = exp(i(phi+lambda)/2) * cos(theta/2)
    theta = 2 * atan(abs(su_mat[2, 1]), abs(su_mat[1, 1]))
    phiplambda = 2 * atan(real(su_mat[2, 2]), imag(su_mat[2, 2]))
    phimlambda = 2 * atan(real(su_mat[2, 1]), imag(su_mat[2, 1]))
    phi = (phiplambda + phimlambda) / 2
    lam = (phiplambda - phimlambda) / 2
    return theta, phi, lam, phase
end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    variation(table)

Return the variation matrix `Τ` of the `table` such that:

- `Τ[i,j] = Var(log(x[i]/x[j]))` for `i, j = 1, ..., D`
"""
function variation(table)
  X = Tables.matrix(table)
  n = Tables.columnnames(table) |> collect
  D = size(X, 2)
  L = log.(X .+ eps())

  T = Matrix{Float64}(undef, D, D)
  for j in 1:D
    for i in j+1:D
      lr = L[:,i] .- L[:,j]
      T[i,j] = var(lr)
    end
    T[j,j] = 0.0
    for i in 1:j-1
      T[i,j] = T[j,i]
    end
  end

  AxisArray(T, row=n, col=n)
end

"""
    alrcov(table)

Return the log-ratio covariance matrix `Σ` of the `table` such that:

- `Σ[i,j] = cov(log(x[i]/x[D]), log(x[j]/x[D]))` for `i, j = 1, ..., d`
"""
function alrcov(table)
  alrtable = table |> ALR()

  Σ = cov(Tables.matrix(alrtable), dims=1)

  vars = Tables.columnnames(alrtable) |> collect
  AxisArray(Σ, row=vars, col=vars)
end

"""
    clrcov(table)

Return the centered log-ratio covariance matrix `Γ` of the `table` such that:

- `Γ[i,j] = cov(log(x[i]/g(x)), log(x[j]/g(x)))` for `i, j = 1, ..., D`,
where `g(x)` is the geometric mean.
"""
function clrcov(table)
  clrtable = table |> CLR()

  Γ = cov(Tables.matrix(clrtable), dims=1)

  vars = Tables.columnnames(clrtable) |> collect
  AxisArray(Γ, row=vars, col=vars)
end

"""
    lrarray(table)

Return the variation array `A` of the `table` such that:

- `A[i,j] = E[log(x[i]/x[j])]` for `i > j`
- `A[i,j] = Var(log(x[i]/x[j]))` for `i < j`
- `A[i,j] = 0` for `i = j`
"""
function lrarray(table)
  X = Tables.matrix(table)
  D = size(X, 2)
  L = log.(X .+ eps())

  A = Matrix{Float64}(undef, D, D)
  for i in 1:D
    for j in i+1:D
      lr = L[:,i] .- L[:,j]
      A[j,i] = mean(lr)
      A[i,j] = var(lr)
    end
    A[i,i] = 0.0
  end

  vars = Tables.columnnames(table) |> collect
  AxisArray(A, row=vars, col=vars)
end


using DifferentialEquations

function parameterized_lorenz(du,u,p,t)
  x,y,z = u
  σ,ρ,β = p
  du[1] = dx = σ*(y-x)
  du[2] = dy = x*(ρ-z) - y
  du[3] = dz = x*y - β*z
end

using ParameterizedFunctions
g = @ode_def begin
  dx = σ*(y-x)
  dy = x*(ρ-z) - y
  dz = x*y - β*z
end σ ρ β

# using LinearAlgebra
# f(du,u,p,t) = mul!(du,A,u)

u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)
p = [10.0,28.0,8/3]
prob = ODEProblem(g,u0,tspan,p)

using StaticArrays, DifferentialEquations
A  = @SMatrix [ 1.0  0.0 0.0 -5.0
                4.0 -2.0 4.0 -3.0
               -4.0  0.0 0.0  1.0
                5.0 -2.0 2.0  3.0]
u0 = @SMatrix rand(4,2)
tspan = (0.0,1.0)
f(u,p,t) = A*u
prob = ODEProblem(f,u0,tspan)
sol = solve(prob)
using Plots;
plot(sol)

# const κ = 1; const tau = 2 * pi
const out = zeros(2) # Define a cache variable
function delay_mathieu_model(du,u,h,p,t)
  δ, ϵ, b, κ = p
  h(out, p, t-tau)
  du[1] = - κ *u[1] - (δ+ϵ*cos(t)) *u[2] -b * out[2]
  du[2] = u[1]
end

p = [1.0,2.0,1,0.1]
tau=2*pi
T=2*pi
lags = [tau]
h(out, p, t) = (out.=1.0)
# h(out, p, t, deriv::Type{Val{i}})
tspan = (0.0,T)
u0 = [1.0,1.0]
prob = DDEProblem(delay_mathieu_model,h,tspan,p; constant_lags=lags)

alg = MethodOfSteps(Tsit5())
# alg = MethodOfSteps(BS3())
# alg = MethodOfSteps(Vern6())
sol = solve(prob,alg)#,reltol=1e-8,abstol=1e-8)

using Plots
plot(sol,vars=(0,2))

# plot(sol,linewidth=5,title="Solution to the linear ODE with a thick line",
#      xaxis="Time (t)",yaxis="u(t) (in μm)",label="My Thick Line!") # legend=false
# plot!(sol.t, t->0.5*exp(1.01t),lw=3,ls=:dash,label="True Solution!")
sol[2][3]
plot(sol,vars=(0,2))


# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module BinaryInteger

export ModuleBinaryInteger
export V001855, V003314, V033156, V054248, V061168, V083652, V097383, V123753
export V295513, BinaryIntegerLength, Bil

"""
For positive n, BinaryIntegerLength is ``⌊ log[2](n) ⌋ + 1``, BinaryIntegerLength(0) = 0.

* BinaryIntegerLength, Bil, V001855, V003314, V033156, V054248, V061168, V083652, V097383, V123753, V295513
"""
const ModuleBinaryInteger = ""

"""
Return the length of the binary extension of an integer ``n``, which is defined as ``0`` if ``n = 0`` and for ``n > 0`` as ``⌊ log[2](n) ⌋ + 1``.
"""
BinaryIntegerLength(n) = n == 0 ? 0 : floor(Int, log2(n)) + 1

"""
Alias for the function BinaryIntegerLength.
"""
Bil(n) = BinaryIntegerLength(n)

"""
Return ``n`` Bil``(n) - 2^{\\text{Bil}(n)}`` where Bil``(n)`` is the binary integer length of ``n``.
"""
V295513(n) = n*Bil(n) - 2^Bil(n)

"""
Maximal number of comparisons for sorting ``n`` elements by binary insertion.
"""
V001855(n) = V295513(n) + 1

"""
Return the sum of lengths of binary expansions of ``0`` through ``n``.
"""
V083652(n) = V295513(n+1) + 2

"""
Recurrence ``a(n) = a(n-1) + ⌊ a(n-1)/(n-1) ⌋ + 2`` for ``m ≥ 2`` and ``a(1) = 1``.
"""
V033156(n) = V295513(n) + 2n

"""
Binary entropy function: ``a(n) = n + `` min ``( a(k) + a(n-k) : 1 ≤ k ≤ n-1 )`` for ``n > 1,`` and ``a(1) = 0``.
"""
V003314(n) = V295513(n) + n

"""
Binary entropy: ``a(n) = n +`` min ``{ a(k) + a(n-k) : 1 ≤ k ≤ n-1 }.``
"""
V054248(n) = V295513(n) + n + rem(n, 2)

"""
Minimum total number of comparisons to find each of the values ``1`` through ``n`` using a binary search with ``3``-way comparisons.
"""
V097383(n) = V295513(n+1) - div(n-1, 2)

"""
Partial sums of the sequence ``⌊ log[2](n) ⌋``.
"""
V061168(n) = V295513(n+1) - n + 1

"""
Partial sums of the sequence of length of the binary expansion of ``2n+1``.
"""
V123753(n) = V295513(n+1) + n + 2

#START-TEST-########################################################

using Test, SeqTests

function test()
    @testset "BinaryInteger" begin

        @test V295513(0) == -1
        @test V295513(1) == -1
        @test V295513(2) == 0
        @test V295513(3) == 2

        if is_oeis_installed()
            V = [V001855, V003314, V033156, V054248, V061168, V097383]
            for v in V SeqTest(v , 'V', 1) end

            V = [V083652, V123753]
            for v in V SeqTest(v , 'V', 0) end
        end
    end
end

function demo()
    println([V295513(n) for n in 0:12])
    println([V123753(n) for n in 0:12])
    println([V001855(n) for n in 1:12])
    println([V083652(n) for n in 1:12])
    println([V033156(n) for n in 1:12])
    println([V003314(n) for n in 1:12])
    println([V054248(n) for n in 1:12])
    println([V097383(n) for n in 1:12])
    println([V061168(n) for n in 1:12])
end

"""
[V295513(k) for k in 0:100000]
    0.014412 seconds (2 allocations: 781.391 KiB)
"""
function perf()
    @time [V295513(k) for k in 0:100000]
end

function main()
    test()
    demo()
    perf()
end

main()

end # module


# this function shows the behavior of the
# likelihood, the log-likelihood and the average
# log-likelihood as the sample size increases
using Distributions
n = 10 # try increasing this to 50, 1000, etc.
p = 0.3
y = rand(Bernoulli(p),n)
f =  pdf.(Bernoulli(p), y)
@show n
@show ℒ = prod(f)
@show lnℒ = sum(log.(f))
@show avglnℒ = (1/n)*sum(log.(f))
nothing


using LPVSpectral, Statistics, Random
export Sinusoidal

"""
    Sinusoidal(fs)
Decompose a timeseries `s` into a **sum** `x + r`, where `x`
are sinusoidal components with the given frequencies `fs` that minimize
coefficients ``A, \\phi`` of the expression
```math
s \\approx A_0 + \\sum_i A_i \\cos(2\\pi f_i t + \\phi_i)
```
with ``A_0 = \\bar{s}`` the mean.

This method uses a new least-squares algorithm in frequency domain using the package
[LPVSpectral.jl](https://github.com/baggepinnen/LPVSpectral.jl), see[^Bagge2017].
It works for non-equispaced `t` axis (and also normal), is generally very accurate
(if choosen frequencies are not too close), but has performance scaling of
O(N^2.4) instead of O(n log(n)) of [`Fourier`](@ref).

Because it can work with arbitrary signal length the method always estimates the
zero-frequency Fourier component, and attributes it to `x`.
The fitted coefficients ``A, \\phi`` are available as fields `.A` and `.φ` of the struct
(first entry is zero-frequency component ``A_0, \\phi_0``).

[^Bagge2017]: F. Bagge Carlson et al., [Linear Parameter-Varying Spectral Decomposition](https://lup.lub.lu.se/search/publication/ac32368e-e199-44ff-b76a-36668ac7d595).
"""
struct Sinusoidal{F<:AbstractArray{<:AbstractFloat}} <: Decomposition
    fs::F
    A::Vector{Float64}
    φ::Vector{Float64}
end
Sinusoidal(fs) = Sinusoidal(fs, zeros(length(fs)+1), zeros(length(fs)+1))

function decompose(t, s, method::Sinusoidal)
    @assert 0 ∉ method.fs
    fs = [0, method.fs...]
    χ, = ls_spectral(s, t, fs)
    As = abs.(χ) ./ sqrt(2*length(fs))
    φs = angle.(χ)
    method.A .= As
    method.φ .= φs
    x = fill(As[1], length(s))
    for i in 2:length(fs)
        @. x += As[i]*cos(2π*fs[i]*t + φs[i])
    end
    return x, s .- x
end


#!/usr/bin/julia

# Trizen
# 15 November 2021
# https://github.com/trizen

# Compute the multiplicative order of `a` modulo `n`: znorder(a, n).
# This is the smallest positive integer k such that a^k == 1 (mod n).

using Primes

function divisors(n)

    d = Int64[1]

    for (p,e) in factor(n)
        t = Int64[]
        r = 1

        for i in 1:e
            r *= p
            for u in d
                push!(t, u*r)
            end
        end

        append!(d, t)
    end

    return sort(d)
end

function znorder(a, n)

    if isprime(n)
        for d in divisors(n-1)
            if (powermod(a, d, n) == 1)
                return d
            end
        end
    end

    f = factor(n)

    if (length(f) == 1)         # is prime power

        p = first(first(f))
        z = znorder(a, p)

        while (powermod(a, z, n) != 1)
            z *= p
        end

        return z
    end

    pp_orders = Int64[]

    for (p,e) in f
        push!(pp_orders, znorder(a, p^e))
    end

    return lcm(pp_orders)
end

isequal(znorder(97, factorial(14)), 25920)          || print("error")
isequal(znorder(53, factorial(15)), 2419200)        || print("error")
isequal(znorder(37, factorial(16)), 116121600)      || print("error")
isequal(znorder(31, factorial(17)), 6220800)        || print("error")
isequal(znorder(89, factorial(18)), 1045094400)     || print("error")
isequal(znorder(101,factorial(19)), 1254113280)     || print("error")
isequal(znorder(97, factorial(20)), 2239488000)     || print("error")

println(znorder(2, 341))        #> 10
println(znorder(97, 5040))      #> 12


# Polynomial kernel function

"""
    Poly <: Kernel

Polynomial kernel (covariance)
```math
k(x,x') = σ²(xᵀx' + c)ᵈ
```
with signal standard deviation ``σ``, additive constant ``c``, and degree ``d``.
"""
mutable struct Poly{T<:Real} <: Kernel
    "Constant"
    c::T
    "Signal variance"
    σ2::T
    "Degree of polynomial"
    deg::Int
    "Priors for kernel parameters"
    priors::Array
end

"""
Polynomial kernel function

    Poly(lc::Real, lσ::Real, deg::Int)

# Arguments
  - `lc::Real`: additive constant (given on log scale)
  - `lσ::Real`: signal standard deviation (given on log scale)
  - `deg::Int`: degree of polynomial
"""
Poly(lc::T, lσ::T, deg::Int) where T = Poly{T}(exp(lc), exp(2 * lσ), deg, [])

function KernelData(k::Poly, X1::AbstractMatrix, X2::AbstractMatrix)
    XtX=X1'*X2
	if X1==X2
    	LinearAlgebra.copytri!(XtX, 'U')
	end
    LinIsoData(XtX)
end
kernel_data_key(k::Poly, X1::AbstractMatrix, X2::AbstractMatrix) = "LinIsoData"

_cov(poly::Poly, xTy) = poly.σ2*(poly.c.+xTy).^poly.deg
function cov(poly::Poly, x::AbstractVector, y::AbstractVector)
    K = _cov(poly, dot(x,y))
end
function cov!(cK::AbstractMatrix, poly::Poly, X::AbstractMatrix, data::LinIsoData)
    cK .= _cov(poly, data.XtX)
end
@inline @inbounds function cov_ij(poly::Poly, X1::AbstractMatrix, X2::AbstractMatrix, i::Int, j::Int, dim::Int)
    return _cov(poly, dotij(X1, X2, i, j, dim))
end
@inline @inbounds function cov_ij(poly::Poly, X1::AbstractMatrix, X2::AbstractMatrix, data::LinIsoData, i::Int, j::Int, dim::Int)
    return _cov(poly, data.XtX[i, j])
end
cov(poly::Poly, X1::AbstractMatrix, X2::AbstractMatrix, data::LinIsoData) = _cov(poly, data.XtX)

get_params(poly::Poly) = [log(poly.c), log(poly.σ2) / 2]
get_param_names(poly::Poly) = [:lc, :lσ]
num_params(poly::Poly) = 2

function set_params!(poly::Poly, hyp::AbstractVector)
    length(hyp) == 2 || throw(ArgumentError("Polynomial function has two parameters"))
    poly.c = exp(hyp[1])
    poly.σ2 = exp(2 * hyp[2])
end

@inline dk_dlc(poly::Poly, xTy::Real) = poly.c*poly.deg*poly.σ2*(poly.c+xTy).^(poly.deg-1)
@inline dk_dlσ(poly::Poly, xTy::Real) = 2 * _cov(poly,xTy)
@inline function dKij_dθp(poly::Poly, X1::AbstractMatrix, X2::AbstractMatrix, i::Int, j::Int, p::Int, dim::Int)
    if p==1
        return dk_dlc(poly, dotij(X1,X2,i,j,dim))
    else
        return dk_dlσ(poly, dotij(X1, X2,i,j,dim))
    end
end
@inline function dKij_dθp(poly::Poly, X1::AbstractMatrix, X2::AbstractMatrix, data::LinIsoData, i::Int, j::Int, p::Int, dim::Int)
    if p==1
        return dk_dlc(poly, data.XtX[i,j])
    else
        return dk_dlσ(poly, data.XtX[i,j])
    end
end


# =========================================================================== #
# Error Functions
# =========================================================================== #

abstract type AbstractErrorFunction
end

struct LeastSquaresErrorFunction <: AbstractErrorFunction
end

function (error_fcn::LeastSquaresErrorFunction)(meas::T, x::T) where {T <: Number}
    dx = -meas + x
    return dx^2
end
function (error_fcn::LeastSquaresErrorFunction)(meas::Vector{T}, x::Vector{T}) where {T}
    dx = sum( (-meas .+ x).^2 )
    return dx
end
function (error_fcn::LeastSquaresErrorFunction)(meas, x::Vector{T}) where {T}
    dx = sum( (-meas .+ x).^2 )
    return dx
end

function differentiate(error_fcn::LeastSquaresErrorFunction, meas, x)
    return -meas + x
end

function eval(error_fcn::LeastSquaresErrorFunction, meas, x)
    return error_fcn(meas, x)
end

function gradient!(error_fcn::LeastSquaresErrorFunction, y_meas, y_out, delta_e)
    for kk = 1:length(delta_e)
        delta_e[kk] = differentiate(error_fcn, y_meas[kk], y_out[kk])
    end
    return delta_e
end

# function error_gradient!(nn::NeuralNetwork, y_meas, y_out, delta_e)
#     return gradient!(nn.error_function, y_meas, y_out, delta_e)
# end


# ------------------------------------------------------------------
# Copyright (c) 2017, Júlio Hoffimann Mendes <juliohm@stanford.edu>
# Licensed under the ISC License. See LICENCE in the project root.
# ------------------------------------------------------------------

"""
    Ellipsoidal(semiaxes, angles)

A distance defined by an ellipsoid with given `semiaxes` and rotation `angles`.

- For 2D ellipsoids, there are two semiaxes and one rotation angle.
- For 3D ellipsoids, there are three semiaxes and three rotation angles.

## Examples

2D ellipsoid making 45ᵒ with the horizontal axis:

```julia
julia> Ellipsoidal([1.0,0.5], [π/2])
```

3D ellipsoid rotated by 45ᵒ in the xy plane:

```julia
julia> Ellipsoidal([1.0,0.5,0.5], [π/2,0.0,0.0])
```
"""
struct Ellipsoidal{N,T} <: Metric
  dist::Mahalanobis{T}

  function Ellipsoidal{N,T}(semiaxes, angles) where {N,T<:Real}
    @assert length(semiaxes) == N "number of semiaxes must match spatial dimension"
    @assert all(semiaxes .> zero(T)) "semiaxes must be positive"
    @assert N ∈ [2,3] "dimension must be either 2 or 3"

    # scaling matrix
    Λ = Diagonal(one(T)./semiaxes.^2)

    # rotation matrix
    if N == 2
      θ = angles[1]

      cosθ = cos(θ)
      sinθ = sin(θ)

      P = [cosθ -sinθ
           sinθ  cosθ]
    end
    if N == 3
      θxy, θyz, θzx = angles

      cosxy = cos(θxy)
      sinxy = sin(θxy)
      cosyz = cos(θyz)
      sinyz = sin(θyz)
      coszx = cos(θzx)
      sinzx = sin(θzx)

      _1 = one(T)
      _0 = zero(T)

      Rxy = [cosxy -sinxy _0
             sinxy  cosxy _0
             _0     _0 _1]

      Ryz = [_1    _0     _0
             _0 cosyz -sinyz
             _0 sinyz  cosyz]

      Rzx = [ coszx _0 sinzx
             _0 _1    _0
             -sinzx _0 coszx]

      P = Rzx*Ryz*Rxy
    end

    # ellipsoid matrix
    Q = P*Λ*P'

    new(Mahalanobis(Q))
  end
end

Ellipsoidal(semiaxes::AbstractVector{T}, angles::AbstractVector{T}) where {T<:Real} =
  Ellipsoidal{length(semiaxes),T}(semiaxes, angles)

function Distances.evaluate(dist::Ellipsoidal{N,T}, a::AbstractVector, b::AbstractVector) where {N,T<:Real}
  evaluate(dist.dist, a, b)
end


using GeometryTypes
using SurfaceTopology
using LinearAlgebra

function quadraticform(vects,vnormal)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        curvatures[v] = H
    end
    return curvatures
end

### Testing
t = ( 1 + sqrt( 5 ) ) / 2;

vertices = Point{3,Float64}[
    [ -1,  t,  0 ], [  1, t, 0 ], [ -1, -t,  0 ], [  1, -t,  0 ],
    [  0, -1,  t ], [  0, 1, t ], [  0, -1, -t ], [  0,  1, -t ],
    [  t,  0, -1 ], [  t, 0, 1 ], [ -t,  0, -1 ], [ -t,  0,  1 ]
] ./ sqrt(1 + t^2)

faces = Face{3,Int64}[
    [1, 12, 6], [1, 6, 2], [1, 2, 8], [1, 8, 11], [1, 11, 12], [2, 6, 10], [6, 12, 5],
    [12, 11, 3], [11, 8, 7], [8, 2, 9], [4, 10, 5], [4, 5, 3], [4, 3, 7], [4, 7, 9],
    [4, 9, 10], [5, 10, 6], [3, 5, 12], [7, 3, 11], [9, 7, 8], [10, 9, 2]
]

### As paraboloid grows faster than sphere the estimated curvature for a coarse mesh would be lower
curvatures = meancurvature(vertices,faces)


#===============================================================================

Using all of the digits 1 through 9 and concatenating them freely to form
decimal integers, different sets can be formed. Interestingly with the set
{2,5,47,89,631}, all of the elements belonging to it are prime.

How many distinct sets containing each of the digits one through nine exactly
once contain only prime elements?

--------------------------------------------------------------------------------

The algorithm proposed is based on the association of a number to its digits.

Let n=d1...dn be a number made of n digits. Then

    encode(n) = 0 if n contains 0 or if a digit appears multiple times.
              = (Σ 2^di)/2

The division by 2 results from the fact the 0 is not allowed, so Σ 2^di is
always odd. Thus the range of encode(n) is 1 to 511 (9 bits set to 1).

Each prime up to 987654321 are encoded as explained above, and each code is
associated with the number of primes it corresponds to, excepted for primes
encoded as 0, which are discarded.

A similar code is also associated to a set of primes, by computing it over all
the digits of all the primes belonging to the set.

With this, a prime p can be added to the set S if encode(p) and encode(S) have
no digits in common, i.e. if encode(p) & encode(S) is null. The code of the new
set is encode(p) | encode(S).

Of course, this applies to all primes whose code is the same as encode(p).

Then the solution is built by exploring recursively the tree of the
possibilities.

===============================================================================#

import Primes

function encode(n::Int)
    result::UInt16 = 0

    while (n != 0)
        n, digit = divrem(n, 10)

        bit = 1 << digit

        if (digit == 0 || result & bit != 0)
            return 0
        end

        result |= bit
    end

    result >> 1
end

primes = Primes.primes(1, 987654321);

function compute_codes(primes)
    result = zeros(Int64, 511)

    for p in primes
        code = encode(p)
        if code != 0
            result[code] += 1
        end
    end

    result;
end

const allcodes = compute_codes(primes)

const fullcode = encode(123456789)

function compute_nbr_sets(code_idx, code_of_set, nbr_of_prime_sets)

    if code_of_set == fullcode
        return nbr_of_prime_sets
    end

    if code_idx == 512
        return 0
    end

    result = compute_nbr_sets(code_idx+1, code_of_set, nbr_of_prime_sets);

    nbr_assoc_primes = allcodes[code_idx]

    if nbr_assoc_primes != 0 && code_of_set & code_idx == 0
        result += compute_nbr_sets(code_idx+1,
                                   code_of_set|code_idx,
                                   nbr_assoc_primes * nbr_of_prime_sets)
    end

    result
end

println("Euler 118: ", compute_nbr_sets(1, 0, 1));


"""
    heapify(array::Array{Int64,1})

The heapify function merges the two subroutines for:
    1) Put elements of the array in heap order
    2) Repair the heap elements `i` whose root element, if is not at the index `largest`
For more information see: https://en.wikipedia.org/wiki/Heapsort

# Arguments
- `array::Array{Int64,1}`: Unsorted array of integers
- `n::Int64`: Length of the array
- `i::Int64`: Current index of the array

"""
function heapify(array::Array{Int64,1}, n::Int64, i::Int64)
    # Initialize largest as root and left and right subarrays
    largest = i
    # Julia notations, which starts with 1
    left = 2 * i
    right = 2 * i + 1

    # See if left child of root exists and is greater than root
    if left <= n && array[i] < array[left]
        largest = left
    end
    # See if right child of root exists and is greater than root
    if right <= n && array[largest] < array[right]
        largest = right
    end
    # Change root, if needed
    if largest != i
        array[i], array[largest] = array[largest], array[i]
        # Heapify the root again
        array = heapify(array, n, largest)
    end
    return array
end

"""
    heap_sorting(array::Array{Int64,1})

As a comparison-based sorting algorithm, the heapsort algorithm (HSA) divides its input
into a sorted and an unsorted region, and it iteratively shrinks the unsorted region by
extracting the largest element from it and inserting it into the sorted region. A specialty
is that the HSA keeps the unsorted region in a heap data structure to find the largest
element in each step more quickly. In more detail, in the first part of the HSA
(while-loop), the largest value has to be found and set to position one. In the second part
of the HSA (while-loop), the array's first and largest value has to be swap to the last
index of the array, and the swapping-procedure starts again for a new interval n-1. For
more information see: [https://en.wikipedia.org/wiki/Heapsort](https://en.wikipedia.org/wiki/Heapsort)

# Arguments
- `array::Array{Int64,1}`: Unsorted array of integers

# Examples
```julia-repl
julia> import ClassicAlgorithmsCollections
julia> arr = [64, 34, 25, 12, 22, 11, 90]
julia> ClassicAlgorithmsCollections.heap_sorting(arr)
[11, 12, 22, 25, 34, 64, 90]
```
"""
function heap_sorting(array::Array{Int64,1})
    # Base case. A array of zero or one elements is sorted, by definition.
    n = length(array)

    # Shrink the active space
    i = fld(n, 2)
    # Build a maxheap
    while i > 0
        array = heapify(array, n, i)
        i -= 1

    end

    # One by one extract elements
    i = n
    while i > 1
        array[i], array[1] = array[1], array[i]
        # Julia specific reducing effectiv index
        array = heapify(array, i - 1, 1)
        i -= 1
    end

    return array
end


##  Eratosthenes' prime number sieve
function primesieve(n::Int64)
    if n <= 1
        return Integer[]
    end

    p = [i for i in 1:2:n]
    q = length(p)
    p[1] = 2
    if n >= 9
        for k in 3:2:isqrt(n)
            if p[(k+1)÷2] != 0
                p[(k*k+1).÷2:k:q] .= 0
            end
        end
    end

    return p[p .> 0]
end

##  Find all prime numbers in  a given interval
function primes2(n::Int64, m::Int64)
    if n > m
        error("Argument 'm' must be bigger than 'n'.")
    end
    if m <= 1000
        P = primesieve(m)
        return P[P .>= n]
    end

    myPrimes = primesieve(isqrt(m))
    N = [n:m]
    l = length(N)  # m-n+1
    A = zeros(Int8, l)
    if n == 1
        A[1] = -1
    end

    for p in myPrimes
        r = n % p
        if r == 0
            i = 1
        else
            i = p - r + 1
        end
        if i <= l && N[i] == p
            i = i + p
        end
        while i <= l
            A[i] = 1
            i = i + p
        end
    end
    return N[A .== 0]
end

##  Return unique prime factors of n sorted
function primefactors(n::Int64)
    if n <= 0
        error("Argument 'n' to be factored must be a positive integer.")
    end

    f = Primes.factor(n)
    return sort(collect(keys(f)))
end

##  Find prime number following n
function nextprime2(n::Int64)
    if n <= 1; return 2; end
    if n == 2; return 3; end
    if iseven(n)
        n += 1
    else
        n += 2
    end
    if isprime(n); return(n); end

    if mod(n, 3) == 1
        a = 4; b = 2
    elseif mod(n, 3) == 2
        a = 2; b = 4
    else
        n += 2
        a = 2; b = 4
    end

    p = n
    while !isprime(p)
        p += a
        if isprime(p); break; end
        p += b
    end

    return p
end

## Find prime number preceeding n
function prevprime2(n::Int64)
    if n <= 2
        return Array(typeof(n), 0)
    elseif n <= 3
        return 2
    end

    if iseven(n)
        n -= 1
    else
        n -= 2
    end
    if isprime(n); return n; end

    if mod(n, 3) == 1
        a = 2; b = 4
    elseif mod(n, 3) == 2
        a = 4; b = 2
    else
        n -= 2
        a = 2; b = 4
    end

    p = n
    while !isprime(p)
        p -= a
        if isprime(p); break; end
        p -= b
    end

    return p
end

##  Find all twin primes
function twinprimes(n::Int64, m::Int64)
    P = primes2(n, m)
    inds = findall(diff(P) .== 2)

    return hcat(P[inds], P[inds.+1])
end

##  Coprimality
function coprime(n::Int64, m::Int64)
    if n == 0 && m == 0
        return false
    end

    if gcd(n, m) > 1
        false
    else
        true
    end
end

##  Solve linear modulo equations
function linmod(a::Int64, b::Int64, m::Int64)
    m > 1 || error("Argument 'm' must be an integer greater 1")
    T = typeof(m)
    g, e, f = gcdx(a, m)

    x = T[]
    if mod(b, g) == 0
        x0 = mod(e * div(b, g), m)
        for i = 0:(g-1)
            x = [x, mod(x0 + i * div(m, g) , m)]
        end
    end

    return x
end

##  Order of the element n (in the ring) modulo m
function ordermod(n::Int64, m::Int64)
    if n <= 0 || m <= 0
        error("Arguments 'n' and 'm' must be positive integers.")
    end
    if m == 1 || gcd(n, m) > 1; return 0; end
    if n == 1; return 1; end

    r = mod(n, m)
    if r == 0; return 0; end

    k = 1
    while r != 1
        r = mod(n*r, m)
        k += 1
    end

    return k
end

##  Find a primitive root modulo m
function primroot(m::Int64)
    if !isprime(m)
        error("Argument 'm' must be a prime number")
    end
    if m == 2; return 1; end

    P = primefactors(m-1)
    for r = 2:(m-1)
        not_found = true
        for p in P
            if powermod(r, div(m-1, p), m) == 1
                not_found = false
            end
        end
        if not_found
            return r
        end
    end

    return 0
end

##  Euler's Phi (or: totient) function
function eulerphi(n::Int64)
    if n <= 0
        error("Argument 'n' must be a (positive) natural number.")
    end

    m = n
    for p in primefactors(n)    # must be unique
        m = m - div(m, p)       # m = m * (1 - 1/p)
    end

    return Integer(round(m))
end


# This program solves the Laplace-Beltrami PDE on a unit sphere.
# In strong form, the problem can be stated as:
# find u s.t.
#   -Δ_S(u) = f on \Omega
# Where Δ_S is the spherical Laplacian defined as Δ_S = ∇_S⋅(∇_S(u)),
# with ∇_S() and  ∇_S⋅ being the spherical gradient and spherical
# divergence, resp.

# Given a known analytical solution u(x), the program manufactures f(x)
# by applying -Δ_S to u, and then performs an error convergence analysis.
# For the discretization of the sphere, it uses the so-called cubed sphere
# mesh

function u(x)
  sin(π*x[1])*cos(π*x[2])*exp(x[3])
end
function uθϕ(θϕ)
  u(θϕ2xyz(θϕ))
end
function f(x)
  -(laplacian_unit_sphere(uθϕ)(xyz2θϕ(x)))
end
function Gridap.∇(::typeof(u))
  gradient_unit_sphere(uθϕ)
end

function solve_laplace_beltrami(model,order,degree,ls=BackslashSolver())
    V = FESpace(model,ReferenceFE(lagrangian,Float64,order); conformity=:H1)
    U = TrialFESpace(V)

    Ω = Triangulation(model)
    dΩ = Measure(Ω,degree)

    a(u,v) = ∫(∇(v)⋅∇(u))dΩ
    b(v)   = ∫(v*f)dΩ

    op = AffineFEOperator(a,b,U,V)
    fels = LinearFESolver(ls)
    uh = solve(fels,op)

    e=u-uh
    ∇e = ∇(uh)-∇(u)∘xyz2θϕ

    # H1-norm
    # sqrt(sum(∫( e*e + (∇e)⋅(∇e) )dΩ))
    sqrt(sum(∫( (∇e)⋅(∇e) )dΩ))#,uh
end


#=
  Two children problem in Turing.jl

  https://robeastaway.com/blog/boy-girl
  """
  I have two children, at least one of whom is a boy.
  What is the chance that the second child is also a boy?

  ...

  The classic, surprising answer is that the chance of a second
  boy is 1/3, not 1/2 as most would expect.  Why?  Because if I
  have two children, there is an equal chance that they will be
  Boy-Boy, Boy-Girl, Girl-Boy or Girl-Girl.  There are three
  equally likely combinations in which at least one child is
  a boy (BB, BG and GB) and only in the first scenario is the
  other child also a boy.
  """

  We assume the inital probability of a boy/girl as 0.5.

  Here we model two similiar but different problems to show
  the difference between these two views of the problem.

  Cf ~/webppl/two_children_problem.wppl

=#

using Turing
include("jl_utils.jl")

#=
  First model:
  """
  I have two children, _at least one of whom_ is a boy.
  What is the chance that the second child is also a boy?
  """

  Probability of 2 boys: 1/3.
  Probability of 1 boy and 1 girl: 2/3.

  Distributions of variable twoBoys (num:0)
  0.00000 =>    6601  (0.660100)
  1.00000 =>    3399  (0.339900)

  Distributions of variable oneBoyOneGirl (num:0)
  1.00000 =>    6601  (0.660100)
  0.00000 =>    3399  (0.339900)

=#
@model function two_children_problem_1()
    boy = 1
    girl = 2

    child1 ~ Categorical([0.5,0.5])
    child2 ~ Categorical([0.5,0.5])

    # Note: it's important that we state the condition
    #   "at least one of whom is a boy"
    # as something like this.
    true ~ Dirac((child1 == boy) + (child2 == boy) >= 1)

    twoBoys ~ Dirac((child1 == boy) + (child2 == boy) == 2)
    oneBoyOneGirl ~ Dirac((child1 == boy) + (child2 == boy) == 1)

end

println("\nModel 1")
model = two_children_problem_1()
# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, IS(), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, NUTS(), 10_000)
# chns = sample(model, HMC(0.1,10), 10_000)

# display(chns)
show_var_dist_pct(chns,:twoBoys)
show_var_dist_pct(chns,:oneBoyOneGirl)

#=
 Second model:
  """
  I have two children, the eldest is a boy.
  What is the chance that the second child is also a boy?
  """

  Probability that the second child is also a boy: 1/2
  Probability that the second child is a girl: 1/2

  Distributions of variable twoBoys (num:0)
  0.00000 =>    5067  (0.506700)
  1.00000 =>    4933  (0.493300)
  Distributions of variable oneBoyOneGirl (num:0)
  1.00000 =>    5067  (0.506700)
  0.00000 =>    4933  (0.493300)

=#
@model function two_children_problem_2()
    boy = 1
    girl = 2

    child1 ~ Categorical([0.5,0.5])
    child2 ~ Categorical([0.5,0.5])

    # "The eldest is a boy"
    true ~ Dirac(child1 == boy)

    twoBoys ~ Dirac((child1 == boy) + (child2 == boy) == 2)
    oneBoyOneGirl ~ Dirac((child1 == boy) + (child2 == boy) == 1)

end

println("\nModel 2")
model = two_children_problem_2()
# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, IS(), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, NUTS(), 10_000)
# chns = sample(model, HMC(0.1,10), 10_000)

# display(chns)
show_var_dist_pct(chns,:twoBoys)
show_var_dist_pct(chns,:oneBoyOneGirl)


"""
Discrete Laplacian in d dimensions and n equidistant discretization points in each direction
"""
function Lap(n::Integer,d::Integer) #returns the tensor of the discrete Laplacian in a box [0,1]^d with n equidistant discretization points in each direction
    A = zeros(n^d,n^d)
    for j in 0:n^d-1
        J = digits(j,base=n,pad=d) #corresponding index as d-tuples
        for k in 0:d-1
            if J[k+1]==0
                A[j+1,j+1] = 2*d
                A[j+1,j+n^k+1] = -1
            elseif J[k+1]==n-1
                A[j+1,j+1] = 2*d
                A[j+1,j-n^k+1] = -1
            else
                A[j+1,j+1] = 2*d
                A[j+1,j-n^k+1] = -1
                A[j+1,j+n^k+1] = -1
            end
        end
    end
    return A
end


"""
    legendre(x, b)

Compute Legendre polynomials.

Legendre polynomials are Jacobi polynomials with parameters a and b equals to zero.

Same interface as `jacobi`. `dlegendre` computes the derivative of the Legendre
polynomial.

# Example
```julia-repl
julia> x = legendre(0.4, 5)
0.27063999999999994

julia> x = legendre(2//5, 5)
3383//12500

julia> x = dlegendre(0.4, 5)
-1.3170000000000002

julia> x = dlegendre(2//5, 5)
-1317//1000
```
"""
function legendre(x, n)

    if n==0
        return one(x)
    elseif n==1
        return(x)
    end

    p0 = one(x)
    p1 = x

    for i = 2:n
        p2 = ( (2i-1)*x*p1 - (i-1)*p0 ) / i
        p0 = p1
        p1 = p2
    end

    return p1
end

dlegendre(x, n) = djacobi(x, n)

legendre_zeros(m) = jacobi_zeros(m)

@doc (@doc legendre) dlegendre
@doc (@doc jacobi_zeros!) legendre_zeros


"""
`prepare_interpolation(nodes::Matrix{T}, kernel::RK = RK_H0()) where {T <: AbstractFloat, RK <: ReproducingKernel_0}`

Prepare the interpolating normal spline by constructing and factoring Gram matrix of the problem.
Initialize the `NormalSpline` object.
# Arguments
- `nodes`: function value interpolation nodes.
           This should be an `n×n_1` matrix, where `n` is dimension of the sampled space and
           `n_1` is the number of function value nodes. It means that each column in the matrix defines one node.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H0` if the spline is constructing as a continuous function,
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: prepared `NormalSpline` object.
"""
function prepare_interpolation(nodes::Matrix{T},
                               kernel::RK = RK_H0()
                              ) where {T <: AbstractFloat, RK <: ReproducingKernel_0}
     spline = _prepare_interpolation(nodes, kernel)
     return spline
end

"""
`construct_interpolation(spline::NormalSpline{T, RK}, values::Vector{T}) where {T <: AbstractFloat, RK <: ReproducingKernel_0}`

Construct the interpolating normal spline by calculating its coefficients and
completely initializing the `NormalSpline` object.
# Arguments
- `spline`: the prepared `NormalSpline` object.
- `values`: function values at interpolation nodes.

Return: constructed `NormalSpline` object.
"""
function construct_interpolation(spline::NormalSpline{T, RK},
                                 values::Vector{T}
                                ) where {T <: AbstractFloat, RK <: ReproducingKernel_0}
    spline = _construct_interpolation(spline, values)
    return spline
end

"""
`interpolate(nodes::Matrix{T}, values::Vector{T}, kernel::RK = RK_H0()) where {T <: AbstractFloat, RK <: ReproducingKernel_0}`

Prepare and construct the interpolating normal spline.
# Arguments
- `nodes`: function value interpolation nodes.
           This should be an `n×n_1` matrix, where `n` is dimension of the sampled space
           and `n_1` is the number of function value nodes.
           It means that each column in the matrix defines one node.
- `values`: function values at interpolation nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H0` if the spline is constructing as a continuous function,
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.
Return: constructed `NormalSpline` object.
"""
function interpolate(nodes::Matrix{T},
                     values::Vector{T},
                     kernel::RK = RK_H0()
                    ) where {T <: AbstractFloat, RK <: ReproducingKernel_0}
     spline = _prepare_interpolation(nodes, kernel)
     spline = _construct_interpolation(spline, values)
     return spline
end

########

"""
`prepare_interpolation(nodes::Matrix{T}, d_nodes::Matrix{T}, es::Matrix{T}, kernel::RK = RK_H1())
                       where {T <: AbstractFloat, RK <: ReproducingKernel_1}`

Prepare the spline by constructing and factoring a Gram matrix of the interpolation problem.
Initialize the `NormalSpline` object.
# Arguments
- `nodes`: function value interpolation nodes.
           This should be an `n×n_1` matrix, where `n` is dimension of the sampled space and
           `n_1` is the number of function value nodes.
            It means that each column in the matrix defines one node.
- `d_nodes`: function directional derivatives nodes.
             This should be an `n×n_2` matrix, where `n` is dimension of the sampled space and
             `n_2` is the number of function directional derivative nodes.
- `es`: Directions of the function directional derivatives.
        This should be an `n×n_2` matrix, where `n` is dimension of the sampled space and
        `n_2` is the number of function directional derivative nodes.
        It means that each column in the matrix defines one direction of the function directional derivative.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: prepared `NormalSpline` object.
"""
function prepare_interpolation(nodes::Matrix{T},
                               d_nodes::Matrix{T},
                               es::Matrix{T},
                               kernel::RK = RK_H1()
                              ) where {T <: AbstractFloat, RK <: ReproducingKernel_1}
     spline = _prepare_interpolation(nodes, d_nodes, es, kernel)
     return spline
end

"""
`construct_interpolation(spline::NormalSpline{T, RK}, values::Vector{T}, d_values::Vector{T})
                         where {T <: AbstractFloat, RK <: ReproducingKernel_1}`

Construct the interpolating normal spline by calculating its coefficients and
completely initializing the `NormalSpline` object.
# Arguments
- `spline`: the prepared `NormalSpline` object.
- `values`: function values at interpolation nodes.
- `d_values`: function directional derivative values at function derivative nodes.

Return: constructed `NormalSpline` object.
"""
function construct_interpolation(spline::NormalSpline{T, RK},
                                 values::Vector{T},
                                 d_values::Vector{T}
                                ) where {T <: AbstractFloat, RK <: ReproducingKernel_1}
     spline = _construct_interpolation(spline, values, d_values)
     return spline
end

"""
`interpolate(nodes::Matrix{T}, values::Vector{T}, d_nodes::Matrix{T}, es::Matrix{T}, d_values::Vector{T}, kernel::RK = RK_H1())
             where {T <: AbstractFloat, RK <: ReproducingKernel_1}`

Prepare and construct the interpolating normal spline.
# Arguments
- `nodes`: function value interpolation nodes.
           This should be an `n×n_1` matrix, where `n` is dimension of the sampled space
           and `n_1` is the number of function value nodes.
           It means that each column in the matrix defines one node.
- `values`: function values at interpolation nodes.
- `d_nodes`: function directional derivative nodes.
            This should be an `n×n_2` matrix, where `n` is dimension of the sampled space and
            `n_2` is the number of function directional derivative nodes.
- `es`: Directions of the function directional derivatives.
       This should be an `n×n_2` matrix, where `n` is dimension of the sampled space and
       `n_2` is the number of function directional derivative nodes.
       It means that each column in the matrix defines one direction of the function directional derivative.
- `d_values`: function directional derivative values at function derivative nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: constructed `NormalSpline` object.
"""
function interpolate(nodes::Matrix{T},
                     values::Vector{T},
                     d_nodes::Matrix{T},
                     es::Matrix{T},
                     d_values::Vector{T},
                     kernel::RK = RK_H1()
                    ) where {T <: AbstractFloat, RK <: ReproducingKernel_1}
     spline = _prepare_interpolation(nodes, d_nodes, es, kernel)
     spline = _construct_interpolation(spline, values, d_values)
     return spline
end

############################## One-dimensional case

"""
`prepare_interpolation(nodes::Vector{T}, kernel::RK = RK_H0()) where {T <: AbstractFloat, RK <: ReproducingKernel_0}`

Prepare the 1D interpolating normal spline by constructing and factoring a Gram matrix of the problem.
Initialize the `NormalSpline` object.
# Arguments
- `nodes`: function value interpolation nodes.
           This should be an `n_1` vector where `n_1` is the number of function value nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H0` if the spline is constructing as a continuous function,
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: prepared `NormalSpline` object.
"""
function prepare_interpolation(nodes::Vector{T},
                               kernel::RK = RK_H0()
                              ) where {T <: AbstractFloat, RK <: ReproducingKernel_0}
     spline = _prepare_interpolation(Matrix(nodes'), kernel)
     return spline
end

"""
`interpolate(nodes::Vector{T}, values::Vector{T}, kernel::RK = RK_H0()) where {T <: AbstractFloat, RK <: ReproducingKernel_0}`

Prepare and construct the 1D interpolating normal spline.
# Arguments
- `nodes`: function value interpolation nodes.
           This should be an `n_1` vector where `n_1` is the number of function value nodes.
- `values`: function values at `n_1` interpolation nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H0` if the spline is constructing as a continuous function,
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: constructed `NormalSpline` object.
"""
function interpolate(nodes::Vector{T},
                     values::Vector{T},
                     kernel::RK = RK_H0()
                    ) where {T <: AbstractFloat, RK <: ReproducingKernel_0}
     spline = _prepare_interpolation(Matrix(nodes'), kernel)
     spline = _construct_interpolation(spline, values)
     return spline
end

"""
`prepare_interpolation(nodes::Vector{T}, d_nodes::Vector{T}, kernel::RK = RK_H1()) where {T <: AbstractFloat, RK <: ReproducingKernel_1}`

Prepare the 1D interpolating normal spline by constructing and factoring a Gram matrix of the problem.
Initialize the `NormalSpline` object.
# Arguments
- `nodes`: function value interpolation nodes.
           This should be an `n_1` vector where `n_1` is the number of function value nodes.
- `d_nodes`: The function derivatives nodes.
             This should be an `n_2` vector where `n_2` is the number of function derivatives nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: prepared `NormalSpline` object.
"""
function prepare_interpolation(nodes::Vector{T},
                               d_nodes::Vector{T},
                               kernel::RK = RK_H1()
                              ) where {T <: AbstractFloat, RK <: ReproducingKernel_1}

     es = ones(T, length(d_nodes))
     spline = _prepare_interpolation(Matrix(nodes'), Matrix(d_nodes'), Matrix(es'), kernel)
     return spline
end

"""
`interpolate(nodes::Vector{T}, values::Vector{T}, d_nodes::Vector{T}, d_values::Vector{T},
             kernel::RK = RK_H1())
             where {T <: AbstractFloat, RK <: ReproducingKernel_1}`

prepare_interpolation and construct the 1D interpolating normal spline.
# Arguments
- `nodes`: function value interpolation nodes.
           This should be an `n_1` vector where `n_1` is the number of function value nodes.
- `values`: function values at `nodes` nodes.
- `d_nodes`: The function derivatives nodes.
             This should be an `n_2` vector where `n_2` is the number of function derivatives nodes.
- `d_values`: function derivative values at function derivative nodes.
- `kernel`: reproducing kernel of Bessel potential space the normal spline is constructed in.
            It must be a struct object of the following type:
              `RK_H1` if the spline is constructing as a differentiable function,
              `RK_H2` if the spline is constructing as a twice differentiable function.

Return: constructed `NormalSpline` object.
"""
function interpolate(nodes::Vector{T},
                     values::Vector{T},
                     d_nodes::Vector{T},
                     d_values::Vector{T},
                     kernel::RK = RK_H1()
                    ) where {T <: AbstractFloat, RK <: ReproducingKernel_1}
     es = ones(T, length(d_nodes))
     spline = _prepare_interpolation(Matrix(nodes'), Matrix(d_nodes'), Matrix(es'), kernel)
     spline = _construct_interpolation(spline, values, d_values)
     return spline
end


#=

# 107: 1D Nonlinear Storage
([source code](SOURCE_URL))

This equation comes from the transformation of the nonlinear diffuision equation.
```math
\partial_t u^\frac{1}{m} -\Delta u = 0
```
in $\Omega=(-1,1)$ with homogeneous Neumann boundary conditions.
We can derive an exact solution from the Barenblatt solution of the previous
example.

=#

module Example107_NonlinearStorage1D
using Printf
using VoronoiFVM
using ExtendableGrids
using GridVisualize

function barenblatt(x,t,m)
    tx=t^(-1.0/(m+1.0))
    xx=x*tx
    xx=xx*xx
    xx=1- xx*(m-1)/(2.0*m*(m+1));
    if xx<0.0
        xx=0.0
    end
    return tx*xx^(1.0/(m-1.0))
end

function main(;n=20,m=2.0,Plotter=nothing,verbose=false, unknown_storage=:sparse,tend=0.01,tstep=0.0001)

    ## Create a one-dimensional discretization
    h=1.0/convert(Float64,n/2)
    X=collect(-1:h:1)
    grid=VoronoiFVM.Grid(X)

    ## Flux function which describes the flux
    ## between neigboring control volumes
    function flux!(f,u0,edge)
        u=unknowns(edge,u0)
        f[1]=u[1,1]-u[1,2]
    end

    ϵ=1.0e-10

    ## Storage term
    ## This needs to be regularized as its derivative
    ## at 0 is infinity
    function storage!(f,u,node)
        f[1]=(ϵ+u[1])^(1.0/m)
    end

    ## Create a physics structure
    physics=VoronoiFVM.Physics(
        flux=flux!,
        storage=storage!)

    ## Create a finite volume system - either
    ## in the dense or  the sparse version.
    ## The difference is in the way the solution object
    ## is stored - as dense or as sparse matrix
    sys=VoronoiFVM.System(grid,physics,unknown_storage=unknown_storage)

    ## Add species 1 to region 1
    enable_species!(sys,1,[1])

    ## Create a solution array
    inival=unknowns(sys)
    solution=unknowns(sys)
    t0=0.001

    ## Broadcast the initial value
    inival[1,:].=map(x->barenblatt(x,t0,m)^m,X)
    solution.=inival

    ## Create solver control info
    control=VoronoiFVM.NewtonControl()
    control.verbose=verbose
    time=t0
    p=GridVisualizer(Plotter=Plotter,layout=(2,1))
    while time<tend
        time=time+tstep
        solve!(solution,inival,sys,control=control,tstep=tstep)
        inival.=solution
        scalarplot!(p[1,1],grid,solution[1,:],title=@sprintf("numerical, t=%.5f",time),clear=true)
        scalarplot!(p[2,1],grid,map(x->barenblatt(x,time,m)^m,grid),title=@sprintf("exact, t=%.4f",time),clear=true)
        sleep(1.0e-5)
#        yield()
    end
    return sum(solution)
end

function test()
    testval=173.84998139534142
    main(unknown_storage=:sparse) ≈ testval && main(unknown_storage=:dense) ≈ testval
end

end