det22(a,b,c,d) = a*d-c*b
det33(a,b,c) = @inbounds a[1]*det22(b[2],c[2],b[3],c[3]) - b[1]*det22(a[2],c[2],a[3],c[3]) + c[1]*det22(a[2],b[2],a[3],b[3])

"""
  solid_angle(topo, P, he, p)

the solid angle for the signed surface area of the triangle of he projected onto unit sphere centred at p
"""
function solid_angle(topo::Topology, P::TA, h::HalfEdgeHandle, p::T) where {T, TA<:Vector{T}}
  if isboundary(topo, h)
    return 0.0
  end

  @inbounds a = P[head(topo, h)] - p
  @inbounds b = P[head(topo, next(topo, h))] - p
  @inbounds c = P[head(topo, next(topo, next(topo, h)))] - p
  aa = norm(a); bb = norm(b); cc = norm(c)
  2.0*atan(det33(a,b,c), aa*bb*cc + a⋅b*cc + b⋅c*aa + c⋅a*bb)
end

function solid_angle_normal(topo::Topology, P::TA, h::HalfEdgeHandle, p::T) where {T, TA<:Vector{T}}
  if isboundary(topo, h)
    return 0.0
  end

  @inbounds a = P[head(topo, h)] - p
  @inbounds b = P[head(topo, next(topo, h))] - p
  @inbounds c = P[head(topo, next(topo, next(topo, h)))] - p
  aa = norm(a); bb = norm(b); cc = norm(c)
  Ω = 2.0*atan(det33(a,b,c), aa*bb*cc + a⋅b*cc + b⋅c*aa + c⋅a*bb)

  if isapprox(Ω, 0.0, atol=1e-16)
    zero(T)
  else
    w = LinearAlgebra.norm(a+b+c)/3.0
    centroidn = normalize((normalize(a) + normalize(b) + normalize(c)))
    centroidn*Ω*w
  end
end

function winding_normal(topo::Topology, P, q::VT) where {T<:Number,
                                                         VT<:AbstractVector{T}}
  mapreduce(i->solid_angle_normal(topo, P, i, q), +, faces(topo))/(4.0π)
end

#=== some tensor stuff ===#
# need this to make approximate winding numbers managable
# works pretty seamlessly with julia's multidimensional arrays.
# pops out an Array quite easily from pretty much any sensible operation
# make sure to convert to Array after precomputation as this isn't very performant for access

struct LazyTensor{T,N} <: AbstractArray{T,N}
  V::Union{NTuple{N}, Vector{LazyTensor{T,N}}}
  dim
end

Base.size(A::LazyTensor) = A.dim
Base.getindex(A::LazyTensor{T,N}, I::Vararg{Int,NN}) where {T,N,NN} = getlazyindex(A.V, I...)

getlazyindex(V::NTuple{N}, I::Vararg{Int,NN}) where {N,NN} = prod(map(i->V[i][I[i]], 1:length(I)))

function getlazyindex(V::Vector{LazyTensor{T,N}}, I::Vararg{Int,NN}) where {T,N,NN}
  sum(map(Aᵢ->getindex(Aᵢ, I...), V))
end

Base.IndexStyle(::Type{LazyTensor}) = IndexCartesian()

cons(a::T, b::NTuple{N,T}) where {N,T} = (a, b...)
cons(a::T, b::Vector{T}) where T = vcat(a, b)
cons(a::NTuple{N,T}, b::T) where {N,T} = (a..., b)
cons(a::Vector{T}, b::T) where T = vcat(a, b)

""" distributes ⊗ into a + b +... """
cons(a::Vector{LazyTensor{T,N}}, b::AbstractVector{T}) where {N,T} = map(aᵢ->aᵢ⊗b, a)
cons(a::AbstractVector{T}, b::Vector{LazyTensor{T,N}}) where {N,T} = map(bᵢ->a⊗bᵢ, b)

⊗(a::T, b::T) where {F<:Real, T<:AbstractVector{F}} = LazyTensor{F,2}((a,b), (length(a), length(b)))

⊗(a::T, b::VT) where {F<:Real, T<:AbstractVector{F}, N, VT<:LazyTensor{F,N}} =
  LazyTensor{F,N+1}(cons(a, b.V), (length(a), b.dim...))

⊗(a::VT, b::T) where {F<:Real, T<:AbstractVector{F}, N, VT<:LazyTensor{F,N}} =
  LazyTensor{F,N+1}(cons(a.V, b), (a.dim..., length(b)))

#!me don't actually need this, so didn't finish implementing
#⊗(a::TA, b::TB) where {F<:Real, NA, TA<:LazyTensor{F,NA}, NB, TB<:LazyTensor{F,NB}} =
#  LazyTensor{F,NA+NB}(cons(a.V, b.V), (a.dim..., b.dim...))

⊗(a) = ⊗(a, a)

Base.:+(a::LT, b::LT) where {T, N, LT<:LazyTensor{T,N}} = LazyTensor{T, N}([a,b], a.dim)

# dipole derivatives from appendix A of "Fast Winding Numbers for Soups and Clouds"
function ∇G(q,x)
  r = x-q
  r/(4.0π*norm(r)^3)
end

function ∇²G(q,x)
  r = x-q
  rlen = norm(r)
  I/(4.0π*rlen^3) - 3.0*(r*r')/(4.0π*rlen^5)
end

# literally from the paper, lots of 0*0 happening and that ⊗ is brutal slow
# likely we wont actually use this third order expansion, so don't worry about optz right now
function ∇³G(q,x)
  r = x-q
  rlen = norm(r)
  Array(sum(map(eᵢ->r⊗eᵢ⊗eᵢ + eᵢ⊗r⊗eᵢ + eᵢ⊗eᵢ⊗r, [SVector(1.0,0.0,0.0),
                                                  SVector(0.0,1.0,0.0),
                                                  SVector(0.0,0.0,1.0)])))/(-4.0π*rlen^5) +
  15.0*Array(r⊗r⊗r)/(4.0π*rlen^7)
end

function centroid(topo, P,  F::Vector{HalfEdgeHandle})
  Pees = unique(mapreduce(partial(vertices, topo), vcat, F; init=[]))
  p̃ = sum(P[Pees])/length(Pees)
end

# triangle integrals from appendix B of "Fast Winding Numbers for Soups and Clouds"

struct ApproxWindingNumber2
  p̃::SVector{3, Float64}
  term1::SVector{3, Float64}
  term2::SMatrix{3,3, Float64}
end

struct ApproxWindingNumber3
  p̃::SVector{3, Float64}
  term1::SVector{3, Float64}
  term2::SMatrix{3,3, Float64}
  term3::Array{Float64, 3}
end

approx_winding_number(awn::ApproxWindingNumber2,
                      q::SVector{3,Float64}) = awn.term1⋅∇G(q, awn.p̃) + awn.term2⋅∇²G(q, awn.p̃)

"""
    approx_winding_number(topo, P, F, order=3)

assumes triangle mesh.  approximate winding number function for a cluster of triangles
"""
function approx_winding_number(topo, P, F::Vector{HalfEdgeHandle}, order=2)

  p̃ = centroid(topo, P, F)

  #!me check https://github.com/alecjacobson/WindingNumber/blob/master/UT_SolidAngle.cpp
  #!me says we need to use NORMALIZED normal to multiply integrals by, paper says otherwise...
  ant = map(partial(weightednormal, topo, P), F)

  term1 = sum(ant)
  term2 = map(zip(ant, F)) do (antᵢ, heh)
    (sum(P[vertices(topo, heh)])/3.0 - p̃)⊗antᵢ
  end |> sum |> SMatrix{3,3}∘Array

  # very inefficient, but order 2 seems good enough, so don't worry about it for now
  if order == 3
    term3 = map(zip(ant, F)) do (antᵢ, heh)
      xᵢ, xⱼ, xk = P[vertices(topo, heh)]
      Ct = ⊗(0.5*(xᵢ+xⱼ)-p̃) +
           ⊗(0.5*(xⱼ+xk)-p̃) +
           ⊗(0.5*(xk+xᵢ)-p̃)
      Ct⊗antᵢ
    end |> sum |> partial(*, 1.0/6.0) |> Array

  #  w̃ = (q)->term1⋅∇G(q, p̃) + term2⋅∇²G(q, p̃) + term3⋅∇³G(q, p̃)
    ApproxWindingNumber3(p̃, term1, term2, term3)
  else
  #  w̃ = (q)->term1⋅∇G(q, p̃) + term2⋅∇²G(q, p̃)
    ApproxWindingNumber2(p̃, term1, term2)
  end

  # is storing a closure reasonably performant?  if not we should just be storing (p̃, term1...)
  #return w̃

end

approx_winding_number(topo, P, F::Vector{FaceHandle}) = approx_winding_number(topo, P, faces(topo)[F])

function approx_winding_number(topo, P, q, node::BT) where {BT<:BinaryTree}
  F = ileaves(node) |> collect |> c->map(value∘data,c)
  approx_winding_number(topo, P, map(i->halfedge(topo, FaceHandle(i)), F))(q)
end

struct WindingNumberCache
  bvh::IndexedBinaryTree{AVLData{AABBNodeData{Float64,Int64}}}
  #W::Vector{Function}
  W::Vector{ApproxWindingNumber2}  #!me ignore order 3 for now
  height_cutoff::Int64
  β::Float64
  F::Vector{Vector{HalfEdgeHandle}}
  WindingNumberCache(a,b,c,d,e) = new(a,b,c,d,e)
  WindingNumberCache(a,b,c,d) = new(a,b,c,d,Vector{Vector{HalfEdgeHandle}})
end

WindingNumberCache(a,b,c) = WindingNumberCache(a, b, c, 2.0)

function winding_number_cache(topo, P, tris_for_brute = 100, order=2)
  aabbs = map((Fᵢ, i)->AABBNodeData(AABB(P[Fᵢ]), i), polygons(topo), 1:nfaces(topo))
  bt = avlitree(map( (Fᵢ, i)->AABBNodeData(AABB(P[Fᵢ]), i), polygons(topo), 1:nfaces(topo)))

  # threshold for when we generate approx w
  height_for_approx = (minheight(tris_for_brute) + maxheight(tris_for_brute))/2 |> Int∘round

  #nnodes = mapreduce(index, max, TraverseWhen(n->height(n) >= height_for_approx, bt, identity))
  nnodes = mapreduce(index, max, itraversewhen(n->height(n) >= height_for_approx, bt, eltype(bt), identity))
  F = polygons(topo)

#  lc = TraverseWhen(n->height(n) >= height_for_approx, bt, Leaves)
  lc = TraverseWhen(n->height(n) >= height_for_approx, bt, ileaves)

  faceheh(n) = n |> data |> value |> i->halfedge(topo, FaceHandle(i))
  W = ApproxWindingNumber2[]
  resize!(W, nnodes)
  W[ TraverseWhen(n->height(n) >= height_for_approx, bt, index) |> collect ] =
    [ approx_winding_number(topo, P, map(faceheh, lcᵢ), order) for lcᵢ in lc ]

  #  W[ itraversewhen(n->height(n) >= height_for_approx, bt, Int, index) ] =
  #    [ approx_winding_number(topo, P, map(faceheh, lcᵢ), order) for lcᵢ in lc ]

  at_height = Iterators.filter(n->height(n) == height_for_approx-1,
                               itraversewhen(n->height(n) >= height_for_approx-1,
                                             bt, eltype(bt), identity))
  iF = map(index, at_height)
  nF = reduce(max, iF)
  F = Vector{HalfEdgeHandle}[]
  resize!(F, nF)
  F[iF] = map(n->faceheh.(ileaves(n)), at_height)

  WindingNumberCache(bt, W, height_for_approx, 2.0, F)
end

"""
"""
function winding_number(topo, P, q, cache::WindingNumberCache)
  node = cache.bvh
  β² = cache.β*cache.β
  winding_number(topo, P, q, node, β², cache)
end

function winding_number(topo, P, q, node::BT, β²::Float64, cache) where {T, BT<:IndexedBinaryTree{T}}
  if height(node) == cache.height_cutoff-1
    return mapreduce(heh->solid_angle(topo, P, heh, q), +, cache.F[index(node)])/(4.0π)
  elseif height(node) < cache.height_cutoff
    return winding_number(topo, P, q, node)
  else
    aabb = key(node)
    p = centre(aabb)
    r = radius_squared(aabb)
    if LinearAlgebra.norm_sqr(q - p) > β²*r
      #println("WHATUP", index(node))
      #return cache.W[index(node)](q)
      return approx_winding_number(cache.W[index(node)], q)
    else
      # descend
      #println("crazy!!!")
      return winding_number(topo, P, q, left(node), β², cache) +
             winding_number(topo, P, q, right(node), β², cache)
    end
  end
end

function iwinding_number(topo, P, q, cache::WindingNumberCache)
  node = cache.bvh
  β² = cache.β*cache.β
  iwinding_number(topo, P, q, node, β², cache)
end

function iwinding_number(topo, P, q, node::BT, β²::Float64, cache) where {T, BT<:IndexedBinaryTree{T}}
  sum = 0.0
  nodestack = [index(node)]

  while !isempty(nodestack)
    node = cache.bvh[pop!(nodestack)]
    #println(height(node), " compare ", cache.height_cutoff-1)
    if height(node) < cache.height_cutoff-1
      sum += winding_number(topo, P, q, node)
    elseif height(node) == cache.height_cutoff-1
      sum += mapreduce(heh->solid_angle(topo, P, heh, q), +, cache.F[index(node)])/(4.0π)
      #  sum += solid_angle(topo, P, halfedge(topo, FaceHandle(value(data(node)))), q)/(4.0π)
      #if height(node) == 0
      #  sum += solid_angle(topo, P, halfedge(topo, FaceHandle(value(data(node)))), q)/(4.0π)
    else
      aabb = key(node)
      p = centre(aabb)
      r = radius_squared(aabb)
      if LinearAlgebra.norm_sqr(q - p) > β²*r
        #println("WHATUP", index(node))
        #return cache.W[index(node)](q)
        sum += approx_winding_number(cache.W[index(node)], q)
      else
        # descend
        #println("crazy!!!")
        push!(nodestack, ileft(node))
        push!(nodestack, iright(node))
      end
    end
  end
  sum
end

function winding_number(topo, P, q, node::BT) where {T, BT<:IndexedBinaryTree{T}}
  #mapreduce(l->solid_angle(topo, P, halfedge(topo, FaceHandle(value(data(l)))), q), +, ileaves(node))/(4.0π)
  wn = 0.0
  nodestack=[node]
  while !isempty(nodestack)
    node = pop!(nodestack)

    if isleaf(node)
      wn += solid_angle(topo, P, halfedge(topo, FaceHandle(value(data(node)))), q)
    else
      l = left(node)
      r = right(node)
      l != BoundingVolumeTrees.EmptyIndexedNode && push!(nodestack, left(node))
      r != BoundingVolumeTrees.EmptyIndexedNode && push!(nodestack, right(node))
    end
  end
  wn/(4.0π)
end

winding_number(topo, P, q) = mapreduce(i->solid_angle(topo, P, i, q), +, faces(topo))/(4.0π)

winding_numbers(topo, P::V, cache; jiggle = eps(T)) where {T,
                                                           PT<:AbstractVector{T},
                                                           V<:AbstractVector{PT}} =
  map( q->iwinding_number(topo, P, q, cache), P+vertexnormals(topo, P)*jiggle)

"""
    winding_numbers(topo, P::V; jiggle = ϵ)

The winding number for each vertex of a mesh.
A vertex is lifted slightly off the surface in the normal direction to give winding number of 0 for vertices on the surface, but otherwise on exterior of the mesh.
The amount of normal lifting is controlled with the jiggle parameter
"""
function winding_numbers(topo, P::V; jiggle = eps(T), approx=false, cacheleaves=50 ) where {T, PT<:AbstractVector{T}, V<:AbstractVector{PT}}
  if approx
    #!me fast winding numbers are anything but atm. some serious performance issue we need to investigate
    cache = winding_number_cache(topo, P, cacheleaves)
    winding_numbers(topo, P, cache; jiggle=jiggle)
    #map( q->iwinding_number(topo, P, q, cache), P+vertexnormals(topo, P)*jiggle)
  else
    map( q->winding_number(topo, P, q), P+vertexnormals(topo, P)*jiggle)
  end
end


using Turing
using Test

xs = [1.5 2.0]
# xx = 1

@model fggibbstest(xs) = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt(s))
  # xx ~ Normal(m, sqrt(s)) # this is illegal

  for i = 1:length(xs)
    xs[i] ~ Normal(m, sqrt(s))
  # for xx in xs
    # xx ~ Normal(m, sqrt(s))
  end
  s, m
