text:
# License for this file: MIT (expat)
# Copyright 2017-2018, DLR Institute of System Dynamics and Control
#
# This file is part of module
#   Modia3D.Composition (Modia3D/Composition/_module.jl)
#

# Utility function that should not be directly called (only to be called from attach(..)
function attachAndReverseParents(newParent::Object3D, obj::Object3D)::Nothing
   @assert(!(newParent ≡ obj))

   # Save elements of obj
   r_rel  = obj.r_rel
   R_rel  = obj.R_rel
   parent = obj.parent

   # Reverse obj, so that newParent is the new parent
   obj.parent = newParent
   push!(newParent.children, obj)

   # Reverse all parents of obj
   obj1::Object3D = obj
   obj2::Object3D = parent
   while !(parent === obj1)
      parent_r_rel = parent.r_rel
      parent_R_rel = parent.R_rel

      if r_rel ≡ Modia3D.ZeroVector3D
         parent.r_rel = Modia3D.ZeroVector3D
         parent.r_abs = obj1.r_abs
      else
         parent.r_rel = -R_rel*r_rel
      end

      if R_rel ≡ Modia3D.NullRotation
         parent.R_rel = Modia3D.NullRotation
         parent.R_abs = obj1.R_abs
      else
         parent.R_rel = R_rel'
      end

      obj2 = parent
      parent = obj2.parent
      r_rel  = parent_r_rel
      R_rel  = parent_R_rel
      obj2.parent = obj1
      removeChild!(obj2, obj1)
      push!(obj1.children, obj2)
      obj1 = obj2
   end

   # If !isnothing(scene), visit all objs and set scene
#=
   if !isnothing(scene)
      stack = scene.stack
      empty!(stack)
      append!(stack, obj.children)
      while length(stack) > 0
         obj1 = pop!(stack)
         obj1.scene = scene
         append!(stack, obj1.children)
      end
   end
=#
   return nothing
end

"""
    (obj1, obj2, cutJoint) = attach(frame_a, frame_b)

If frame_a and frame_b do not have the same root, they are attached to each other
and cutJoint = false is returned.

If they have the same root, the tree is not modified and cutJoint=true is returned.
"""
function attach(obj1::Object3D, obj2::Object3D)
   root1 = rootObject3D(obj1)
   root2 = rootObject3D(obj2)
   #println("attach: obj1 = ", Modia3D.fullName(obj1), ", root = ", Modia3D.fullName(root1))
   #println("attach: obj2 = ", Modia3D.fullName(obj2), ", root = ", Modia3D.fullName(root2))

   if root1 ≡ root2
      # Compute absolute positions
      # updatePosition!(root1)
      return (obj1,obj2,true)
   end

   if isWorld(root2)
      attachAndReverseParents(obj2, obj1)
      return (obj2,obj1,false)
   else
      attachAndReverseParents(obj1, obj2)
      return (obj1,obj2,false)
   end
end

#=
function connect(obj1::Object3D, obj2::Object3D;
                 r::AbstractVector = Modia3D.ZeroVector3D,
                 R::Union{Frames.RotationMatrix,Nothing} = nothing,
                 q::Union{Frames.Quaternion,Nothing} = nothing,
                 fixed::Bool = true)::Nothing
   if !isnothing(R) && !isnothing(q)
      error("Modia3D.connect(...): either R or q must be nothing but both have a value.")
   end
   if !isnothing(R)
      Modia3D.assertRotationMatrix(R)
   elseif !isnothing(q)
      Modia3D.assertQuaternion(q)
   end

   (parentObject3D, obj, cutJoint) = attach(obj1, obj2)
  # println("... connect, fixed = ", fixed, ", obj1=",Modia3D.fullName(obj1), ", obj2 = ", Modia3D.fullName(obj2), ", obj = ", Modia3D.fullName(obj))

   if cutJoint
      error("Error from Modia3D.Composition.connect(", obj1.name, ",", obj2.name, "):\n",
            "Not yet supported to rigidly connect two objs that have the same root.")
   end

   r_rel = SVector{3,Float64}(r)
   R_rel = !isnothing(R) ? R                   :
   !isnothing(q) ? Modia3D.from_q(q) : Modia3D.NullRotation

   obj.r_rel = obj===obj2 ? r_rel : -R_rel*r_rel
   obj.R_rel = obj===obj2 ? R_rel :  R_rel'

   #r_abs = parent.r_abs + r_rel
   #R_abs = R_rel*parent.R_abs

   if fixed
      obj.joint = fixedJoint
   else
      obj1.hasChildJoint = true
      q_start = !isnothing(R) ? Modia3D.from_R(R) :
      !isnothing(q) ? q                   : Modia3D.NullQuaternion
      q_start = obj===obj2 ? q_start : Modia3D.inverseRotation(q_start)

      obj.joint = FreeMotion(obj, r_start = obj.r_rel, q_start = q_start)
   end
   return nothing
end
=#

function updatePosition!(obj::Object3D)::Nothing
   stack = Object3D[]
   # Push initial children on stack
   append!(stack, obj.children)
   while length(stack) > 0
      obj = pop!(stack)

      # computeKinematics!(obj.joint, obj, Modia3D.KinematicAnalysis, 0.0)

      parent = obj.parent

      if obj.r_rel ≡ Modia3D.ZeroVector3D
         obj.r_abs = parent.r_abs
      else
         obj.r_abs = parent.r_abs + parent.R_abs'*obj.r_rel
      end

      if obj.R_rel ≡ Modia3D.NullRotation
         obj.R_abs = parent.R_abs
      else
         obj.R_abs = obj.R_rel*parent.R_abs
      end

      if length(obj.visualizationFrame) == 1
         obj.visualizationFrame[1].r_abs = obj.r_abs
         obj.visualizationFrame[1].R_abs = obj.R_abs
      end

      append!(stack, obj.children)
   end
   return nothing
end

"""
    R = Rfromrot123(rot123::AbstractVector)

Return rotation matrix `R` from Cardan angles (rotation sequence x-y-z).
`rot123` are the Cardan angles (rotation sequence x-y-z) of rotation from frame `1` to frame `2`.
"""
function Rfromrot123(rot123::AbstractVector)

    (sal, cal) = sincos(rot123[1])
    (sbe, cbe) = sincos(rot123[2])
    (sga, cga) = sincos(rot123[3])
    return @SMatrix[ cbe*cga   sal*sbe*cga+cal*sga  -cal*sbe*cga+sal*sga ;
                    -cbe*sga  -sal*sbe*sga+cal*cga   cal*sbe*sga+sal*cga ;
                     sbe      -sal*cbe               cal*cbe             ]

end

"""
    w = wfromrot123(rot123::AbstractVector, derrot123::AbstractVector)

Return relative rotational velocity Vector3D `w` from frame `1` to frame `2` resolved in frame `2`.

`rot123` are the Cardan angles (rotation sequence x-y-z) of rotation from frame `1` to frame `2`.
`derrot123` are the time derivatives of `rot123`.
"""
function wfromrot123(rot123::AbstractVector, derrot123::AbstractVector)

    # 2_w_12 = (0, 0, gad) + Rga21*[(0, bed, 0) + Rbe21*(ald, 0, 0)]
    #        = (cbe*cga*ald + sga*bed, -cbe*sga*ald + cga*bed, sbe*ald + gad)
    (sbe, cbe) = sincos(rot123[2])
    (sga, cga) = sincos(rot123[3])
    return @SVector[ cbe*cga*derrot123[1] + sga*derrot123[2]                ,
                    -cbe*sga*derrot123[1] + cga*derrot123[2]                ,
                         sbe*derrot123[1]                    + derrot123[3] ]

end

# Next function only for backwards compatibility (do not use for new model)
computeKinematics!(scene::Scene, joint::Modia3D.AbstractJoint, obj::Object3D, analysis::Modia3D.AnalysisType, time::Float64)::Nothing =
    computeKinematics!(scene, [obj], time)

"""
    computeKinematics!(scene::Scene, tree::Vector{Object3D}, time)

Compute position, velocity, acceleration variables of the Object3Ds that are connected
in form of a tree. Variable `tree` contains the Object3Ds in a traversal order (e.g. pre-order traversal).
`tree[1]` is the root object. It is assumed that the kinematic
variables of tree[1].parent have a meaningful value.
"""
function computeKinematics!(scene::Scene, tree::Vector{Object3D}, time)::Nothing
    for obj in tree
        parent    = obj.parent
        jointKind = obj.jointKind

        if jointKind == FixTranslationKind
            obj.r_abs = parent.r_abs + parent.R_abs'*obj.r_rel
            obj.R_abs = parent.R_abs

            obj.v0 = parent.v0 + parent.R_abs'*cross(parent.w, obj.r_rel)
            obj.a0 = parent.a0 + parent.R_abs'*(cross(parent.z, obj.r_rel) + cross(parent.w, cross(parent.w, obj.r_rel)))

            obj.w = parent.w
            obj.z = parent.z

        elseif jointKind == FixKind
            obj.r_abs = parent.r_abs + parent.R_abs'*obj.r_rel
            obj.R_abs = obj.R_rel*parent.R_abs

            obj.v0 = parent.v0 + parent.R_abs'*cross(parent.w, obj.r_rel)
            obj.a0 = parent.a0 + parent.R_abs'*(cross(parent.z, obj.r_rel) + cross(parent.w, cross(parent.w, obj.r_rel)))

            obj.w = obj.R_rel*parent.w
            obj.z = obj.R_rel*parent.z

        elseif jointKind == RevoluteKind
            revolute  = scene.revolute[obj.jointIndex]

            obj.r_abs = parent.r_abs
            obj.R_rel = Frames.rotAxis(revolute.posAxis, revolute.posMovement, revolute.phi)
            obj.R_abs = obj.R_rel*parent.R_abs

            obj.v0 = parent.v0
            obj.a0 = parent.a0

            w_rel = Frames.axisValue(revolute.posAxis, revolute.posMovement, revolute.w)
            z_rel = Frames.axisValue(revolute.posAxis, revolute.posMovement, revolute.a)

            obj.w = obj.R_rel*(parent.w + w_rel)
            obj.z = obj.R_rel*(parent.z + z_rel + cross(parent.w, w_rel))

        elseif jointKind == PrismaticKind
            prismatic = scene.prismatic[obj.jointIndex]

            obj.r_rel = prismatic.eAxis*prismatic.s
            obj.r_abs = parent.r_abs + parent.R_abs'*obj.r_rel
            obj.R_abs = parent.R_abs

            v_rel  = prismatic.eAxis*prismatic.v
            a_rel  = prismatic.eAxis*prismatic.a
            obj.v0 = parent.v0 + parent.R_abs'*(v_rel + cross(parent.w, obj.r_rel))
            obj.a0 = parent.a0 + parent.R_abs'*(a_rel + cross(parent.z, obj.r_rel) + cross(parent.w, v_rel + cross(parent.w, obj.r_rel)) )

            obj.w  = parent.w
            obj.z  = parent.z

        elseif jointKind == AbsoluteFreeMotionKind
            freeMotion  = scene.freeMotion[obj.jointIndex]

            obj.r_rel = freeMotion.r
            obj.R_rel = Rfromrot123(freeMotion.rot)

            obj.r_abs = obj.r_rel
            obj.R_abs = obj.R_rel

            obj.v0 = freeMotion.v
            obj.a0 = freeMotion.a

            obj.w  = freeMotion.w
            obj.z  = freeMotion.z

        elseif jointKind == FreeMotionKind
            freeMotion  = scene.freeMotion[obj.jointIndex]

            obj.r_rel = freeMotion.r
            obj.R_rel = Rfromrot123(freeMotion.rot)

            obj.r_abs = parent.r_abs + parent.R_abs'*obj.r_rel
            obj.R_abs = obj.R_rel*parent.R_abs

            obj.v0 = parent.v0 + parent.R_abs'*(freeMotion.v + cross(parent.w, obj.r_rel))
            obj.a0 = parent.a0 + parent.R_abs'*(freeMotion.a + cross(parent.z, obj.r_rel) +
                                                    cross(parent.w, cross(parent.w, obj.r_rel)))
            obj.w  = parent.R_abs*parent.w + freeMotion.w
            obj.z  = parent.R_abs*parent.z + freeMotion.z

        else
            error("Bug in Modia3D/src/Composition/specifics/specifics.jl (computeKinematics!): jointKind = $jointKind is not known")
        end
    end
    return nothing
end

"""
    computeKinematics_for_leq_mode_pos!(scene::Scene, tree::Vector{Object3D}, time)

Compute accelerations that are only a function of qdd, but not of q and qd.
of the Object3Ds that are connected in form of a tree.
"""
function computeKinematics_for_leq_mode_pos!(scene::Scene, tree::Vector{Object3D}, time)::Nothing
    for obj in tree
        parent    = obj.parent
        jointKind = obj.jointKind

        if jointKind == FixTranslationKind
            obj.a0 = parent.a0 + parent.R_abs'*(cross(parent.z, obj.r_rel))
            obj.z  = parent.z

        elseif jointKind == FixKind
            obj.a0 = parent.a0 + parent.R_abs'*(cross(parent.z, obj.r_rel))
            obj.z  = obj.R_rel*parent.z

        elseif jointKind == RevoluteKind
            revolute = scene.revolute[obj.jointIndex]

            obj.a0 = parent.a0
            z_rel  = Frames.axisValue(revolute.posAxis, revolute.posMovement, revolute.a)
            obj.z  = obj.R_rel*(parent.z + z_rel)

        elseif jointKind == PrismaticKind
            prismatic = scene.prismatic[obj.jointIndex]

            eAxis  = prismatic.posMovement ? prismatic.posAxis : -prismatic.posAxis
            a_rel  = prismatic.eAxis*prismatic.a
            obj.a0 = parent.a0 + parent.R_abs'*(a_rel + cross(parent.z, obj.r_rel))
            obj.z  = parent.z

        elseif jointKind == AbsoluteFreeMotionKind
            freeMotion = scene.freeMotion[obj.jointIndex]

            obj.a0 = freeMotion.a
            obj.z  = freeMotion.z

        elseif jointKind == FreeMotionKind
            freeMotion = scene.freeMotion[obj.jointIndex]

            obj.a0 = parent.a0 + parent.R_abs'*(freeMotion.a + cross(parent.z, obj.r_rel))
            obj.z  = parent.R_abs*parent.z + freeMotion.z

        else
            error("Bug in Modia3D/src/Composition/specifics/specifics.jl (computeKinematics_for_leq_mode_pos!): jointKind = $jointKind is not known")
        end
    end
    return nothing
end

"""
    computeForcesTorquesAndResiduals!(scene::Scene, tree::Vector{Object3D}, time)

Compute forces/torques and residuals in a backward recursion from tree[end] to tree[1].
Variable `tree` contains the Object3Ds in a traversal order (e.g. pre-order traversal).
It is assumed that all force/torque variables are initialized (e.g. to zero), including
tree[1].parent.
"""
function computeForcesTorquesAndResiduals!(scene::Scene, tree::Vector{Object3D}, time)::Nothing
    for i = length(tree):-1:1
        obj       = tree[i]
        parent    = obj.parent
        jointKind = obj.jointKind

        if jointKind == FixTranslationKind
            parent.f += obj.f
            parent.t += obj.t + cross(obj.r_rel, obj.f)

        elseif jointKind == FixKind
            parent.f += obj.R_rel'*obj.f
            parent.t += obj.R_rel'*obj.t + cross(obj.r_rel, obj.f)

        elseif jointKind == RevoluteKind
            revolute  = scene.revolute[obj.jointIndex]
            parent.f += obj.R_rel'*obj.f
            parent.t += obj.R_rel'*obj.t
            obj.f     = -obj.f
            obj.t     = -obj.t
            revolute.residue = -revolute.tau + (revolute.posMovement ? obj.t[revolute.posAxis] : -obj.t[revolute.posAxis])

        elseif jointKind == PrismaticKind
            prismatic = scene.prismatic[obj.jointIndex]
            parent.f += obj.f
            parent.t += obj.t + cross(obj.r_rel, obj.f)
            obj.f     = -obj.f
            obj.t     = -obj.t
            prismatic.residue = -prismatic.f + dot(prismatic.eAxis,obj.f)

        elseif jointKind == AbsoluteFreeMotionKind || jointKind == FreeMotionKind
            freeMotion           = scene.freeMotion[obj.jointIndex]
            freeMotion.residue_f = obj.f
            freeMotion.residue_t = obj.t

        else
            error("Bug in Modia3D/src/Composition/specifics/specifics.jl (computeForcesTorquesAndResiduals!): jointKind = $jointKind is not known")
        end
    end
    return nothing
end

"""
    setJointVariables!(scene::Scene, objects::Vector{Object3D}, args, qdd, startIndex, ndof)

Copy specific variables into the corresponding Object3Ds.
"""
function setJointVariables!(scene::Scene, objects::Vector{Object3D}, args, qdd, startIndex, ndof)::Nothing
    for (i,obj) in enumerate(objects)
        jointKind = obj.jointKind
        args_i    = args[i]
        beg       = startIndex[i]

        if jointKind == RevoluteKind
            @assert(ndof[i] == 1)
            revolute     = scene.revolute[obj.jointIndex]
            revolute.phi = args_i[1]
            revolute.w   = args_i[2]
            revolute.a   = qdd[beg]
            revolute.tau = args_i[3]

        elseif jointKind == PrismaticKind
            @assert(ndof[i] == 1)
            prismatic   = scene.prismatic[obj.jointIndex]
            prismatic.s = args_i[1]
            prismatic.v = args_i[2]
            prismatic.a = qdd[beg]
            prismatic.f = args_i[3]

        elseif jointKind == AbsoluteFreeMotionKind || jointKind == FreeMotionKind
            @assert(ndof[i] == 6)
            freeMotion     = scene.freeMotion[obj.jointIndex]
            freeMotion.r   = SVector{3,Float64}(args_i[1])
            freeMotion.rot = SVector{3,Float64}(args_i[2])
            freeMotion.v   = SVector{3,Float64}(args_i[3])
            freeMotion.w   = SVector{3,Float64}(args_i[4])
            freeMotion.a   = qdd[beg+0:beg+2]
            freeMotion.z   = qdd[beg+3:beg+5]

        else
           error("Bug in Modia3D/src/Composition/specifics/specifics.jl (setJointVariables!): jointKind = $jointKind is not known")
        end
    end
    return nothing
end

"""
    getJointResiduals_for_leq_mode_0!(scene::Scene, objects::Vector{Object3D}, residuals, startIndex::Int, ndof::Int, cache_h)

Copy specific variables into their objects for leq_mode = 0.
"""
function getJointResiduals_for_leq_mode_0!(scene::Scene, objects::Vector{Object3D}, residuals, startIndex::Vector{Int}, ndof::Vector{Int}, cache_h)::Nothing
    for (i,obj) in enumerate(objects)
        jointKind = obj.jointKind
        beg       = startIndex[i]

        if jointKind == RevoluteKind
            @assert(ndof[i] == 1)
            revolute       = scene.revolute[obj.jointIndex]
            cache_h[  beg] = revolute.residue + revolute.tau
            residuals[beg] = revolute.residue

        elseif jointKind == PrismaticKind
            @assert(ndof[i] == 1)
            prismatic      = scene.prismatic[obj.jointIndex]
            cache_h[  beg] = prismatic.residue + prismatic.f
            residuals[beg] = prismatic.residue

        elseif jointKind == AbsoluteFreeMotionKind || jointKind == FreeMotionKind
            @assert(ndof[i] == 6)
            freeMotion             = scene.freeMotion[obj.jointIndex]
            cache_h[  beg+0:beg+2] = freeMotion.residue_f
            cache_h[  beg+3:beg+5] = freeMotion.residue_t

            residuals[beg+0:beg+2] = freeMotion.residue_f
            residuals[beg+3:beg+5] = freeMotion.residue_t

        else
           error("Bug in Modia3D/src/Composition/specifics/specifics.jl (getJointResiduals_for_leq_mode_0!): jointKind = $jointKind is not known")
        end
    end
    return nothing
end

"""
    getJointResiduals_for_leq_mode_pos!(scene::Scene, objects::Vector{Object3D}, residuals, startIndex::Int, ndof::Int, cache_h)

Copy specific variables into their objects for leq_mode > 0.
"""
function getJointResiduals_for_leq_mode_pos!(scene::Scene, objects::Vector{Object3D}, residuals, startIndex::Vector{Int}, ndof::Vector{Int}, cache_h)::Nothing
    for (i,obj) in enumerate(objects)
        jointKind = obj.jointKind
        beg       = startIndex[i]

        if jointKind == RevoluteKind
            @assert(ndof[i] == 1)
            residuals[beg] = scene.revolute[obj.jointIndex].residue + cache_h[beg]

        elseif jointKind == PrismaticKind
            @assert(ndof[i] == 1)
            residuals[beg] = scene.prismatic[obj.jointIndex].residue + cache_h[beg]

        elseif jointKind == AbsoluteFreeMotionKind || jointKind == FreeMotionKind
            @assert(ndof[i] == 6)
            freeMotion             = scene.freeMotion[obj.jointIndex]
            residuals[beg+0:beg+2] = freeMotion.residue_f + cache_h[beg+0:beg+2]
            residuals[beg+3:beg+5] = freeMotion.residue_t + cache_h[beg+3:beg+5]

        else
           error("Bug in Modia3D/src/Composition/specifics/specifics.jl (getJointResiduals_for_leq_mode_0!): jointKind = $jointKind is not known")
        end
    end
    return nothing
end

# For backwards compatibility (do not use for new models)
function setAngle!(revolute::Revolute, phi::Float64)
   obj          = revolute.obj2
   revolute.phi = phi
   obj.R_rel    = Frames.rotAxis(revolute.posAxis, revolute.posMovement, phi)
   obj.R_abs    = obj.R_rel*obj.parent.R_abs
   return revolute
end

# For backwards compatibility (do not use for new models)
function setDistance!(prismatic::Prismatic, s::Float64)
   obj         = prismatic.obj2
   prismatic.s = s
   obj.r_rel   = prismatic.eAxis*s
   obj.r_abs   = parent.r_abs + parent.R_abs'*obj.r_rel
   return prismatic
end


text:
#=
# MNIST

We begin by importing the required packages.
We load MNIST via the MLDatasets.jl package.
=#

#=
Some references

<https://www.jmlr.org/beta/papers/v17/14-237.html>,

<http://www.cs.toronto.edu/~tang/papers/RbmZM.pdf>,

<https://doi.org/10.1007/978-3-642-35289-8_3>

=#

import CairoMakie
import Makie
import MLDatasets
import Flux
import RestrictedBoltzmannMachines as RBMs
import WhitenedRBMs as WhiteRBMs
using Statistics: mean, var, std
using Random: bitrand
using ValueHistories: MVHistory
using RestrictedBoltzmannMachines: BinaryRBM
using WhitenedRBMs: whiten, blacken, whiten_visible_from_data
nothing #hide

#=
Useful function to plot MNIST digits.
=#

"""
    imggrid(A)

Given a four dimensional tensor `A` of size `(width, height, ncols, nrows)`
containing `width x height` images in a grid of `nrows x ncols`, this returns
a matrix of size `(width * ncols, height * nrows)`, that can be plotted in a heatmap
to display all images.
"""
function imggrid(A::AbstractArray{<:Any,4})
    reshape(permutedims(A, (1,3,2,4)), size(A,1)*size(A,3), size(A,2)*size(A,4))
end

#=
Now load the MNIST dataset.
=#

Float = Float32
train_x, train_y = MLDatasets.MNIST.traindata()
digit = 0 # the digit we work with
train_x = Array{Float}(train_x[:, :, train_y .== digit] .> 0.5)
train_y = train_y[train_y .== digit]
println(length(train_y), " train samples")

#=
Initialize and train a normal RBM
=#

rbm = BinaryRBM(Float, (28,28), 200)
RBMs.initialize!(rbm, train_x)
rbm = whiten_visible_from_data(whiten(rbm), train_x, WhiteRBMs.Identity(); ϵ=1f-3)
rbm.hidden.θ .= 0
nothing #hide

# Train init

batchsize = 256
optim = Flux.ADAM()
vm = bitrand(28, 28, batchsize) # fantasy chains
history = MVHistory()
push!(history, :lpl, mean(RBMs.log_pseudolikelihood(rbm, train_x)))
nothing #hide

# Pseudolikelihood before training

mean(@time RBMs.log_pseudolikelihood(rbm, train_x))

# Train

@time for epoch in 1:100 # track pseudolikelihood every 5 epochs
    WhiteRBMs.pcd!(
        rbm, train_x; epochs=5, vm, history, batchsize, optim, ϵv=1f-3,
        transform_v=WhiteRBMs.Identity(), transform_h=WhiteRBMs.Identity()
    )
    push!(history, :lpl, mean(RBMs.log_pseudolikelihood(rbm, train_x)))
end
nothing #hide

# Convert to equivalent RBM (without affine transforms)

rbm = WhiteRBMs.blacken(rbm)
nothing #hide

# Plot log-pseudolikelihood of train data during learning.

fig = Makie.Figure(resolution=(600, 300))
ax = Makie.Axis(fig[1,1], xlabel="epochs", ylabel="pseudolikelihood")
Makie.lines!(ax, get(history, :lpl)...)
fig

# Seconds per epoch.

fig = Makie.Figure(resolution=(600, 300))
ax = Makie.Axis(fig[1,1], xlabel="epoch", ylabel="seconds")
Makie.lines!(ax, get(history, :Δt)...)
fig

# Log-pseudolikelihood vs. computation time instead of epoch count.

fig = Makie.Figure(resolution=(600, 300))
ax = Makie.Axis(fig[1,1], xlabel="seconds", ylabel="pseudolikelihood")
Makie.lines!(ax, cumsum([0; get(history, :Δt)[2]])[1:5:end], get(history, :lpl)[2])
fig

# Now we do the Gibbs sampling to generate RBM digits.

nrows, ncols = 10, 15
nsteps = 5000
fantasy_F = zeros(nrows*ncols, nsteps)
fantasy_x = bitrand(28,28,nrows*ncols)
fantasy_F[:,1] .= RBMs.free_energy(rbm, fantasy_x)
@time for t in 2:nsteps
    fantasy_x .= RBMs.sample_v_from_v(rbm, fantasy_x)
    fantasy_F[:,t] .= RBMs.free_energy(rbm, fantasy_x)
end
nothing #hide

# Check equilibration of sampling

fig = Makie.Figure(resolution=(400,300))
ax = Makie.Axis(fig[1,1], xlabel="sampling time", ylabel="free energy")
fantasy_F_μ = vec(mean(fantasy_F; dims=1))
fantasy_F_σ = vec(std(fantasy_F; dims=1))
Makie.band!(ax, 1:nsteps, fantasy_F_μ - fantasy_F_σ/2, fantasy_F_μ + fantasy_F_σ/2)
Makie.lines!(ax, 1:nsteps, fantasy_F_μ)
fig

# Plot the resulting samples.

fig = Makie.Figure(resolution=(40ncols, 40nrows))
ax = Makie.Axis(fig[1,1], yreversed=true)
Makie.image!(ax, imggrid(reshape(fantasy_x, 28, 28, ncols, nrows)), colorrange=(1,0))
Makie.hidedecorations!(ax)
Makie.hidespines!(ax)
fig

ext: jl


text:
module SIMeasurementsBenchmark
using DataFrames
using BenchmarkTools
using SIMeasurements: Meter, Kilogram, Unit, Dimensions
import SIMeasurements: conversion_factors

const Klurdge = Unit{:Koala, Dimensions(length=2)}()
conversion_factors(
    ::Unit{:Koala, Dimensions(length=1)}, ::typeof(Meter)) = 1//2, 0