end

gibbs = Gibbs(2, PG(10, 2, :s), HMC(1, 0.4, 8, :m))
chain = sample(fggibbstest(xs), gibbs);

#
#
# using Turing
# using Test
#
# @model ttt(a) = begin
#   a ~ Normal(0, 1)
#   bb ~ Normal(0, 1)
#   bb
# end
#
# a = 1
#
# sample(ttt(a), SMC(10));


# Metatheory.options[:printiter] = true
using Metatheory
using Metatheory.EGraphs
using Metatheory.Classic
using Test

or_alg = @theory begin
    ((p ∨ q) ∨ r)       ==  (p ∨ (q ∨ r))
    (p ∨ q)             ==  (q ∨ p)
    (p ∨ p)             =>  p
    (p ∨ true)          =>  true
    (p ∨ false)         =>  p
end

and_alg = @theory begin
    ((p ∧ q) ∧ r)       ==  (p ∧ (q ∧ r))
    (p ∧ q)             ==  (q ∧ p)
    (p ∧ p)             =>  p
    (p ∧ true)          =>  p
    (p ∧ false)         =>  false
end

comb = @theory begin
    # DeMorgan
    ¬(p ∨ q)            ==  (¬p ∧ ¬q)
    ¬(p ∧ q)            ==  (¬p ∨ ¬q)
    # distrib
    (p ∧ (q ∨ r))       ==  ((p ∧ q) ∨ (p ∧ r))
    (p ∨ (q ∧ r))       ==  ((p ∨ q) ∧ (p ∨ r))
    # absorb
    (p ∧ (p ∨ q))       =>  p
    (p ∨ (p ∧ q))       =>  p
    # complement
    (p ∧ (¬p ∨ q))      =>  p ∧ q
    (p ∨ (¬p ∧ q))      =>  p ∨ q
end

negt = @theory begin
    (p ∧ ¬p)            =>  false
    (p ∨ ¬(p))          =>  true
    ¬(¬p)               ==  p
end

impl = @theory begin
    (p == ¬p)           =>  false
    (p == p)            =>  true
    (p == q)            =>  (¬p ∨ q) ∧ (¬q ∨ p)
    (p => q)            =>  (¬p ∨ q)
end

fold = @theory begin
    (p::Bool == q::Bool)   |> (p == q)
    (p::Bool ∨ q::Bool)    |> (p || q)
    (p::Bool => q::Bool)   |> ((p || q) == q)
    (p::Bool ∧ q::Bool)    |> (p && q)
    ¬(p::Bool)             |> (!p)
end

t = or_alg ∪ and_alg ∪ comb ∪ negt ∪ impl ∪ fold


symmetrize(P::Matrix) = (P+P')/2 #+ 1e-15eye(P) #TODO: Hacky. Square root filter?

function likelihood{T}(y::Vector{T}, C::Matrix{T}, σϵ²::T, Q::Matrix{T})

    n   = length(y)
    nₓ  = size(C, 1)

    x  = zeros(nₓ)
    P∞ = eye(nₓ)
    P = zeros(nₓ, nₓ)

    ll = -n*log(2π)

    d = 0
    diffuseP = true

    while diffuseP

        d += 1
        Cₜ = C[:,d]

        ν⁰ = y[d] - dot(Cₜ, x)

        F∞  = dot(Cₜ, P∞*Cₜ)
        F   = dot(Cₜ, P*Cₜ) + σϵ²
        F¹  = 1/F∞
        F²  = - F¹ * F * F¹
        posdefF∞ = !isapprox(0, F∞, atol=1e-10)

        M∞  = P∞ * Cₜ
        M   = P * Cₜ

        K⁰  = posdefF∞ ? M∞*F¹ : M/F
        K¹  = M * F¹ + M∞ * F²

        L⁰  = I - K⁰ * Cₜ'
        L¹  = -K¹ * Cₜ'

        # Calculate loglikelihood contribution
        ll-= posdefF∞ ? log(F∞) : log(F) + ν⁰*ν⁰/F

        # Calculate next values
        x   = x + K⁰*ν⁰
        P   = posdefF∞ ? symmetrize(P∞*L¹' + P*L⁰' + Q) : symmetrize(P*L⁰' + Q)
        P∞  = posdefF∞ ? symmetrize(P∞*L⁰') : P∞

        diffuseP = findfirst(round(P∞,10)) > 0

        d == n && error("Failed to substantiate diffuse initial conditions")

    end #while

    for t in d+1:n-1

        Cₜ = C[:,t]

        νₜ = y[t] - dot(Cₜ, x)
        Fₜ = dot(Cₜ, P*Cₜ) + σϵ²
        Kₜ = P * Cₜ / Fₜ

        # Calculate loglikelihood contribution
        ll -= Cₜ[1] == 0 ? 0 : log(Fₜ) + νₜ*νₜ/Fₜ

        # Predict next values
        x += Kₜ * νₜ
        P = P * (I - Cₜ * Kₜ') + Q

    end #for

    Cₜ = C[:,n]
    νₜ = y[n] - dot(Cₜ, x)
    Fₜ = dot(Cₜ, P*Cₜ) + σϵ²
    ll -= log(Fₜ) + νₜ*νₜ/Fₜ

    return -ll/2

end #likelihood

function smooth{T}(y::Vector{T}, C::Matrix{T}, σϵ²::T, σμ²::T)

    n   = length(y)
    nₓ  = size(C, 1)
    Q = zeros(nₓ, nₓ)
    Q[1] = σμ²

    K   = Array{T}(nₓ, n)
    F⁻¹ = Array{T}(n)
    ν   = Array{T}(n)

    P∞  = Matrix{T}[]
    L⁰  = Matrix{T}[]
    L¹  = Matrix{T}[]
    F∞  = T[]
    F²  = T[]

    x   = Array{T}(nₓ, n)
    P   = Array{T}(nₓ, nₓ, n)
    ϵ   = Array{T}(n)
    η   = Array{T}(nₓ, n)

    xₜ  = zeros(nₓ)
    P∞ₜ = eye(nₓ)
    Pₜ  = zeros(nₓ, nₓ)

    ll = -n*log(2π)

    d           = 0
    x[:,1]      = xₜ
    P[:,:,1]    = Pₜ
    push!(P∞, P∞ₜ)
    diffuseP    = true

    while diffuseP

        d += 1
        Cₜ = C[:,d]
        ν⁰ = y[d] - dot(Cₜ, xₜ)

        F∞ₜ = dot(Cₜ, P∞ₜ * Cₜ)
        Fₜ  = dot(Cₜ, Pₜ * Cₜ) + σϵ²
        F¹ₜ = 1/F∞ₜ
        F²ₜ = - F¹ₜ * Fₜ * F¹ₜ

        M∞  = P∞ₜ * Cₜ
        M   = Pₜ * Cₜ

        K⁰  = F∞ₜ == 0 ? M/Fₜ : M∞*F¹ₜ
        K¹ = M * F¹ₜ + M∞ * F²ₜ

        L⁰ₜ = I - K⁰ * Cₜ'
        L¹ₜ = -K¹ * Cₜ'

        ν[d]    = ν⁰
        F⁻¹[d]  = 1/Fₜ
        K[:, d] = K⁰
        push!(F∞, F∞ₜ)
        push!(F², F²ₜ)
        push!(L⁰, L⁰ₜ)
        push!(L¹, L¹ₜ)

        # Calculate next values
        xₜ  = xₜ + K⁰ * ν⁰
        Pₜ  = F∞ₜ == 0 ? Pₜ*L⁰ₜ' + Q : P∞ₜ*L¹ₜ' + Pₜ*L⁰ₜ' + Q
        P∞ₜ = F∞ₜ == 0 ? P∞ₜ         : P∞ₜ*L⁰ₜ'

        x[:,d+1]    = xₜ
        P[:,:,d+1]  = Pₜ
        push!(P∞, P∞ₜ)

        diffuseP = findfirst(round(P∞ₜ,12)) > 0

    end #while

    for t in d+1:n-1

        Cₜ = C[:,t]
        νₜ = y[t] - dot(Cₜ, xₜ)
        Fₜ = dot(Cₜ, Pₜ*Cₜ) + σϵ²
        Kₜ = Pₜ * Cₜ / Fₜ

        ν[t]    = νₜ
        F⁻¹[t]  = 1/Fₜ
        K[:,t]  = Kₜ

        # Predict next values
        xₜ += Kₜ * νₜ
        Pₜ = Pₜ * (I - Cₜ * Kₜ') + Q

        x[:,t+1]   = xₜ
        P[:,:,t+1] = Pₜ

    end #for

    Cₜ = C[:,n]
    νₜ = y[n] - dot(Cₜ, xₜ)
    Fₜ = dot(Cₜ, Pₜ*Cₜ) + σϵ²
    Kₜ = Pₜ * Cₜ / Fₜ

    ν[n]    = νₜ
    F⁻¹[n]  = 1/Fₜ
    K[:,n]  = Kₜ

    rₜ, Nₜ = zeros(nₓ), zeros(nₓ, nₓ)
    uₜ = νₜ/Fₜ - dot(Kₜ, rₜ)
    Dₜ = 1/Fₜ + dot(Kₜ, Nₜ * Kₜ)

    # Calculate smoothed disturbances
    ϵ[n] = σϵ² * uₜ
    η[:, n] = Q * rₜ

    # Calculate helper recursions
    rₜ = Cₜ * uₜ + rₜ
    Θ = Cₜ * Kₜ' * Nₜ
    Nₜ = Cₜ * Dₜ * Cₜ' + Nₜ - Θ - Θ'

    for t in n-1:-1:d+1

        Cₜ    = C[:,t]
        F⁻¹ₜ  = F⁻¹[t]
        νₜ    = ν[t]
        Kₜ    = K[:,t]
        Pₜ    = P[:,:,t]

        uₜ = F⁻¹ₜ * νₜ - dot(Kₜ, rₜ)
        Dₜ = F⁻¹ₜ + dot(Kₜ, Nₜ * Kₜ)

        # Calculate smoothed disturbances
        ϵ[t] = σϵ² * uₜ
        η[:, t] = Q * rₜ

        # Calculate helper recursions
        rₜ = Cₜ * uₜ + rₜ
        Θ = Cₜ * Kₜ' * Nₜ
        Nₜ = Cₜ * Dₜ * Cₜ' + Nₜ - Θ - Θ'

        # Calculate smoothed states
        x[:, t] += Pₜ * rₜ
        P[:,:,t] -= Pₜ * Nₜ * Pₜ

    end #for

    r⁰ₜ, r¹ₜ      = rₜ, zeros(rₜ)
    N⁰ₜ, N¹ₜ, N²ₜ = Nₜ, zeros(Nₜ), zeros(Nₜ)

    for t in d:-1:1

        Cₜ        = C[:,t]
        F∞ₜ       = F∞[t]
        #F¹ₜ       = F⁻¹[t]
        F²ₜ       = F²[t]
        F⁻¹ₜ      = F∞ₜ == 0 ? F⁻¹[t] : 0
        ν⁰ₜ       = ν[t]
        K⁰ₜ       = K[:,t]
        L⁰ₜ, L¹ₜ  = L⁰[t], L¹[t]
        Pₜ, P∞ₜ   = P[:,:,t], P∞[t]

        uₜ        = F⁻¹ₜ*ν⁰ₜ - dot(K⁰ₜ, r⁰ₜ)
        Dₜ        = F⁻¹ₜ + dot(K⁰ₜ, N⁰ₜ * K⁰ₜ)

        # Calculate smoothed disturbances
        ϵ[t] = σϵ² * uₜ
        η[:, t] = Q * rₜ

        # Calculate helper recursions
        r¹ₜ = Cₜ*ν⁰ₜ/F∞ₜ + L⁰ₜ'*r¹ₜ + L¹ₜ'*r⁰ₜ
        r⁰ₜ += Cₜ * uₜ
        Δ = L⁰ₜ' * N¹ₜ * L¹ₜ
        N²ₜ = Cₜ*F²ₜ*Cₜ' + L⁰ₜ'*N²ₜ*L⁰ₜ + Δ + Δ' + L¹ₜ'*N⁰ₜ*L¹ₜ
        N¹ₜ = Cₜ*Cₜ'/F∞ₜ + L⁰ₜ'*N¹ₜ*L⁰ₜ + L¹ₜ'*N⁰ₜ*L⁰ₜ
        Θ = Cₜ*K⁰ₜ'*N⁰ₜ
        N⁰ₜ += Cₜ*Dₜ*Cₜ' - Θ - Θ'

        # Calculate smoothed states
        Γ = P∞ₜ * N¹ₜ * Pₜ
        x[:,t] += Pₜ*r⁰ₜ + P∞ₜ*r¹ₜ
        P[:,:,t] -= Pₜ*N⁰ₜ*Pₜ + P∞ₜ*N²ₜ*P∞ₜ + Γ + Γ'

    end #for

    return ϵ, x, P, η

end #statesmooth


# This file defines the const varibles

const ρ = 1.225 #空气密度 量纲kg*m^-3
const ν = 1.46e-4 #运动黏性系数 无量纲
const g = 9.8 #重力加速度  量纲kg*m*s^-2
const v_sound = 343.2 #音速 量纲m/s
const NR = 2  #桨叶片数
const Nbe = 10 #旋翼分段数量
const R = 1.05  #旋翼半径 量纲m
const ecut = 0.1 #桨叶根切比例 无量纲
const m_ = 0.35/R #桨叶质量密度 量纲kg/m
const Ω = 100  #旋翼转速 量纲rad/s
const αs = -2*π/180.0  #旋翼轴倾角  量纲rad
const Kβ = 0.0  #桨叶根部挥舞弹簧刚度 量纲？？？
const Iβ = m_/3*(1-ecut^3)*R^3  #桨叶挥舞惯量 量纲kg*m^2
# const R0 = 0.382  #lenght of root cut 量纲m
# const R1 = 0.136  #length of tip loss 量纲m
const vair = 50 # 来流速度 量纲 m/s
const v_air = [vair*cos(αs),0.0,vair*sin(αs)] #forward wind speed  量纲m/s
const μ_air = v_air[1]/(Ω*R)  #来流在桨盘edgewise方向分量  无量纲
const λ_air = v_air[3]/(Ω*R)  #来流在桨盘轴向分量  无量纲
const A = π*R^2 #参考面积 量纲为m^2
const fnonc = ρ*A*Ω*R^2*Ω #力的无量纲化参数 量纲kg*m*s^-2
const mnonc = ρ*A*Ω^2*R^3 #力矩的无量纲化参数 量纲kg*m^2/s^2
const T = 250 #飞行器重量 (量纲为kg*m*s^-2)
const CT = T/fnonc  #无量纲重量系数
const dpsideg = 10  # 方位角不仅长度（量纲为deg）
const dψ = dpsideg*π/180 #方位角步进步长 (量纲为rad)
const npsi = 360/dpsideg # 周向分割步数
const βang0 = 0.0 # 挥舞角初值
const dβ0 = 0.0 # 挥舞角速度初值
const θ7 = 2.0/180*π # 旋翼总距
const θtw = -5./180*π # 桨叶扭转角
const thelat = -0/180*π # 横向周期变距
const thelon = 0/180*π # 纵向周期变距

ch = Array{Float64}(NR,Nbe) #桨叶弦长
rb = Array{Float64}(NR,Nbe)  #桨叶叶素点分段径向长度值
dr = Array{Float64}(NR,Nbe)  #桨叶叶素分段长度值

for k in 1:NR
  for i in 1:Nbe
    ch[k,i] = 0.35
    dr[k,i] = R/Nbe
    rb[k,i] = R/Nbe*(i-1)+dr[k,i]/2
  end
end


using NearestNeighbors
using Distances
using LinearAlgebra
data = rand(3, 10^4);
point = rand(3);
#println("Our point")
#println(point)
tree = BallTree(data);

t = @elapsed minimum.(data)
t = @elapsed minimum.(data)
println("Reference time: ",t)
@elapsed tree = BallTree(data);
t = @elapsed tree = BallTree(data);
println("BallTree time: ",t)

@elapsed kdtree = KDTree(data);

k = 10;
t = @elapsed idxs, dists = knn(kdtree, point, k, true);
t = @elapsed idxs, dists = knn(kdtree, point, k, true);
#println("time:",t)

# display(idxs)
# println("")
# display(dists)
# println("")

@elapsed idxs, dists = knn(tree, point, k, true)
t = @elapsed idxs, dists = knn(tree, point, k, true)
#println("time:",t)

# display(idxs)
# println("")
# display(dists)
# println("")

r = 0.02;
@elapsed idxs = inrange(tree, point, r)
t = @elapsed idxs = inrange(tree, point, r)
# println("time:",t)
# display(data[:,idxs])
#println("**************************************************************************************")
A = randn(3,3);
A = A'*A + 0.1*Matrix(1.0I,3,3);
t = @elapsed tree = BallTree(data,Mahalanobis{Float64}(A));
t = @elapsed tree = BallTree(data,Mahalanobis{Float64}(A));
println("time:",t)

t = @elapsed idxs, dists = knn(tree, point, k, true)

t = @elapsed idxs, dists = knn(tree, point, k, true)
# println("time:",t)

# display(idxs)
# println("")
# display(dists)
# println("")

r = 0.02;
@elapsed idxs = inrange(tree, point, r)
t = @elapsed idxs = inrange(tree, point, r)
# println("time:",t)
# display(data[:,idxs])


immutable MultivariateTestFunction{T <: Function}
    f::T
    name::String
    x1s::Vector{Float64}
    x2s::Vector{Float64}
end

# Test cases generating as follows
# using Distributions
# repr(rand(Cauchy(0, 10), 5))

test_funcs = (
    MultivariateTestFunction(
        x -> sin(x[1]) + cos(x[2]),
        "x -> sin(x[1]) + cos(x[2])",
        [62.1908, 7.19278, 32.4944, -10.0597, 1.23525],
        [1.91818, -96.7718, -54.5914, 105.894, -2.74696],
    ),
    MultivariateTestFunction(
        x -> sin(x[1]) + 2 * sin(x[2]),
        "x -> sin(x[1]) + 2 * sin(x[2])",
        [-3.54442, 7.14114, -50.0388 , -5.73865, -8.88167],
        [-5.33968, -5.38714, -24.4391, -16.1525, -18.9814],
    ),
    MultivariateTestFunction(
        x -> x[1]^2 + x[2]^2,
        "x -> x[1]^2 + x[2]^2",
        [5.50783, -3.14069, 5.6383, 188.917, 98.7357],
        [6.71689, -15.1161, -19.7469, -2.62965, -14.9729],
    ),
    MultivariateTestFunction(
        x -> (1.0 - x[1])^2 + (2.0 - x[2])^2,
        "x -> (1.0 - x[1])^2 + (2.0 - x[2])^2",
        [5.66663, 22.7718, -13.5606, 9.80007, 0.0533083],
        [-11.0918, 1.55486, -133.914, -45.432, -19.3287],
    ),
    MultivariateTestFunction(
        x -> (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2,
        "x -> (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2",
        [25.9188, 35.6369, -17.8868, 47.76, 18.5091],
        [19.2642, -0.886276, 9.51653, 47.9309, -3.34237],
    ),
    MultivariateTestFunction(
        x -> 1 / (1 + exp(-(x[1] + x[2]))),
        "x -> 1 / (1 + exp(-(x[1] + x[2])))",
        [7.88823, -98.6984, 3.51154, -39.7024, 6.60017],
        [-55.2978, -212.439, 5.52012, 44.009, -10.9622],
    ),
    MultivariateTestFunction(
        x -> log(x[1]) + log(x[2]),
        "x -> log(x[1]) + log(x[2])",
        [7.88823, 98.6984, 3.51154, 39.7024, 6.60017],
        [55.2978, 212.439, 5.52012, 44.009, 10.9622],
    ),
)


# function tcm_evaluateKL(N, v, transaction_counts)
# 	M = size(transaction_counts,1)

# 	empirical_probs = transaction_counts ./ sum(transaction_counts, 2)

# 	KL = 0.0;

# 	for m in 1:M
# 		for i in 1:N
# 			if (transaction_counts[m,i] > 0)
# 				temp = transaction_counts[m,i] * log( v[m,i] / empirical_probs[m,i] )
# 				KL += temp
# 			end
# 		end
# 	end

# 	KL_avg = KL / sum(transaction_counts)

# 	return -KL_avg
# end

function tcm_evaluateKL(N, v, transaction_counts)
	M = size(transaction_counts,1)

	empirical_probs = transaction_counts ./ sum(transaction_counts, 2)

	KL_noinf = 0.0
	total_noinf_transactions = 0;

	for m in 1:M
		for i in 1:N
			if (transaction_counts[m,i] > 0)
				temp = transaction_counts[m,i] * log( v[m,i] / empirical_probs[m,i] )
				if (!isinf(temp))
					KL_noinf += temp
					total_noinf_transactions += transaction_counts[m,i]
				end
			end
		end
	end
	KL_noinf = -KL_noinf

	KL_noinf_avg = KL_noinf / total_noinf_transactions * 100

	return KL_noinf, total_noinf_transactions, KL_noinf_avg
end

# function tcm_evaluateKL_debug(N, v, transaction_counts)
# 	M = size(transaction_counts,1)

# 	empirical_probs = transaction_counts ./ sum(transaction_counts, 2)

# 	KL = 0.0;

# 	KL_matrix = zeros(M, N);

# 	for m in 1:M
# 		for i in 1:N
# 			if (transaction_counts[m,i] > 0)
# 				KL += transaction_counts[m,i] * log( v[m,i] / empirical_probs[m,i] )
# 				KL_matrix[m,i] = transaction_counts[m,i] * log( v[m,i] / empirical_probs[m,i] )
# 			end
# 		end
# 	end
# 	return -KL, KL_matrix
# end


### A Pluto.jl notebook ###
# v0.18.0

using Markdown
using InteractiveUtils

# ╔═╡ 6f9be298-5374-11ec-3b9d-fd74a32d6b53
using Colors, PlutoUI, Manopt, Manifolds

# ╔═╡ e701b513-4ad9-474c-9d79-60a24f545ebd
md"""
# Bezier curves and their acceleration

This tutorial illustrates how Bézier curves are generalized to manifolds and how to
minimize their acceleration, i.e. how to get a curve that is as straight or as geodesic
while fulfilling constraints

This example also illustrates how to apply the minimization on the corresponding [`PowerManifold`](https://juliamanifolds.github.io/Manifolds.jl/latest/manifolds/power.html) manifold using a [`gradient_descent`](https://manoptjl.org/stable/solvers/gradient_descent.html) with [`ArmijoLinesearch`](https://manoptjl.org/stable/plans/index.html#Manopt.ArmijoLinesearch).
"""

# ╔═╡ 087fb979-d68d-4ee7-9be5-fad992fdb792
md"and we define some colors from [Paul Tol](https://personal.sron.nl/~pault/)"

# ╔═╡ 85d8985b-f5ee-4cba-bcc5-528049798330
begin
    black = RGBA{Float64}(colorant"#000000")
    TolVibrantBlue = RGBA{Float64}(colorant"#0077BB")
    TolVibrantOrange = RGBA{Float64}(colorant"#EE7733")
    TolVibrantMagenta = RGBA{Float64}(colorant"#EE3377")
    TolVibrantCyan = RGBA{Float64}(colorant"#33BBEE")
    TolVibrantTeal = RGBA{Float64}(colorant"#009988")
end;

# ╔═╡ e62e36e9-19ce-4a7d-8bda-5634813c1ee9
md"Let's also set up a few parameters"

# ╔═╡ 18d378dd-abc6-4fa6-a2c9-1cc09e4cec46
begin
    geo_pts = collect(range(0.0, 1.0; length=101))
    bezier_pts = collect(range(0.0, 3.0; length=201))
    camera_position = (-1.0, -0.7, 0.3)
    #render asy yes/no. If not, images included w/ markdown are assumed to be prerendered
    render_asy = false
    localpath = join(splitpath(@__FILE__)[1:(end - 1)], "/") # remove file to get files folder
    image_prefix = localpath * "/bezier"
    @info image_prefix
end;

# ╔═╡ dee4eb65-145c-4b2b-98a7-ae3a7045f24b
md"""
We finally load our data, see [`artificial_S2_composite_bezier_curve`](https://manoptjl.org/stable/helpers/data.html#Manopt.artificial_S2_composite_bezier_curve-Tuple{}), a composite Bezier curve consisting of 3 segments on the Sphere. The middle segment consists of the control points
"""

# ╔═╡ 2b485956-459d-40e7-9364-79ea77073804
begin
    B = artificial_S2_composite_bezier_curve()
    b = B[2].pts
end

# ╔═╡ 98c925ac-9ac8-42f3-9855-a01e5eaac350
M = Sphere(2)

# ╔═╡ d22e825e-d8ff-44fe-8b06-b14c09765180
md"""
On Euclidean spaces Bézier curves of these ``n=4`` so called control points like this segment yield polynomials of degree``3``.
The resulting curve ``γ: [0,1] → ℝ^m`` is called [Bezier curve](https://en.wikipedia.org/wiki/Bézier_curve) or Bézier spline and is named after [Piérre Bezier](https://en.wikipedia.org/wiki/Pierre_Bézier) (1910–1999).
They can be evaluated by the de Casteljau algorithm by evaluating line segments between points.
While it is not easy to evaluate polynomials on a manifold, evaluating line segments generalizes to the evaluation of [`shortest_geodesic`](https://juliamanifolds.github.io/Manifolds.jl/latest/interface.html#ManifoldsBase.shortest_geodesic-Tuple{Manifold,Any,Any})s.
We will illustrate this using these points ``b=(b_1,b_2,b_3,b_4)`` on the [`Sphere`](https://juliamanifolds.github.io/Manifolds.jl/latest/manifolds/sphere.html) ``\mathbb S^2``.
Let's evaliuate this at the point ``t=\frac{1}{4}∈[0,1]``. We first compute
"""

# ╔═╡ 157d13c8-7311-4832-a9ff-da22125b2bbf
begin
    t = 0.66
    pts1 = shortest_geodesic.(Ref(M), b[1:3], b[2:4], Ref(t))
end

# ╔═╡ 5af58241-50c7-464e-bfa3-53819e605793
md"""
We obtain 3 points on the geodesics connecting the control points. Repeating this again twice
"""

# ╔═╡ b372fcfc-cad0-48bc-b759-54d15bd2e28b
begin
    pts2 = shortest_geodesic.(Ref(M), pts1[1:2], pts1[2:3], Ref(t))
    p = shortest_geodesic(M, pts2[1], pts2[2], t)
end

# ╔═╡ 62115e3f-cdfe-464c-91a1-eb80e674a610
md"""we obtain the point on the Bézier curve ``c(t)``."""

# ╔═╡ 8653ca13-21e9-4049-a81e-02b6047d2c66
begin
    curves1 = [shortest_geodesic(M, b[i], b[i + 1], geo_pts) for i in 1:3]
    curves2 = [shortest_geodesic(M, pts1[i], pts1[i + 1], geo_pts) for i in 1:2]
    bezier_curve = [shortest_geodesic(M, pts2[1], pts2[2], geo_pts)]
end;

# ╔═╡ 05612a12-1080-4b79-b505-198e079a4b2a
render_asy && begin
    asymptote_export_S2_signals(
        image_prefix * "/Casteljau-illustr.asy";
        points=[b, pts1, pts2, [p]],
        curves=[curves1..., curves2..., bezier_curve..., de_casteljau(M, B[2], geo_pts)],
        colors=Dict(
            :points =>
                [TolVibrantBlue, TolVibrantCyan, TolVibrantTeal, TolVibrantOrange],
            :curves => [
                TolVibrantBlue,
                TolVibrantBlue,
                TolVibrantBlue,
                TolVibrantCyan,
                TolVibrantCyan,
                TolVibrantTeal,
                black,
            ],
        ),
        dot_size=3.5,
        line_widths=[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.5],
        camera_position=camera_position,
    )
    render_asymptote(image_prefix * "/Casteljau-illustr.asy"; render=2)
end;

# ╔═╡ bdb85fb8-0707-47d1-95ee-9bf2e56500f5
md"""This image summarizes Casteljaus algorithm:"""

# ╔═╡ 38b8f8af-ba81-4d3b-83b7-56b0d41ea708
PlutoUI.LocalResource(image_prefix * "/Casteljau-illustr.png")

# ╔═╡ d6759a8e-a705-4d95-a8da-e172ada644d8
md"""
From the control points (blue) and their geodesics, ont evaluation per geodesic yields three interims points (cyan), their two successive geodeics another two points (teal) and at its geodesic at ``t=0.66`` we obtain the point on the curve.

In Manopt.jl, to evaluate a Bézier curve knowing its [`BezierSegment`](https://manoptjl.org/stable/functions/bezier.html#Manopt.BezierSegment), use [`de_casteljau`](https://manoptjl.org/stable/functions/bezier.html#Manopt.de_casteljau-Tuple{AbstractManifold,%20Vararg{Any,%20N}%20where%20N}).

There are a few nice observations to make, that hold also for these Bézier curves on manifolds:
* The curve starts in the first controlpoint ``b_0`` and ends in the last controlpoint ``b_3``
* The tangent vector to the curve at the start ``\dot c(0)`` is equal to ``\frac{1}{3}\log_{b_0}b_1 = \dot γ_{b_0,b_1}(0)``, where ``γ_{a,b}`` denotes the shortest geodesic between ``a`` and ``b``.
* The tangent vector to the curve at the end ``\dot c(1)`` is equal to ``-\frac{1}{3}\log_{b_3}b_2 = -γ_{b_3,b_2}(0) = \dot γ_{b_2,b_3}(1)``.
* the curve is differentiable.

For more details on these properties, see for example [^PopielNoakes2007].
"""

# ╔═╡ 7c6947a0-7d0b-4459-b999-58ea4f2b0169
md"""
## Composite Bézier curves

With the properties of a single Bézier curve, also called Bézier segment, we can “stitch” curves together. Let ``a_0,…,a_n`` and ``b_0,…,b_m`` be two sets of controlpoints for the Bézier segments ``c(t)`` and ``d(t)``, respectively.
We define the composite Bézier curve by ``B(t) = \begin{cases} c(t) & \text{ if } 0\leq t < 1, \\ d(t-1) & \text{ if } 1\leq t \leq 2,\end{cases}`` where ``t∈[0,2]``.
This can of course be generalised straight forward to more than two cases.
With the properties from the previous section we can now state that

* the curve ``B(t)`` is continuous if ``c(1)=d(0)`` or in other words ``a_n=b_0``
* the curve ``B(t)`` is differentiable if additionally ``\dot c(1)=\dot d(0)`` or in
other words ``-\log_{a_n}a_{n-1} = \log_{b_0}b_1``. This is equivalent to ``a_n=b_0 = \gamma_{a_{n-1}b_1}(\tfrac{1}{2})``.

One nice interpretation of the last characterization is, that the tangents ``\log_{a_n}a_{n-1}`` and ``\log_{b_0}b_1`` point into opposite directions.
For a continuous curve, the first point of every segment (except for the first segment) can be ommitted, for a differentiable curve the first two points (except for the first segment) can be ommitted.
You can reduce storage by calling [`get_bezier_points`](https://manoptjl.org/stable/functions/bezier.html#Manopt.get_bezier_points), though for econstruciton with [`get_bezier_segments`](https://manoptjl.org/stable/functions/bezier.html#Manopt.get_bezier_segments-Union{Tuple{P},%20Tuple{AbstractManifold,%20Vector{P},%20Any},%20Tuple{AbstractManifold,%20Vector{P},%20Any,%20Symbol}}%20where%20P) you also need [`get_bezier_degrees`](https://manoptjl.org/stable/functions/bezier.html#Manopt.get_bezier_degrees-Tuple{AbstractManifold,%20AbstractVector{var%22#s56%22}%20where%20var%22#s56%22%3C:BezierSegment}).
The reduced storage is represented as an array of points, i.e. an element of the corresponding [`PowerManifold`](https://juliamanifolds.github.io/Manifolds.jl/latest/manifolds/power.html).

For the three segment example from the beginning this looks as follows
"""

# ╔═╡ 3f1d0989-2f97-4945-82f8-22e9c7bda6d2
render_asy && begin
    asymptote_export_S2_signals(
        image_prefix * "/Bezier-composite-curve.asy";
        curves=[de_casteljau(M, B, bezier_pts)],
        points=[get_bezier_junctions(M, B), get_bezier_inner_points(M, B)],
        tangent_vectors=[[
            Tuple(a) for a in zip(
                get_bezier_junctions(M, B, true),
                get_bezier_junction_tangent_vectors(M, B),
            )
        ]],
        colors=Dict(
            :curves => [black],
            :points => [TolVibrantBlue, TolVibrantTeal],
            :tvectors => [TolVibrantCyan],
        ),
        camera_position=camera_position,
        arrow_head_size=10.0,
        line_widths=[1.5, 1.5],
        dot_size=4.0,
    )
    render_asymptote(image_prefix * "/Bezier-composite-curve.asy"; render=2)
end;

# ╔═╡ 41ba8ac2-5a81-46f5-9269-404cc6a1d287
PlutoUI.LocalResource(image_prefix * "/Bezier-composite-curve.png")

# ╔═╡ e7fc05b3-a198-4a64-b62a-43cdf8f56238
md"""
## Minimizing the acceleration of a composite Bézier curve
The motivation to minimize the acceleration of the composite Bézier curve is, that the curve should get “straighter” or more geodesic like.
If we discretize the curve ``B(t)`` with its control points denoted by ``b_{i,j}`` for the ``j``th note in the ``i``th segment, the discretized model for equispaced ``t_i``, ``i=0,…,N`` in the domain of ``B`` reads[^BergmannGousenbourger2018]

```math
A(b) :eqq\sum_{i=1}^{N-1}\frac{\mathrm{d}^2_2 \bigl[ B(t_{i-1}), B(t_{i}), B(t_{i+1}) \bigr]}{\Delta_t^3},
```

where $\mathrm{d}_2$ denotes the second order finite difference using the mid point approach, see [`costTV2`](@ref)[^BacakBergmannSteidlWeinmann2016],

```math
d_2(x,y,z) := \min_{c ∈ \mathcal C_{x,z}} d_{\mathcal M}(c,y),\qquad x,y,z∈\mathcal M.
```

Another model is based on logarithmic maps, see [^BoumalAbsil2011], but that is not considered here.
An advantage of the model considered here is, that it only consist of the evaluation of geodesics.
This yields a gradient of ``A(b)`` with respect to ``b`` [`adjoint_Jacobi_field`](https://manoptjl.org/stable/functions/Jacobi_fields.html#Manopt.adjoint_Jacobi_field)s. The following image shows the negative gradient (scaled)
"""

# ╔═╡ 7b545a1e-b9c3-4c57-a5ee-3a629407b67a
render_asy && begin
    gradFullB = Manopt._grad_acceleration_bezier(
        M,
        get_bezier_points(M, B, :differentiable),
        [3, 3, 3],
        collect(range(0.0, 3.0; length=151)),
    )
    asymptote_export_S2_signals(
        image_prefix * "/Bezier-composite-curve-gradient.asy";
        curves=[de_casteljau(M, B, bezier_pts)],
        points=[get_bezier_junctions(M, B), get_bezier_inner_points(M, B)],
        tangent_vectors=[[
            Tuple(a) for a in zip(
                get_bezier_points(M, B, :continuous),
                -0.05 .* get_bezier_points(M, gradFullB, :continuous),
            )
        ]],
        colors=Dict(
            :curves => [black],
            :points => [TolVibrantBlue, TolVibrantTeal],
            :tvectors => [TolVibrantOrange],
        ),
        camera_position=camera_position,
        arrow_head_size=10.0,
        line_widths=[1.5, 1.5],
        dot_size=4.0,
    )
    render_asymptote(image_prefix * "/Bezier-composite-curve-gradient.asy"; render=2)
end;

# ╔═╡ 5157ac76-d359-4f18-883d-8afbe78d3b4d
PlutoUI.LocalResource(image_prefix * "/Bezier-composite-curve-gradient.png")

# ╔═╡ 7640628b-e232-4ba3-94aa-d8a00218ff6a
md"""
In the following we consider two cases: Interpolation, which fixes the junction and end points of ``B(t)``
and approximation, where a weight and a dataterm are additionally introduced.
"""

# ╔═╡ da970042-8326-4d17-b440-ab1960f10283
md"""
### Interpolation

For interpolation, the junction points are fixed and their gradient entries are hence set to zero.
After transferring to the already mentioned [`PowerManifold`](https://juliamanifolds.github.io/Manifolds.jl/latest/manifolds/power.html), we can then perform a  [`gradient_descent`](https://manoptjl.org/stable/solvers/gradient_descent.html)  as follows
"""

# ╔═╡ 7db3e138-6ec1-4a0c-8f3b-279f7251b4c3
begin
    #number of sampling points refers to exactness of approximating d^2
    curve_samples = collect(range(0.0, 3.0; length=151))
    pB = get_bezier_points(M, B, :differentiable)
    N = PowerManifold(M, NestedPowerRepresentation(), length(pB))
    function F(M, pB)
        return cost_acceleration_bezier(
            M.manifold, pB, get_bezier_degrees(M.manifold, B), curve_samples
        )
    end
    function gradF(M, pB)
        return grad_acceleration_bezier(
            M.manifold, pB, get_bezier_degrees(M.manifold, B), curve_samples
        )
    end
    x0 = get_bezier_points(M, B, :differentiable)
    pB_opt_ip = gradient_descent(
        N,
        F,
        gradF,
        x0;
        stepsize=ArmijoLinesearch(1.0, ExponentialRetraction(), 0.5, 0.0001),
        stopping_criterion=StopWhenChangeLess(5 * 10.0^(-7)),
    )
end;

# ╔═╡ 5e8e98f6-8896-462e-b610-4b379c727a72
md"and the resut looks like"

# ╔═╡ 50854504-af68-4d7f-999e-f3054c088fb6
render_asy && begin
    B_opt_ip = get_bezier_segments(M, pB_opt_ip, get_bezier_degrees(M, B), :differentiable)
    asymptote_export_S2_signals(
        image_prefix * "/Bezier-IP-Min.asy";
        curves=[de_casteljau(M, B_opt_ip, bezier_pts), de_casteljau(M, B, bezier_pts)],
        points=[get_bezier_junctions(M, B_opt_ip), get_bezier_inner_points(M, B_opt_ip)],
        tangent_vectors=[[
            Tuple(a) for a in zip(
                get_bezier_junctions(M, B_opt_ip, true),
                get_bezier_junction_tangent_vectors(M, B_opt_ip),
            )
        ]],
        colors=Dict(
            :curves => [TolVibrantBlue, black],
            :points => [TolVibrantBlue, TolVibrantTeal],
            :tvectors => [TolVibrantCyan],
        ),
        camera_position=camera_position,
        arrow_head_size=10.0,
        line_widths=[1.5, 0.75, 1.5],
        dot_size=4.0,
    )
    render_asymptote(image_prefix * "/Bezier-IP-Min.asy"; render=2)
end;

# ╔═╡ 4f35b3e8-90ca-4d21-a377-007f93eb0efd
PlutoUI.LocalResource(image_prefix * "/Bezier-IP-Min.png")

# ╔═╡ 203a0538-d608-4d3e-b57d-edec20756300
md"""
Where the original courve is shown in black and the interpolating curve with minimized (discretized) acceleration is shown in blue including its junction points (also blue), tangent vectors (light blue) and control points (teal).
"""

# ╔═╡ 3e43f3fe-5c9a-4fb5-9033-db4a34bac5a5
md"""
### Approximation
Similarly if we introduce the junction points as data fixed given ``d_i`` and set (for simplicity) ``p_i=b_{i,0}`` and ``p_{n+1}=b_{n,4}`` and ``λ=3`` in

```math
\frac{λ}{2}\sum_{k=0}^3 d_{\mathcal M}(d_i,p_i)^2 + A(b),
```

then ``λ`` models how important closeness to the data ``d_i`` is.

Then we obtain the folowing code to minimize the acceleration while approximating the original curve
"""

# ╔═╡ 2f9ef252-5046-4af7-aa69-449a04bab26e
begin
    λ = 3.0
    d = get_bezier_junctions(M, B)
    function F2(M, pB)
        return cost_L2_acceleration_bezier(
            M.manifold, pB, get_bezier_degrees(M.manifold, B), curve_samples, λ, d
        )
    end
    function gradF2(M, pB)
        return grad_L2_acceleration_bezier(
            M.manifold, pB, get_bezier_degrees(M.manifold, B), curve_samples, λ, d
        )
    end
    x1 = get_bezier_points(M, B, :differentiable)
    pB_opt_appr = gradient_descent(
        N,
        F2,
        gradF2,
        x1;
        stepsize=ArmijoLinesearch(1.0, ExponentialRetraction(), 0.5, 0.001),
        stopping_criterion=StopWhenChangeLess(10.0^(-5)),
    )
end

# ╔═╡ 6cfb38a8-c29f-42a6-b642-31c047dd6cb8
render_asy && begin
    B_opt_appr = get_bezier_segments(M, pB_opt_appr, get_bezier_degrees(M, B), :differentiable)
    asymptote_export_S2_signals(
        image_prefix * "/Bezier-Appr-Min.asy";
        curves=[de_casteljau(M, B_opt_appr, bezier_pts), de_casteljau(M, B, bezier_pts)],
        points=[
            d,
            get_bezier_junctions(M, B_opt_appr),
            get_bezier_inner_points(M, B_opt_appr),
        ],
        tangent_vectors=[[
            Tuple(a) for a in zip(
                get_bezier_junctions(M, B_opt_appr, true),
                get_bezier_junction_tangent_vectors(M, B_opt_appr),
            )
        ]],
        colors=Dict(
            :curves => [TolVibrantBlue, black],
            :points => [TolVibrantOrange, TolVibrantBlue, TolVibrantTeal],
            :tvectors => [TolVibrantCyan],
        ),
        camera_position=camera_position,
        arrow_head_size=10.0,
        line_widths=[1.5, 0.75, 1.5],
        dot_size=4.0,
    )
    render_asymptote(image_prefix * "/Bezier-Appr-Min.asy"; render=2)
end;

# ╔═╡ db8110e0-d271-424c-b920-817dc1aab0c9
PlutoUI.LocalResource(image_prefix * "/Bezier-Appr-Min.png")

# ╔═╡ 0855f66e-f9af-4c63-b041-424c9f3fd655
md"""Additionally to the last image, the data points ``d_i`` (junction points of the original curve) are shown in orange, the distance between these and the blue junction points is part of the cost function here."""

# ╔═╡ 67e8e794-a16a-4d39-804a-6bfb350f8247
md"""
The role of ``λ`` can be interpreted as follows: for large values of $λ$, the
minimizer, i.e. the resulting curve, is closer to the original Bézier junction points.
For small ``λ`` the resting curve is closer to a geodesic and the control points are closer to the curve.
For ``λ=0`` _any_ (not necessarily shortest) geodesic is a solution and the problem is ill-posed.
To illustrate the effect of ``λ``, the following image contains 1000 runs for ``λ=10`` in dark currant to ``λ=0.01`` in bright yellow.

![Approximation min Acc](https://manoptjl.org/stable/assets/images/tutorials/Bezier_Approximation_video-Summary-result.png)

The effect of the data term can also be seen in the following video, which starts a little slow and takes about 40 seconds.

![Video of the effect of lambda, the weight of the dataterm](https://manoptjl.org/stable/assets/videos/tutorials/Bezier_Approximation_video-movie.mp4)
"""

# ╔═╡ 446f7276-7ed8-4b5c-a3f1-fdde4fab1479
md"""
## Literature

[^BacakBergmannSteidlWeinmann2016]:
    > Bačák, M., Bergmann, R., Steidl, G. and Weinmann, A.: _A second order nonsmooth
    > variational model for restoring manifold-valued images_,
    > SIAM Journal on Scientific Computations, Volume 38, Number 1, pp. A567–597,
    > doi: [10.1137/15M101988X](https://doi.org/10.1137/15M101988X),
    > arXiv: [1506.02409](https://arxiv.org/abs/1506.02409)
[^BergmannGousenbourger2018]:
    > Bergmann, R. and Gousenbourger, P.-Y.: _A variational model for data fitting on manifolds
    > by minimizing the acceleration of a Bézier curve_.
    > Frontiers in Applied Mathematics and Statistics, 2018.
    > doi: [10.3389/fams.2018.00059](https://dx.doi.org/10.3389/fams.2018.00059),
    > arXiv: [1807.10090](https://arxiv.org/abs/1807.10090)
[^BoumalAbsil2011]:
    > Boumal, N. and Absil, P.-A.: _A discrete regression method on manifolds and its application to data on SO(n)._
    > In: IFAC Proceedings Volumes (IFAC-PapersOnline). Vol. 18. Milano (2011). p. 2284–89.
    > doi: [10.3182/20110828-6-IT-1002.00542](https://dx.doi.org/10.3389/10.3182/20110828-6-IT-1002.00542),
    > web: [www](https://web.math.princeton.edu/~nboumal/papers/Boumal_Absil_A_discrete_regression_method_on_manifolds_and_its_application_to_data_on_SOn.pdf)
[^Casteljau1959]:
    > de Casteljau, P.: _Outillage methodes calcul_, Enveloppe Soleau 40.040 (1959),
    > Institute National de la Propriété Industrielle, Paris.
[^Casteljau1963]:
    > de Casteljau, P.: _Courbes et surfaces à pôles_, Microfiche P 4147-1,
    > André Citroën Automobile SA, Paris, (1963).
[^PopielNoakes2007]:
    > Popiel, T. and Noakes, L.: _Bézier curves and ``C^2`` interpolation in Riemannian
    > manifolds_. Journal of Approximation Theory (2007), 148(2), pp. 111–127.-
    > doi: [10.1016/j.jat.2007.03.002](https://doi.org/10.1016/j.jat.2007.03.002).
"""

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"

[compat]
Colors = "~0.12.8"
Manifolds = "~0.7.1"
Manopt = "~0.3.14"
PlutoUI = "~0.7.22"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

julia_version = "1.7.2"
manifest_format = "2.0"

[[deps.AbstractFFTs]]
deps = ["LinearAlgebra"]
git-tree-sha1 = "485ee0867925449198280d4af84bdb46a2a404d0"
uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
version = "1.0.1"

[[deps.AbstractPlutoDingetjes]]
deps = ["Pkg"]
git-tree-sha1 = "abb72771fd8895a7ebd83d5632dc4b989b022b5b"
uuid = "6e696c72-6542-2067-7265-42206c756150"
version = "1.1.2"

[[deps.ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[deps.ArnoldiMethod]]
deps = ["LinearAlgebra", "Random", "StaticArrays"]
git-tree-sha1 = "62e51b39331de8911e4a7ff6f5aaf38a5f4cc0ae"
uuid = "ec485272-7323-5ecc-a04f-4719b315124d"
version = "0.2.0"

[[deps.ArrayInterface]]
deps = ["Compat", "IfElse", "LinearAlgebra", "Requires", "SparseArrays", "Static"]
git-tree-sha1 = "265b06e2b1f6a216e0e8f183d28e4d354eab3220"
uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
version = "3.2.1"

[[deps.Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[deps.Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[deps.ChainRulesCore]]
deps = ["Compat", "LinearAlgebra", "SparseArrays"]
git-tree-sha1 = "f885e7e7c124f8c92650d61b9477b9ac2ee607dd"
uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
version = "1.11.1"

[[deps.ChangesOfVariables]]
deps = ["LinearAlgebra", "Test"]
git-tree-sha1 = "9a1d594397670492219635b35a3d830b04730d62"
uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
version = "0.1.1"

[[deps.ColorSchemes]]
deps = ["ColorTypes", "Colors", "FixedPointNumbers", "Random"]
git-tree-sha1 = "a851fec56cb73cfdf43762999ec72eff5b86882a"
uuid = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
version = "3.15.0"

[[deps.ColorTypes]]
deps = ["FixedPointNumbers", "Random"]
git-tree-sha1 = "024fe24d83e4a5bf5fc80501a314ce0d1aa35597"
uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
version = "0.11.0"

[[deps.Colors]]
deps = ["ColorTypes", "FixedPointNumbers", "Reexport"]
git-tree-sha1 = "417b0ed7b8b838aa6ca0a87aadf1bb9eb111ce40"
uuid = "5ae59095-9a9b-59fe-a467-6f913c188581"
version = "0.12.8"

[[deps.Compat]]
deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
git-tree-sha1 = "dce3e3fea680869eaa0b774b2e8343e9ff442313"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "3.40.0"

[[deps.CompilerSupportLibraries_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae"

[[deps.CovarianceEstimation]]
deps = ["LinearAlgebra", "Statistics", "StatsBase"]
git-tree-sha1 = "bc3930158d2be029e90b7c40d1371c4f54fa04db"
uuid = "587fd27a-f159-11e8-2dae-1979310e6154"
version = "0.2.6"

[[deps.DataAPI]]
git-tree-sha1 = "cc70b17275652eb47bc9e5f81635981f13cea5c8"
uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
version = "1.9.0"

[[deps.DataStructures]]
deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
git-tree-sha1 = "7d9d316f04214f7efdbb6398d545446e246eff02"
uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
version = "0.18.10"

[[deps.Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[deps.DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[deps.DensityInterface]]
deps = ["InverseFunctions", "Test"]
git-tree-sha1 = "80c3e8639e3353e5d2912fb3a1916b8455e2494b"
uuid = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
version = "0.4.0"

[[deps.Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[deps.Distributions]]
deps = ["ChainRulesCore", "DensityInterface", "FillArrays", "LinearAlgebra", "PDMats", "Printf", "QuadGK", "Random", "SparseArrays", "SpecialFunctions", "Statistics", "StatsBase", "StatsFuns", "Test"]
git-tree-sha1 = "7f3bec11f4bcd01bc1f507ebce5eadf1b0a78f47"
uuid = "31c24e10-a181-5473-b8eb-7969acd0382f"
version = "0.25.34"

[[deps.DocStringExtensions]]
deps = ["LibGit2"]
git-tree-sha1 = "b19534d1895d702889b219c382a6e18010797f0b"
uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
version = "0.8.6"

[[deps.Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[deps.Einsum]]
deps = ["Compat"]
git-tree-sha1 = "4a6b3eee0161c89700b6c1949feae8b851da5494"
uuid = "b7d42ee7-0b51-5a75-98ca-779d3107e4c0"
version = "0.4.1"

[[deps.FillArrays]]
deps = ["LinearAlgebra", "Random", "SparseArrays", "Statistics"]
git-tree-sha1 = "8756f9935b7ccc9064c6eef0bff0ad643df733a3"
uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
version = "0.12.7"

[[deps.FixedPointNumbers]]
deps = ["Statistics"]
git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc"
uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
version = "0.8.4"

[[deps.Graphs]]
deps = ["ArnoldiMethod", "DataStructures", "Distributed", "Inflate", "LinearAlgebra", "Random", "SharedArrays", "SimpleTraits", "SparseArrays", "Statistics"]
git-tree-sha1 = "92243c07e786ea3458532e199eb3feee0e7e08eb"
uuid = "86223c79-3864-5bf0-83f7-82e725a168b6"
version = "1.4.1"

[[deps.HybridArrays]]
deps = ["LinearAlgebra", "Requires", "StaticArrays"]
git-tree-sha1 = "c1d5b1dcdf2140644e1c6beb9ca09fbed601c241"
uuid = "1baab800-613f-4b0a-84e4-9cd3431bfbb9"
version = "0.4.9"

[[deps.Hyperscript]]
deps = ["Test"]
git-tree-sha1 = "8d511d5b81240fc8e6802386302675bdf47737b9"
uuid = "47d2ed2b-36de-50cf-bf87-49c2cf4b8b91"
version = "0.0.4"

[[deps.HypertextLiteral]]
git-tree-sha1 = "2b078b5a615c6c0396c77810d92ee8c6f470d238"
uuid = "ac1192a8-f4b3-4bfe-ba22-af5b92cd3ab2"
version = "0.9.3"

[[deps.IOCapture]]
deps = ["Logging", "Random"]
git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a"
uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
version = "0.2.2"

[[deps.IfElse]]
git-tree-sha1 = "debdd00ffef04665ccbb3e150747a77560e8fad1"
uuid = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
version = "0.1.1"

[[deps.Inflate]]
git-tree-sha1 = "f5fc07d4e706b84f72d54eedcc1c13d92fb0871c"
uuid = "d25df0c9-e2be-5dd7-82c8-3ad0b3e990b9"
version = "0.1.2"

[[deps.InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[deps.InverseFunctions]]
deps = ["Test"]
git-tree-sha1 = "a7254c0acd8e62f1ac75ad24d5db43f5f19f3c65"
uuid = "3587e190-3f89-42d0-90ee-14403ec27112"
version = "0.1.2"

[[deps.IrrationalConstants]]
git-tree-sha1 = "7fd44fd4ff43fc60815f8e764c0f352b83c49151"
uuid = "92d709cd-6900-40b7-9082-c6be49f344b6"
version = "0.1.1"

[[deps.JLLWrappers]]
deps = ["Preferences"]
git-tree-sha1 = "642a199af8b68253517b80bd3bfd17eb4e84df6e"
uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
version = "1.3.0"

[[deps.JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[deps.Kronecker]]
deps = ["LinearAlgebra", "NamedDims", "SparseArrays", "StatsBase"]
git-tree-sha1 = "a51f46415c844dee694cb8b20a3fcbe6dba342c2"
uuid = "2c470bb0-bcc8-11e8-3dad-c9649493f05e"
version = "0.5.1"

[[deps.LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[deps.LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[deps.LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[deps.LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[deps.Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[deps.LinearAlgebra]]
deps = ["Libdl", "libblastrampoline_jll"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[deps.LogExpFunctions]]
deps = ["ChainRulesCore", "ChangesOfVariables", "DocStringExtensions", "InverseFunctions", "IrrationalConstants", "LinearAlgebra"]
git-tree-sha1 = "be9eef9f9d78cecb6f262f3c10da151a6c5ab827"
uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
version = "0.3.5"

[[deps.Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[deps.MacroTools]]
deps = ["Markdown", "Random"]
git-tree-sha1 = "3d3e902b31198a27340d0bf00d6ac452866021cf"
uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
version = "0.5.9"

[[deps.Manifolds]]
deps = ["Colors", "Distributions", "Einsum", "Graphs", "HybridArrays", "Kronecker", "LinearAlgebra", "ManifoldsBase", "Markdown", "Random", "RecipesBase", "RecursiveArrayTools", "Requires", "SimpleWeightedGraphs", "SpecialFunctions", "StaticArrays", "Statistics", "StatsBase"]
git-tree-sha1 = "3811935fd2549b0f5f9b365d6e7173fcdb98ec9a"
uuid = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
version = "0.7.1"

[[deps.ManifoldsBase]]
deps = ["LinearAlgebra", "Markdown"]
git-tree-sha1 = "77a5949567437d185ee929c405e3c6c0768118ea"
uuid = "3362f125-f0bb-47a3-aa74-596ffd7ef2fb"
version = "0.12.9"

[[deps.Manopt]]
deps = ["ColorSchemes", "ColorTypes", "Colors", "DataStructures", "Dates", "LinearAlgebra", "ManifoldsBase", "Markdown", "Random", "Requires", "SparseArrays", "StaticArrays", "Statistics", "Test"]
git-tree-sha1 = "e87c818e19f79444fff1c77880cdb43b51974dfe"
uuid = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
version = "0.3.14"

[[deps.Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[deps.MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[deps.Missings]]
deps = ["DataAPI"]
git-tree-sha1 = "bf210ce90b6c9eed32d25dbcae1ebc565df2687f"
uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
version = "1.0.2"

[[deps.Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[deps.MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[deps.NamedDims]]
deps = ["AbstractFFTs", "ChainRulesCore", "CovarianceEstimation", "LinearAlgebra", "Pkg", "Requires", "Statistics"]
git-tree-sha1 = "4f34e3ff2fa7f2a1c03fb2c4fdd637380d760bbc"
uuid = "356022a1-0364-5f58-8944-0da4b18d706f"
version = "0.2.42"

[[deps.NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[deps.OpenBLAS_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"

[[deps.OpenLibm_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "05823500-19ac-5b8b-9628-191a04bc5112"

[[deps.OpenSpecFun_jll]]
deps = ["Artifacts", "CompilerSupportLibraries_jll", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "13652491f6856acfd2db29360e1bbcd4565d04f1"
uuid = "efe28fd5-8261-553b-a9e1-b2916fc3738e"
version = "0.5.5+0"

[[deps.OrderedCollections]]
git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.4.1"

[[deps.PDMats]]
deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
git-tree-sha1 = "ee26b350276c51697c9c2d88a072b339f9f03d73"
uuid = "90014a1f-27ba-587c-ab20-58faa44d9150"
version = "0.11.5"

[[deps.Parsers]]
deps = ["Dates"]
git-tree-sha1 = "ae4bbcadb2906ccc085cf52ac286dc1377dceccc"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.1.2"

[[deps.Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[deps.PlutoUI]]
deps = ["AbstractPlutoDingetjes", "Base64", "Dates", "Hyperscript", "HypertextLiteral", "IOCapture", "InteractiveUtils", "JSON", "Logging", "Markdown", "Random", "Reexport", "UUIDs"]
git-tree-sha1 = "565564f615ba8c4e4f40f5d29784aa50a8f7bbaf"
uuid = "7f904dfe-b85e-4ff6-b463-dae2292396a8"
version = "0.7.22"

[[deps.Preferences]]
deps = ["TOML"]
git-tree-sha1 = "00cfd92944ca9c760982747e9a1d0d5d86ab1e5a"
uuid = "21216c6a-2e73-6563-6e65-726566657250"
version = "1.2.2"

[[deps.Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[deps.QuadGK]]
deps = ["DataStructures", "LinearAlgebra"]
git-tree-sha1 = "78aadffb3efd2155af139781b8a8df1ef279ea39"
uuid = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
version = "2.4.2"

[[deps.REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[deps.Random]]
deps = ["SHA", "Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[deps.RecipesBase]]
git-tree-sha1 = "6bf3f380ff52ce0832ddd3a2a7b9538ed1bcca7d"
uuid = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
version = "1.2.1"

[[deps.RecursiveArrayTools]]
deps = ["ArrayInterface", "ChainRulesCore", "DocStringExtensions", "FillArrays", "LinearAlgebra", "RecipesBase", "Requires", "StaticArrays", "Statistics", "ZygoteRules"]
git-tree-sha1 = "c944fa4adbb47be43376359811c0a14757bdc8a8"
uuid = "731186ca-8d62-57ce-b412-fbd966d074cd"
version = "2.20.0"

[[deps.Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[deps.Requires]]
deps = ["UUIDs"]
git-tree-sha1 = "4036a3bd08ac7e968e27c203d45f5fff15020621"
uuid = "ae029012-a4dd-5104-9daa-d747884805df"
version = "1.1.3"

[[deps.Rmath]]
deps = ["Random", "Rmath_jll"]
git-tree-sha1 = "bf3188feca147ce108c76ad82c2792c57abe7b1f"
uuid = "79098fc4-a85e-5d69-aa6a-4863f24498fa"
version = "0.7.0"

[[deps.Rmath_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "68db32dff12bb6127bac73c209881191bf0efbb7"
uuid = "f50d1b31-88e8-58de-be2c-1cc44531875f"
version = "0.3.0+0"

[[deps.SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[deps.Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[deps.SharedArrays]]
deps = ["Distributed", "Mmap", "Random", "Serialization"]
uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"

[[deps.SimpleTraits]]
deps = ["InteractiveUtils", "MacroTools"]
git-tree-sha1 = "5d7e3f4e11935503d3ecaf7186eac40602e7d231"
uuid = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
version = "0.9.4"

[[deps.SimpleWeightedGraphs]]
deps = ["Graphs", "LinearAlgebra", "Markdown", "SparseArrays", "Test"]
git-tree-sha1 = "a6f404cc44d3d3b28c793ec0eb59af709d827e4e"
uuid = "47aef6b3-ad0c-573a-a1e2-d07658019622"
version = "1.2.1"

[[deps.Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[deps.SortingAlgorithms]]
deps = ["DataStructures"]
git-tree-sha1 = "b3363d7460f7d098ca0912c69b082f75625d7508"
uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
version = "1.0.1"

[[deps.SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[deps.SpecialFunctions]]
deps = ["ChainRulesCore", "IrrationalConstants", "LogExpFunctions", "OpenLibm_jll", "OpenSpecFun_jll"]
git-tree-sha1 = "f0bccf98e16759818ffc5d97ac3ebf87eb950150"
uuid = "276daf66-3868-5448-9aa4-cd146d93841b"
version = "1.8.1"

[[deps.Static]]
deps = ["IfElse"]
git-tree-sha1 = "e7bc80dc93f50857a5d1e3c8121495852f407e6a"
uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
version = "0.4.0"

[[deps.StaticArrays]]
deps = ["LinearAlgebra", "Random", "Statistics"]
git-tree-sha1 = "3c76dde64d03699e074ac02eb2e8ba8254d428da"
uuid = "90137ffa-7385-5640-81b9-e52037218182"
version = "1.2.13"

[[deps.Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[deps.StatsAPI]]
git-tree-sha1 = "0f2aa8e32d511f758a2ce49208181f7733a0936a"
uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0"
version = "1.1.0"

[[deps.StatsBase]]
deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"]
git-tree-sha1 = "2bb0cb32026a66037360606510fca5984ccc6b75"
uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
version = "0.33.13"

[[deps.StatsFuns]]
deps = ["ChainRulesCore", "InverseFunctions", "IrrationalConstants", "LogExpFunctions", "Reexport", "Rmath", "SpecialFunctions"]
git-tree-sha1 = "bedb3e17cc1d94ce0e6e66d3afa47157978ba404"
uuid = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
version = "0.9.14"

[[deps.SuiteSparse]]
deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"]
uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"

[[deps.TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[deps.Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[deps.Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[deps.UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[deps.Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[deps.Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[deps.ZygoteRules]]
deps = ["MacroTools"]
git-tree-sha1 = "8c1a8e4dfacb1fd631745552c8db35d0deb09ea0"
uuid = "700de1a5-db45-46bc-99cf-38207098b444"
version = "0.2.2"

[[deps.libblastrampoline_jll]]
deps = ["Artifacts", "Libdl", "OpenBLAS_jll"]
uuid = "8e850b90-86db-534c-a0d3-1478176c7d93"

[[deps.nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[deps.p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─e701b513-4ad9-474c-9d79-60a24f545ebd
# ╠═6f9be298-5374-11ec-3b9d-fd74a32d6b53
# ╠═087fb979-d68d-4ee7-9be5-fad992fdb792
# ╠═85d8985b-f5ee-4cba-bcc5-528049798330
# ╟─e62e36e9-19ce-4a7d-8bda-5634813c1ee9
# ╠═18d378dd-abc6-4fa6-a2c9-1cc09e4cec46
# ╟─dee4eb65-145c-4b2b-98a7-ae3a7045f24b
# ╠═2b485956-459d-40e7-9364-79ea77073804
# ╟─98c925ac-9ac8-42f3-9855-a01e5eaac350
# ╟─d22e825e-d8ff-44fe-8b06-b14c09765180
# ╠═157d13c8-7311-4832-a9ff-da22125b2bbf
# ╟─5af58241-50c7-464e-bfa3-53819e605793
# ╠═b372fcfc-cad0-48bc-b759-54d15bd2e28b
# ╟─62115e3f-cdfe-464c-91a1-eb80e674a610
# ╠═8653ca13-21e9-4049-a81e-02b6047d2c66
# ╠═05612a12-1080-4b79-b505-198e079a4b2a
# ╟─bdb85fb8-0707-47d1-95ee-9bf2e56500f5
# ╠═38b8f8af-ba81-4d3b-83b7-56b0d41ea708
# ╟─d6759a8e-a705-4d95-a8da-e172ada644d8
# ╟─7c6947a0-7d0b-4459-b999-58ea4f2b0169
# ╠═3f1d0989-2f97-4945-82f8-22e9c7bda6d2
# ╟─41ba8ac2-5a81-46f5-9269-404cc6a1d287
# ╟─e7fc05b3-a198-4a64-b62a-43cdf8f56238
# ╠═7b545a1e-b9c3-4c57-a5ee-3a629407b67a
# ╟─5157ac76-d359-4f18-883d-8afbe78d3b4d
# ╟─7640628b-e232-4ba3-94aa-d8a00218ff6a
# ╟─da970042-8326-4d17-b440-ab1960f10283
# ╠═7db3e138-6ec1-4a0c-8f3b-279f7251b4c3
# ╟─5e8e98f6-8896-462e-b610-4b379c727a72
# ╠═50854504-af68-4d7f-999e-f3054c088fb6
# ╟─4f35b3e8-90ca-4d21-a377-007f93eb0efd
# ╟─203a0538-d608-4d3e-b57d-edec20756300
# ╟─3e43f3fe-5c9a-4fb5-9033-db4a34bac5a5
# ╠═2f9ef252-5046-4af7-aa69-449a04bab26e
# ╠═6cfb38a8-c29f-42a6-b642-31c047dd6cb8
# ╠═db8110e0-d271-424c-b920-817dc1aab0c9
# ╟─0855f66e-f9af-4c63-b041-424c9f3fd655
# ╟─67e8e794-a16a-4d39-804a-6bfb350f8247
# ╟─446f7276-7ed8-4b5c-a3f1-fdde4fab1479
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


# # Bayesian Neural Network

# We borrow this tutorial from the [official Turing Docs](https://turing.ml/dev/tutorials/03-bayesian-neural-network/). We will show how the explicit parameterization of Lux enables first-class composability with packages which expect flattened out parameter vectors.

# We will use [Turing.jl](https://turing.ml) with [Lux.jl](https://lux.csail.mit.edu/stable) to implement implementing a classification algorithm. Lets start by importing the relevant libraries

## Import libraries.
using Lux
using Pkg #hide
Pkg.activate(joinpath(dirname(pathof(Lux)), "..", "examples")) #hide
using Turing, Plots, Random, ReverseDiff, NNlib, Functors

## Hide sampling progress.
Turing.setprogress!(false);

## Use reverse_diff due to the number of parameters in neural networks.
Turing.setadbackend(:reversediff)

# ## Generating data

# Our goal here is to use a Bayesian neural network to classify points in an artificial dataset. The code below generates data points arranged in a box-like pattern and displays a graph of the dataset we'll be working with.

## Number of points to generate.
N = 80
M = round(Int, N / 4)
rng = Random.default_rng()
Random.seed!(rng, 1234)

## Generate artificial data.
x1s = rand(rng, Float32, M) * 4.5f0;
x2s = rand(rng, Float32, M) * 4.5f0;
xt1s = Array([[x1s[i] + 0.5f0; x2s[i] + 0.5f0] for i in 1:M])
x1s = rand(rng, Float32, M) * 4.5f0;
x2s = rand(rng, Float32, M) * 4.5f0;
append!(xt1s, Array([[x1s[i] - 5.0f0; x2s[i] - 5.0f0] for i in 1:M]))

x1s = rand(rng, Float32, M) * 4.5f0;
x2s = rand(rng, Float32, M) * 4.5f0;
xt0s = Array([[x1s[i] + 0.5f0; x2s[i] - 5.0f0] for i in 1:M])
x1s = rand(rng, Float32, M) * 4.5f0;
x2s = rand(rng, Float32, M) * 4.5f0;
append!(xt0s, Array([[x1s[i] - 5.0f0; x2s[i] + 0.5f0] for i in 1:M]))

## Store all the data for later.
xs = [xt1s; xt0s]
ts = [ones(2 * M); zeros(2 * M)]

## Plot data points.
function plot_data()
    x1 = first.(xt1s)
    y1 = last.(xt1s)
    x2 = first.(xt0s)
    y2 = last.(xt0s)

    plt = Plots.scatter(x1, y1; color="red", clim=(0, 1))
    Plots.scatter!(plt, x2, y2; color="blue", clim=(0, 1))

    return plt
end

plot_data()

# ## Building the Neural Network

# The next step is to define a feedforward neural network where we express our parameters as distributions, and not single points as with traditional neural networks. For this we will use `Dense` to define liner layers and compose them via `Chain`, both are neural network primitives from `Lux`. The network `nn` we will creat will have two hidden layers with `tanh` activations and one output layer with `sigmoid` activation, as shown below.

# The `nn` is an instance that acts as a function and can take data, parameters and current state as inputs and output predictions. We will define distributions on the neural network parameters.

## Construct a neural network using Lux
nn = Chain(Dense(2, 3, tanh), Dense(3, 2, tanh), Dense(2, 1, sigmoid))

## Initialize the model weights and state
ps, st = Lux.setup(rng, nn)

Lux.parameterlength(nn) # number of paraemters in NN

# The probabilistic model specification below creates a parameters variable, which has IID normal variables. The parameters represents all parameters of our neural net (weights and biases).

## Create a regularization term and a Gaussian prior variance term.
alpha = 0.09
sig = sqrt(1.0 / alpha)

# Construct named tuple from a sampled parameter vector. We could also use ComponentArrays here and simply broadcast to avoid doing this. But let's do it this way to avoid dependencies.
function vector_to_parameters(ps_new::AbstractVector, ps::NamedTuple)
    @assert length(ps_new) == Lux.parameterlength(ps)
    i = 1
    function get_ps(x)
        z = reshape(view(ps_new, i:(i + length(x) - 1)), size(x))
        i += length(x)
        return z
    end
    return fmap(get_ps, ps)
end

## Specify the probabilistic model.
@model function bayes_nn(xs, ts)
    global st

    ## Sample the parameters
    nparameters = Lux.parameterlength(nn)
    parameters ~ MvNormal(zeros(nparameters), sig .* ones(nparameters))

    ## Forward NN to make predictions
    preds, st = nn(xs, vector_to_parameters(parameters, ps), st)

    ## Observe each prediction.
    for i in 1:length(ts)
        ts[i] ~ Bernoulli(preds[i])
    end
end

# Inference can now be performed by calling sample. We use the HMC sampler here.

## Perform inference.
N = 5000
ch = sample(bayes_nn(hcat(xs...), ts), HMC(0.05, 4), N)

# Now we extract the parameter samples from the sampled chain as theta (this is of size `5000 x 20` where `5000` is the number of iterations and `20` is the number of parameters). We'll use these primarily to determine how good our model's classifier is.

## Extract all weight and bias parameters.
theta = MCMCChains.group(ch, :parameters).value;

# ## Prediction Visualization

## A helper to run the nn through data `x` using parameters `theta`
nn_forward(x, theta) = nn(x, vector_to_parameters(theta, ps), st)[1]

## Plot the data we have.
plot_data()

## Find the index that provided the highest log posterior in the chain.
_, i = findmax(ch[:lp])

## Extract the max row value from i.
i = i.I[1]

## Plot the posterior distribution with a contour plot
x1_range = collect(range(-6; stop=6, length=25))
x2_range = collect(range(-6; stop=6, length=25))
Z = [nn_forward([x1, x2], theta[i, :])[1] for x1 in x1_range, x2 in x2_range]
contour!(x1_range, x2_range, Z)

# The contour plot above shows that the MAP method is not too bad at classifying our data. Now we can visualize our predictions.

# $p(\tilde{x} | X, \alpha) = \int_{\theta} p(\tilde{x} | \theta) p(\theta | X, \alpha) \approx \sum_{\theta \sim p(\theta | X, \alpha)}f_{\theta}(\tilde{x})$

# The `nn_predict` function takes the average predicted value from a network parameterized by weights drawn from the MCMC chain.

## Return the average predicted value across multiple weights.
function nn_predict(x, theta, num)
    return mean([nn_forward(x, view(theta, i, :))[1] for i in 1:10:num])
end

# Next, we use the `nn_predict` function to predict the value at a sample of points where the x1 and x2 coordinates range between -6 and 6. As we can see below, we still have a satisfactory fit to our data, and more importantly, we can also see where the neural network is uncertain about its predictions much easier---those regions between cluster boundaries.

# Plot the average prediction.
plot_data()

n_end = 1500
x1_range = collect(range(-6; stop=6, length=25))
x2_range = collect(range(-6; stop=6, length=25))
Z = [nn_predict([x1, x2], theta, n_end)[1] for x1 in x1_range, x2 in x2_range]
contour!(x1_range, x2_range, Z)

# Suppose we are interested in how the predictive power of our Bayesian neural network evolved between samples. In that case, the following graph displays an animation of the contour plot generated from the network weights in samples 1 to 1,000.

## Number of iterations to plot.
n_end = 1000

anim = @gif for i in 1:n_end
    plot_data()
    Z = [nn_forward([x1, x2], theta[i, :])[1] for x1 in x1_range, x2 in x2_range]
    contour!(x1_range, x2_range, Z; title="Iteration $i", clim=(0, 1))
end every 5


abstract type UStatistics end

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

struct OneSampleUstatistics{N} <: UStatistics
  h
end

struct TwoSampleUstatistics{N, M} <: UStatistics
  h
end

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

compute(s::UStatistics, X::AbstractArray) = compute(s, s.h, X)
compute(s::UStatistics, X::AbstractArray, Y::AbstractArray) = compute(s, s.h, X, Y)

function compute(::OneSampleUstatistics{1}, h, X::AbstractVector)
  v = zero(eltype(X))
  n = 0
  for x in X
    n += 1
    v += h(x)
  end

  v / n
end

function compute(::OneSampleUstatistics{1}, h, X::AbstractMatrix)
  n = size(X, 1)
  v = 0.
  @inbounds for i=1:n
    v += h(view(X, i, :))
  end

  v / n
end

function compute(::OneSampleUstatistics{2}, h, X::AbstractVector)
  n = length(X)
  v = 0.
  @inbounds for i1=1:n
    for i2=i1+1:n
      v += h( X[i1], X[i2] )::Float64
    end
  end

  2. * v / (n * (n-1))
end

function compute(::OneSampleUstatistics{2}, h, X::AbstractMatrix)
  n = size(X, 1)
  v = 0.
  @inbounds for i1=1:n
    for i2=i1+1:n
      v += h( view(X, i1, :),
              view(X, i2, :) )
    end
  end

  2. * v / n / (n-1)
end

function compute(::TwoSampleUstatistics{1, 1}, h, X::AbstractVector, Y::AbstractVector)
  nx = length(X)
  ny = length(Y)

  v = 0.
  for i=1:nx
    for j=1:ny
      v += h(X[i], Y[j])
    end
  end
  v / nx / ny
end

function compute(s::TwoSampleUstatistics{1, 1}, X::AbstractMatrix, Y::AbstractMatrix)
  h = s.h
  nx = size(X, 1)
  ny = size(Y, 1)

  v = 0.
  for i=1:nx
    for j=1:ny
      v += h(view(X, i, :), view(Y, j, :))
    end
  end
  v / nx / ny
end

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

variance(s::UStatistics, X::AbstractArray) = variance(s, s.h, X)
variance(s::UStatistics, X::AbstractArray, Y::AbstractArray) = variance(s, s.h, X, Y)

function variance(::OneSampleUstatistics{1}, h, X::AbstractVector)
  n = length(X)
  q = map(h, X)
  t = mean(q)
  (sum(abs2, q) - n * t^2) / (n - 1)
end

function variance(::OneSampleUstatistics{1}, h, X::AbstractMatrix)
  n = size(X, 1)
  q = mapslices(h, X; dims = 2)
  t = mean(q)
  (sum(abs2, q) - n * t^2) / (n - 1)
end

function variance(::TwoSampleUstatistics{1, 1}, h, X::AbstractMatrix, Y::AbstractMatrix)
  nx = size(X, 1)
  ny = size(Y, 1)

  q = zeros(nx)
  @inbounds for k=1:nx
    vx = view(X, k, :)
    for j=1:ny
      q[k] += h(vx, view(Y, j, :))
    end
    q[k] /= ny
  end

  r = zeros(ny)
  @inbounds for k=1:ny
    vy = view(Y, k, :)
    for j=1:nx
      r[k] += h(view(X, j, :), vy)
    end
    r[k] /= nx
  end

  t = mean(q)

  σ1 = sum(abs2, q) / (nx - 1) - nx / (nx - 1) * t^2
  σ2 = sum(abs2, r) / (ny - 1) - ny / (ny - 1) * t^2

  return σ1/nx + σ2/ny
end

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

### examples

function createVarianceStatistic()
  @inline function h(x, y)
    (x - y)^2 / 2.
  end
  OneSampleUstatistics{2}(h)
end

function createCovarianceStatistic(row::Int, col::Int)
  @inline function h(X, Y)
      @inbounds (X[row] - Y[row]) * (X[col] - Y[col]) / 2.
  end
  OneSampleUstatistics{2}(h)
end

function createSecondMomentStatistic()
  h(X) = X*X
  OneSampleUstatistics{1}(h)
end

function createSecondMomentStatistic(row::Int, col::Int)
  @inbounds h(X) = X[row] * X[col]
  OneSampleUstatistics{1}(h)
end


# # Free-vibration solution for a clamped plate in fluid

# ## Description

# Free-vibration solution for a clamped plate fully submerged in inviscid
# fluid. The plate is modeled with energy-stabilized NICE tetrahedra.

# For the dry plate, these are the natural frequencies:

# | Mode | Dry frequencies [rad/s] |
# | ----- | ---------------------- |
# | 1 | 12.45   |
# | 2 | 29.44   |
# | 3 | 75.05   |
# | 4 | 94.27   |
# | 5 | 106.79 |

# For the immersed (wet) plate, these are the natural frequencies:

# | Mode | Wet frequencies [rad/s] |
# | ----- | ---------------------- |
# | 1 | 7.35   |
# | 2 | 20.2   |
# | 3 | 50.45   |
# | 4 | 70.41   |
# | 5 | 78.85 |

# ## References

# [1] Fu, Y., and Price, W. G., 1987, “Interactions Between a Partially
# or Totally Immersed Vibrating Cantilever Plate and the Surrounding Fluid,”
# J. Sound Vib., 118(3), pp. 495–513.

# ## Goals

# - Define model data for an algorithm to solve the free vibration problem.

##
# ## Definitions

# Bring in needed basic facilities.
using LinearAlgebra
using Arpack

# Bring in the finite element tools, and the linear deformation support.
using FinEtools
using FinEtoolsDeforLinear
# The dry-vibration solver is provided by this algorithm module.
using FinEtoolsDeforLinear.AlgoDeforLinearModule

# The wet-vibration solver is provided by this algorithm module.
using FinEtoolsVibInFluids.AlgoVibInFluidsModule

# The basic data corresponds to a steel plate.
E = 206000*phun("MPa");# Young's modulus
nu = 0.3;# Poisson ratio
rho = 7850*phun("KG*M^-3");# mass density
# The fluid properties correspond to water.
rhow = 1000*phun("KG*M^-3");
# Dimensions of the plate: it is quite massive.
Length= 10e3*phun("MM"); Width= 10e3*phun("MM"); Height= 0.238e3*phun("MM");
# The mesh is quite coarse, but still provides engineering accuracy for the
# angular frequencies.
nHeight = 4; nLength  = 2*10; nWidth = nLength;
tolerance = Height/nHeight/100;

# The reference angular frequencies in radians per second. The dry frequencies:
dreffs = [12.45153,   29.44122,   75.04792,   94.26498,  106.7881]*phun("rad/s")
# The reference angular frequencies in radians per second. The wet frequencies:
wreffs = [7.35, 20.2, 50.45, 70.41, 78.85]*phun("rad/s")
# Solve for this many natural frequencies.
neigvs = length(wreffs);

# Generate a tetrahedral-mesh block.
fens,fes  = T4block(Length, Width, Height, nLength, nWidth, nHeight)

# Define the material model.
MR = DeforModelRed3D
material = MatDeforElastIso(MR, rho, E, nu, 0.0)

# The finite element machine for the interior of the domain: we used to
# energy-stabilized tetrahedra which are quite effective for modeling thin
# structures.
femm = FEMMDeforLinearESNICET4(MR, IntegDomain(fes, NodalSimplexRule(3)), material)

# The plate is clamped along $X = 0$. Select the nodes using the box criterion.
# Set all the displacements components to zero.
nl1 = selectnode(fens, box=[0.0 0.0 -Inf Inf -Inf Inf], inflate=tolerance)
ebc1 = FDataDict("node_list"=>nl1, "component"=>1, "displacement"=>0.0)
ebc2 = FDataDict("node_list"=>nl1, "component"=>2, "displacement"=>0.0)
ebc3 = FDataDict("node_list"=>nl1, "component"=>3, "displacement"=>0.0)

# Make the region
region1 = FDataDict("femm"=>femm, "femm_mass"=>femm)

# Make model data
modeldata =  FDataDict("fens"=> fens, "regions"=>  [region1], "essential_bcs"=>[ebc1 ebc2 ebc3], "omega_shift"=>0.0, "neigvs"=>neigvs)

# Solve for the in-vacuo free vibration modes
modeldata = AlgoDeforLinearModule.modal(modeldata)

println("Dry angular frequencies: $(modeldata["omega"]) [rad/s]")
println("Reference dry ang. fre.: $(dreffs) [rad/s]")

# Determine the wet boundary faces. Both square faces of the plate are exposed
# to the fluid as it is fully immersed. We select the faces inside the boxes.
bfes = meshboundary(fes)
rl = vcat(selectelem(fens, bfes, box = [-Inf Inf -Inf Inf Height Height], inflate = tolerance),
	selectelem(fens, bfes, box = [-Inf Inf -Inf Inf 0.0 0.0], inflate = tolerance));
# These are the wet surface panels. Note that the algorithm assumes that these
# are three-node triangular finite elements.
modeldata["wet_boundary_fes"] = subset(bfes, rl);
# The only property of the fluid that matters is the mass density.
modeldata["rhow"] = rhow

# Solve for the wet-vibration modes. The solution is returned in the model data.
modeldata = AlgoVibInFluidsModule.modal(modeldata)

println("Wet angular frequencies: $(modeldata["wet_omega"]) [rad/s]")
println("Reference wet ang. fre.: $(wreffs) [rad/s]")

# Export the fluid pressure modes, and the vibration modes of the structure. The
# names of the VTK files are left up to the algorithms. We should be able to
# locate them in the current (working) folder.
modeldata = AlgoVibInFluidsModule.exportfluidpressuremode(modeldata)
modeldata["postprocessing"]["mode"] = 1:neigvs
modeldata = AlgoVibInFluidsModule.exportmode(modeldata)

true


x = range(0, stop=1, length=20)
plot(x, x.^2, x.^3, lw=0.05, ls=["--", "-"], marker=["wo", "none"],
     keys=["quadratic", "cubic"],)
scatter!(x, sin.(x), key=["sine"])
legend(; position="top-left")


"""
Returns an array of tuples of directions.

```julia
    get_directions(dim::Integer; include_zero::Bool = false) -> Array{NTuple{N, Integer}, 1}

julia> get_directions(3)
27-element Array{Tuple{Int64,Int64,Int64},1}:
 (-1, -1, -1)
 (0, -1, -1)
 (1, -1, -1)
 (-1, 0, -1)
 (0, 0, -1)
 (1, 0, -1)
 (-1, 1, -1)
 (0, 1, -1)
 (1, 1, -1)
 (-1, -1, 0)
 (0, -1, 0)
 (1, -1, 0)
 (-1, 0, 0)
 (0, 0, 0)
 (1, 0, 0)
 (-1, 1, 0)
 (0, 1, 0)
 (1, 1, 0)
 (-1, -1, 1)
 (0, -1, 1)
 (1, -1, 1)
 (-1, 0, 1)
 (0, 0, 1)
 (1, 0, 1)
 (-1, 1, 1)
 (0, 1, 1)
 (1, 1, 1)
```
"""
function get_directions(dim::T; include_zero::Bool = false) where T <: Integer
    D = reshape(collect(Iterators.product([-one(T):one(T) for _ in 1:dim]...)), :)
    return include_zero ? D : filter(d -> d ≠ 𝟘(dim), D)
end
get_directions(M::AbstractArray{T, N}; include_zero::Bool = false) where {T, N} = get_directions(ndims(M), include_zero = include_zero)
# get_directions(M::AbstractArray{T, N}; include_zero::Bool = false) where {T, N} = get_directions(N, include_zero = include_zero)

function _expand_as_required(M::AbstractArray{T, N}, expand_by::T, inds::AbstractIndexOrIndices{N}...) where {T, N}
    indices = Vector{Union{T, Nothing}}()
    ind_extrema = extrema_indices(inds...)

    for d in 1:ndims(M)
        ith_extrema = ind_extrema[d]
        if ! all(map(m -> m ≤ size(M, d), ith_extrema))
            for invalid_idx in filter(i -> i > size(M, i) || i < 1, ith_extrema)
                difference = invalid_idx - size(M, d)
                M = sign(difference) == 1 ? append_n_times(M, abs(difference), expand_by, dims = d) : append_n_times_backwards(M, abs(difference), expand_by, dims = d)
            end
        end
    end

    return M
end

"""
    expand_as_required(M::AbstractArray{T, N}, expand_by::T, inds::AbstractIndexOrIndices{N})

Given indices, `expand_as_required` will fill in a matrix with `expand_by` where needed if such indices are not currently accessible.  If no `expand_by` is given, and the element type of the matrix is a number, fills with zero.

```julia
julia> A = rand(Int8, 2, 2)
2×2 Array{Int8,2}:
 -96   83
 -94  -39

julia> expand_as_required(A, zero(Int8), [(1, 1), (2, 2), (3, 3)])
3×3 Array{Int8,2}:
 -96   83  0
 -94  -39  0
   0    0  0
```
"""
expand_as_required(M::AbstractArray{T, N}, expand_by::T, inds::AbstractIndices{N}) where {T, N} = _expand_as_required(M, expand_by, inds)
expand_as_required(M::AbstractArray{T, N}, expand_by::T, inds::AbstractIndex{N}...) where {T, N} = _expand_as_required(M, expand_by, inds)
expand_as_required(M::AbstractArray{T, N}, inds::AbstractIndices{N}) where {T <: Number, N} = _expand_as_required(M, zero(T), inds)
expand_as_required(M::AbstractArray{T, N}, inds::AbstractIndex{N}...) where {T <: Number, N} = _expand_as_required(M, zero(T), inds)

"""
    adjacencies(M::AbstractArray{T, N}, idx::AbstractIndex{N})

`adjacencies` will get all adjacent element to an index in a given matrix.
"""
function adjacencies(M::AbstractArray{T, N}, idx::AbstractIndex{N}) where {T, N}
    direction_multipliers = get_directions(M, include_zero = false)
    adjacent_indices = AbstractIndex{N}[idx .+ j for j in direction_multipliers]
    return T[k for k in tryindex(M, adjacent_indices) if ! isnothing(k)]
end

"""
    expanded_adjacencies(M::Array{T, N}, expand_by::T, idx::AbstractIndex{N})

`expanded_adjacencies` will get all adjacencies of an index in the matrix M, given that the matrix is infinitely expanding by a single element beyond the specified M.  See also `adjacencies` and `expand_as_required`.  If no `expand_by` is given, and the element type of the matrix is a number, fills with zero.
"""
function expanded_adjacencies(M::Array{T, N}, expand_by::T, idx::AbstractIndex{N}) where {T, N}
    𝟎 = ntuple(_ -> zero(Int), N)
    idx = Tuple(idx)
    direction_multipliers = get_directions(M, include_zero = false)
    adjacent_indices = AbstractIndex{N}[idx .+ j for j in direction_multipliers]
    M = expand_as_required(M, expand_by, adjacent_indices)
    idx_shift = abs.(size(M) .- idx) # in case we have had to shift the array for non-positive indices
    return T[M[(i .+ idx_shift .- 1)...] for i in adjacent_indices]
end
expanded_adjacencies(M::Array{T, N}, idx::AbstractIndex{N}) where {T <: Number, N} = expanded_adjacencies(M, zero(T), idx)

"""
    n_adjacent_to(M::AbstractArray{T, N}, idx::AbstractIndex{N}, adj_elem::T)

Given a matrix M, counts the number of adjacent elements to index `idx` that are exactly the `adj_elem`.
"""
function n_adjacent_to(M::AbstractArray{T, N}, idx::AbstractIndex{N}, adj_elem::T) where {N, T}
    n = 0

    for i in adjacencies(M, idx)
        if i == adj_elem
            n += 1
        end
    end

    return n
end

"""
    expanded_n_adjacent_to(M::AbstractArray{T, N}, expand_by::T, idx::AbstractIndex{N}, adj_elem::T)

Given a matrix M, counts the number of adjacent elements to index `idx` that are exactly the `adj_elem`, expanding the matrix if needed (see `expanded_adjacencies`).  If no `expand_by` is given, and the element type of the matrix is a number, fills with zero.
"""
function expanded_n_adjacent_to(M::AbstractArray{T, N}, expand_by::T, idx::AbstractIndex{N}, adj_elem::T) where {N, T}
    n = 0

    for i in expanded_adjacencies(M, expand_by, idx)
        if i == adj_elem
            n += 1
        end
    end

    return n
end
expanded_n_adjacent_to(M::AbstractArray{T, N}, idx::AbstractIndex{N}, adj_elem::T) where {N, T <: Number} = expanded_n_adjacent_to(M, zero(T), idx, adj_elem)

"""
    global_adjacencies(M::AbstractArray{T, N}, idx::AbstractIndex{N}, ignored_elem::T)

Returns the "globally" adjacent indices in arbitrary positions in the cardinal directions from a specified index in matrix `M`, ignoring certain adjacent elements (i.e., skipping over them).
"""
function global_adjacencies_indices(M::AbstractArray{T, N}, idx::AbstractIndex{N}, ignored_elem::T) where {T, N}
    adjacent_indices = AbstractIndex{N}[]
    directional_shifts = get_directions(M, include_zero = false)
    n_adjacent, adjacent_count = n_adjacencies(M), 0

    while adjacent_count < n_adjacent
        for directional_shift in directional_shifts
            while true
                adj_index = idx .+ directional_shift

                if nothing ∈ tryindex(M, adj_index)
                    n_adjacent -= 1
                    break
                end

                if M[adj_index...] ≠ ignored_elem
                    adjacent_count += 1
                    push!(adjacent_indices, adj_index)
                    break
                else
                    directional_shift = (abs.(directional_shift) .+ 1) .* sign.(directional_shift)
                end
            end
        end
    end

    return adjacent_indices
end

"""
    global_adjacencies(M::AbstractArray{T}, idx::AbstractIndex{N}, ignored_elem::T)

Using `global_adjacencies_indices`, returns the elements of each globally adjacent index.  This function is mainly used for `global_n_adjacent_to`.
"""
function global_adjacencies(M::AbstractArray{T, N}, idx::AbstractIndex{N}, ignored_elem::T) where {N, T}
    return T[M[i...] for i in global_adjacencies_indices(M, idx, ignored_elem)]
end

"""
    global_n_adjacent_to(M::AbstractArray{T, N},idx::AbstractIndex{N}, adj_elem::T, ignored_elem::T) where {N, T}

Given a matrix M, counts the number of adjacent elements to index `idx` that are exactly the `adj_elem`, skipping over certain elements if needed (see `expanded_adjacencies`).
"""
function global_n_adjacent_to(M::AbstractArray{T, N},idx::AbstractIndex{N}, adj_elem::T, ignored_elem::T) where {N, T}
    n = 0

    for i in global_adjacencies(M, idx, ignored_elem)
        if i == adj_elem
            n += 1
        end
    end

    return n
end


_hashpowers(u::UInt) = u
function _hashpowers(u::UInt, power::Tuple, powers...)
    if iszero(power[2])
        _hashpowers(u, powers...)
    else
        _hashpowers(hash(power, u), powers...)
    end
end
function Base.hash(m::AbstractMonomial, u::UInt)
    nnz = count(!iszero, exponents(m))
    if iszero(nnz)
        hash(1, u)
    elseif isone(nnz) && isone(degree(m))
        hash(variable(m), u)
    else
        _hashpowers(u, powers(m)...)
    end
end

function Base.hash(t::AbstractTerm, u::UInt)
    if iszero(t)
        hash(0, u)
    elseif coefficient(t) == 1
        hash(monomial(t), u)
    else
        hash(monomial(t), hash(coefficient(t), u))
    end
end

function Base.hash(p::AbstractPolynomial, u::UInt)
    if iszero(p)
        hash(0, u)
    else
        reduce((u, t) -> hash(t, u), terms(p); init=u)
    end
end


# ----------------------
# -- read             --
# ----------------------
function read_nodenum(skipnum)
    """
    xmax : 仮想セルも含めたnodeのxの数
    ymax : 仮想セルも含めたnodeのyの数
    """
    fff=[]
    open("grid/nodesnum", "r") do f
        fff=read(f,String)
    end
    fff=split(fff,"\n",keepempty=false)
    num_nodes=length(fff)-skipnum

    for i in 1+skipnum:length(fff)
        fff[i]=replace(fff[i]," \r" => "")
    end

    temp = split(fff[2]," ")
    xmax = parse(Int64,temp[1])
    ymax = parse(Int64,temp[2])

    return xmax, ymax
end

function read_nodes(skipnum,xmax,ymax)
    """
    nodes[i][j][k]
    i : x点の番号
    j : y点の番号
    k=1 : 点のx座標
    k=2 : 点のy座標
    """

    fff=[]
    open("grid/nodes", "r") do f
        fff=read(f,String)
    end
    fff=split(fff,"\n",keepempty=false)
    num_nodes=length(fff)-skipnum

    for i in 1+skipnum:length(fff)
        fff[i]=replace(fff[i]," \r" => "")
    end

    nodes=zeros(xmax,ymax,2)
    for i in 1:num_nodes
        temp=split(fff[i+skipnum]," ")

        x = parse(Int64,temp[1])
        y = parse(Int64,temp[2])
        nodes[x,y,1]=parse(Float64,temp[3])
        nodes[x,y,2]=parse(Float64,temp[4])
    end
    return nodes
end

function read_vecA(skipnum,xmax,ymax)
    """ vecAx[i][j][k]
        i : x番号
        j : y番号
        k=1 : x vec
        k=2 : y vec
    """
    fff=[]
    open("grid/vecAx", "r") do f
        fff=read(f,String)
    end
    fff=split(fff,"\n",keepempty=false)
    num_cell=length(fff)-skipnum

    for i in 1+skipnum:length(fff)
        fff[i]=replace(fff[i]," \r" => "")
    end

    vecAx=zeros(xmax,ymax-1,2)
    for i in 1:num_cell
        temp=split(fff[i+skipnum]," ")

        x = parse(Int64,temp[1])
        y = parse(Int64,temp[2])
        vecAx[x,y,1]=parse(Float64,temp[3])
        vecAx[x,y,2]=parse(Float64,temp[4])
    end

    fff=[]
    open("grid/vecAy", "r") do f
        fff=read(f,String)
    end
    fff=split(fff,"\n",keepempty=false)
    num_cell=length(fff)-skipnum

    for i in 1+skipnum:length(fff)
        fff[i]=replace(fff[i]," \r" => "")
    end

    vecAy=zeros(xmax-1,ymax,2)
    for i in 1:num_cell
        temp=split(fff[i+skipnum]," ")

        x = parse(Int64,temp[1])
        y = parse(Int64,temp[2])
        vecAy[x,y,1]=parse(Float64,temp[3])
        vecAy[x,y,2]=parse(Float64,temp[4])
    end
    return vecAx, vecAy
end

function read_allgrid()
    skip=1
    xmax,ymax=read_nodenum(skip)
    nodes=read_nodes(skip,xmax,ymax)
    vecAx, vecAy=read_vecA(skip,xmax,ymax)
    println("fin read grid")
    return xmax,ymax,nodes,vecAx,vecAy
end