suite = BenchmarkGroup()
suite[:scalar] = BenchmarkGroup(["scalar"])
suite[:array] = BenchmarkGroup(["array"])
suite[:scalar][:same] = BenchmarkGroup(["same units"])
for op in (:*, :/, ://)
    suite[:scalar][op] = BenchmarkGroup([string(op)])
    suite[:scalar][op][:nothing] =
        @benchmarkable $op($(rand(1:10)), $(rand(1:10)))
    suite[:scalar][op][:same] =
        @benchmarkable $op($(rand(1:10)Meter), $(rand(1:10)Meter))
    suite[:scalar][op][:different] =
        @benchmarkable $op($(rand(1:10)Meter), $(rand(1:10)Kilogram))
    suite[:scalar][op][:inv] =
        @benchmarkable $op($(rand(1:10)Meter), $(rand(1:10)Meter^-1))
    suite[:scalar][op][:systems] =
        @benchmarkable $op($(rand(1:10)Klurdge), $(rand(1:10)Meter^-1))
end

for op in (:+, :-)
    suite[:scalar][op] = BenchmarkGroup([string(op)])
    suite[:scalar][op][:nothing] =
        @benchmarkable $op($(rand(1:10)), $(rand(1:10)))
    suite[:scalar][op][:same] =
        @benchmarkable $op($(rand(1:10)Meter), $(rand(1:10)Meter))
    suite[:scalar][op][:systems] =
        @benchmarkable $op($(rand(1:10)Klurdge), $(rand(1:10)Meter^2))
end

tune!(suite)
results = run(suite)
for op in keys(suite[:scalar])
    for key in keys(suite[:scalar][op])
        key == :nothing && continue
        bench = results[:scalar][op]
        print(op, ", ", key, " ")
        println(ratio(median(bench[key]), median(bench[:nothing])))
    end
end

end


text:
workspace()

include("generate_random_unit_quaternion.jl")
include("characteristic_matrix_ellipsoid.jl")
include("characteristic_matrix_ellipse.jl")
include("rotation_matrix.jl")
include("relax_system.jl")
include("equilibrate_system.jl")
include("compress_system.jl")
include("generate_proposal_position.jl")
include("generate_proposal_orientation.jl")
include("overlap_ellipsoid.jl")
include("overlap_ellipse.jl")
include("overlap_cuboid.jl")
include("signed_distance_mod.jl")
include("quaternion_mult.jl")

function run()
	# Inititalization of random number generation device.
	random_seed::Int64 = 1#convert(Int64, time_ns())
	srand(random_seed)

	# Type of particles.
	#particle_type::String = "sphere"
	#particle_type::String = "ellipse"
	#particle_type::String = "ellipsoid"
	particle_type::String = "cuboid"

	particle_type_index::Int64 = 0
	if particle_type == "sphere"
		particle_type_index = 1
	elseif particle_type == "ellipse"
		particle_type_index = 2
	elseif particle_type == "ellipsoid"
		particle_type_index = 3
	elseif particle_type == "cuboid"
		particle_type_index = 4
	end

	# Number of particles.
	number_of_particles::Int64 = 125

	# Particle property matrix (i.e. radii).
	number_of_properties::Int64 = 0
	if particle_type == "sphere"
		number_of_properties = 1
	elseif particle_type == "ellipse"
		number_of_properties = 2
	elseif particle_type == "ellipsoid"
		number_of_properties = 3
	elseif particle_type == "cuboid"
		number_of_properties = 3
	end

	R::Array{Float64, 2} = zeros(number_of_particles, number_of_properties)
	R[:, 1] = 1.0 * ones(number_of_particles, 1)
	R[:, 2] = 1.0 * ones(number_of_particles, 1)
	R[:, 3] = 1.0 * ones(number_of_particles, 1)
	#R[:, 1] = 0.5 + 1.0 * rand(number_of_particles, 1)
	#R[:, 2] = 0.5 + 1.0 * rand(number_of_particles, 1)
	#R[:, 3] = 0.5 + 1.0 * rand(number_of_particles, 1)

	# Simulation domain dimensions.
	phi_start::Float64 = 0.30#0.40
	Lx::Float64 = ( sum(8.0 * R[:, 1] .* R[:, 2] .* R[:, 3]) / phi_start )^(1/3)
	#Lx::Float64 = ( sum(4.0 * pi / 3.0 * R[:, 1] .* R[:, 2] .* R[:, 3]) / phi_start )^(1/3)
	Ly::Float64 = Lx
	Lz::Float64 = Lx

	# Positions.
	X::Array{Float64, 1} = Lx * rand(number_of_particles)
	Y::Array{Float64, 1} = Ly * rand(number_of_particles)
	Z::Array{Float64, 1} = Lz * rand(number_of_particles)
	# X::Array{Float64, 1} = zeros(number_of_particles)
	# Y::Array{Float64, 1} = zeros(number_of_particles)
	# Z::Array{Float64, 1} = zeros(number_of_particles)
	# count::Int64 = 0
	# for i = 1:8
		# for j = 1:8
			# for k = 1:8
				# count += 1
				# x = Lx/8*(i-1) + Lx/16
				# y = Ly/8*(j-1) + Ly/16
				# z = Lz/8*(k-1) + Lz/16
				# X[count] = x
				# Y[count] = y
				# Z[count] = z
			# end
		# end
	# end

	# Quaternions. Generate random orientations.
	Q0::Array{Float64, 1} = zeros(number_of_particles)#zeros(number_of_particles)
	Q1::Array{Float64, 1} = zeros(number_of_particles)
	Q2::Array{Float64, 1} = zeros(number_of_particles)
	Q3::Array{Float64, 1} = zeros(number_of_particles)
	q0::Float64 = 0.0
	q1::Float64 = 0.0
	q2::Float64 = 0.0
	q3::Float64 = 0.0
	if particle_type != "sphere"
		for current_particle = 1:number_of_particles
			(q0, q1, q2, q3) = generate_random_unit_quaternion()
			Q0[current_particle] = q0
			Q1[current_particle] = q1
			Q2[current_particle] = q2
			Q3[current_particle] = q3
		end
	end

	# Simulation parameters.
	sigma_translation_max::Float64 = 0.05#1.0
	sigma_rotation_max::Float64 = 0.01#5#0.01
	sigma_translation::Float64 = sigma_translation_max
	sigma_rotation::Float64 = sigma_rotation_max

	# Characteristic/rotation matrix entries.
	A11::Array{Float64, 1} = zeros(number_of_particles)
	A12::Array{Float64, 1} = zeros(number_of_particles)
	A13::Array{Float64, 1} = zeros(number_of_particles)
	A21::Array{Float64, 1} = zeros(number_of_particles)
	A22::Array{Float64, 1} = zeros(number_of_particles)
	A23::Array{Float64, 1} = zeros(number_of_particles)
	A31::Array{Float64, 1} = zeros(number_of_particles)
	A32::Array{Float64, 1} = zeros(number_of_particles)
	A33::Array{Float64, 1} = zeros(number_of_particles)

	a11::Float64 = 0.0
	a12::Float64 = 0.0
	a13::Float64 = 0.0
	a21::Float64 = 0.0
	a22::Float64 = 0.0
	a23::Float64 = 0.0
	a31::Float64 = 0.0
	a32::Float64 = 0.0
	a33::Float64 = 0.0
	if particle_type == "ellipse"
		for current_particle = 1:number_of_particles
			(a11, a12, a13, a21, a22, a23, a31, a32, a33) = characteristic_matrix_ellipse(Q0[current_particle], Q1[current_particle], Q2[current_particle], Q3[current_particle], R[current_particle, 1], R[current_particle, 2])
			A11[current_particle] = a11
			A12[current_particle] = a12
			A13[current_particle] = a13
			A21[current_particle] = a21
			A22[current_particle] = a22
			A23[current_particle] = a23
			A31[current_particle] = a31
			A32[current_particle] = a32
			A33[current_particle] = a33
		end
	elseif particle_type == "ellipsoid"
		for current_particle = 1:number_of_particles
			(a11, a12, a13, a21, a22, a23, a31, a32, a33) = characteristic_matrix_ellipsoid(Q0[current_particle], Q1[current_particle], Q2[current_particle], Q3[current_particle], R[current_particle, 1], R[current_particle, 2], R[current_particle, 3])
			A11[current_particle] = a11
			A12[current_particle] = a12
			A13[current_particle] = a13
			A21[current_particle] = a21
			A22[current_particle] = a22
			A23[current_particle] = a23
			A31[current_particle] = a31
			A32[current_particle] = a32
			A33[current_particle] = a33
		end
	elseif particle_type == "cuboid"
		for current_particle = 1:number_of_particles
			(a11, a12, a13, a21, a22, a23, a31, a32, a33) = rotation_matrix(Q0[current_particle], Q1[current_particle], Q2[current_particle], Q3[current_particle])
			A11[current_particle] = a11
			A12[current_particle] = a12
			A13[current_particle] = a13
			A21[current_particle] = a21
			A22[current_particle] = a22
			A23[current_particle] = a23
			A31[current_particle] = a31
			A32[current_particle] = a32
			A33[current_particle] = a33
		end
	end

	# Relax system until zero energy is reached.
	(X, Y, Z, Q0, Q1, Q2, Q3, A11, A12, A13, A21, A22, A23, A31, A32, A33, sigma_translation, sigma_rotation) = relax_system(Lx, Ly, Lz, particle_type, R, X, Y, Z, Q0, Q1, Q2, Q3, A11, A12, A13, A21, A22, A23, A31, A32, A33, sigma_translation, sigma_translation_max, sigma_rotation, sigma_rotation_max)

	# Equilibrate system.
	number_of_equilibration_sweeps::Int64 = 1000
	(X, Y, Z, Q0, Q1, Q2, Q3, A11, A12, A13, A21, A22, A23, A31, A32, A33, sigma_translation, sigma_rotation) = equilibrate_system(Lx, Ly, Lz, particle_type, R, X, Y, Z, Q0, Q1, Q2, Q3, A11, A12, A13, A21, A22, A23, A31, A32, A33, sigma_translation, sigma_translation_max, sigma_rotation, sigma_rotation_max, number_of_equilibration_sweeps)

	# Compress system.
	delta_phi::Float64 = 1e-3
	phi_target::Float64 = 1.0
	number_of_sweeps_max::Int64 = 1000#1000
	(Lx, Ly, Lz, X, Y, Z, Q0, Q1, Q2, Q3, A11, A12, A13, A21, A22, A23, A31, A32, A33, sigma_translation, sigma_rotation) = compress_system(Lx, Ly, Lz, particle_type, R, X, Y, Z, Q0, Q1, Q2, Q3, A11, A12, A13, A21, A22, A23, A31, A32, A33, sigma_translation, sigma_translation_max, sigma_rotation, sigma_rotation_max, delta_phi, phi_target, number_of_sweeps_max)

	# Verify non-overlap.
	println("OVERLAPS")
	test_energy = 0
	for currentA = 1:number_of_particles
		for currentB = [1:currentA-1 ; currentA+1:number_of_particles]

			xAB = signed_distance_mod(X[currentA], X[currentB], Lx)
			yAB = signed_distance_mod(Y[currentA], Y[currentB], Ly)
			zAB = signed_distance_mod(Z[currentA], Z[currentB], Lz)
			overlapfun = overlap_cuboid(xAB, yAB, zAB, A11[currentA], A12[currentA], A13[currentA], A21[currentA], A22[currentA], A23[currentA], A31[currentA], A32[currentA], A33[currentA], A11[currentB], A12[currentB], A13[currentB], A21[currentB], A22[currentB], A23[currentB], A31[currentB], A32[currentB], A33[currentB], R[currentA, 1], R[currentA, 2], R[currentA, 3], R[currentB, 1], R[currentB, 2], R[currentB, 3])
			if overlapfun > 0.0
				println((currentA, currentB))
			end
			test_energy += overlapfun
			#overlapfun = overlap_ellipsoid(xAB, yAB, zAB, A11[currentA], A12[currentA], A13[currentA], A21[currentA], A22[currentA], A23[currentA], A31[currentA], A32[currentA], A33[currentA], A11[currentB], A12[currentB], A13[currentB], A21[currentB], A22[currentB], A23[currentB], A31[currentB], A32[currentB], A33[currentB], R[currentA, 1]^2 * R[currentA, 2]^2 * R[currentA, 3]^2)
			#if overlapfun < 1.0
			#	println((currentA, currentB))
			#end
		end
	end
	println("============")
	println(test_energy)
	println("============")
	# Write result to file.
	file_name_output::String = "output.dat"
	file_stream_output::IOStream = open(file_name_output, "w")
	#writedlm(file_stream_output, particle_type_index, ",")
	writedlm(file_stream_output, [Lx Ly Lz], ",")
	writedlm(file_stream_output, [R X Y Z Q0 Q1 Q2 Q3], ",")
	close(file_stream_output)

	println(random_seed)

	nothing
end

@time run()

max_stars_repo_name: roding/whitefish

max_issues_repo_name: roding/whitefish

max_forks_repo_name: roding/whitefish


text:
using Pkg
 Pkg.activate(".")

verbose = true
 if verbose && !isdefined(Main,:RvSpectML)  println("# Loading RvSpecML")    end
 using RvSpectML
include("shared/scripts.jl")
 if verbose   println("# Loading other packages")    end
 using DataFrames, Query, Statistics, Dates

include("read_expres_data_101501.jl")

#order_list_timeseries = extract_orders(all_spectra,pipeline_plan, orders_to_use = [EXPRES.orders_to_use_default(all_spectra[1].inst)[2]], recalc=true )
order_list_timeseries = extract_orders(all_spectra,pipeline_plan, recalc=true )

line_list_df = prepare_line_list_pass1(linelist_for_ccf_filename, all_spectra, pipeline_plan,  v_center_to_avoid_tellurics=ccf_mid_velocity, Δv_to_avoid_tellurics = RvSpectMLBase.max_bc)

(ccfs, v_grid) = ccf_total(order_list_timeseries, line_list_df, pipeline_plan,  mask_scale_factor=10.0, ccf_mid_velocity=ccf_mid_velocity, recalc=true)

line_width = RvSpectML.calc_line_width(v_grid,view(ccfs,:,1),frac_depth=0.05)

line_list_df = prepare_line_list_pass1(linelist_for_ccf_filename, all_spectra, pipeline_plan,  v_center_to_avoid_tellurics=ccf_mid_velocity, Δv_to_avoid_tellurics = line_width)

#chunk_list_timeseries = make_chunk_list_timeseries(all_spectra,line_list_df)
EXPRES.continuum_normalize_spectra!(all_spectra)

(ccfs, v_grid) = ccf_total(order_list_timeseries, line_list_df, pipeline_plan,  mask_type=:gaussian, mask_scale_factor=10.0, ccf_mid_velocity=ccf_mid_velocity, recalc=true)
 rvs_ccf = calc_rvs_from_ccf_total(ccfs, pipeline_plan, v_grid=v_grid, times = order_list_timeseries.times, recalc=true)

(ccfs_expr, v_grid_expr) = ccf_total(order_list_timeseries, line_list_df, pipeline_plan, mask_type=:gaussian, mask_scale_factor=10.0, ccf_mid_velocity=ccf_mid_velocity, recalc=true, use_pixel_vars=true)
 rvs_ccf_expr = calc_rvs_from_ccf_total(ccfs_expr, pipeline_plan, v_grid=v_grid_expr, times = order_list_timeseries.times, recalc=true)
 println(" mean Δχ²= ",mean(minimum(length(order_list_timeseries.chunk_list[1][1].flux).*log.(1 .- ccfs_expr[:,t_idx].^2 ),dims=1)))

if make_plot(pipeline_plan, :ccf_total)
   using Plots
   t_idx = 1:40
   plt = plot()
   #plot!(plt,v_grid,ccfs[:,t_idx]./maximum(ccfs[:,t_idx],dims=1),label=:none)
   #if isdefined(Main,:ccfs_expr)   scatter!(plt,v_grid_expr,-ccfs_expr[:,t_idx]./maximum(ccfs_expr[:,t_idx],dims=1),markersize=1.2,label=:none)  end
   #lot!(plt,v_grid,ccfs[:,t_idx],label=:none)
   #if isdefined(Main,:ccfs_expr)   plot!(plt,v_grid_expr,ccfs_expr[:,t_idx],markersize=1.2,label=:none)   end
   if isdefined(Main,:ccfs_expr)   plot!(plt,v_grid_expr,length(order_list_timeseries.chunk_list[1][1].flux).*log.(1 .- ccfs_expr[:,t_idx].^2 ) ,markersize=1.2,label=:none)   end
   #if isdefined(Main,:ccfs_expr)   scatter!(plt,v_grid_expr,ccfs_expr[:,t_idx] .-maximum(ccfs_expr[:,t_idx],dims=1),markersize=1.2,label=:none)   end
   xlabel!("v (m/s)")
   ylabel!("CCF")
   if save_plot(pipeline_plan,:ccf_total)   savefig(plt,joinpath(output_dir,target_subdir * "_ccf_sum.png"))   end
   display(plt)
end

if make_plot(pipeline_plan, :ccf_total)
   include("../scripts/plots/spectra.jl")
   zvals = ccfs./maximum(ccfs,dims=1).-mean(ccfs./maximum(ccfs,dims=1),dims=2)
   colorscale = cgrad(:balance)
   plt = heatmap(v_grid,collect(1:size(ccfs,2)),zvals', c=colorscale, clims=(-maximum(abs.(zvals)),maximum(abs.(zvals))) )
   add_time_gap_lines(plt,order_list_timeseries.times)
   xlabel!("v (m/s)")
   ylabel!("Observation #")
   title!("CCF(v,t)-<CCF>(v) vs time")
   if save_plot(pipeline_plan,:ccf_total)   savefig(plt,joinpath(output_dir,target_subdir * "_ccf_sum_vs_time_heatmap.png"))   end
   display(plt)
end

rvs_ccf = calc_rvs_from_ccf_total(ccfs, pipeline_plan, v_grid=v_grid, times = order_list_timeseries.times, recalc=true)
# Store estimated RVs in metadata for use when making template
map(i->order_list_timeseries.metadata[i][:rv_est] = rvs_ccf[i]-mean(rvs_ccf), 1:length(rvs_ccf) )
rvs_ccf_expr = calc_rvs_from_ccf_total(ccfs_expr, pipeline_plan, v_grid=v_grid_expr, times = order_list_timeseries.times, recalc=true)

if need_to(pipeline_plan,:scalpels)
   rvs_scalpels = map(n->Scalpels.clean_rvs_scalpels(rvs_ccf, ccfs, num_basis=n), 1:5)
   println("RMS RVs cleaned by Scalpels: ",std.(rvs_scalpels) )
   dont_need_to!(pipeline_plan,:scalpels)
end

make_plot!(pipeline_plan,:scalpels)
if make_plot(pipeline_plan, :scalpels)
   @assert !need_to(pipeline_plan, :rvs_ccf_total)
   @assert !need_to(pipeline_plan, :ccf_total)
   include("../scripts/plots/scalpels_code.jl")
   plt = Scalpels.make_plots_scalpels(rvs_ccf, ccfs, max_num_basis=2, v_grid=v_grid, times=order_list_timeseries.times, output_path="examples/output/figures")
   display(plt)
end

if make_plot(pipeline_plan, :rvs_ccf_total)
   using Plots
   rvs_ccf .-= mean(rvs_ccf)
   plt = scatter(rvs_ccf,markersize=3,label="RVs CCF Std", legend=:bottomright)
   if isdefined(Main, :rvs_ccf_expr)
      rvs_ccf_expr .-= mean(rvs_ccf_expr)
      scatter!(plt,rvs_ccf_expr,markersize=3,label="RVs CCF Expr")
   end
   if isdefined(Main, :rvs_scalpels)
      for i in 1:length(rvs_scalpels)
         rvs_scalpels[i] .-= mean(rvs_scalpels[i])
         global plt
         scatter!(plt,rvs_scalpels[i],markersize=3,color=i+2,label="RVs post-Scalpels " * string(i+1))
         plot!(plt,rvs_scalpels[i],markersize=3,color=i+2,label=:none)
      end
   end
   ylabel!("v (m/s)")
   xlabel!("Time (#)")
   if save_plot(pipeline_plan,:rvs_ccf_total)   savefig(plt,joinpath(output_dir,target_subdir * "_rvs_ccf_sum.png"))   end
   display(plt)
   #=
   df_yale_resutls = CSV.read(joinpath(homedir(),"Data/EXPRES/inputs/101501/101501_activity.csv"))
   rvs_yale_ccf = df_yale_resutls["CCF RV [m/s]"]
   rvs_yale_ccf .-= mean(rvs_yale_ccf)
   rvs_yale_cbc = df_yale_resutls["CBC RV [m/s]"]
   rvs_yale_cbc .-= mean(rvs_yale_cbc)
   scatter!(rvs_yale_cbc,markersize=2,label="Yale CBC")
   scatter!(rvs_yale_ccf,markersize=2,label="Yale CCF")
   =#
   #diff = rvs_yale_ccf.-rvs_yale_cbc
   #diff = rvs_ccf.-rvs_yale_cbc
   #diff = rvs_ccf2.-rvs_yale_cbc
   #diff = rvs_ccf.-rvs_yale_ccf
   #diff = rvs_ccf2.-rvs_yale_ccf
end
if make_plot(pipeline_plan, :rvs_ccf_total)
   if isdefined(Main, :rvs_ccf_expr)
      diff = rvs_ccf.-rvs_ccf_expr
      println(std(diff))
      plt = scatter(order_list_timeseries.times,diff,markersize=4,label="Delta RV Expr")
   end
   if isdefined(Main, :rvs_scalpels)
      diff = rvs_ccf.-rvs_scalpels[max(2,length(rvs_scalpels))]
      println(std(diff))
      plt = scatter(order_list_timeseries.times,diff,markersize=4,label="Delta RV Scalpels")
   end
   ylabel!("Δv (m/s) (Two mask shapes)")
   xlabel!("Time (d)")
   if save_plot(pipeline_plan,:rvs_ccf_total)   savefig(plt,joinpath(output_dir,target_subdir * "_rvs_ccf_sum.png"))   end
   display(plt)
end

#need_to!(pipeline_plan, :ccf_orders)
(order_ccfs, v_grid_order_ccfs) = ccf_orders(order_list_timeseries, line_list_df, pipeline_plan)

if make_plot(pipeline_plan, :ccf_orders)
   # order ccfs averaged over observations at multiple times
   obs = 1:length(order_list_timeseries.times)
   order_labels = map(c->order_list_timeseries.chunk_list[1].data[c].λ.indices[2], 1:size(order_ccfs,2))
   orders_to_plot = findall(c->sum(order_ccfs[:,c,obs])>0, 1:size(order_ccfs,2))
   zvals =  reshape(sum(order_ccfs[:,orders_to_plot,obs],dims=3)./maximum(sum(order_ccfs[:,orders_to_plot,obs],dims=3),dims=1)   ,size(order_ccfs,1),size(order_ccfs[:,orders_to_plot,obs],2)) .-
            reshape(sum(order_ccfs[:,orders_to_plot,obs],dims=(2,3))./maximum(sum(order_ccfs[:,orders_to_plot,obs],dims=(2,3))),size(order_ccfs,1) )
   plt = heatmap(v_grid,order_labels[orders_to_plot], zvals',c = cgrad(:balance), clims=(-maximum(abs.(zvals)),maximum(abs.(zvals))) )

   xlabel!("v (m/s)")
   ylabel!("Order ID")
   title!("CCF-<CCF> for obs ID=" * string(obs))
   if save_plot(pipeline_plan,:ccf_orders)   savefig(plt,joinpath(output_dir,target_subdir * " _ccf_orders.png"))   end
   display(plt)
end

#RvSpectML.make_plot!(pipeline_plan, :movie)
if make_plot(pipeline_plan, :ccf_orders) && make_plot(pipeline_plan, :movie)
   # order ccfs averaged over observations at multiple times
   anim = @animate for obs ∈  1:length(order_list_timeseries.times)
      local order_labels = map(c->order_list_timeseries.chunk_list[1].data[c].λ.indices[2], 1:size(order_ccfs,2))
      local orders_to_plot = findall(c->sum(order_ccfs[:,c,obs])>0, 1:size(order_ccfs,2))
      local zvals =  reshape(sum(order_ccfs[:,orders_to_plot,obs],dims=3)./maximum(sum(order_ccfs[:,orders_to_plot,obs],dims=3),dims=1)   ,size(order_ccfs,1),size(order_ccfs[:,orders_to_plot,obs],2)) .-
               reshape(sum(order_ccfs[:,orders_to_plot,obs],dims=(2,3))./maximum(sum(order_ccfs[:,orders_to_plot,obs],dims=(2,3))),size(order_ccfs,1) )
      local plt = heatmap(v_grid,order_labels[orders_to_plot], zvals',c = cgrad(:balance), clims=(-maximum(abs.(zvals)),maximum(abs.(zvals))) )

      xlabel!("v (m/s)")
      ylabel!("Order ID")
      title!("CCF-<CCF> for obs ID=" * string(obs))
   end
   gif(anim, joinpath(output_dir,target_subdir * "_ccf_order_movie_obs.gif"), fps = 5)
end

if make_plot(pipeline_plan, :ccf_orders)
   ord = 3
   zvals = reshape(sum(order_ccfs[:,ord,:],dims=3)./maximum(sum(order_ccfs[:,ord,:],dims=3),dims=1),size(order_ccfs,1),size(order_ccfs[:,ord,:],2)).-
   reshape(sum(order_ccfs[:,ord,:],dims=(2,3))./maximum(sum(order_ccfs[:,ord,:],dims=(2,3))),size(order_ccfs,1) )
   plt = heatmap(v_grid,1:size(order_ccfs,3), zvals',c = cgrad(:balance), clims=(-maximum(abs.(zvals)),maximum(abs.(zvals))) )
   add_time_gap_lines(plt,order_list_timeseries.times)
   xlabel!("v (m/s)")
   ylabel!("Observation ID")
   title!("CCF-<CCF> for order=" * string(ord))
   if save_plot(pipeline_plan,:ccf_orders)   savefig(plt,joinpath(output_dir,target_subdir * "_ccf_obs_order=" * string(ord) * ".png"))   end
   display(plt)
end

if make_plot(pipeline_plan, :ccf_orders) && make_plot(pipeline_plan, :movie)
   orders_to_plot = findall(c->sum(order_ccfs[:,c,obs])>0, 1:size(order_ccfs,2))
   min_z_val = Inf
   max_z_val = -Inf
   for ord ∈ orders_to_plot
      local zvals = reshape(sum(order_ccfs[:,ord,:],dims=3)./maximum(sum(order_ccfs[:,ord,:],dims=3),dims=1),size(order_ccfs,1),size(order_ccfs[:,ord,:],2)) .-
                    reshape(sum(order_ccfs[:,ord,:],dims=(2,3))./maximum(sum(order_ccfs[:,ord,:],dims=(2,3))),size(order_ccfs,1) )
      lo, hi = extrema(zvals)
      #println(" ord = ", ord, " lo = ", lo, " hi = ", hi)
      global min_z_val = min(min_z_val,lo)
      global max_z_val = max(max_z_val,hi)
   end
   #println(min_z_val, " - ", max_z_val)
   if -min_z_val > max_z_val
         max_z_val = -min_z_val
   else
         min_z_val = -max_z_val
   end
   #println(min_z_val, " - ", max_z_val)
   anim = @animate for ord ∈ orders_to_plot
      local zvals = reshape(sum(order_ccfs[:,ord,:],dims=3)./maximum(sum(order_ccfs[:,ord,:],dims=3),dims=1),size(order_ccfs,1),size(order_ccfs[:,ord,:],2)).-
      reshape(sum(order_ccfs[:,ord,:],dims=(2,3))./maximum(sum(order_ccfs[:,ord,:],dims=(2,3))),size(order_ccfs,1) )
      local plt = heatmap(v_grid,1:size(order_ccfs,3), zvals',c = cgrad(:balance), clims=(min_z_val,max_z_val) )
      add_time_gap_lines(plt,order_list_timeseries.times)
      xlabel!("v (m/s)")
      ylabel!("Observation ID")
      title!("CCF-<CCF> for order=" * string(order_labels[ord]))
   end
   gif(anim, joinpath(output_dir,target_subdir * "_ccf_order_movie=.gif"), fps = 5)
end


text:
cd(@__DIR__) do
    run(`make`)
end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule FriBidi_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("FriBidi")
JLLWrappers.@generate_main_file("FriBidi", UUID("559328eb-81f9-559d-9380-de523a88c83c"))
end  # module FriBidi_jll


text:
using GeneralizedGammaDistribution
using Test

@testset "generalized gamma" begin
    @test GeneralizedGamma() == GeneralizedGamma(1.0, 1.0, 1.0)
end


text:
function LazyLoadedFrames(arg0::LazyLoadedEop, arg1::TimeScales, arg2::CelestialBodies)
    return LazyLoadedFrames((LazyLoadedEop, TimeScales, CelestialBodies), arg0, arg1, arg2)
end

function add_default_eop1980_history_loaders(obj::LazyLoadedFrames, arg0::JString, arg1::JString, arg2::JString, arg3::JString, arg4::JString)
    return jcall(obj, "addDefaultEOP1980HistoryLoaders", void, (JString, JString, JString, JString, JString), arg0, arg1, arg2, arg3, arg4)
end

function add_default_eop2000_history_loaders(obj::LazyLoadedFrames, arg0::JString, arg1::JString, arg2::JString, arg3::JString, arg4::JString)
    return jcall(obj, "addDefaultEOP2000HistoryLoaders", void, (JString, JString, JString, JString, JString), arg0, arg1, arg2, arg3, arg4)
end

function add_eop_history_loader(obj::LazyLoadedFrames, arg0::IERSConventions, arg1::EOPHistoryLoader)
    return jcall(obj, "addEOPHistoryLoader", void, (IERSConventions, EOPHistoryLoader), arg0, arg1)
end

function clear_eop_history_loaders(obj::LazyLoadedFrames)
    return jcall(obj, "clearEOPHistoryLoaders", void, ())
end

function get_eop_history(obj::LazyLoadedFrames, arg0::IERSConventions, arg1::jboolean)
    return jcall(obj, "getEOPHistory", EOPHistory, (IERSConventions, jboolean), arg0, arg1)
end

function set_eop_continuity_threshold(obj::LazyLoadedFrames, arg0::jdouble)
    return jcall(obj, "setEOPContinuityThreshold", void, (jdouble,), arg0)
end


text:
using LinearAlgebra, Statistics

import Statistics: cov
import LinearAlgebra: logdet, AbstractTriangular, Adjoint, transpose, adjoint,
    AbstractMatrix, cholesky, Cholesky
import Base: size, getindex, IndexStyle, ==, isapprox, +, -, *, map, broadcast, \

"""
    logdet(U::UpperTriangular)

Compute the log determinant by summing the logdet of the diagonal of an `UpperTriangular`
matrix. Implementing this here is a horrible bit of type piracy, but it's necessary.
"""
logdet(T::AbstractTriangular) = sum(logdet, view(T, diagind(T)))
logdet(A::Adjoint{T, <:AbstractTriangular{T}}) where T = logdet(A')

"""
    LazyPDMat{T<:Real, TΣ<:AbstractMatrix{T}} <: AbstractMatrix{T}

A positive definite matrix which evaluates its Cholesky lazily and caches the result.
Please don't mutate it this object: `setindex!` isn't defined for a reason.
"""
mutable struct LazyPDMat{T<:Real, TΣ<:AbstractMatrix{T}} <: AbstractMatrix{T}
    Σ::TΣ
    C::Union{Nothing, Cholesky}
    ϵ::T
    LazyPDMat(Σ::TΣ) where {T<:Real, TΣ<:AbstractMatrix{T}} = new{T, TΣ}(Σ, nothing, 1e-6)
    function LazyPDMat(Σ::TΣ, ϵ::Real) where {T<:Real, TΣ<:AbstractMatrix{T}}
        return new{T, TΣ}(Σ, nothing, ϵ)
    end
end
@inline LazyPDMat(Σ::LazyPDMat) = Σ
@inline LazyPDMat(σ::Real) = σ
@inline unbox(Σ::LazyPDMat) = Σ.Σ
@inline unbox(A) = A
@inline AbstractMatrix(Σ::LazyPDMat) = unbox(Σ)

size(Σ::LazyPDMat) = size(unbox(Σ))
@inline getindex(Σ::LazyPDMat, i::Int...) = getindex(unbox(Σ), i...)
IndexStyle(::Type{<:LazyPDMat}) = IndexLinear()
==(Σ1::LazyPDMat, Σ2::LazyPDMat) = unbox(Σ1) == unbox(Σ2)
isapprox(Σ1::LazyPDMat, Σ2::LazyPDMat) = isapprox(unbox(Σ1), unbox(Σ2))

# Unary functions.
logdet(Σ::LazyPDMat) = logdet(cholesky(Σ))
function cholesky(Σ::LazyPDMat{T, <:AbstractMatrix{T}} where T<:Real)
    if Σ.C == nothing
        Σ.C = cholesky(Symmetric(unbox(Σ) + Σ.ϵ * I))
    end
    return Σ.C
end
chol(Σ) = cholesky(Σ).U

# Binary functions.
+(Σ1::LazyPDMat, Σ2::LazyPDMat) = LazyPDMat(unbox(Σ1) + unbox(Σ2))
+(Σ1::LazyPDMat, Σ2::UniformScaling) = (Σ2.λ > 0 ? LazyPDMat : identity)(unbox(Σ1) + Σ2)
-(Σ1::LazyPDMat, Σ2::LazyPDMat) = unbox(Σ1) - unbox(Σ2)
*(Σ1::LazyPDMat, Σ2::LazyPDMat) = LazyPDMat(unbox(Σ1) * unbox(Σ2))
map(::typeof(+), Σs::LazyPDMat...) = LazyPDMat(map(+, map(unbox, Σs)...))
map(::typeof(*), Σs::LazyPDMat...) = LazyPDMat(map(*, map(unbox, Σs)...))

broadcast(::typeof(*), Σ1::LazyPDMat, Σ2::LazyPDMat) = LazyPDMat(unbox(Σ1) .* unbox(Σ2))

# Specialised operations to exploit the Cholesky.
@noinline function Xt_A_X(A::LazyPDMat, X::AbstractMatrix)
    return LazyPDMat(Symmetric(X' * unbox(unbox(A)) * X))
end
@noinline Xt_A_X(A::LazyPDMat, x::AbstractVector) = sum(abs2, chol(A) * x)
@noinline Xt_A_Y(X::AVM, A::LazyPDMat, Y::AVM) = (chol(A) * X)' * (chol(A) * Y)
@noinline function Xt_invA_X(A::LazyPDMat, X::AVM)
    V = chol(A)' \ X
    return LazyPDMat(Symmetric(V'V))
end
@noinline Xt_invA_X(A::LazyPDMat, x::AbstractVector) = sum(abs2, chol(A)' \ x)
# @noinline Xt_invA_X(A::LazyPDMat{<:Real, <:SymmetricToeplitz}, x::AV) = sum(abs2, chol(A) \ x)
@noinline Xt_invA_Y(X::AVM, A::LazyPDMat, Y::AVM) = (chol(A)' \ X)' * (chol(A)' \ Y)
@noinline \(Σ::LazyPDMat, X::Union{AM, AV}) = chol(Σ) \ (chol(Σ)' \ X)

# Some generic operations that are useful for operations involving covariance matrices.
diag_AᵀA(A::AbstractMatrix) = vec(sum(abs2, A, dims=1))
function diag_AᵀB(A::AbstractMatrix, B::AbstractMatrix)
    @assert size(A) == size(B)
    return vec(sum(A .* B, dims=1))
end

diag_Xᵀ_A_X(A::LazyPDMat, X::AbstractMatrix) = diag_AᵀA(chol(A) * X)
diag_Xᵀ_A_Y(X::AM, A::LazyPDMat, Y::AM) = diag_AᵀB(chol(A) * X, chol(A) * Y)

diag_Xᵀ_invA_X(A::LazyPDMat, X::AbstractMatrix) = diag_AᵀA(chol(A)' \ X)
diag_Xᵀ_invA_Y(X::AM, A::LazyPDMat, Y::AM) = diag_AᵀB(chol(A)' \ X, chol(A)' \ Y)

# Specialised solver routine, especially useful for Titsias conditionals.
function Xtinv_A_Xinv(A::LazyPDMat, X::LazyPDMat)
    @assert size(A) == size(X)
    C = chol(X) \ (chol(X)' \ chol(A)')
    return LazyPDMat(Symmetric(C * C'))
end

transpose(A::LazyPDMat) = A
adjoint(A::LazyPDMat) = A


text:
##
using ACE, Printf, Test, LinearAlgebra
using ACE: evaluate, evaluate_d, read_dict, write_dict
using ACE.Testing

verbose = false
maxdeg = 25

using BenchmarkTools
pin = 0; pcut = 2
trans = PolyTransform(2, 1.0)
B = transformed_jacobi(maxdeg, trans, 3.0, 0.5, pin = pin, pcut = pcut)
Rn_new = ACE.Rn1pBasis_new(B.J; trans=trans)
X = State( rr = ACE.rand_sphere() * ACE.rand_radial(B) )
evaluate(Rn_new, X)

function runn(f, Rn, X, N)
   for n = 1:N
      B = f(Rn, X)
      ACE.release!(B)
   end
   return nothing
end

@btime runn($evaluate, $Rn_new, $X, 1_000)
@btime runn($evaluate, $(Rn_new.basis), $(X.rr), 1_000)

function runn!(B, Rn, X, N, trans)
   for n = 1:N
      B = ACE.evaluate!(B, Rn, trans(X))
      ACE.release!(B)
   end
   return nothing
end

tmpc = evaluate(Rn_new, X)
tmp = copy(tmpc.A)
@btime runn!($tmp, $(B.J), $(norm(X.rr)), 1_000, $trans)
@btime runn!($tmpc, $(B.J), $(norm(X.rr)), 1_000, $trans)

##
@profview runn(evaluate, Rn_new.basis, X.rr, 100_000_000)

@profview runn!(tmp, B.J, norm(X.rr), 100_000_000, trans)

##
using Profile
Profile.clear()
@profile runn(evaluate, Rn_new.basis, X.rr, 10_000_000)
Profile.print()

##
Rn_old = ACE.chain(trans, B.J)

let B0 = Rn_old, B1 = Rn_new.basis, X=X, B2 = B.J, tmpc=tmpc, tmp=tmp, trans=trans
   N = 10_000_000
   @time runn(evaluate, B0, norm(X.rr), N)
   @time runn(evaluate, B0, norm(X.rr), N)
   @time runn(evaluate, B1, X.rr, N)
   @time runn(evaluate, B1, X.rr, N)
   @time runn!(tmp, B2, norm(X.rr), N, trans)
   @time runn!(tmp, B2, norm(X.rr), N, trans)
end


text:
using DSGE, ClusterManagers, HDF5

##########################################################################################
## SETUP
##########################################################################################

# What do you want to do?
run_estimation     = false
run_modal_forecast = false
run_full_forecast  = true

# Initialize model object
# Note that the default for m1010 uses 6 anticipated shocks
m = Model1010("ss20")

# Settings for data, paths, etc.
dataroot = joinpath(dirname(@__FILE__()), "input_data")
saveroot = dirname(@__FILE__())
m <= Setting(:dataroot, dataroot, "Input data directory path")
m <= Setting(:saveroot, saveroot, "Output data directory path")
m <= Setting(:data_vintage, "161223")
m <= Setting(:use_population_forecast, false)

# Settings for estimation
# set to false => will load pre-computed mode and hessian before MCMC
m <= Setting(:reoptimize, true)
m <= Setting(:calculate_hessian, true)

# Settings for forecast dates
m <= Setting(:date_forecast_start,  quartertodate("2016-Q4"))
m <= Setting(:date_conditional_end, quartertodate("2016-Q4"))
m <= Setting(:shockdec_startdate,   Nullable(date_mainsample_start(m)))

# Parallelization
m <= Setting(:forecast_block_size,  500)
nworkers = 1
addprocsfcn = addprocs_sge # choose to work with your scheduler; see ClusterManagers.jl

##########################################################################################
## RUN
##########################################################################################

# Run estimation
if run_estimation

    if reoptimize(m)
        # Start from ss18 mode
        mode_file = rawpath(m, "estimate", "paramsmode.h5")
        mode_file = replace(mode_file, "ss20", "ss18")
        update!(m, h5read(mode_file, "params"))
    else
        # Use calculated ss18 mode
        mode_file = joinpath(dataroot, "user", "paramsmode_vint=161223.h5")
        specify_mode!(m, mode_file)
    end

    # Use calculated ss18 hessian
    if !calculate_hessian(m)
        hessian_file = joinpath(dataroot, "user", "hessian_vint=161223.h5")
        specify_hessian(m, hessian_file)
    end

    estimate(m)

    # Print tables of estimated parameter moments
    groupings = DSGE.parameter_groupings(m)
    moment_tables(m, groupings = groupings)
end

# Forecast step: produces smoothed histories and shock decompositions
if run_modal_forecast || run_full_forecast

    # what do we want to produce?
    output_vars = [:histpseudo, :forecastpseudo, :shockdecpseudo]

    # conditional type
    cond_type = :none

    # Forecast label: all forecast output filenames will contain this string
    forecast_string = ""

    # Modal forecast
    if run_modal_forecast
        # run modal forecasts and save all draws
        forecast_one(m, :mode, cond_type, output_vars; verbose = :high)

        # compute means and bands
        compute_meansbands(m, :mode, cond_type, output_vars)
    end

    # Full-distribution forecast
    if run_full_forecast
        my_procs = addprocsfcn(nworkers)
        @everywhere using DSGE

        forecast_one(m, :full, cond_type, output_vars; verbose = :high, forecast_string = forecast_string)
        rstar_bands = [0.68, 0.95]
        compute_meansbands(m, :full, cond_type, output_vars; verbose = :high, density_bands = rstar_bands,
                           forecast_string = forecast_string)
        rmprocs(my_procs)

        meansbands_to_matrix(m, :full, cond_type, output_vars; forecast_string = forecast_string)

        # print history means and bands tables to csv
        table_vars = [:ExAnteRealRate, :Forward5YearRealRate, :Forward10YearRealRate,
                      :RealNaturalRate, :Forward5YearRealNaturalRate,
                      :Forward10YearRealNaturalRate, :Forward20YearRealNaturalRate,
                      :Forward30YearRealNaturalRate]
        write_meansbands_tables_all(m, :full, cond_type, [:histpseudo], forecast_string = forecast_string,
                                    vars = table_vars)

        # print shockdec means and bands tables to csv
        if any(x->contains(string(x), "shockdec"), output_vars)
            shockdec_vars = [:RealNaturalRate, :Forward30YearRealNaturalRate]

            write_meansbands_tables_all(m, :full, cond_type, [:shockdecpseudo, :trendpseudo, :dettrendpseudo],
                                        vars = shockdec_vars,
                                        forecast_string = forecast_string)

        end
    end
end

nothing


text:
using MonteCarloMeasurements, Measurements, Plots, BenchmarkTools, OrdinaryDiffEq, PrettyTables, ChangePrecision, LinearAlgebra
# pgfplots()
default(size=(600,400))
color_palette = [
    RGB(0.1399999, 0.1399999, 0.4),
    RGB(1.0, 0.7075, 0.35),
    RGB(0.414999, 1.0, 1.0),
    RGB(0.6, 0.21, 0.534999),
    RGB(0,0.6,0),
]

function sim((±)::F, tspan, plotfun=plot!, args...; kwargs...) where F
    @changeprecision Float32 begin
        g = 9.79 ± 0.02; # Gravitational constant
        L = 1.00 ± 0.01; # Length of the pendulum
        u₀ = [0.0 ± 0.0, π / 3.0 ± 0.02] # Initial speed and initial angle
        # @show typeof(u₀)
        gL = g/L
        #Define the problem
        function simplependulum(du,u,p,t)
            θ  = u[1]
            dθ = u[2]
            du[1] = dθ
            du[2] = -gL * sin(θ)
        end

        prob = ODEProblem(simplependulum, u₀, tspan)
        sol = solve(prob, Tsit5(), reltol = 1e-6)

        plotfun(sol.t, getindex.(sol.u, 2), args...; kwargs...)
    end
end

## Special function needed to construct the parameters as sigmapoints, as they all have to be constructed in one single call
function sigmasim(tspan, plotfun=plot!, args...; kwargs...) where F
    @changeprecision Float32 begin
        g,L,u02 = StaticParticles(sigmapoints([9.79, 1.0, pi/3.0], diagm([0.02, 0.01, 0.02].^2)))
        u₀ = [0, u02] # Initial speed and initial angle
        # @show typeof(u₀)
        gL = g/L
        #Define the problem
        function simplependulum(du,u,p,t)
            θ  = u[1]
            dθ = u[2]
            du[1] = dθ
            du[2] = -gL * sin(θ)
        end

        prob = ODEProblem(simplependulum, u₀, tspan)
        sol = solve(prob, Tsit5(), reltol = 1e-6)

        plotfun(sol.t, getindex.(sol.u, 2), args...; kwargs...)
    end
end

##
tspan = (0.0f0, 0.5f0)
plot()
sim(Measurements.:±, tspan, label = "Linear", xlims=(tspan[2]-2,tspan[2]), color=color_palette[4])
sim(MonteCarloMeasurements.:±, tspan, errorbarplot!, 0.8413, label = "MCM", xlims=(tspan[2]-0.5,tspan[2]), l=(:dot,), color=color_palette[2], xlabel="Time [s]", ylabel="\$\\theta\$")

# savefig("/home/fredrikb/mcm_paper/figs/0-2.pdf")
##
tspan = (0.0f0, 200)
plot()
sim(Measurements.:±, tspan, label = "Linear", xlims=(tspan[2]-5,tspan[2]), color=color_palette[4])
sim(MonteCarloMeasurements.:±, tspan, label = "Monte Carlo", xlims=(tspan[2]-5,tspan[2]), l=(:dot,), color=color_palette[2], xlabel="Time [s]", ylabel="\$\\theta\$")

##
# We now integrated over 200 seconds and look at the last 5 seconds. This result maybe looks a bit confusing, the linear uncertainty propagation is very sure about the amplitude at certain points but not at others, whereas the Monte-Carlo approach is completely unsure. Furthermore, the linear approach thinks that the amplitude at some points is actually much higher than the starting amplitude, implying that energy somehow has been added to the system! The picture might become a bit more clear by plotting the individual trajectories of the particles

tspan = (0.0f0, 200)
plot()
sim(Measurements.:±, tspan, label = "Linear", xlims=(tspan[2]-5,tspan[2]), l=(5,), color=color_palette[4])
sim(MonteCarloMeasurements.:∓, tspan, mcplot!, xlims=(tspan[2]-5,tspan[2]), l=(color_palette[2],0.3), xlabel="Time [s]", ylabel="\$\\theta\$", label="MCM")
sigmasim(tspan, mcplot!, xlims=(tspan[2]-5,tspan[2]), l=(color_palette[3],0.9,2), xlabel="Time [s]", ylabel="\$\\theta\$", label="MCM \\Sigma")
# savefig("/home/fredrikb/mcm_paper/figs/mcplot.pdf")

# It now becomes clear that each trajectory has a constant amplitude (although individual trajectories amplitudes vary slightly due to the uncertainty in the initial angle), but the phase is all mixed up due to the slightly different frequencies!

# These problems grow with increasing uncertainty and increasing integration time. In fact, the uncertainty reported by Measurements.jl goes to infinity as the integration time does the same.

# Of course, the added accuracy from using MonteCarloMeasurements does not come for free, as it costs some additional computation. We have the following timings for integrating the above system 100 seconds using three different uncertainty representations

##
function naive_mc(tspan)
    for i = 1:100
        sim(certain, tspan, (args...;kwargs...)->nothing)
    end
end

tspan = (0.0f0, 100f0)
certain = (x,y)->x+y*randn()
table = Matrix{Any}(undef,3,6)
t1 = @benchmark sim($certain, $tspan, (args...;kwargs...)->nothing)
t2 = @benchmark sim($Measurements.:±, $tspan, (args...;kwargs...)->nothing) samples=500
t3 = @benchmark sim($MonteCarloMeasurements.:∓, $tspan, (args...;kwargs...)->nothing) samples=500
t4 = @benchmark sigmasim($tspan, (args...;kwargs...)->nothing) samples=500
t5 = @benchmark naive_mc($tspan) samples=500

# table[1,1] = ""
table[1,1] = "Time [ms]"
table[2,1] = "Memory [MiB]"
table[3,1] = "k Allocations"

for (i,t) in enumerate((t1,t2,t3,t4,t5))
    table[1,i+1] = time(t)/1000_000
    table[2,i+1] = memory(t)/1000_000
    table[3,i+1] = allocs(t)/1000
end

# pretty_table(table, ["" "Float32" "Linear" "MCM" "MCM \\Sigma" "Naive MC"], backend=:latex, formatters=ft_printf("%5.1f"))
pretty_table(table, ["" "Float32" "Linear" "MCM" "MCM \\Sigma" "Naive MC"], backend=:text, tf=markdown, formatters=ft_printf("%5.1f"))


text:
module flint_rock3_examples
using FinEtools
using FinEtools.MeshImportModule
using FinEtools.MeshExportModule
using FinEtoolsDeforLinear
using FinEtoolsDeforLinear.AlgoDeforLinearModule
using FinEtoolsVibInFluids.LaplBEM
using FinEtoolsVibInFluids.AlgoVibInFluidsModule
using FinEtoolsVoxelMesher
using LinearAlgebra
using TimerOutputs
using Arpack

# Free-vibration solution for a clamped plate partially or fully submerged in inviscid fluid.
#
# Reference: 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.

E = 185000*phun("MPa");# Young's modulus
nu = 0.27;# Poisson ratio
rho = 2300*phun("KG*M^-3");# mass density
rhow = 1000*phun("KG*M^-3");

tolerance = 30.0*phun("MM")/100;
OmegaShift = (2*pi*10.0); # to resolve rigid body modes
neigvs = 10;

sigdig(n) = round(n * 1000) / 1000

function volumeetc(fens, fes)
	geom  =  NodalField(fens.xyz)
	femm  =  FEMMBase(IntegDomain(fes, TetRule(1)))
	V = integratefunction(femm, geom, (x) ->  1.0)
	Sx = integratefunction(femm, geom, (x) ->  x[1])
	Sy = integratefunction(femm, geom, (x) ->  x[2])
	Sz = integratefunction(femm, geom, (x) ->  x[3])
	centroid = [Sx, Sy, Sz] ./ V
	return V, centroid
end

function free_vibration_solver()
	# Free-vibration solution for a clamped plate
	# Reference: 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.
	# Dry Natural frequencies [radians per second] (nLength=10,nWidth=10)
	# 12.45153      29.44122      75.04792      94.26498      106.7881
	to = TimerOutput()

	println("Loading mesh")
    @timeit to "Loading mesh" begin
	    MR = DeforModelRed3D
	    output = MeshImportModule.import_ABAQUS("./flint-rock3-1_5mm.inp")
	    fens, fes = output["fens"], output["fesets"][1]
	    fens.xyz .*= phun("MM")
	end

    println("Number of nodes: $(count(fens))")
    println("Interior mesh: $(count(fes)) tets")
    println("Surface mesh: $(count(meshboundary(fes))) triangles")

    geom = NodalField(fens.xyz)
    u = NodalField(zeros(size(fens.xyz,1),3)) # displacement field
    # nl = selectnode(fens, box=[0.0 0.0 -Inf Inf -Inf Inf], inflate=tolerance)
    # setebc!(u, nl, true, collect(1:3))
    applyebc!(u)
    numberdofs!(u)

    material = MatDeforElastIso(MR, rho, E, nu, 0.0)

    @timeit to "Set up FE machine" begin
	    femm = FEMMDeforLinearESNICET4(MR, IntegDomain(fes, NodalSimplexRule(3)), material)
	    associategeometry!(femm,  geom)
	end

	println("Stiffness matrix")
    @timeit to "Stiffness matrix" begin
	    K = stiffness(femm, geom, u)
	end
	println("Mass matrix")
	@timeit to "Mass matrix" begin
    	M = mass(femm, geom, u)
    end

    println("Eigenvalue problem")
    @timeit to "Eigenvalue problem" begin
	    d,v,nev,nconv = eigs(K+OmegaShift^2*M, M; nev=neigvs, which=:SM,  explicittransform=:none)
	    d = d .- OmegaShift^2;
	end

	println("Done")

	return Dict("fens"=>fens, "femm"=>femm, "geom"=>geom, "u"=>u, "eigenvectors"=>v, "eigenvalues"=>d, "timeroutput" =>to)
end # free_vibration_solver

function flint_rock3_dry()

	model = free_vibration_solver()

	@show model["timeroutput"]

	fens = model["fens"]
	femm = model["femm"]
	u = model["u"]
	d = model["eigenvalues"]
	v = model["eigenvectors"]

	fs = real(sqrt.(complex(d)))/(2*pi)
	println("Eigenvalues: $(sigdig.(2*pi*fs)) [radians per second]")
	println("Natural frequencies: $(sigdig.(fs)) [Hz]")

	vectors = []
	for mode in 1:length(fs)
		scattersysvec!(u, v[:,mode])
		push!(vectors, ("mode$mode", deepcopy(u.values)))
	end
	File =  "flint_rock3_dry.vtk"
	vtkexportmesh(File, fens, femm.integdomain.fes; vectors=vectors)
	@async run(`"paraview.exe" $File`)

    true
end # flint_rock3_dry

function flint_rock3_wet()
	to = TimerOutput()

	model = free_vibration_solver()

	@show model["timeroutput"]

	fens = model["fens"]
	femm = model["femm"]
	d = model["eigenvalues"]
	v = model["eigenvectors"]
	u = model["u"]

	fs = real(sqrt.(complex(d)))/(2*pi)
	println("Dry Eigenvalues: $(sigdig.(2*pi*fs)) [radians per second]")
	println("Dry Natural frequencies: $(sigdig.(fs)) [Hz]")

	neigvs = size(v, 2)

	fes = femm.integdomain.fes

	bfes = meshboundary(fes)
	# These are the wet surface panels
	wbfes = bfes;
	nnperel = nodesperelem(wbfes)

	pnormals = LaplBEM.allnormals(LaplBEM.allelxs(fens.xyz, connasarray(wbfes)))
	vn = fill(0.0, count(wbfes), neigvs)
	uk = fill(0.0, 3)
	for mode in 1:neigvs
		for i in 1:count(wbfes)
			vn[i, mode] = 0.0
			for k in wbfes.conn[i]
				for m in 1:3
					dof = u.dofnums[k, m]
					uk[m] = (dof != 0 ? v[dof, mode] : 0.0)
				end
				vn[i, mode] += dot(pnormals[i], uk)
			end
			vn[i, mode] /= nnperel
		end
	end

	scalars = []
	for mode in 1:neigvs
		push!(scalars, ("vn$mode", deepcopy(vn[:, mode])))
	end
	# File =  "flint_rock3_wet-vn.vtk"
	# vtkexportmesh(File, fens, wbfes; scalars=scalars)
	# @async run(`"paraview.exe" $File`)

	println("Layer potential matrices")
	@timeit to "Layer potential matrices" begin
		A = fill(0.0, count(wbfes), count(wbfes))
		A = LaplBEM.doublelayer!(A, fens.xyz, connasarray(wbfes), TriRule(3))
		B = fill(0.0, count(wbfes), count(wbfes))
		B = LaplBEM.singlelayer!(B, fens.xyz, connasarray(wbfes), TriRule(3), TriRule(3))
	end
	println("Solve BEM equations")
	@timeit to "Solve BEM equations" begin
		pp = A \ (B * vn)
	end

	areas = fill(0.0, count(wbfes))
	n = fill(0.0, 3)
	for i in 1:count(wbfes)
		ix = collect(wbfes.conn[i])
		LaplBEM.trinml!(n, @view fens.xyz[ix, :])
		areas[i] = norm(n) / 2
	end

	rK = diagm(vec(d))
	rM = I
	raM = fill(0.0, neigvs, neigvs);
	println("Added mass matrix")
	@timeit to " " begin
		for i in 1:neigvs
		    for j in 1:neigvs
		        raM[i,j] = rhow*dot(vn[:,i].*areas, pp[:,j]);
		    end
		end
		raM = (raM +raM')/2;# it should be symmetric, but make sure
	end

	decomp = eigen(rK, rM+raM)
	wv = v*decomp.vectors;
	d = decomp.values

	@show to

	fs = real(sqrt.(complex(d)))/(2*pi)
	println("Wet Eigenvalues: $(sigdig.(2*pi*fs)) [radians per second]")
	println("Wet Natural frequencies: $(sigdig.(fs)) [Hz]")

	true
end # flint_rock3_wet

function flint_rock3_wet_algo()
	output = MeshImportModule.import_ABAQUS("./flint-rock3-1_5mm.inp")
	fens, fes = output["fens"], output["fesets"][1]
	fens.xyz .*= phun("MM")

	MR = DeforModelRed3D
	material = MatDeforElastIso(MR, rho, E, nu, 0.0)

	femm = FEMMDeforLinearESNICET4(MR, IntegDomain(fes, NodalSimplexRule(3)), material)

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

	# Make model data
	modeldata =  FDataDict("fens"=> fens, "regions"=>  [region1],
		"omega_shift"=>OmegaShift^2, "neigvs"=>neigvs)

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

	d = modeldata["omega"]
	fs = real.(complex.(d))/(2*pi)
	println("Dry angular frequencies: $(sigdig.(d)) [rad/s]")
	println("Dry Natural frequencies: $(sigdig.(fs)) [Hz]")

	# Determine the wet boundary faces
	bfes = meshboundary(fes)
	# These are the wet surface panels
	modeldata["wet_boundary_fes"] = bfes;
	modeldata["rhow"] = rhow

	# Solve for the wet-vibration modes
	modeldata = AlgoVibInFluidsModule.modal(modeldata)

	d = modeldata["wet_omega"]
	fs = real.(complex.(d))/(2*pi)
	println("Wet angular frequencies: $(sigdig.(d)) [rad/s]")
	println("Wet Natural frequencies: $(sigdig.(fs)) [Hz]")

	modeldata = AlgoVibInFluidsModule.exportfluidpressuremode(modeldata)
	modeldata["postprocessing"]["mode"] = 1:neigvs
	modeldata = AlgoVibInFluidsModule.exportmode(modeldata)
    true
end # flint_rock3_wet_algo

function free_vibration_solver_w_remeshing(nunref = 6, nref = 2)
	to = TimerOutput()

	println("Loading mesh")
    @timeit to "Loading mesh" begin
	    MR = DeforModelRed3D
	    output = MeshImportModule.import_ABAQUS("./flint-rock3-1_5mm.inp")
	    fens, fes = output["fens"], output["fesets"][1]
	    fens.xyz .*= phun("MM")
	end

    println("Before unrefinement: ")
    println("Number of nodes: $(count(fens))")
    println("Interior mesh: $(count(fes)) tets")
    println("Surface mesh: $(count(meshboundary(fes))) triangles")

    File = "Original.vtk"
    vtkexportmesh(File, fens, fes)

   	V0, centroid0 = volumeetc(fens, fes)

    remesher = Remesher(fens.xyz, connasarray(fes), [1 for idx in 1:count(fes)], 0.0)
    for pass in 1:nunref
    	remesh!(remesher)
    	t, v, tmid = meshdata(remesher)
    	fens.xyz = v
    	fes = fromarray!(fes, t)
    	setlabel!(fes, tmid)
    	V, centroid = volumeetc(fens, fes)
    	stretch = (V0 / V)^(1/3)
    	for i in 1:size(fens.xyz, 1)
    		for k in 1:size(fens.xyz, 2)
    			fens.xyz[i, k] = (fens.xyz[i, k] - centroid[k]) * stretch + centroid[k]
    		end
    	end
    	File = "Unref-$(pass).vtk"
    	vtkexportmesh(File, fens, fes)
    	println("After unrefinement: ")
    	println("Number of nodes: $(count(fens))")
    	println("Interior mesh: $(count(fes)) tets")
    	println("Surface mesh: $(count(meshboundary(fes))) triangles")
    end

    for pass in 1:nref
    	fens, fes = T4refine(fens, fes)
    end

    File = "Ref.vtk"
    vtkexportmesh(File, fens, fes)
    # @async run(`"paraview.exe" $File`)

    println("After refinement: ")
    println("Number of nodes: $(count(fens))")
    println("Interior mesh: $(count(fes)) tets")
    println("Surface mesh: $(count(meshboundary(fes))) triangles")

    geom = NodalField(fens.xyz)
    u = NodalField(zeros(size(fens.xyz,1),3)) # displacement field
    # nl = selectnode(fens, box=[0.0 0.0 -Inf Inf -Inf Inf], inflate=tolerance)
    # setebc!(u, nl, true, collect(1:3))
    applyebc!(u)
    numberdofs!(u)

    material = MatDeforElastIso(MR, rho, E, nu, 0.0)

    @timeit to "Set up FE machine" begin
	    femm = FEMMDeforLinearESNICET4(MR, IntegDomain(fes, NodalSimplexRule(3)), material)
	    associategeometry!(femm,  geom)
	end

	println("Stiffness matrix")
    @timeit to "Stiffness matrix" begin
	    K = stiffness(femm, geom, u)
	end
	println("Mass matrix")
	@timeit to "Mass matrix" begin
    	M = mass(femm, geom, u)
    end

    println("Eigenvalue problem")
    @timeit to "Eigenvalue problem" begin
	    d,v,nev,nconv = eigs(K+OmegaShift^2*M, M; nev=neigvs, which=:SM,  explicittransform=:none)
	    d = d .- OmegaShift^2;
	end

	println("Done")

	return Dict("fens"=>fens, "femm"=>femm, "geom"=>geom, "u"=>u, "eigenvectors"=>v, "eigenvalues"=>d, "timeroutput" =>to)
end # free_vibration_solver_w_remeshing

function free_vibration_solver_w_remeshing_ren(nunref = 6, nref = 2)
	# As free_vibration_solver_w_remeshing, but with renumbering of the nodes to
	# minimize the skyline (fill-in).
	to = TimerOutput()

	println("Loading mesh")
    @timeit to "Loading mesh" begin
	    MR = DeforModelRed3D
	    output = MeshImportModule.import_ABAQUS("./flint-rock3-1_5mm.inp")
	    fens, fes = output["fens"], output["fesets"][1]
	    fens.xyz .*= phun("MM")
	end

    println("Before unrefinement: ")
    println("Number of nodes: $(count(fens))")
    println("Interior mesh: $(count(fes)) tets")
    println("Surface mesh: $(count(meshboundary(fes))) triangles")

    File = "Original.vtk"
    vtkexportmesh(File, fens, fes)

   	V0, centroid0 = volumeetc(fens, fes)

    remesher = Remesher(fens.xyz, connasarray(fes), [1 for idx in 1:count(fes)], 0.0)
    for pass in 1:nunref
    	remesh!(remesher)
    	t, v, tmid = meshdata(remesher)
    	fens.xyz = v
    	fes = fromarray!(fes, t)
    	setlabel!(fes, tmid)
    	V, centroid = volumeetc(fens, fes)
    	stretch = (V0 / V)^(1/3)
    	for i in 1:size(fens.xyz, 1)
    		for k in 1:size(fens.xyz, 2)
    			fens.xyz[i, k] = (fens.xyz[i, k] - centroid[k]) * stretch + centroid[k]
    		end
    	end
    	# @show V, centroid = volumeetc(fens, fes)
    	File = "Unref-$(pass).vtk"
    	vtkexportmesh(File, fens, fes)
    	println("After unrefinement: ")
    	println("Number of nodes: $(count(fens))")
    	println("Interior mesh: $(count(fes)) tets")
    	println("Surface mesh: $(count(meshboundary(fes))) triangles")
    end

    for pass in 1:nref
    	fens, fes = T4refine(fens, fes)
    end

    numbering = let
    	C = connectionmatrix(FEMMBase(IntegDomain(fes, TetRule(1))), count(fens))
    	ag = adjgraph(C)
    	nd = nodedegrees(ag)
    	numbering = revcm(ag, nd)
    end

    File = "Ref.vtk"
    vtkexportmesh(File, fens, fes)
    # @async run(`"paraview.exe" $File`)

    println("After refinement: ")
    println("Number of nodes: $(count(fens))")
    println("Interior mesh: $(count(fes)) tets")
    println("Surface mesh: $(count(meshboundary(fes))) triangles")

    geom = NodalField(fens.xyz)
    u = NodalField(zeros(size(fens.xyz,1),3)) # displacement field
    # nl = selectnode(fens, box=[0.0 0.0 -Inf Inf -Inf Inf], inflate=tolerance)
    # setebc!(u, nl, true, collect(1:3))
    applyebc!(u)
    numberdofs!(u, numbering)

    material = MatDeforElastIso(MR, rho, E, nu, 0.0)

    @timeit to "Set up FE machine" begin
	    femm = FEMMDeforLinearESNICET4(MR, IntegDomain(fes, NodalSimplexRule(3)), material)
	    associategeometry!(femm,  geom)
	end

	println("Stiffness matrix")
    @timeit to "Stiffness matrix" begin
	    K = stiffness(femm, geom, u)
	end
	println("Mass matrix")
	@timeit to "Mass matrix" begin
    	M = mass(femm, geom, u)
    end

    println("Eigenvalue problem")
    @timeit to "Eigenvalue problem" begin
	    d,v,nev,nconv = eigs(K+OmegaShift^2*M, M; nev=neigvs, which=:SM)
	    d = d .- OmegaShift^2;
	end

	println("Done")

	return Dict("fens"=>fens, "femm"=>femm, "geom"=>geom, "u"=>u, "eigenvectors"=>v, "eigenvalues"=>d, "timeroutput" =>to)
end # free_vibration_solver_w_remeshing_ren

function flint_rock3_wet_unrefine()
	to = TimerOutput()

	model = free_vibration_solver_w_remeshing()

	@show model["timeroutput"]

	fens = model["fens"]
	femm = model["femm"]
	d = model["eigenvalues"]
	v = model["eigenvectors"]
	u = model["u"]

	fs = real(sqrt.(complex(d)))/(2*pi)
	println("Dry Eigenvalues: $(sigdig.(2*pi*fs)) [radians per second]")
	println("Dry Natural frequencies: $(sigdig.(fs)) [Hz]")

	neigvs = size(v, 2)

	fes = femm.integdomain.fes

	bfes = meshboundary(fes)
	# These are the wet surface panels
	wbfes = bfes;
	nnperel = nodesperelem(wbfes)

	pnormals = LaplBEM.allnormals(LaplBEM.allelxs(fens.xyz, connasarray(wbfes)))
	vn = fill(0.0, count(wbfes), neigvs)
	uk = fill(0.0, 3)
	for mode in 1:neigvs
		for i in 1:count(wbfes)
			vn[i, mode] = 0.0
			for k in wbfes.conn[i]
				for m in 1:3
					dof = u.dofnums[k, m]
					uk[m] = (dof != 0 ? v[dof, mode] : 0.0)
				end
				vn[i, mode] += dot(pnormals[i], uk)
			end
			vn[i, mode] /= nnperel
		end
	end

	scalars = []
	for mode in 1:neigvs
		push!(scalars, ("vn$mode", deepcopy(vn[:, mode])))
	end
	# File =  "flint_rock3_wet-vn.vtk"
	# vtkexportmesh(File, fens, wbfes; scalars=scalars)
	# @async run(`"paraview.exe" $File`)

	println("Layer potential matrices")
	@timeit to "Layer potential matrices" begin
		A = fill(0.0, count(wbfes), count(wbfes))
		A = LaplBEM.doublelayer!(A, fens.xyz, connasarray(wbfes), TriRule(3))
		B = fill(0.0, count(wbfes), count(wbfes))
		B = LaplBEM.singlelayer!(B, fens.xyz, connasarray(wbfes), TriRule(3), TriRule(3))
	end
	println("Solve BEM equations")
	@timeit to "Solve BEM equations" begin
		pp = A \ (B * vn)
	end

	areas = fill(0.0, count(wbfes))
	n = fill(0.0, 3)
	for i in 1:count(wbfes)
		ix = collect(wbfes.conn[i])
		LaplBEM.trinml!(n, @view fens.xyz[ix, :])
		areas[i] = norm(n) / 2
	end

	rK = diagm(vec(d))
	rM = I
	raM = fill(0.0, neigvs, neigvs);
	println("Added mass matrix")
	@timeit to " " begin
		for i in 1:neigvs
		    for j in 1:neigvs
		        raM[i,j] = rhow*dot(vn[:,i].*areas, pp[:,j]);
		    end
		end
		raM = (raM +raM')/2;# it should be symmetric, but make sure
	end

	decomp = eigen(rK, rM+raM)
	wv = v*decomp.vectors;
	d = decomp.values

	@show to

	fs = real(sqrt.(complex(d)))/(2*pi)
	println("Wet Eigenvalues: $(sigdig.(2*pi*fs)) [radians per second]")
	println("Wet Natural frequencies: $(sigdig.(fs)) [Hz]")

	true
end # flint_rock3_wet_unrefine

function allrun()
    println("#####################################################")
    println("# flint_rock3_wet_unrefine ")
    flint_rock3_wet_unrefine()
    println("#####################################################")
    println("# flint_rock3_dry ")
    flint_rock3_dry()
    println("#####################################################")
    println("# flint_rock3_wet ")
    flint_rock3_wet()
    println("#####################################################")
    println("# flint_rock3_wet_algo ")
    flint_rock3_wet_algo()
    return true
end # function allrun

end # module flint_rock3_examples

max_stars_repo_name: PetrKryslUCSD/FinEtoolsVibInFluids

max_issues_repo_name: PetrKryslUCSD/FinEtoolsVibInFluids

max_forks_repo_name: PetrKryslUCSD/FinEtoolsVibInFluids


text:
export Convolution

#############
# Convolution implements the convolution operator:
#
# \int_Γ G(|x-y|) u(y) dy or dΓ(y), for x ∈ Γ,
#
# where op qualifies the integral operator, and
# where G is supplied as a Fun.
#############

abstract type Convolution{S,BT,T} <: Operator{T} end

struct ConcreteConvolution{S<:Space,BT,T} <: Convolution{S,BT,T}
    op::BT
    G::Fun{S,T}
end

Convolution(op::BT,G::Fun{S,T}) where {S,BT,T} = ConcreteConvolution(op,G)

domain(C::ConcreteConvolution) = domain(C.op)
domainspace(C::ConcreteConvolution) = domainspace(C.op)
rangespace(C::ConcreteConvolution) = UnsetSpace()
bandwidths(C::ConcreteConvolution) = error("No range space attached to Convolution")
getindex(C::ConcreteConvolution,k::Integer,j::Integer) = error("No range space attached to Convolution")

rangespace(C::ConcreteConvolution{Laurent{D},BT}) where {D,BT} = space(C.G)
rangespace(C::ConcreteConvolution{Fourier{D},BT}) where {D,BT} = space(C.G)
rangespace(C::ConcreteConvolution{CosSpace{D},BT}) where {D,BT} = Fourier(domain(C.G))

bandwidths(C::ConcreteConvolution{Laurent{D},BT}) where {D,BT} = 0,0
bandwidths(C::ConcreteConvolution{Fourier{D},BT}) where {D,BT} = 0,0
bandwidths(C::ConcreteConvolution{CosSpace{D},BT}) where {D,BT} = 0,0

function getindex(C::ConcreteConvolution{Fourier{D},ConcreteDefiniteLineIntegral{Fourier{D},T1},T2},k::Integer,j::Integer) where {D,T1,T2}
    T = promote_type(T1,T2)
    if k == j
        if k == 1 && k ≤ ncoefficients(C.G)
            (C.op[1]*C.G.coefficients[k])::T
        elseif isodd(k) && k ≤ ncoefficients(C.G)
            (C.op[1]*C.G.coefficients[k]/2)::T
        elseif iseven(k) && k ≤ ncoefficients(C.G)-1
            (C.op[1]*C.G.coefficients[k+1]/2)::T
        else
            zero(T)
        end
    else
        zero(T)
    end
end

function getindex(C::ConcreteConvolution{CosSpace{D},ConcreteDefiniteLineIntegral{Fourier{D},T1},T2},k::Integer,j::Integer) where {D,T1,T2}
    T = promote_type(T1,T2)
    if k == j
        if k == 1 && k ≤ ncoefficients(C.G)
            (C.op[1]*C.G.coefficients[k])::T
        elseif k ≤ 2ncoefficients(C.G)-1
            (C.op[1]*C.G.coefficients[k÷2+1]/2)::T
        else
            zero(T)
        end
    else
        zero(T)
    end
end

function getindex(C::ConcreteConvolution{Laurent{D},ConcreteDefiniteLineIntegral{Laurent{D},T1},T2},k::Integer,j::Integer) where {D,T1,T2}
    T = promote_type(T1,T2)
    if k == j && k ≤ ncoefficients(C.G)
        (C.op[1]*C.G.coefficients[k])::T
    else
        zero(T)
    end
end

function getindex(C::ConcreteConvolution{Laurent{D},ConcreteSingularIntegral{Laurent{D},OT,T1},T2},k::Integer,j::Integer) where {D,OT,T1,T2}
    T = promote_type(T1,T2)
    c = C.G.coefficients
    Op = C.op
    N = ncoefficients(C.G)
    ret = zero(T)
    if k == j
        if isodd(k)
            for ℓ = 1:2:N
                ret += c[ℓ]*Op[ℓ+k-1,ℓ+k-1]
            end
            for ℓ = 2:2:min(k-1,N)
                ret += c[ℓ]*Op[k-ℓ,k-ℓ]
            end
            for ℓ = (k+1):2:N
                ret += c[ℓ]*Op[ℓ-k+1,ℓ-k+1]
            end
        else
            for ℓ = 2:2:N
                ret += c[ℓ]*Op[ℓ+k,ℓ+k]
            end
            for ℓ = 1:2:min(k-1,N)
                ret += c[ℓ]*Op[k-ℓ+1,k-ℓ+1]
            end
            for ℓ = (k+1):2:N
                ret += c[ℓ]*Op[ℓ-k,ℓ-k]
            end
        end
    end
    ret
end


text:
#
# Solve the 2d wave equation using finite differencing
#
# @author arbennett
#

# Begin timing
#t_0 = time()

# Parameters for solving
const N = 50
const M = 50
const T = 6000
const dt = 0.1
const plot_frequency = 10
u = zeros(3,N,M)
f = open("solution.dat", "w")

# Set the initial conditions
for i=1:N
    for j=1:M
        u[2:3,i,j] = 1 + cos(2 * pi * (i/N)) + cos(2 * pi * (j/M))
        #u[3,i,j] = sin(2 * pi * ((i+1)/N)) + cos(2 * pi * ((j+1)/M))
    end
end

# Enforce boundary conditions
u[:,:,1] = u[:,:,2]
u[:,:,end] = u[:,:,end-1]
u[:,1,:] = u[:,2,:]
u[:,end,:] = u[:,end-1,:]

# Do some timestepping
for t=1:int(T)
    # Calculate the new u
    u[1,2:end-1,2:end-1] = (2*u[2,2:end-1,2:end-1] - u[3,2:end-1,2:end-1] +
        dt^2 * (u[2,2:end-1,1:end-2] + u[2,2:end-1,3:end] + u[2,1:end-2,2:end-1]
        + u[2,3:end,2:end-1] - 4*u[2,2:end-1,2:end-1]))

    # Enforce boundary conditions
    u[:,:,1] = u[:,:,2]
    u[:,:,end] = u[:,:,end-1]
    u[:,1,:] = u[:,2,:]
    u[:,end,:] = u[:,end-1,:]

    # Move the time forwards
    u = circshift(u,[1,0,0])
    if mod(t, plot_frequency) == 0
        writedlm("out/solution.dat."*lpad(string(t),length(string(int(T))),"0"), u[1,:,:])
    end
end

# Report the time taken
#t_1 = time()
#println("Time elapsed: ", t_1 - t_0, "s")

max_stars_repo_name: arbennett/Simple-Simulations-In-Julia

max_issues_repo_name: arbennett/Simple-Simulations-In-Julia

max_forks_repo_name: arbennett/Simple-Simulations-In-Julia


text:
using DataFrames

function mungdata(filename)
    lines = readlines(filename)
    numlines = length(lines)
    dates = Array{DateTime, 1}(numlines)
    means = zeros(Float64, numlines)
    numvalid = zeros(Int, numlines)
    invalidlength = zeros(Int, numlines)
    invalidpos = zeros(Int, numlines)
    datamatrix = Array{Float64,2}(numlines, 24)
    datamatrix .= NaN
    totalsum = 0.0
    totalgood = 0
    for (linenum,line) in enumerate(lines)
        data = split(line)
        validcount = badlength = 0
        validsum = 0.0
        for i in 2:2:length(data)-1
            if parse(Int, data[i+1]) >= 0
                validsum += (datamatrix[linenum, Int(i/2)] = parse(Float64, data[i]))
                validcount += 1
                badlength = 0
            else
                badlength += 1
                if badlength > invalidlength[linenum]
                    invalidlength[linenum] = badlength
                    invalidpos[linenum] = Int(i/2) - invalidlength[linenum] + 1
                end
            end
        end
        dates[linenum] = DateTime(data[1], "y-m-d")
        means[linenum] = validsum / validcount
        numvalid[linenum] = validcount
        totalsum += validsum
        totalgood += validcount
    end
    dt = DataFrame(Date = dates, Mean = means, ValidValues = numvalid,
                   MaximumGap = invalidlength, GapPosition = invalidpos)
    for i in 1:size(datamatrix)[2]
        dt[Symbol("$(i-1):00")] = datamatrix[:,i]
    end
    dt, totalsum/totalgood
end

datafilename = "data.txt" # this is taken from the example listed on the task, since the actual text file is not available
df, dmean = mungdata(datafilename)
println(df)
println("The overall mean is $dmean")
maxbadline = indmax(df[:MaximumGap])
maxbadval = df[:MaximumGap][maxbadline]
maxbadtime = df[:GapPosition][maxbadline] - 1
maxbaddate = replace("$(df[:Date][maxbadline])", r"T.+$", "")
println("The largest run of bad values is $(maxbadval), on $(maxbaddate) beginning at $(maxbadtime):00 hours.")

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
# test g-distribution
@testset "g-dist" begin
    a = 3.5
    samples2 = [KissMCMC.sample_g(a) for i=1:50000];

    @test all(1/a .<= samples2 .<= a)
    @test KissMCMC.cdf_g_inv(1, a) ≈ a
    @test KissMCMC.cdf_g_inv(0, a) ≈ 1/a
    z = 1/a:0.01:a
    meang = sum(z.*KissMCMC.g_pdf.(z, a)) *step(z)
    @test isapprox(mean(samples2), meang, atol=1e-2)
    stdg = sqrt(sum((meang.-z).^2 .* KissMCMC.g_pdf.(z, a)) *step(z))
    @test isapprox(std(samples2), stdg, atol=1e-2)
end

# test MCMC
@testset "emcee" begin
    for tc in testcases
        pdf = make_lpdf(tc.dist)
        @inferred pdf(tc.theta0)
        theta0s = make_theta0s(tc.theta0, tc.ball_radius,
                               pdf, tc.nwalkers,
                               hasblob=tc.hasblob)
        samples = emcee(pdf, theta0s;
                        niter=tc.niter,
                        hasblob=tc.hasblob,
                        use_progress_meter=false,
                        tc.otherkws...);
        @test length.(samples[1:3]) == (tc.nwalkers, tc.nwalkers, tc.nwalkers)
        if tc.hasblob
            @test length(samples[4]) == tc.nwalkers
        else
            @test samples[4]==nothing
        end
        @test length(samples[1][1]) == tc.niter÷tc.nwalkers÷2 # ÷2 because of burnin
        thetas, accept_ratio, logdensities, blobs = squash_walkers(samples...;
                                                                   verbose=false,
                                                                   tc.squash_walkers_kws...)

        !tc.hasblob && @test blobs==nothing
        @test length(thetas) == tc.niter÷2
        @test length(logdensities) == tc.niter÷2
        @test accept_ratio>0.1
        test_mean_std(thetas, tc, tc.tole)
        test_blobs(tc.blob_truths, blobs)
        #test_logdensities(tc.truths, logdensities, pdf)
    end
end


text:
module MachineLearningMetrics

##############################################################################
##
## Dependencies
##
##############################################################################

##############################################################################
##
## Exported methods and types (in addition to everything reexported above)
##
##############################################################################

export @~,
    check_args,
    absolute_error,
    percent_error,
    log_error,
    squared_error,
    squared_log_error,
    absolute_percent_error,
    mean_error,
    mean_absolute_error,
    median_absolute_error,
    mean_percent_error,
    median_percent_error,
    mean_squared_error,
    median_squared_error,
    sum_squared_error,
    mean_squared_log_error,
    mean_absolute_percent_error,
    median_absolute_percent_error,
    symmetric_mean_absolute_percent_error,
    symmetric_median_absolute_percent_error,
    mean_absolute_scaled_error,
    total_variance_score,
    explained_variance_score,
    unexplained_variance_score,
    r2_score,
    total_population,
    true_positive,
    true_negative,
    false_positive,
    false_negative,
    condition_positive,
    condition_negative,
    predicted_condition_positive,
    predicted_condition_negative,
    accuracy_score,
    accuracy,
    prevalence,
    positive_predictive_value,
    precision,
    false_discovery_rate,
    negative_predictive_value,
    false_omission_rate,
    true_positive_rate,
    sensitivity,
    recall,
    false_positive_rate,
    false_negative_rate,
    true_negative_rate,
    specificity,
    positive_likelihood_ratio,
    negative_likelihood_ratio,
    diagnostic_odds_ratio,
    f1_score,
    matthews_corrcoef

##############################################################################
##
## Load files
##
##############################################################################

filenames = [
    "classification.jl",
    "common.jl",
    "regression.jl"
  ]

for filename in filenames
    include(filename)
end

end # module MachineLearningMetrics


text:
const cuda_memstep = 1024
const cuda_ptrs = Dict{Int,Vector{Ptr{Void}}}()

type CudaArray{T,N}
  ptr::Ptr{T}
  dims::NTuple{N,Int}
  dev::Int
end

function allocate!{T}(a::CudaArray{T})
  id = length(a) * sizeof(T) ÷ cuda_memstep + 1
  if haskey(cuda_ptrs, id)
    ptrs = cuda_ptrs[id]
  else
    ptrs = Ptr{Void}[]
    cuda_ptrs[id] = ptrs
  end
  if length(ptrs) == 0
    p = Ptr{Void}[0]
    cudaMalloc(p, id * cuda_memstep)
    a.ptr = p[1]
  else
    a.ptr = convert(Ptr{T}, pop!(ptrs))
  end
end

function release{T}(a::CudaArray{T})
  id = length(a) * sizeof(T) ÷ cuda_memstep + 1
  p = convert(Ptr{Void}, a.ptr)
  push!(cuda_ptrs[id], p)
end

function device_gc()
  for ptrs in cuda_ptrs
    for p in ptrs
      cudaFree(p)
    end
  end
end

function CudaArray{T,N}(::Type{T}, dims::NTuple{N,Int})
  a = CudaArray(Ptr{T}(0), dims, 1)
  allocate!(a)
  finalizer(a, release)
  a
end
CudaArray{T}(::Type{T}, dims...) = CudaArray(T, dims)

typealias CudaVector{T} CudaArray{T,1}
typealias CudaMatrix{T} CudaArray{T,2}

Base.length(a::CudaArray) = prod(a.dims)
Base.size(a::CudaArray) = a.dims
Base.ndims{T,N}(a::CudaArray{T,N}) = N
Base.eltype{T}(a::CudaArray{T}) = T
Base.stride(a::CudaArray, dim::Int) = prod(size(a)[1:dim-1])
Base.pointer(a::CudaArray) = a.ptr

Base.similar{T}(a::CudaArray{T}) = CudaArray(T, size(a))

Base.unsafe_convert(::Type{Ptr{Void}}, a::CudaArray) = convert(Ptr{Void}, a.ptr)

Base.Array{T}(a::CudaArray{T}) = copy!(Array(T, size(a)), a)

# To limit the likelihood of ambiguity warnings with other packages,
# these should be the only two-argument definitions of copy!
#copy!(dst::Array, src::HostOrDevArray; stream=null_stream) = _copy!(dst, src, stream)
#copy!(dst::HostOrDevArray, src::HostOrDevArray; stream=null_stream) = _copy!(dst, src, stream)

function copy!{T}(dest::Array{T}, src::CudaArray{T}, stream=nullstream)
  if length(dest) != length(src)
    throw(ArgumentError("Inconsistent array length."))
  end
  nbytes = length(src) * sizeof(T)
  cudaMemcpyAsync(dest, src, nbytes, cudaMemcpyDeviceToHost, stream)
  dest
end
#_copy!{T}(dst::ContiguousArray{T}, src::ContiguousArray, stream) = _copy!(dst, to_eltype(T, src), stream)
#_copy!{T}(dst::AbstractCudaArray{T}, src, stream) = _copy!(dst, copy!(Array(T, size(src)), src), stream)


text:
using Flux
using Flux: onehot, chunk, batchseq, throttle, logitcrossentropy
using StatsBase: wsample
using Base.Iterators: partition
using Parameters: @with_kw

# Hyperparameter arguments
@with_kw mutable struct Args
    lr::Float64 = 1e-2	# Learning rate
    seqlen::Int = 50	# Length of batchseqences
    nbatch::Int = 50	# number of batches text is divided into
    throttle::Int = 30	# Throttle timeout
end

function getdata(args)
    # Download the data if not downloaded as 'input.txt'
    isfile("input.txt") ||
        download("https://cs.stanford.edu/people/karpathy/char-rnn/shakespeare_input.txt","input.txt")

    text = collect(String(read("input.txt")))

    # an array of all unique characters
    alphabet = [unique(text)..., '_']

    text = map(ch -> onehot(ch, alphabet), text)
    stop = onehot('_', alphabet)

    N = length(alphabet)

    # Partitioning the data as sequence of batches, which are then collected as array of batches
    Xs = collect(partition(batchseq(chunk(text, args.nbatch), stop), args.seqlen))
    Ys = collect(partition(batchseq(chunk(text[2:end], args.nbatch), stop), args.seqlen))

    return Xs, Ys, N, alphabet
end

# Function to construct model
function build_model(N)
    return Chain(
            LSTM(N, 128),
            LSTM(128, 128),
            Dense(128, N))
end

function train(; kws...)
    # Initialize the parameters
    args = Args(; kws...)

    # Get Data
    Xs, Ys, N, alphabet = getdata(args)

    # Constructing Model
    m = build_model(N)

    function loss(xs, ys)
      l = sum(logitcrossentropy.(m.(xs), ys))
      return l
    end

    ## Training
    opt = ADAM(args.lr)
    tx, ty = (Xs[5], Ys[5])
    evalcb = () -> @show loss(tx, ty)

    Flux.train!(loss, params(m), zip(Xs, Ys), opt, cb = throttle(evalcb, args.throttle))
    return m, alphabet
end

# Sampling
function sample(m, alphabet, len; seed="")
    m = cpu(m)
    Flux.reset!(m)
    buf = IOBuffer()
    if seed == ""
        seed = string(rand(alphabet))
    end
    write(buf, seed)
    c = wsample(alphabet, softmax(m.(map(c -> onehot(c, alphabet), collect(seed)))[end]))
    for i = 1:len
        write(buf, c)
        c = wsample(alphabet, softmax(m(onehot(c, alphabet))))
    end
    return String(take!(buf))
end

cd(@__DIR__)
m, alphabet = train()
sample(m, alphabet, 1000) |> println

max_stars_repo_name: shawn-kern/model-zoo

max_issues_repo_name: shawn-kern/model-zoo

max_forks_repo_name: shawn-kern/model-zoo