using JuMP
import NLopt
using Ipopt

function D2(v0::Vector)

    function ExplicitFunction(uw::Vector, grad::Vector, v::Vector)
        println("called my own OFunction")
        if length(grad) > 0
            grad[1] = 0
            grad[2] = 0.5 / sqrt(x[2])
        end
        w1, w2, u1, u2, u3, u4 = uw
        return u3 * (u2 + u4 - u1 - u3 + (u1 + u2 - u3 - u4 - 2w2 - 6w1) / (2w1 + 2w2)) +
               u1 * (u2 + u4 - u1 - u3 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2)) +
               u4 * (-4 + u1 + u3 - u2 - u4 + (u1 + u2 - u3 - u4 - 2w2 - 6w1) / (2w1 + 2w2)) +
               u2 * (-4 + u1 + u3 - u2 - u4 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2)) +
               w1 * (-3 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2))^2 +
               (1 / 2) * (-u1 + u2 - u3 + u4)^2 +
               w2 * (-1 + (u3 + u4 - u1 - u2 + 2w2 + 6w1) / (2w1 + 2w2))^2 - w1 * v[1] - w2 * v[2]
    end
    function ImplicitFunction(uw::Vector, grad::Vector, v::Vector)

        if length(grad) > 0
            grad[1] = 0
            grad[2] = 0.5 / sqrt(x[2])
        end
        model2 = Model(Ipopt.Optimizer)
        JuMP.set_silent(model2)
        JuMP.@variable(model2, x1)
        JuMP.@variable(model2, x2)
        JuMP.@NLexpression(model2, my_expr, uw[1] * (x1 - 3)^2 + (x2 - 1)^2 + uw[2] * (x1 - 1)^2 + (x2 - 1)^2)
        JuMP.@NLobjective(model2, Min,
            uw[3] * (x1 + 2 * x2 - 2) +
            uw[4] * (x1 - 2 * x2 - 2) +
            uw[5] * (-x1 + 2 * x2 - 2) +
            uw[6] * (-x1 - 2 * x2 - 2) +
            my_expr
        )
        JuMP.optimize!(model2)
        return JuMP.objective_value(model2) - uw[1] * v[1] - uw[2] * v[2]
    end

    function constraint(uw::Vector, _grad::Vector)
        uw[1] + uw[2] - 1
    end

    opt = NLopt.Opt(:LN_COBYLA, 6) # LN_COBYLA LD_MMA
    opt.lower_bounds = repeat([0], 6)
    # opt.xtol_rel = 1e-2

    opt.max_objective = (x, g) -> ImplicitFunction(x, g, v0)
    NLopt.equality_constraint!(opt, constraint, 1e-8)

    (minf, minx, ret) = NLopt.optimize(opt, repeat([0.9], 6))
    numevals = opt.numevals # the number of function evaluations
    w1, w2, u1, u2, u3, u4 = minx
    return (value.([w1, w2]), value.([u1, u2, u3, u4]))
end


mutable struct UIO{V<:Real}
    F::Matrix{V}
    T::Matrix{V}
    B::Matrix{V}
    H::Matrix{V}
    C::OutputMatrix{V}
    A1::Matrix{V}
    K1::Matrix{V}
    K::Matrix{V}
    z::Vector{V}

    function UIO(A::U, B::U, C::OutputMatrix{V}, E::U) where {V<:Real, U<:Matrix{V}}
        (ny, feasible, lrE) = checkUIOFeasibility(A,C,E);
        if feasible
            H = lrE*pinv(C*lrE)
            T = I - H*C
            A1 = T*A
            nz = size(A1, 1)
            # K2 = F*H;
            new{V}(
                zeros(V, nz, nz),
                T,
                B,
                H,
                C,
                A1,
                zeros(V, nz, ny),
                zeros(V, nz, ny),
                zeros(V, nz));
        else
            error("unfeasible UIO");
        end
    end
end

function UIO(A::U, B::U, C::OutputMatrix{V}, E::U, K1::U) where {V<:Real, U<:Matrix{V}}
    uio = UIO(A, B, C, E)
    setF(uio, K1)
    uio
end

function checkUIOFeasibility(A::Matrix{V},
        C::OutputMatrix{V},
        E::Matrix{V}) where {V<:Real}
    rE = rank(E)
    nx = size(A,1)
    if rE < nx
        SVD = svd(E);
        lrE = SVD.U[:,1:rE] * sqrt(Diagonal(SVD.S[1:rE]))
    else
        lrE = E
    end
    matC = C == I ? 1.0 * Matrix(I, nx, nx) : C
    size(matC, 1), rank(matC*lrE) == rE, lrE
end

function setF(o::UIO{T}, K1::AbstractMatrix{T}) where {T}
    o.F = o.A1 - K1*o.C
    o.K1 = K1;
    o.K = K1 + o.F*o.H
end

function updateState(o::UIO{T}, u::Vector{T}, y::Vector{T}) where {T}
    o.z = o.F*o.z + o.T*o.B*u + o.K*y
end

updateState(o::UIO{T}, u::T, y::Vector{T}) where {T} = updateState(o, [u], y)

outputMap(o::UIO{T}, y::Vector{T}) where {T} = o.z + o.H*y

function Base.getproperty(sys::UIO, s::Symbol)
    if s === :K2
        return sys.K - sys.K1
    else
        return getfield(sys, s)
    end
end

function Base.show(io::IO, sys::UIO)
    # Compose the name vectors
    #inputs = format_names(s.inputnames, "u", "?")
    #outputs = format_names(s.outputnames, "y", "?")
    println(io, "UIO:")
    # println(io, typeof(sys))
    println(io, "F = \n", _string_mat_with_headers(sys.F))
    println(io, "T = \n", _string_mat_with_headers(sys.T))
    println(io, "B = \n", _string_mat_with_headers(sys.B))
    println(io, "K = \n", _string_mat_with_headers(sys.K))
    println(io, "H = \n", _string_mat_with_headers(sys.H))
end


module matrix

    using LinearAlgebra

    # simple function for building a
    # strickly increasing matrix,
    # which is useful for testing.
    function get_inc_mat(T, n::Integer)::Matrix
        mat = Matrix{T}(undef, n, n)

        for row in 1:n
            for col in 1:n
                mat[row, col] = (row - 1)*n + col
            end
        end

        return mat
    end

    #= Rotate Matrix.

    Given an NxN image represented by an NxN matrix,
    where each pixel in the image is 4 bytes, write a method
    to rotate the image by 90°.

    Can you do it in place?
    =#

    export rotate_matrix1

    # long way, basically a brute force solution
    # with O(n²) run-time and O(N) memory.
    # Which, is bad.
    function rotate_matrix1(mat::Array)::Array

        length(mat) <= 1 && return mat

        width, height = size(mat)
        result = Matrix(undef, width, height)

        for row in 1:height
            for col in 1:width
                result[col, height - row + 1] = mat[row, col]
            end
        end

        return result
    end

    export rotate_matrix2

    # Rotate 90+ = transpose, rev
    # in place sln
    function rotate_matrix2(mat::Array)::Array
        length(mat) <= 1 && return mat

        return reverse(transpose(Matrix(mat)), dims=2)
    end

    #= Zero Matrix.

    Write an algorithm such that if an element in an NxM matrix is
    0, its entire row and column are set to 0.
    =#

    export zero_matrix1

    # O(N) solution
    function zero_matrix1(mat::Matrix)::Array

        function zero_row(mat::Matrix, row::Integer)

            height, width = size(mat)

            for index in 1:width
                mat[row, index] = 0
            end
        end

        function zero_column(mat::Matrix, col::Integer)

            height, width = size(mat)

            for index in 1:height
                mat[index, col] = 0
            end
        end

        length(mat) <= 1 && return mat

        height, width = size(mat)

        zero_rows = Set{Integer}()
        zero_cols = Set{Integer}()

        for row in 1:height
            for col in 1:width
                if mat[row, col] == 0
                    zero_rows = union(zero_rows, row)
                    zero_cols = union(zero_cols, col)
                end
            end
        end

        for row in zero_rows
            zero_row(mat, row)
        end

        for col in zero_cols
            zero_column(mat, col)
        end

        return mat
    end

end # end module


using JuliaPerformanceVerification
using Base.Test

# write your own tests here
@test 1 == 1

iteration_cnt_range = [ 10^i for i in 1:7 ]
vector_size_range   = [ 10^i for i in 1:7 ]
vector_wins = trues(length(vector_size_range), length(iteration_cnt_range))
vector_faster = zeros(length(vector_size_range), length(iteration_cnt_range))
for (i, current_vector_size) in enumerate(vector_size_range), (j, iteration_cnt) in enumerate(iteration_cnt_range)
	if current_vector_size * iteration_cnt >= 10^10
	    continue
	end
	a = rand(current_vector_size)
	b = rand(current_vector_size)
	c_vec = zeros(current_vector_size)
	c_iter = zeros(current_vector_size)
	c_vec_oped = zeros(current_vector_size)
	c_iter_oped = zeros(current_vector_size)
	tic()
	vecaddbyvec!(a, b, c_vec, iteration_cnt)
	vector_time = toq();
	tic()
	vecaddbyiter!(a, b, c_iter, iteration_cnt)
	iteration_time = toq();
	@test c_vec == c_iter
	tic()
	vecaddbyvec_oped!(a, b, c_vec_oped, iteration_cnt)
	vector_oped_time = toq();
	tic()
	vecaddbyiter_oped!(a, b, c_iter_oped, iteration_cnt)
	iteration_oped_time = toq();
	@test c_vec == c_iter
	@test c_vec_oped == c_iter
	@test c_vec == c_iter_oped
	vector_wins[i, j] = iteration_time > vector_time ? true: false
	vector_faster[i, j] = iteration_time / vector_time
	if vector_wins[i, j]
		@printf "Two %d dim vectors adding for %d times, vector add wins %.03f faster!\n" current_vector_size iteration_cnt vector_faster[i, j]
	else
		@printf "Two %d dim vectors adding for %d times, iteration add wins %.03f faster!\n" current_vector_size iteration_cnt 1/vector_faster[i, j]
	end
	@printf "  @inbounds and @simd accelerate ratio Vector: %.03f; Iteration: %.03f\n\n" vector_time/vector_oped_time iteration_time/iteration_oped_time
end

@printf "The vector add wins:\n %s \n" vector_wins
@printf "The vector add is x times fast as iteration:\n %s \n" vector_faster


#=
  https://dtai.cs.kuleuven.be/problog/tutorial/basic/03_dice.html
  """
  The following example illustrates the use of a logic program with recursion and lists.
  We start by rolling the first die. Every roll determines the next die to roll, but we stop
  if we have used that die before. We query for the possible sequences of rolled dice.
  We use three-sided dice instead of the regular six-sided ones simply to restrict the number
  of possible outcomes (and thus inference time).
  """

  Cf ~/blog/rolling_dice5.blog
     ~/webppl/rolling_dice5.wppl
=#

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

@model function rolling_dice5(n=3)

    function roll(a)
        # t ~ DiscreteUniform(1,n)
        t = rand(DiscreteUniform(1,n))
        if t in a
            return a
        else
            return roll(vcat(a,t))
        end
    end

    a = roll([])
    len ~ DiscreteUniform(1,n)
    len ~ Dirac(length(a))

    # In the example, the first roll is 1.
    true ~ Dirac(a[1] == 1)

end

function run_model(len=0)
    model = rolling_dice5(len)

    num_chns = 4

    # chns = sample(model, Prior(), 10_000)

    # chns = sample(model, MH(), 10_000)
    chns = sample(model, MH(), 100_000)

    # chns = sample(model, PG(15), MCMCThreads(), 1_000, num_chns)

    # chns = sample(model, SMC(1000), MCMCThreads(), 10_000, num_chns)
    # chns = sample(model, SMC(1000), 10_000)
    # chns = sample(model, IS(), 10_000)
    #
    display(chns)
    show_var_dist_pct(chns,:len,1000)

end

genq = undef
for val in 1:6
    global genq
    println("\nval:$val")
    @time run_model(val)
end

println("\nval=10")
# println("mean:", mean(run_model(10)))
run_model(10)
println("\nval=100")
# println("mean:", mean(run_model(100)))
run_model(100)


export Tensor

struct Tensor <: AbstractTensor
dim::Int          # "dimension"
T::SparseVector{Float64,Int}
get::Function
op::AbstractOrthoPoly

    # inner constructors
    function Tensor(dim::Int,mop::MultiOrthoPoly)
        tensorEntries = computeTensorizedSP(dim, mop)
        getfun(ind) = getentry(ind, tensorEntries, mop.ind, dim)
        new(dim, tensorEntries, getfun, mop)
    end

    function Tensor(dim::Int,opq::AbstractOrthoPoly)
        tensorEntries = computeTensorizedSP(dim, opq)
        getfun(ind) = getentry(ind, tensorEntries, calculateMultiIndices(1, opq.deg), dim)
        new(dim, tensorEntries, getfun, opq)
    end
end


@defstruct FMmodel (
  w0::Float64                = 0.0,
  w::Vector{Float64}         = Vector{Float64}(),
  v::Array{Float64}          = Array{Float64}(),

  # controls
  (SOLVER::String            = "TDAP", SOLVER in ["ALS" "MCMC" "SGD" "BGD" "FTRL" "TDAP"]),
  (TASK::String              = "CLASSIFICATION", TASK in ["CLASSIFICATION" "REGRESSION"]),
  (nthreads::Int             = 1, nthreads >= 1),
  (num_attribute::Int        = 1, num_attribute > 0),
  k0::Bool                   = true,
  k1::Bool                   = true,
  (num_factor::Int           = 2, num_factor >= 0),
  (l1_regw::Float64          = 0.0, l1_regw >= 0),
  (l1_regv::Float64          = 0.0, l1_regv >= 0),
  (l2_regw::Float64          = 0.0, l2_regw >= 0),
  (l2_regv::Float64          = 0.0, l2_regv >= 0),
  (l2_reg0::Float64          = 0.0, l2_reg0 >= 0),
  init_mean::Float64         = 0.0,
  (init_stdev::Float64       = 0.0, init_stdev >= 0),

  # temporary
  m_sum::Vector{Float64}     = Vector{Float64}(),
  m_sum_sqr::Vector{Float64} = Vector{Float64}()
)

function InitModel(model::FMmodel)
  model.w = zeros(Float64, model.num_attribute)
  model.v = zeros(Float64, (model.num_factor, model.num_attribute))
  model.m_sum = zeros(Float64, model.num_factor)
  model.m_sum_sqr = zeros(Float64, model.num_factor)
end

function SaveModel(model::FMmodel, file::String)
  checkPath(file)
  open(file, "w+") do io
    write(io, "FMmodel\n")
    for field in fieldnames(model)
      writeKV(io, string(field), getfield(model, field))
    end
  end
end

function LoadModel(file::String)
  checkPath(file)
  isfile(file) || error(file * " does not exist")
  open(file, "r") do io
    ftype = readline(io)
    @assert(ftype == "FMmodel\n", "the file read in is not a FMmodel output")
    for line in eachline(io)
      key, value = readKV(line)
      set(model, convert(String, key), value)
    end
  end
  return model
end

function predict(model::FMmodel, vec::SparseVector)
  pred::Float64 = 0.0
  model.k0 && pred += model.w0
  model.m_sum = zeros(model.m_sum)
  model.m_sum_sqr = zeros(model.m_sum_sqr)

  for (_idx, _val) in zip(vec.nzind, vec.nzval):
    model.k1 &&　pred += model.w[_idx] * _val
    for i in 1:model.num_factor
      _tmp = model.v[i, _idx] * _val
      model.m_sum[i] += _tmp
      model.m_sum_sqr[i] += _tmp^2
    end
  end

  for i in 1:model.num_factor
    pred += 0.5 * (model.m_sum[i] ^ 2 - model.m_sum_sqr[i])
  end

  return pred
end


@testset "Testing Kalman Filter..." begin

using LinearAlgebra, Statistics, Plots

Random.seed!(0)

F = Matrix(I,4,4) * 0.9
V = random_cov(4)
G = [1.0 0.0 0.5 0.1;
	 0.0 0.5 1.0 1.0]
W = random_cov(2)

m = LinearGaussianSSM(F, V, G, W)
x0 = MvNormal(randn(4), Matrix(I,4,4) .* 100.0)

x1 = predict(m, x0)
println(mean(x1))

y1 = m.G(1) * mean(x1) + randn(2) / 10
u1 = update(m, x1, y1)
println(mean(u1))

xx, yy = simulate(m, 100, x0)
yy[1, 50] = NaN  # throw in a missing value

plot(xx')
savefig("xx.png")
plot(yy')
savefig("yy.png")

fs = filter(m, yy, x0)
 print(fs)

# plot(xx', "k")
# plot(mean(fs)', "r")

y_new = fs.observations[:, end] + randn(2) / 10
update!(m, fs, y_new)

ss = smooth(m, fs)

@test loglikelihood(fs) < loglikelihood(ss)

end


# key words in forward flight
V #forward velocity
  #thus the rotor blade loading and motion are periodic with a fundamental
  #frequency equal to the rotor speed Ω
A = π*R^2 #the area of a circle with 1/2 diameter
          #also: Cdi = CL^2/(π*A*R)
vi #the induced velocity of the helicopter rotor in high-speed forward flight
vi = T/(2*ρ*A*V)
AR = R/c = (N/π)*σ #aspect ratio should be high
U = sqrt(uT^2+uP^2) #
Φ = atan(uP/UT)
α = θ - Φ #the aerodynamic angle-of-attack of the blade element
L = 1/2*ρ*U^2*c*cl
D = 1/2*ρ*U^2*c*cd
Fz = L*cos(Φ)-D*sin(Φ)
Fx = L*sin(Φ)+D*cos(Φ)
dT = N*Fz*dr
dQ = N*Fx*r*dr
dP = Ω*dQ = N*Fx*Ω*r*dr


py"""
def DiffBasicWeightConvexSet(D, wealth, lower_bound, upper_bound, max_risk, μ, sqrtΣ):
    w = cp.Variable(D)
    c = [
        cp.sum(w) <= wealth,
        w >= lower_bound,
        w <= upper_bound,
        cp.SOC(max_risk, sqrtΣ @ w),
    ]
    return w, c
"""

"""
    DiffBasicWeightConvexSet

Basic convex set for weights.

Atributes:
- `D::Int`: dimension (i.e. number of assets)
- `wealths::Float64`: total wealth to invest
- `lower_bound::Float64`: lower bound of weights
- `upper_bound::Float64`: upper bound of weights
- `max_risk::Float64`: maximum risk
- `μ::Array{Float64,1}`: predicted mean vector
- `sqrtΣ::Array{Float64,2}`: predicted upper triangular Cholesky factor of covariance matrix
"""
DiffBasicWeightConvexSet = py"DiffBasicWeightConvexSet"


function boys_function(n::Int, x::Float64)::Float64
    x < 1e-6 && return 1/(2n+1)
    n == 0 && return 0.5*√(π/x)*erf(√x)
    return ((2n-1)*boys_function(n-1, x) - exp(-x))/(2x)
end

"S(gs, gs)"
function overlap_integral(g1::Gaussian_s, g2::Gaussian_s)::Float64
    p = (g1.α + g2.α)
    μ = (g1.α * g2.α) / p
    γ = 4μ / p
    return γ^(3/4) * exp(-μ * abs2(g1.R - g2.R))
end

"S(gs, gp)"
function overlap_integral(g1::Gaussian_s, g2::Gaussian_p)::Float64
    p = g1.α + g2.α
    μ = (g1.α * g2.α) / p
    γ = 4μ / p
    Rab = g1.R - g2.R
    return √(g1.α) * γ^(5/4) * exp(-μ * abs2(Rab)) * (Rab * g2.n)
end

"S(gp, gs)"
function overlap_integral(g1::Gaussian_p, g2::Gaussian_s)::Float64
    overlap_integral(g2, g1)
end

"S(gp, gp)"
function overlap_integral(g1::Gaussian_p, g2::Gaussian_p)::Float64
    p = g1.α + g2.α
    μ = (g1.α * g2.α) / p
    γ = 4μ / p
    Rab = g1.R - g2.R
    return γ^(5/4) * ((g1.n * g2.n) -√(g1.α*g2.α*γ) * (Rab * g1.n) * (Rab * g2.n)) * exp(-μ * abs2(Rab))
end

"T(gs, gs)"
function kinetic_integral(g1::Gaussian_s, g2::Gaussian_s)::Float64
    p = g1.α + g2.α
    μ = (g1.α * g2.α) / p
    γ = 4μ / p
    Rab2 = abs2(g1.R - g2.R)
    return γ^(3/4) * μ * (3-2μ*Rab2) * exp(-μ * Rab2)
end

"T(gp, gs)"
function kinetic_integral(g1::Gaussian_p, g2::Gaussian_s)::Float64
    p = g1.α + g2.α
    μ = g1.α * g2.α / p
    γ = 4μ / p
    Rab = g1.R - g2.R
    Rab2 = abs2(Rab)
    return γ^(5/4) * exp(-μ * Rab2) * μ * (2μ*Rab2-5) * √g2.α * (Rab * g1.n)
end

"T(gs, gp)"
function kinetic_integral(g1::Gaussian_s, g2::Gaussian_p)::Float64
    kinetic_integral(g2, g1)
end

"T(gp, gp)"
function kinetic_integral(g1::Gaussian_p, g2::Gaussian_p)::Float64
    p = g1.α + g2.α
    μ = g1.α * g2.α / p
    γ = 4μ / p
    Rab = g1.R - g2.R
    dotnab = g1.n * g2.n
    μRab2 = μ * abs2(Rab)
    M = (g1.n * Rab) * (g2.n * Rab)
    return γ^(5/4) * exp(-μRab2) * μ * (5dotnab - 14μ * M - 2μRab2 * dotnab + 4μ * μRab2 * M)
end

"Vn(gs, gs)"
function nuclear_acctaction_integral(Zc::Int, Rc::Vec3, g1::Gaussian_s, g2::Gaussian_s)::Float64
    p = g1.α + g2.α
    μ = (g1.α * g2.α) / p
    γ = 4μ / p
    Rμ = (g1.α / p) * g1.R + (g2.α / p) * g2.R
    return -2Zc * γ^(3/4) * √(p/π) * exp(-μ * abs2(g1.R - g2.R)) * boys_function(0, p*abs2(Rμ-Rc))
end

"Vn(gp,gs)"
function nuclear_acctaction_integral(Zc::Int, Rc::Vec3, g1::Gaussian_p, g2::Gaussian_s)::Float64
    p = g1.α + g2.α
    μ = (g1.α * g2.α) / p
    γ = 4μ / p
    Rμ = (g1.α / p) * g1.R + (g2.α / p) * g2.R
    pR2 = p*abs2(Rμ - Rc)
    return -4Zc * γ^(3/4) * √(g1.α * p / π) * exp(-μ * abs2(g1.R-g2.R)) *
        (-(g1.n * (Rμ-Rc)) * boys_function(1, pR2) + (g1.n * (Rμ-g1.R)) * boys_function(0, pR2))
end

"Vn(gs,gp)"
function nuclear_acctaction_integral(Zc::Int, Rc::Vec3, g1::Gaussian_s, g2::Gaussian_p)::Float64
    nuclear_acctaction_integral(Zc, Rc, g2, g1)
end

"Vn(gp, gp)"
function nuclear_acctaction_integral(Zc::Int, Rc::Vec3, g1::Gaussian_p, g2::Gaussian_p)::Float64
    p = g1.α + g2.α
    μ = g1.α * g2.α / p
    γ = 4μ / p
    Rab = g1.R - g2.R
    Rμ = (g1.α / p) * g1.R + (g2.α / p) * g2.R
    Rμc = Rμ - Rc
    Dac = g1.n * Rμc
    Dbc = g2.n * Rμc
    Daa = g1.n * (Rμ - g1.R)
    Dbb = g2.n * (Rμ - g2.R)
    x = p * abs2(Rμc)
    F0x = boys_function(0, x)
    F1x = boys_function(1, x)
    F2x = boys_function(2, x)
    return -2Zc * γ^(5/4) * √(p/π) * exp(-μ*abs2(g1.R - g2.R)) * (
        (g1.n * g2.n) * (F0x - F1x) + 2p * Dac * Dbc * F2x -
        2p * (Dac * Dbb + Dbc * Daa) * F1x + 2p * Daa * Dbb * F0x
    )
end

function electron_repulsion_integral(g1::Gaussian_s, g2::Gaussian_s, g3::Gaussian_s, g4::Gaussian_s)
    p1 = g1.α + g4.α
    p2 = g2.α + g3.α
    μ1 = (g1.α * g4.α) / p1
    μ2 = (g2.α * g3.α) / p2
    γ1 = 4μ1 / p1
    γ2 = 4μ2 / p2
    Rμ1 = (g1.α / p1) * g1.R + (g4.α / p1) * g4.R
    Rμ2 = (g2.α / p2) * g2.R + (g3.α / p2) * g3.R
    λ = p1 * p2 / (p1 + p2)
    R = abs(Rμ1 - Rμ2)
    if R < 1e-6
        return (γ1*γ2)^(3/4) * exp(-μ1 * abs2(g1.R-g4.R)-μ2 * abs2(g2.R-g3.R)) * 2 * √(λ / π)
    else
        return (γ1*γ2)^(3/4) * exp(-μ1 * abs2(g1.R-g4.R)-μ2 * abs2(g2.R-g3.R)) * erf(√λ * R) / R
    end
end

function two_center_integral(s1::STONG{N}, s2::STONG{N}, integral::Function) where N
    total = 0.0
    for i = 1:N
        for j = 1:N
            total += s1.d[i] * s2.d[j] * integral(s1.g[i], s2.g[j])
        end
    end
    return total
end

function four_center_integral(s1::STONG{N}, s2::STONG{N}, s3::STONG{N}, s4::STONG{N}, integral::Function) where N
    total = 0.0
    for i = 1:N
        for j = 1:N
            for k = 1:N
                for l = 1:N
                    total += s1.d[i]*s2.d[j]*s3.d[k]*s4.d[l] * integral(s1.g[i], s2.g[j], s3.g[k], s4.g[l])
                end
            end
        end
    end
    return total
end

overlap_integral(s1::STONG, s2::STONG) = two_center_integral(s1, s2, overlap_integral)
kinetic_integral(s1::STONG, s2::STONG) = two_center_integral(s1, s2, kinetic_integral)
nuclear_acctaction_integral(Zc::Int, Rc::Vec3, s1::STONG, s2::STONG) = two_center_integral(
    s1, s2,
    (g1,g2)->nuclear_acctaction_integral(Zc, Rc, g1, g2)
)
electron_repulsion_integral(s1::STONG, s2::STONG, s3::STONG, s4::STONG) = four_center_integral(
    s1, s2, s3, s4,
    electron_repulsion_integral
)


""" Linear Kalman Filter Structure """

abstract type AbstractFilter end

mutable struct KalmanFilter{a<:AbstractMatrix, b<:AbstractMatrix, q<:Symmetric, h<:AbstractMatrix, r<:Symmetric} <: AbstractFilter
    A::a # process matrix
    B::b # control matrix
    Q::q # process zero mean noise covariance
    H::h # measurement matrix
    R::r # measurement zero mean noise covariance
end

function KalmanFilter(A::AbstractMatrix, B::AbstractMatrix, Q::AbstractMatrix,
    H::AbstractMatrix, R::AbstractMatrix)
    return KalmanFilter(A, B, Symmetric(Q), H, Symmetric(R))
end

function KalmanFilter(A::AbstractMatrix, Q::AbstractMatrix,
    H::AbstractMatrix, R::AbstractMatrix)
    return KalmanFilter(A, zeros(size(A)), Symmetric(Q), H, Symmetric(R))
end

""" Extended Kalman Filter Structure """

mutable struct ExtendedKalmanFilter{a<:Function, q<:Symmetric, b<:Function, r<:Symmetric} <: AbstractFilter
    f::a # non linear process function f(x, u)
    Q::q # process zero mean noise covariance
    h::b # non linear observation function h(x)
    R::r# measurement zero mean noise covariance
end

function ExtendedKalmanFilter(f::Function, Q::AbstractMatrix, h::Function, R::AbstractMatrix)
    return ExtendedKalmanFilter(f, Symmetric(Q), h, Symmetric(R))
end

""" Parametrized Kalman and Extended Kalman Filters """

abstract type AbstractParamFilter end

mutable struct ParamKalmanFilter{a, b, q, h, r} <: AbstractParamFilter
    # parametrized kalman filter where each matrix is a function of θ
    A::a # process matrix A(θ)
    B::b # control matrix B(θ)
    Q::q # process zero mean noise covariance Q(θ)
    H::h # measurement matrix H(θ)
    R::r # measurement zero mean noise covariance R(θ)
end

mutable struct ExtendedParamKalmanFilter{a, q, b, r} <: AbstractParamFilter
    f::a # non linear process function f(x, u)
    Q::q # process zero mean noise covariance
    h::b # non linear observation function h(x)
    R::r# measurement zero mean noise covariance
end

""" State: mean and cov """
struct State{a<:Number, b<:Number}
    x::AbstractVector{a}
    P::Symmetric{b}
end

function State(x::AbstractVector, P::AbstractMatrix)
    return State(x, Symmetric(P))
end


##################################################
# Compute the NR local diffusion coefficients in (E,L)
##################################################

include("../sources/julia/Main.jl")

r = 1.0 # Radius
E = -0.5 # Energy
L = 0.1 # Angular momentum
m = 0.1 # Mass of field star

println("--------------------")

println("Plummer potential")

println("q = ", qCalc)
println("m = ", m)
println("r = ", r)
println("E = ", E)
println("L = ", L)

println("--------------------")

println("Lc(E) = ", Lc(_tE(E)))
println("Coulomb log = ", logCoulomb)

println("--------------------")

@time dE, dE2, dL, dEdL, dL2 = localOrbitChange(r,E,L,m,100,m)

println("D_E  =  ", dE)
println("D_EE =  ", dE2)
println("D_L  =  ", dL)
println("D_LL =  ", dEdL)
println("D_EL =  ", dL2)

println("--------------------")


@doc raw"""
    AbstractGroupOperation

Abstract type for smooth binary operations $∘$ on elements of a Lie group $\mathcal{G}$:
```math
∘ : \mathcal{G} × \mathcal{G} → \mathcal{G}
```
An operation can be either defined for a specific [`AbstractGroupManifold`](@ref) over
number system `𝔽` or in general, by defining for an operation `Op` the following methods:

    identity!(::AbstractGroupManifold{𝔽,Op}, q, q)
    identity(::AbstractGroupManifold{𝔽,Op}, p)
    inv!(::AbstractGroupManifold{𝔽,Op}, q, p)
    inv(::AbstractGroupManifold{𝔽,Op}, p)
    compose(::AbstractGroupManifold{𝔽,Op}, p, q)
    compose!(::AbstractGroupManifold{𝔽,Op}, x, p, q)

Note that a manifold is connected with an operation by wrapping it with a decorator,
[`AbstractGroupManifold`](@ref). In typical cases the concrete wrapper
[`GroupManifold`](@ref) can be used.
"""
abstract type AbstractGroupOperation end

@doc raw"""
    AbstractGroupManifold{𝔽,O<:AbstractGroupOperation} <: AbstractDecoratorManifold{𝔽}

Abstract type for a Lie group, a group that is also a smooth manifold with an
[`AbstractGroupOperation`](@ref), a smooth binary operation. `AbstractGroupManifold`s must
implement at least [`inv`](@ref), [`identity`](@ref), [`compose`](@ref), and
[`translate_diff`](@ref).
"""
abstract type AbstractGroupManifold{𝔽,O<:AbstractGroupOperation,T<:AbstractEmbeddingType} <:
              AbstractEmbeddedManifold{𝔽,T} end

"""
    GroupManifold{𝔽,M<:Manifold{𝔽},O<:AbstractGroupOperation} <: AbstractGroupManifold{𝔽,O}

Decorator for a smooth manifold that equips the manifold with a group operation, thus making
it a Lie group. See [`AbstractGroupManifold`](@ref) for more details.

Group manifolds by default forward metric-related operations to the wrapped manifold.

# Constructor

    GroupManifold(manifold, op)
"""
struct GroupManifold{𝔽,M<:Manifold{𝔽},O<:AbstractGroupOperation} <:
       AbstractGroupManifold{𝔽,O,TransparentIsometricEmbedding}
    manifold::M
    op::O
end

Base.show(io::IO, G::GroupManifold) = print(io, "GroupManifold($(G.manifold), $(G.op))")

const GROUP_MANIFOLD_BASIS_DISAMBIGUATION =
    [AbstractDecoratorManifold, ValidationManifold, VectorBundle]

"""
    base_group(M::Manifold) -> AbstractGroupManifold

Un-decorate `M` until an `AbstractGroupManifold` is encountered.
Return an error if the [`base_manifold`](@ref) is reached without encountering a group.
"""
base_group(M::AbstractDecoratorManifold) = base_group(decorated_manifold(M))
function base_group(M::Manifold)
    return error("base_group: no base group found.")
end
base_group(G::AbstractGroupManifold) = G

base_manifold(G::GroupManifold) = G.manifold

decorator_group_dispatch(::Manifold) = Val(false)
function decorator_group_dispatch(M::AbstractDecoratorManifold)
    return decorator_group_dispatch(decorated_manifold(M))
end
decorator_group_dispatch(::AbstractGroupManifold) = Val(true)

function is_group_decorator(M::Manifold)
    return _extract_val(decorator_group_dispatch(M))
end

default_decorator_dispatch(::AbstractGroupManifold) = Val(false)

# piping syntax for decoration
if VERSION ≥ v"1.3"
    (op::AbstractGroupOperation)(M::Manifold) = GroupManifold(M, op)
    (::Type{T})(M::Manifold) where {T<:AbstractGroupOperation} = GroupManifold(M, T())
end

function decorator_transparent_dispatch(
    ::typeof(get_coordinates!),
    ::AbstractGroupManifold{<:Any,<:AbstractGroupOperation,TransparentIsometricEmbedding},
    args...,
)
    return Val(:transparent)
end
function decorator_transparent_dispatch(
    ::typeof(get_vector!),
    ::AbstractGroupManifold{<:Any,<:AbstractGroupOperation,TransparentIsometricEmbedding},
    args...,
)
    return Val(:transparent)
end

###################
# Action directions
###################

"""
    ActionDirection

Direction of action on a manifold, either [`LeftAction`](@ref) or [`RightAction`](@ref).
"""
abstract type ActionDirection end

"""
    LeftAction()

Left action of a group on a manifold.
"""
struct LeftAction <: ActionDirection end

"""
    RightAction()

Right action of a group on a manifold.
"""
struct RightAction <: ActionDirection end

"""
    switch_direction(::ActionDirection)

Returns a [`RightAction`](@ref) when given a [`LeftAction`](@ref) and vice versa.
"""
switch_direction(::ActionDirection)
switch_direction(::LeftAction) = RightAction()
switch_direction(::RightAction) = LeftAction()

##################################
# General Identity element methods
##################################

@doc raw"""
    Identity(G::AbstractGroupManifold, p)

The group identity element $e ∈ \mathcal{G}$ represented by point `p`.
"""
struct Identity{G<:AbstractGroupManifold,PT}
    group::G
    p::PT
end

Identity(M::AbstractDecoratorManifold, p) = Identity(decorated_manifold(M), p)
function Identity(M::Manifold, p)
    return error("Identity not implemented for manifold $(M) and point $(p).")
end

function Base.:(==)(e1::Identity, e2::Identity)
    return e1.p == e2.p && e1.group == e2.group
end

make_identity(M::Manifold, p) = Identity(M, identity(M, p))

Base.show(io::IO, e::Identity) = print(io, "Identity($(e.group), $(e.p))")

# To ensure allocate_result_type works
number_eltype(e::Identity) = Bool

Base.copyto!(e::TE, ::TE) where {TE<:Identity} = e
Base.copyto!(p, ::TE) where {TE<:Identity} = copyto!(p, e.p)
Base.copyto!(p::AbstractArray, e::TE) where {TE<:Identity} = copyto!(p, e.p)

Base.isapprox(p, e::Identity; kwargs...) = isapprox(e::Identity, p; kwargs...)
Base.isapprox(e::Identity, p; kwargs...) = isapprox(e.group, e, p; kwargs...)
Base.isapprox(e::E, ::E; kwargs...) where {E<:Identity} = true

function allocate_result(
    M::Manifold,
    f::typeof(get_coordinates),
    e::Identity,
    X,
    B::AbstractBasis,
)
    T = allocate_result_type(M, f, (e.p, X))
    return allocate(e.p, T, Size(number_of_coordinates(M, B)))
end

function decorator_transparent_dispatch(
    ::typeof(get_vector),
    ::AbstractGroupManifold,
    args...,
)
    return Val(:parent)
end
function allocate_result(M::Manifold, f::typeof(get_vector), e::Identity, Xⁱ)
    is_group_decorator(M) && return allocate_result(base_group(M), f, e, Xⁱ)
    return error("allocate_result not implemented for manifold $(M), function $(f), point $(e), and vector $(Xⁱ).")
end
function allocate_result(M::AbstractGroupManifold, f::typeof(get_vector), e::Identity, Xⁱ)
    return error("allocate_result not implemented for group manifold $(M), function $(f), $(e), and vector $(Xⁱ).")
end
function allocate_result(
    G::GT,
    ::typeof(get_vector),
    ::Identity{GT},
    Xⁱ,
) where {GT<:AbstractGroupManifold}
    B = VectorBundleFibers(TangentSpace, G)
    return allocate(Xⁱ, Size(representation_size(B)))
end

function decorator_transparent_dispatch(
    ::typeof(get_coordinates),
    ::AbstractGroupManifold,
    args...,
)
    return Val(:parent)
end
function allocate_result(
    M::AbstractDecoratorManifold,
    f::typeof(get_coordinates),
    e::Identity,
    X,
)
    is_group_decorator(M) && return allocate_result(base_group(M), f, e, X)
    return error("allocate_result not implemented for manifold $(M), function $(f), point $(e), and vector $(X).")
end
function allocate_result(
    M::AbstractGroupManifold,
    f::typeof(get_coordinates),
    e::Identity,
    X,
)
    return error("allocate_result not implemented for group manifold $(M), function $(f), $(e), and vector $(X).")
end
function allocate_result(
    G::GT,
    ::typeof(get_coordinates),
    ::Identity{GT},
    X,
) where {GT<:AbstractGroupManifold}
    return allocate(X, Size(manifold_dimension(G)))
end

function get_vector(M::AbstractGroupManifold, e::Identity, X, B::VeeOrthogonalBasis)
    M != e.group && error("On $(M) the identity $(e) does not match to perform get_vector.")
    return get_vector(decorated_manifold(M), e.p, X, B)
end
function get_vector(M::Manifold, e::Identity, X, B::VeeOrthogonalBasis)
    M != e.group.manifold &&
        error("On $(M) the identity $(e) does not match to perform get_vector.")
    return get_vector(M, e.p, X, B)
end
for MT in GROUP_MANIFOLD_BASIS_DISAMBIGUATION
    eval(
        quote
            @invoke_maker 1 Manifold get_vector(
                M::$MT,
                e::Identity,
                X,
                B::VeeOrthogonalBasis,
            )
        end,
    )
end
function get_vector!(M::AbstractGroupManifold, Y, e::Identity, X, B::VeeOrthogonalBasis)
    M != e.group && error("On $(M) the identity $(e) does not match to perform get_vector!")
    return get_vector!(decorated_manifold(M), Y, e.p, X, B)
end
function get_vector!(M::Manifold, Y, e::Identity, X, B::VeeOrthogonalBasis)
    M != e.group.manifold &&
        error("On $(M) the identity $(e) does not match to perform get_vector!")
    return get_vector!(M, Y, e.p, X, B)
end
for MT in GROUP_MANIFOLD_BASIS_DISAMBIGUATION
    eval(
        quote
            @invoke_maker 1 Manifold get_vector!(
                M::$MT,
                Y,
                e::Identity,
                X,
                B::VeeOrthogonalBasis,
            )
        end,
    )
end

function get_coordinates(M::AbstractGroupManifold, e::Identity, X, B::VeeOrthogonalBasis)
    M != e.group &&
        error("On $(M) the identity $(e) does not match to perform get_coordinates")
    return get_coordinates(decorated_manifold(M), e.p, X, B)
end
function get_coordinates(M::Manifold, e::Identity, X, B::VeeOrthogonalBasis)
    M != e.group.manifold &&
        error("On $(M) the identity $(e) does not match to perform get_coordinates")
    return get_coordinates(M, e.p, X, B)
end
for MT in GROUP_MANIFOLD_BASIS_DISAMBIGUATION
    eval(
        quote
            @invoke_maker 1 Manifold get_coordinates(
                M::$MT,
                e::Identity,
                X,
                B::VeeOrthogonalBasis,
            )
        end,
    )
end

function get_coordinates!(
    M::AbstractGroupManifold,
    Y,
    e::Identity,
    X,
    B::VeeOrthogonalBasis,
)
    M != e.group &&
        error("On $(M) the identity $(e) does not match to perform get_coordinates!")
    return get_coordinates!(decorated_manifold(M), Y, e.p, X, B)
end
function get_coordinates!(M::Manifold, Y, e::Identity, X, B::VeeOrthogonalBasis)
    M != e.group.manifold &&
        error("On $(M) the identity $(e) does not match to perform get_coordinates!")
    return get_coordinates!(M, Y, e.p, X, B)
end
for MT in GROUP_MANIFOLD_BASIS_DISAMBIGUATION
    eval(
        quote
            @invoke_maker 1 Manifold get_coordinates!(
                M::$MT,
                Y,
                e::Identity,
                X,
                B::VeeOrthogonalBasis,
            )
        end,
    )
end

manifold_dimension(G::GroupManifold) = manifold_dimension(G.manifold)

function check_manifold_point(G::AbstractGroupManifold, e::Identity; kwargs...)
    return DomainError(e, "The identity element $(e) does not belong to $(G).")
end
function check_manifold_point(
    G::GT,
    e::Identity{GT};
    kwargs...,
) where {GT<:AbstractGroupManifold}
    return nothing
end
##########################
# Group-specific functions
##########################

@doc raw"""
    inv(G::AbstractGroupManifold, p)

Inverse $p^{-1} ∈ \mathcal{G}$ of an element $p ∈ \mathcal{G}$, such that
$p \circ p^{-1} = p^{-1} \circ p = e ∈ \mathcal{G}$, where $e$ is the [`identity`](@ref)
element of $\mathcal{G}$.
"""
inv(::AbstractGroupManifold, ::Any...)
@decorator_transparent_function :intransparent function Base.inv(
    G::AbstractGroupManifold,
    p,
)
    q = allocate_result(G, inv, p)
    return inv!(G, q, p)
end

@decorator_transparent_function :intransparent function inv!(G::AbstractGroupManifold, q, p)
    return inv!(G.manifold, q, p)
end

@doc raw"""
    identity(G::AbstractGroupManifold, p)

Identity element $e ∈ \mathcal{G}$, such that for any element $p ∈ \mathcal{G}$,
$p \circ e = e \circ p = p$.
The returned element is of a similar type to `p`.
"""
identity(::AbstractGroupManifold, ::Any)
@decorator_transparent_signature Base.identity(G::AbstractDecoratorManifold, p)
function Base.identity(G::AbstractGroupManifold, p)
    y = allocate_result(G, identity, p)
    return identity!(G, y, p)
end

@decorator_transparent_signature identity!(G::AbstractDecoratorManifold, q, p)
function decorator_transparent_dispatch(::typeof(identity!), ::AbstractGroupManifold, q, p)
    return Val(:intransparent)
end

function Base.isapprox(
    G::GT,
    e::Identity{GT},
    p;
    kwargs...,
) where {GT<:AbstractGroupManifold}
    return isapprox(G, e.p, p; kwargs...)
end
function Base.isapprox(
    G::GT,
    p,
    e::Identity{GT};
    kwargs...,
) where {GT<:AbstractGroupManifold}
    return isapprox(G, e, p; kwargs...)
end
function Base.isapprox(
    ::GT,
    ::E,
    ::E;
    kwargs...,
) where {GT<:AbstractGroupManifold,E<:Identity{GT}}
    return true
end

function decorator_transparent_dispatch(
    ::typeof(isapprox),
    M::AbstractDecoratorManifold,
    e::E,
    p,
) where {E<:Identity{<:AbstractGroupManifold}}
    return Val(:transparent)
end
function decorator_transparent_dispatch(
    ::typeof(isapprox),
    M::AbstractDecoratorManifold,
    p,
    e::E,
) where {E<:Identity{<:AbstractGroupManifold}}
    return Val(:transparent)
end
function decorator_transparent_dispatch(
    ::typeof(isapprox),
    M::AbstractDecoratorManifold,
    e1::E1,
    e2::E2,
) where {E1<:Identity{<:AbstractGroupManifold},E2<:Identity{<:AbstractGroupManifold}}
    return Val(:transparent)
end

@doc raw"""
    compose(G::AbstractGroupManifold, p, q)

Compose elements $p,q ∈ \mathcal{G}$ using the group operation $p \circ q$.
"""
compose(::AbstractGroupManifold, ::Any...)
@decorator_transparent_signature compose(M::AbstractDecoratorManifold, p, q)
function compose(G::AbstractGroupManifold, p, q)
    x = allocate_result(G, compose, p, q)
    return compose!(G, x, p, q)
end
function decorator_transparent_dispatch(
    ::typeof(compose),
    M::AbstractGroupManifold,
    x,
    p,
    q,
)
    return Val(:intransparent)
end

@decorator_transparent_signature compose!(M::AbstractDecoratorManifold, x, p, q)
function decorator_transparent_dispatch(
    ::typeof(compose!),
    M::AbstractGroupManifold,
    x,
    p,
    q,
)
    return Val(:intransparent)
end

_action_order(p, q, conv::LeftAction) = (p, q)
_action_order(p, q, conv::RightAction) = (q, p)

@doc raw"""
    translate(G::AbstractGroupManifold, p, q)
    translate(G::AbstractGroupManifold, p, q, conv::ActionDirection=LeftAction()])

Translate group element $q$ by $p$ with the translation $τ_p$ with the specified
`conv`ention, either left ($L_p$) or right ($R_p$), defined as
```math
\begin{aligned}
L_p &: q ↦ p \circ q\\
R_p &: q ↦ q \circ p.
\end{aligned}
```
"""
translate(::AbstractGroupManifold, ::Any...)
@decorator_transparent_signature translate(M::AbstractDecoratorManifold, p, q)
function translate(G::AbstractGroupManifold, p, q)
    return translate(G, p, q, LeftAction())
end
@decorator_transparent_signature translate(
    M::AbstractDecoratorManifold,
    p,
    q,
    conv::ActionDirection,
)
function translate(G::AbstractGroupManifold, p, q, conv::ActionDirection)
    return compose(G, _action_order(p, q, conv)...)
end

function decorator_transparent_dispatch(
    ::typeof(translate),
    ::AbstractGroupManifold,
    args...,
)
    return Val(:intransparent)
end

@decorator_transparent_signature translate!(M::AbstractDecoratorManifold, X, p, q)
function translate!(G::AbstractGroupManifold, X, p, q)
    return translate!(G, X, p, q, LeftAction())
end
@decorator_transparent_signature translate!(
    M::AbstractDecoratorManifold,
    X,
    p,
    q,
    conv::ActionDirection,
)
function translate!(G::AbstractGroupManifold, X, p, q, conv::ActionDirection)
    return compose!(G, X, _action_order(p, q, conv)...)
end

function decorator_transparent_dispatch(
    ::typeof(translate!),
    ::AbstractGroupManifold,
    args...,
)
    return Val(:intransparent)
end

@doc raw"""
    inverse_translate(G::AbstractGroupManifold, p, q)
    inverse_translate(G::AbstractGroupManifold, p, q, conv::ActionDirection=LeftAction())

Inverse translate group element $q$ by $p$ with the inverse translation $τ_p^{-1}$ with the
specified `conv`ention, either left ($L_p^{-1}$) or right ($R_p^{-1}$), defined as
```math
\begin{aligned}
L_p^{-1} &: q ↦ p^{-1} \circ q\\
R_p^{-1} &: q ↦ q \circ p^{-1}.
\end{aligned}
```
"""
inverse_translate(::AbstractGroupManifold, ::Any...)
@decorator_transparent_function :intransparent function inverse_translate(
    G::AbstractGroupManifold,
    p,
    q,
)
    return inverse_translate(G, p, q, LeftAction())
end
@decorator_transparent_function :intransparent function inverse_translate(
    G::AbstractGroupManifold,
    p,
    q,
    conv::ActionDirection,
)
    return translate(G, inv(G, p), q, conv)
end

@decorator_transparent_function :intransparent function inverse_translate!(
    G::AbstractGroupManifold,
    X,
    p,
    q,
)
    return inverse_translate!(G, X, p, q, LeftAction())
end
@decorator_transparent_function :intransparent function inverse_translate!(
    G::AbstractGroupManifold,
    X,
    p,
    q,
    conv::ActionDirection,
)
    return translate!(G, X, inv(G, p), q, conv)
end

@doc raw"""
    translate_diff(G::AbstractGroupManifold, p, q, X)
    translate_diff(G::AbstractGroupManifold, p, q, X, conv::ActionDirection=LeftAction())

For group elements $p, q ∈ \mathcal{G}$ and tangent vector $X ∈ T_q \mathcal{G}$, compute
the action of the differential of the translation $τ_p$ by $p$ on $X$, with the specified
left or right `conv`ention. The differential transports vectors:
```math
(\mathrm{d}τ_p)_q : T_q \mathcal{G} → T_{τ_p q} \mathcal{G}\\
```
"""
translate_diff(::AbstractGroupManifold, ::Any...)
@decorator_transparent_signature translate_diff(M::AbstractDecoratorManifold, p, q, X)
function translate_diff(G::AbstractGroupManifold, p, q, X)
    return translate_diff(G, p, q, X, LeftAction())
end
@decorator_transparent_signature translate_diff(
    M::AbstractDecoratorManifold,
    p,
    q,
    X,
    conv::ActionDirection,
)
function translate_diff(G::AbstractGroupManifold, p, q, X, conv::ActionDirection)
    Y = allocate_result(G, translate_diff, X, p, q)
    translate_diff!(G, Y, p, q, X, conv)
    return Y
end

@decorator_transparent_signature translate_diff!(M::AbstractDecoratorManifold, Y, p, q, X)
function translate_diff!(G::AbstractGroupManifold, Y, p, q, X)
    return translate_diff!(G, Y, p, q, X, LeftAction())
end
@decorator_transparent_signature translate_diff!(
    M::AbstractDecoratorManifold,
    Y,
    p,
    q,
    X,
    conv::ActionDirection,
)
function decorator_transparent_dispatch(
    ::typeof(translate_diff!),
    G::AbstractGroupManifold,
    Y,
    p,
    q,
    X,
    conv,
)
    return Val(:intransparent)
end

@doc raw"""
    inverse_translate_diff(G::AbstractGroupManifold, p, q, X)
    inverse_translate_diff(G::AbstractGroupManifold, p, q, X, conv::ActionDirection=LeftAction())

For group elements $p, q ∈ \mathcal{G}$ and tangent vector $X ∈ T_q \mathcal{G}$, compute
the action on $X$ of the differential of the inverse translation $τ_p$ by $p$, with the
specified left or right `conv`ention. The differential transports vectors:
```math
(\mathrm{d}τ_p^{-1})_q : T_q \mathcal{G} → T_{τ_p^{-1} q} \mathcal{G}\\
```
"""
inverse_translate_diff(::AbstractGroupManifold, ::Any...)
@decorator_transparent_function :intransparent function inverse_translate_diff(
    G::AbstractGroupManifold,
    p,
    q,
    X,
)
    return inverse_translate_diff(G, p, q, X, LeftAction())
end
@decorator_transparent_function :intransparent function inverse_translate_diff(
    G::AbstractGroupManifold,
    p,
    q,
    X,
    conv::ActionDirection,
)
    return translate_diff(G, inv(G, p), q, X, conv)
end

@decorator_transparent_function :intransparent function inverse_translate_diff!(
    G::AbstractGroupManifold,
    Y,
    p,
    q,
    X,
)
    return inverse_translate_diff!(G, Y, p, q, X, LeftAction())
end
@decorator_transparent_function :intransparent function inverse_translate_diff!(
    G::AbstractGroupManifold,
    Y,
    p,
    q,
    X,
    conv::ActionDirection,
)
    return translate_diff!(G, Y, inv(G, p), q, X, conv)
end

@doc raw"""
    group_exp(G::AbstractGroupManifold, X)

Compute the group exponential of the Lie algebra element `X`.

Given an element $X ∈ 𝔤 = T_e \mathcal{G}$, where $e$ is the [`identity`](@ref) element of
the group $\mathcal{G}$, and $𝔤$ is its Lie algebra, the group exponential is the map

````math
\exp : 𝔤 → \mathcal{G},
````
such that for $t,s ∈ ℝ$, $γ(t) = \exp (t X)$ defines a one-parameter subgroup with the
following properties:

````math
\begin{aligned}
γ(t) &= γ(-t)^{-1}\\
γ(t + s) &= γ(t) \circ γ(s) = γ(s) \circ γ(t)\\
γ(0) &= e\\
\lim_{t → 0} \frac{d}{dt} γ(t) &= X.
\end{aligned}
````

!!! note
    In general, the group exponential map is distinct from the Riemannian exponential map
    [`exp`](@ref).

```
group_exp(G::AbstractGroupManifold{𝔽,AdditionOperation}, X) where {𝔽}
```

Compute $q = X$.

    group_exp(G::AbstractGroupManifold{𝔽,MultiplicationOperation}, X) where {𝔽}

For `Number` and `AbstractMatrix` types of `X`, compute the usual numeric/matrix
exponential,

````math
\exp X = \operatorname{Exp} X = \sum_{n=0}^∞ \frac{1}{n!} X^n.
````
"""
group_exp(::AbstractGroupManifold, ::Any...)
@decorator_transparent_signature group_exp(G::AbstractDecoratorManifold, X)
function group_exp(G::AbstractGroupManifold, X)
    q = allocate_result(G, group_exp, X)
    return group_exp!(G, q, X)
end

@decorator_transparent_signature group_exp!(G::AbstractDecoratorManifold, q, X)

function decorator_transparent_dispatch(::typeof(group_exp), ::AbstractGroupManifold, q, X)
    return Val(:intransparent)
end
function decorator_transparent_dispatch(::typeof(group_exp!), ::AbstractGroupManifold, q, X)
    return Val(:intransparent)
end

@doc raw"""
    group_log(G::AbstractGroupManifold, q)

Compute the group logarithm of the group element `q`.

Given an element $q ∈ \mathcal{G}$, compute the right inverse of the group exponential map
[`group_exp`](@ref), that is, the element $\log q = X ∈ 𝔤 = T_e \mathcal{G}$, such that
$q = \exp X$

!!! note
    In general, the group logarithm map is distinct from the Riemannian logarithm map
    [`log`](@ref).

```
group_log(G::AbstractGroupManifold{𝔽,AdditionOperation}, q) where {𝔽}
```

Compute $X = q$.

    group_log(G::AbstractGroupManifold{𝔽,MultiplicationOperation}, q) where {𝔽}

For `Number` and `AbstractMatrix` types of `q`, compute the usual numeric/matrix logarithm:

````math
\log q = \operatorname{Log} q = \sum_{n=1}^∞ \frac{(-1)^{n+1}}{n} (q - e)^n,
````

where $e$ here is the [`identity`](@ref) element, that is, $1$ for numeric $q$ or the
identity matrix $I_m$ for matrix $q ∈ ℝ^{m × m}$.
"""
group_log(::AbstractGroupManifold, ::Any...)
@decorator_transparent_signature group_log(M::AbstractDecoratorManifold, q)
function group_log(G::AbstractGroupManifold, q)
    X = allocate_result(G, group_log, q)
    return group_log!(G, X, q)
end
function decorator_transparent_dispatch(::typeof(group_log), ::AbstractGroupManifold, X, q)
    return Val(:intransparent)
end

@decorator_transparent_signature group_log!(M::AbstractDecoratorManifold, X, q)
function decorator_transparent_dispatch(::typeof(group_log!), ::AbstractGroupManifold, X, q)
    return Val(:intransparent)
end
############################
# Group-specific Retractions
############################

"""
    GroupExponentialRetraction{D<:ActionDirection} <: AbstractRetractionMethod

Retraction using the group exponential [`group_exp`](@ref) "translated" to any point on the
manifold.

For more details, see
[`retract`](@ref retract(::GroupManifold, p, X, ::GroupExponentialRetraction)).

# Constructor

    GroupExponentialRetraction(conv::ActionDirection = LeftAction())
"""
struct GroupExponentialRetraction{D<:ActionDirection} <: AbstractRetractionMethod end

function GroupExponentialRetraction(conv::ActionDirection = LeftAction())
    return GroupExponentialRetraction{typeof(conv)}()
end

"""
    GroupLogarithmicInverseRetraction{D<:ActionDirection} <: AbstractInverseRetractionMethod

Retraction using the group logarithm [`group_log`](@ref) "translated" to any point on the
manifold.

For more details, see
[`inverse_retract`](@ref inverse_retract(::GroupManifold, p, q ::GroupLogarithmicInverseRetraction)).

# Constructor

    GroupLogarithmicInverseRetraction(conv::ActionDirection = LeftAction())
"""
struct GroupLogarithmicInverseRetraction{D<:ActionDirection} <:
       AbstractInverseRetractionMethod end

function GroupLogarithmicInverseRetraction(conv::ActionDirection = LeftAction())
    return GroupLogarithmicInverseRetraction{typeof(conv)}()
end

direction(::GroupExponentialRetraction{D}) where {D} = D()
direction(::GroupLogarithmicInverseRetraction{D}) where {D} = D()

@doc raw"""
    retract(
        G::AbstractGroupManifold,
        p,
        X,
        method::GroupExponentialRetraction{<:ActionDirection},
    )

Compute the retraction using the group exponential [`group_exp`](@ref) "translated" to any
point on the manifold.
With a group translation ([`translate`](@ref)) $τ_p$ in a specified direction, the
retraction is

````math
\operatorname{retr}_p = τ_p \circ \exp \circ (\mathrm{d}τ_p^{-1})_p,
````

where $\exp$ is the group exponential ([`group_exp`](@ref)), and $(\mathrm{d}τ_p^{-1})_p$ is
the action of the differential of inverse translation $τ_p^{-1}$ evaluated at $p$ (see
[`inverse_translate_diff`](@ref)).
"""
function retract(G::AbstractGroupManifold, p, X, method::GroupExponentialRetraction)
    conv = direction(method)
    Xₑ = inverse_translate_diff(G, p, p, X, conv)
    pinvq = group_exp(G, Xₑ)
    q = translate(G, p, pinvq, conv)
    return q
end

function retract!(G::AbstractGroupManifold, q, p, X, method::GroupExponentialRetraction)
    conv = direction(method)
    Xₑ = inverse_translate_diff(G, p, p, X, conv)
    pinvq = group_exp(G, Xₑ)
    return translate!(G, q, p, pinvq, conv)
end

@doc raw"""
    inverse_retract(
        G::AbstractGroupManifold,
        p,
        X,
        method::GroupLogarithmicInverseRetraction{<:ActionDirection},
    )

Compute the inverse retraction using the group logarithm [`group_log`](@ref) "translated"
to any point on the manifold.
With a group translation ([`translate`](@ref)) $τ_p$ in a specified direction, the
retraction is

````math
\operatorname{retr}_p^{-1} = (\mathrm{d}τ_p)_e \circ \log \circ τ_p^{-1},
````

where $\log$ is the group logarithm ([`group_log`](@ref)), and $(\mathrm{d}τ_p)_e$ is the
action of the differential of translation $τ_p$ evaluated at the identity element $e$
(see [`translate_diff`](@ref)).
"""
function inverse_retract(G::GroupManifold, p, q, method::GroupLogarithmicInverseRetraction)
    conv = direction(method)
    pinvq = inverse_translate(G, p, q, conv)
    Xₑ = group_log(G, pinvq)
    return translate_diff(G, p, Identity(G, p), Xₑ, conv)
end

function inverse_retract!(
    G::AbstractGroupManifold,
    X,
    p,
    q,
    method::GroupLogarithmicInverseRetraction,
)
    conv = direction(method)
    pinvq = inverse_translate(G, p, q, conv)
    Xₑ = group_log(G, pinvq)
    return translate_diff!(G, X, p, Identity(G, p), Xₑ, conv)
end

#################################
# Overloads for AdditionOperation
#################################

"""
    AdditionOperation <: AbstractGroupOperation

Group operation that consists of simple addition.
"""
struct AdditionOperation <: AbstractGroupOperation end

const AdditionGroup = AbstractGroupManifold{𝔽,AdditionOperation} where {𝔽}

Base.:+(e::Identity{G}) where {G<:AdditionGroup} = e
Base.:+(p::Identity{G}, ::Identity{G}) where {G<:AdditionGroup} = p
Base.:+(::Identity{G}, p) where {G<:AdditionGroup} = p
Base.:+(p, ::Identity{G}) where {G<:AdditionGroup} = p
Base.:+(e::E, ::E) where {G<:AdditionGroup,E<:Identity{G}} = e

Base.:-(e::Identity{G}) where {G<:AdditionGroup} = e
Base.:-(e::Identity{G}, ::Identity{G}) where {G<:AdditionGroup} = e
Base.:-(::Identity{G}, p) where {G<:AdditionGroup} = -p
Base.:-(p, ::Identity{G}) where {G<:AdditionGroup} = p
Base.:-(e::E, ::E) where {G<:AdditionGroup,E<:Identity{G}} = e

Base.:*(e::Identity{G}, p) where {G<:AdditionGroup} = e
Base.:*(p, e::Identity{G}) where {G<:AdditionGroup} = e
Base.:*(e::E, ::E) where {G<:AdditionGroup,E<:Identity{G}} = e

Base.zero(e::Identity{G}) where {G<:AdditionGroup} = e

Base.identity(::AdditionGroup, p) = zero(p)

identity!(::AdditionGroup, q, p) = fill!(q, 0)

Base.inv(::AdditionGroup, p) = -p

inv!(::AdditionGroup, q, p) = copyto!(q, -p)

compose(::AdditionGroup, p, q) = p + q

function compose!(::GT, x, p, q) where {GT<:AdditionGroup}
    p isa Identity{GT} && return copyto!(x, q)
    q isa Identity{GT} && return copyto!(x, p)
    x .= p .+ q
    return x
end

translate_diff(::AdditionGroup, p, q, X, ::ActionDirection) = X

translate_diff!(::AdditionGroup, Y, p, q, X, ::ActionDirection) = copyto!(Y, X)

inverse_translate_diff(::AdditionGroup, p, q, X, ::ActionDirection) = X

function inverse_translate_diff!(::AdditionGroup, Y, p, q, X, ::ActionDirection)
    return copyto!(Y, X)
end

group_exp(::AdditionGroup, X) = X

group_exp!(::AdditionGroup, q, X) = copyto!(q, X)

group_log(::AdditionGroup, q) = q

group_log!(::AdditionGroup, X, q) = copyto!(X, q)

#######################################
# Overloads for MultiplicationOperation
#######################################

"""
    MultiplicationOperation <: AbstractGroupOperation

Group operation that consists of multiplication.
"""
struct MultiplicationOperation <: AbstractGroupOperation end

const MultiplicationGroup = AbstractGroupManifold{𝔽,MultiplicationOperation} where {𝔽}

Base.:*(e::Identity{G}) where {G<:MultiplicationGroup} = e
Base.:*(::Identity{G}, p) where {G<:MultiplicationGroup} = p
Base.:*(p, ::Identity{G}) where {G<:MultiplicationGroup} = p
Base.:*(e::E, ::E) where {G<:MultiplicationGroup,E<:Identity{G}} = e
Base.:*(::Identity{<:MultiplicationGroup}, e::Identity{<:AdditionGroup}) = e

Base.:/(p, ::Identity{G}) where {G<:MultiplicationGroup} = p
Base.:/(::Identity{G}, p) where {G<:MultiplicationGroup} = inv(p)
Base.:/(e::E, ::E) where {G<:MultiplicationGroup,E<:Identity{G}} = e

Base.:\(p, ::Identity{G}) where {G<:MultiplicationGroup} = inv(p)
Base.:\(::Identity{G}, p) where {G<:MultiplicationGroup} = p
Base.:\(e::E, ::E) where {G<:MultiplicationGroup,E<:Identity{G}} = e

Base.inv(e::Identity{G}) where {G<:MultiplicationGroup} = e

Base.one(e::Identity{G}) where {G<:MultiplicationGroup} = e

Base.transpose(e::Identity{G}) where {G<:MultiplicationGroup} = e

LinearAlgebra.det(::Identity{<:MultiplicationGroup}) = 1

LinearAlgebra.mul!(q, e::Identity{G}, p) where {G<:MultiplicationGroup} = copyto!(q, p)
LinearAlgebra.mul!(q, p, e::Identity{G}) where {G<:MultiplicationGroup} = copyto!(q, p)
function LinearAlgebra.mul!(q, e::E, ::E) where {G<:MultiplicationGroup,E<:Identity{G}}
    return identity!(e.group, q, e)
end

Base.identity(::MultiplicationGroup, p) = one(p)

function identity!(G::GT, q, p) where {GT<:MultiplicationGroup}
    isa(p, Identity{GT}) || return copyto!(q, one(p))
    return error("identity! not implemented on $(typeof(G)) for points $(typeof(q)) and $(typeof(p))")
end
identity!(::MultiplicationGroup, q::AbstractMatrix, p) = copyto!(q, I)

Base.inv(::MultiplicationGroup, p) = inv(p)

inv!(G::MultiplicationGroup, q, p) = copyto!(q, inv(G, p))

compose(::MultiplicationGroup, p, q) = p * q

# TODO: x might alias with p or q, we might be able to optimize it if it doesn't.
compose!(::MultiplicationGroup, x, p, q) = copyto!(x, p * q)

inverse_translate(::MultiplicationGroup, p, q, ::LeftAction) = p \ q
inverse_translate(::MultiplicationGroup, p, q, ::RightAction) = q / p

function inverse_translate!(G::MultiplicationGroup, x, p, q, conv::ActionDirection)
    return copyto!(x, inverse_translate(G, p, q, conv))
end

function group_exp!(G::MultiplicationGroup, q, X)
    X isa Union{Number,AbstractMatrix} && return copyto!(q, exp(X))
    return error("group_exp! not implemented on $(typeof(G)) for vector $(typeof(X)) and element $(typeof(q)).")
end


# This file was generated, do not modify it. # hide
using Soss
using MeasureTheory
using TupleVectors

m = @model n begin
    α ~ Normal()
    β ~ Normal()
    x ~ Normal() |> iid(n)
    σ ~ Exponential(λ=1)
    y ~ For(x) do xj
        Normal(α + β * xj, σ)
    end
    return y
end


using Distributions, HCubature, Printf, LinearAlgebra, Random, ProgressMeter

function test_inside_logarithm()

    rg = MersenneTwister(1)

    N = 5 # number of observations

    y = randn(rg, N)
    f = randn(rg, N)
    A = Diagonal(rand(rg, N))
    b = rand(rg, N)
    S = Diagonal(rand(rg, N)) + 0.2*I
    K = 0.1*randn(rg, N, N) ; K= K*K' + 1.1*I

    # Exact
    a1 = logpdf(MvNormal(A*f + b, S), y)
    a2 = logpdf(MvNormal(zeros(N), K), f)

    # My implementation

    b1 = -0.5*N*log(2π) - 0.5*logdet(S)

    c1 = -0.5*(y .- A*f .- b)'*(S\( y.- A*f .- b))

    b2 = -0.5*N*log(2π)

    c2 = -0.5*logdet(K) - 0.5*f'*(K\f)

    # a1, b1+c1, a2, b2+c
    a1+a2, b1+c1+b2+c2
end


function pWaves(ngrid,H,Np,Br,m,T)

# Evaluates the frequency and spatial structure of
# polar magnetic Rossby waves in a spherical geometry.
# The wave equation is formulated in terms of the
# meridional component of the magnetic field.
# The meridional coordinate is y = cos(colat), although
# a second grid array is stored as latitude for plotting.
# The radial magnetic field is assumed to be constant
# over the surface of the CMB. The same governing equation
# describes equatorial MAC waves in a spherical geometry. The
# two types of waves are distinguished by the sign of the frequency.
# The equatorial waves travel eastward (positive freq.) and the
# high latitude waves travel westward (negative freq.)
#
# input
#
# ngrid - grid used to represent the waves
# H    - layer thickness (km)
# Np   - dimensionless buoyancy frequency (N/Omega)
# Br   - radial magnetic field  (mT)
# m    - angular order m
# T    - initial estimate of period  (years, -'ve' for westward)
#
# output
#
# lat  -  array of latitude for plotting eigenfunction by(y)
# y    -  array of coordinates "y"
# byp   -  eigenfunction by(y)  (modified according to Eq. 13)
# omega - frequency of wave
# pd    - period of wave (years)
# q     - quality factor

# initial estimate of frequency
sectoyear = 365.25 * 24 * 60 * 60;
omega0 = 2 * pi / (T*sectoyear) + 0.0im;
B = Br * 1.0e-3;    # convert to Tesla

# set up grid
y,lat = mgrid(ngrid);

# solve eigenvalue problem for byp
omega,byp = modes_y(y,H,Np,B,m,omega0);

#  compute b_x
bxp = modes_x(y,byp,H,Np,B,m,omega);

# transform back to original variables
ynew,bx,by = transform(y,bxp,byp);

# evaluate period and quality factor of wave
pd = 2 * pi / (sectoyear * real(omega));
q = abs(0.5*real(omega)/imag(omega));

return pd,q,bx,by,ynew

end


#Copyright (c) 2015 Rex Computing and Isaac Yonemoto

#see LICENSE.txt

#this work was supported in part by DARPA Contract D15PC00135

#frag-mult-test.jl
#test fragmented multiplication

function frag_mult(a::Array{UInt64,1}, b::Array{UInt64,1})
  #note that frag_mult fails for absurdly high length integer arrays.
  length(a) != length(b) && throw(ArgumentError("mismatched arrays"))

  #take these two UInt64 arrays and reinterpret them as UInt32 arrays
  a_32 = reinterpret(UInt32, a)
  b_32 = reinterpret(UInt32, b)

  scratchpad = zeros(UInt32, length(a) * 4)
  carries    = zeros(UInt32, length(a) * 4)

  for idx = 1:length(a_32)
    for jdx = 1:length(b_32)
      #multiply the two values.
      multres::UInt64 = a_32[idx] * b_32[jdx]
      #bin these into a high word and a lwo word
      multresbins = reinterpret(UInt32, [multres])
      #add the low word to the scratchpad
      scratchpad[idx + jdx - 1] += multresbins[1]
      #save the carry, if need be.
      (scratchpad[idx + jdx - 1] < multresbins[1]) && (carries[idx + jdx] += 1)
      #add the high word to the scratchpad
      scratchpad[idx + jdx] += multresbins[2]
      #save the carry, if need be.
      (scratchpad[idx + jdx] < multresbins[2]) && (carries[idx + jdx + 1] += 1)
    end
  end
  #go through and resolve the carries.
  for idx = 1:length(carries)
    scratchpad[idx] += carries[idx]
    (scratchpad[idx] < carries[idx]) && (carries[idx + 1] += 1)
  end
  reinterpret(UInt64, scratchpad)
end

function calculate(a::Array{UInt64,1})
  sum = big(0)
  for idx = 1:length(a)
    sum += big(a[idx]) << (64 * (idx - 1))
  end
  sum
end

<<<<<<< HEAD:oddsandends/frag-mult-test.jl
function test_mults(cells::Integer)
=======
function test_mults(cells::Int)
>>>>>>> directory-reorg:src/oddsandends/frag-mult-test.jl
  x = rand(UInt64, cells)
  y = rand(UInt64, cells)

  x_big = calculate(x)
  y_big = calculate(y)

  #calculate using fragment multiplication
  res_frag = calculate(frag_mult(x,y))
  res_big = x_big * y_big

  println(res_frag)
  println(res_big)
end

#test that fragment multiplication in general works.
test_mults(1)
test_mults(2)
test_mults(4)
test_mults(8)

#next up:  truncated fragment multiplication.
function trunc_frag_mult(a::Array{UInt64,1}, b::Array{UInt64,1})
  #note that frag_mult fails for absurdly high length integer arrays.
  length(a) != length(b) && throw(ArgumentError("mismatched arrays"))

  #take these two UInt64 arrays and reinterpret them as UInt32 arrays
  a_32 = reinterpret(UInt32, a)
  b_32 = reinterpret(UInt32, b)
  l = length(a_32)

  scratchpad = zeros(UInt32, l + 1)
  carries    = zeros(UInt32, l)

  #first indexsum is length(a_32)
  indexsum = l
  for (aidx = 1:(indexsum - 1))
    temp_res::UInt64 = a_32[aidx] * b_32[indexsum - aidx]
    temp_res_high::UInt32 = (temp_res >> 32)
    scratchpad[1] += temp_res_high
    (scratchpad[1] < temp_res_high) && (carries[1] += 1)
  end
  #now proceed with the rest of the additions.
  for aidx = 1:l
    for bidx = (l + 1 - aidx):l
      temp_res = a_32[aidx] * b_32[bidx]
      temp_res_bins = reinterpret(UInt32, [temp_res])
      temp_res_low::UInt32 = temp_res
      temp_res_high = (temp_res >> 32)

      scratchindex = aidx + bidx - l

      scratchpad[scratchindex] += temp_res_low
      (temp_res_low > scratchpad[scratchindex]) && (carries[scratchindex] += 1)

      scratchpad[scratchindex + 1] += temp_res_high
      (temp_res_high > scratchpad[scratchindex + 1]) && (carries[scratchindex + 1] += 1)
    end
  end

  #go through and resolve the carries.
  for idx = 1:length(carries) - 1
    scratchpad[idx + 1] += carries[idx]
    (scratchpad[idx + 1] < carries[idx]) && (carries[idx + 1] += 1)
  end
  reinterpret(UInt64, scratchpad[2:end])
end

<<<<<<< HEAD:oddsandends/frag-mult-test.jl
function test_truncs(cells::Integer)
=======
function test_truncs(cells::Int)
>>>>>>> directory-reorg:src/oddsandends/frag-mult-test.jl
  x = rand(UInt64, cells)
  y = rand(UInt64, cells)

  z = frag_mult(x, y)[cells + 1:2 *cells]
  w = trunc_frag_mult(x, y)

  println(z)
  println(w)
end

test_truncs(1)
test_truncs(2)
test_truncs(4)
test_truncs(8)


using ZChop
function squeezing_param(N, ΔJ1, J2m, J3m)
    """
        ξ2 = squeezing_param(N, ΔJ1, J2m, J3m)

    Returns the squeezing parameter defined, e.g., as
    the inverse of Eq. (1) in Phys. Rev. A 65, 061801 (2002).

    The state is squeezed if greater than one.
    """
    return (J2m.^2 + J3m.^2) ./ (N * ΔJ1)
end

function squeezing_param(rho::QContinuousMeasurement.FixedjState)
    """
        ξ2 = squeezing_param(rho)

    Returns the squeezing parameter defined, e.g., as
    the inverse of Eq. (1) in Phys. Rev. A 65, 061801 (2002).

    The state is squeezed if greater than one.
    """
    Nj = nspins(rho)
    firstblock = block_sizes(Nj)[1]
    (Jx, Jy, Jz) =  map(x -> sparse(x[1:firstblock, 1:firstblock]), jspin(Nj))

    Jy2 = Jy^2
    Jzm = expectation_value!(rho, Jz)
    Jxm =  expectation_value!(rho, Jx)
    Jym =  expectation_value!(rho, Jy)
    Jy2m = expectation_value!(rho, Jy2)
    return (Jzm.^2 + Jxm.^2) ./ (Nj * (Jy2m - Jym.^2))
end

function squeezing_param(rho::QContinuousMeasurement.State)
    """
        ξ2 = squeezing_param(rho)

    Returns the squeezing parameter defined, e.g., as
    the inverse of Eq. (1) in Phys. Rev. A 65, 061801 (2002).

    The state is squeezed if greater than one.
    """
    Nj = nspins(rho)
    (Jx, Jy, Jz) = jspin(Nj)
    Jy2 = Jy^2
    Jzm = expectation_value!(rho, Jz)
    Jxm =  expectation_value!(rho, Jx)
    Jym =  expectation_value!(rho, Jy)
    Jy2m = expectation_value!(rho, Jy2)
    return (Jzm.^2 + Jxm.^2) ./ (Nj * (Jy2m - Jym.^2))
end

"""
    density(M)

Returns the density of a sparse matrix
"""
function density(s::SparseMatrixCSC)
    return length(s.nzval) / (s.n * s.m)
end

function Unconditional_QFI_Dicke(Nj::Int64, Tfinal::Real, dt::Real;
    κ::Real=1.,                    # Independent noise strength
    κcoll::Real=1.,                # Collective noise strength
    ω::Real=0.0                   # Frequency of the Hamiltonian
    )
    return Eff_QFI_HD_Dicke(Nj, 1, Tfinal, dt; κ=κ, κcoll=κcoll, ω=ω, η=0.0)
end

# Specialize zchop! to BlockDiagonal
function ZChop.zchop!(A::BlockDiagonal)
    for b in blocks(A)
        zchop!(b)
    end
end


@testset "update CIR Tests" begin
    tmp = DualOptimalFiltering.update_CIR_params([0.5, 0.5], 1., 1.5, 1., [0, 1], [5])
    @test tmp[1] == 2.5
    @test tmp[2] == [5, 6]
    @test tmp[3][1] ≈ 0.131579 atol=10.0^(-5)
    @test tmp[3][2] ≈ 0.868421 atol=10.0^(-5)
end;

@testset "predict CIR Tests" begin
    tmp = DualOptimalFiltering.predict_CIR_params([1.], 1., 2., 1., 1., [5], 1.)
    @test tmp[1] ≈ 1.072578883495754 atol=10.0^(-5)
    @test tmp[2] == 0:5
    [@test tmp[3][k] ≈ [0.686096, 0.268465, 0.0420196, 0.0032884, 0.000128673, 2.01396e-6][k] atol=10.0^(-5) for k in 1:6]
end

@testset "CIR filtering tests" begin
    Random.seed!(0)
    X = DualOptimalFiltering.generate_CIR_trajectory(range(0, stop = 2, length = 20), 3, 3., 0.5, 1);
    Y = map(λ -> rand(Poisson(λ),10), X);
    data = Dict(zip(range(0, stop = 2, length = 20), Y))
    Λ_of_t, wms_of_t, θ_of_t = DualOptimalFiltering.filter_CIR(3., 0.5, 1.,1.,data);
    [@test isreal(k) for k in Λ_of_t |> keys]
    [@test isinteger(sum(k)) for k in Λ_of_t |> values]
    [@test isreal(sum(k)) for k in wms_of_t |> values]
    [@test isreal(k) for k in θ_of_t |> values]
end;


macro linklibrary_add(n, es)
  path = normpath(Pkg.dir("FastSigmoid"),"c-src","libfastposit.so")
  fnname = QuoteNode(Symbol(:p, n, :e, es, :_add_e))

  posittype = Symbol(:UInt, n)

  esc(:(padd = (a, b) -> ccall( ($fnname, $path), $posittype, ($posittype, $posittype), a, b )))
end

#spot unit testing
@linklibrary_add(8,0)
@test padd(0x00, 0x00) == 0x00  #0 + 0 == 0
@test padd(0x00, 0x40) == 0x40  #0 + 1 == 1
@test padd(0x40, 0x40) == 0x60  #1 + 1 == 2
@test padd(0x40, 0xc0) == 0x00  #1 + -1 == 0
@test padd(0x80, 0xc0) == 0x80  #Inf + 1 == Inf

#comprehensive integration testing
@linklibrary_add(8,0)
for x = 0x00:0xFF, y = 0x00:0xFF
  if (x != 0x80 || y != 0x80) #ignore the one corner case that doesn't work.
    @test (x,y,padd(x, y)) == (x,y,reinterpret(UInt8, (reinterpret(Posit{8,0}, x) + reinterpret(Posit{8,0}, y))))
  end
end
@linklibrary_add(8,1)
for x = 0x00:0xFF, y = 0x00:0xFF
  if (x != 0x80 || y != 0x80) #ignore the one corner case that doesn't work.
    @test (x,y,padd(x, y)) == (x,y,reinterpret(UInt8, (reinterpret(Posit{8,1}, x) + reinterpret(Posit{8,1}, y))))
  end
end

@linklibrary_add(8,2)
for x = 0x00:0xFF, y = 0x00:0xFF
  if (x != 0x80 || y != 0x80) #ignore the one corner case that doesn't work.
    @test (x,y,padd(x, y)) == (x,y,reinterpret(UInt8, (reinterpret(Posit{8,2}, x) + reinterpret(Posit{8,2}, y))))
  end
end

# testing 16-bit number types.

macro test_big_bittype_add(N, ES)
  quote
    @linklibrary_add($N, $ES)

    println("testing addition on p$($N)e$($ES)")

    UType = UIntLookup[$N]

    for idx = 1:10000
      x = rand(UType)
      y = rand(UType)
      if (x != topbits($N) || y != topbits($N)) #ignore the one corner case that doesn't work.
        @test (x,y,padd(x, y)) == (x,y,reinterpret(UType, (reinterpret(Posit{$N,$ES}, x) + reinterpret(Posit{$N,$ES}, y))))
      end
    end
  end
end

@test_big_bittype_add(16,0)
@test_big_bittype_add(16,1)
@test_big_bittype_add(16,2)

@test_big_bittype_add(32,0)
@test_big_bittype_add(32,1)
@test_big_bittype_add(32,2)
@test_big_bittype_add(32,3)


using JLD2, Plots, Printf, Oceananigans

# ENV["GKSwstype"] = "100"

name = "veron_and_melville_Nz128_β1.0e-05_waves"

file = jldopen(name * "_averages.jld2")

Nx = file["grid/Nx"]
Ny = file["grid/Ny"]
Nz = file["grid/Nz"]

Lx = file["grid/Lx"]
Ly = file["grid/Ly"]
Lz = file["grid/Lz"]

grid = RegularRectilinearGrid(size = (Nx, Ny, Nz), halo = (3, 3, 3),
                              x = (0, Lx),
                              y = (0, Ly),
                              z = (-Lz, 0),
                              topology = (Periodic, Bounded, Bounded))

xc, yc, zc = nodes((Center, Center, Center), grid)

iterations = parse.(Int, keys(file["timeseries/t"]))

t = [file["timeseries/t/$iter"] for iter in iterations]

C = zeros(length(t), Nz)

for (i, iter) in enumerate(iterations)
    C[i, :] .= file["timeseries/c/$iter"][1, 1, :]
end

clim = 0.2
clevels = vcat(range(-0.01, clim, length=30), [1.0])

p = contourf(t, zc, C';
             linewidth = 0,
             # aspectratio = :equal,
             xlims = (0, t[end]),
             ylims = (-grid.Lz, 0),
             xlabel = "Time (s)",
             ylabel = "z (m)",
             levels = clevels,
             clims = (0, clim),
             color = :thermal)

display(p)


# ---
# title: 78. Subsets
# id: problem78
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Array, Backtracking, Bit Manipulation
# link: <https://leetcode.com/problems/subsets/description/>
# hidden: true
# ---
#
# Given an integer array `nums`, return _all possible subsets (the power set)_.
#
# The solution set must not contain duplicate subsets.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [1,2,3]
#     Output: [[],[1],[2],[1,2],[3],[1,3],[2,3],[1,2,3]]
#
#
# **Example 2:**
#
#
#
#     Input: nums = [0]
#     Output: [[],[0]]
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 10`
#   * `-10 <= nums[i] <= 10`
#
#
## @lc code=start
using LeetCode

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


# # Langmuir turbulence
#
# This script reproduces the Langmuir turbulence simulation reported in section 4 of
#
# [McWilliams, J. C. et al., "Langmuir Turbulence in the ocean," Journal of Fluid Mechanics (1997)](https://www.cambridge.org/core/journals/journal-of-fluid-mechanics/article/langmuir-turbulence-in-the-ocean/638FD0E368140E5972144348DB930A38).

using Pkg

Pkg.instantiate()

using Printf
using JLD2

## https://discourse.julialang.org/t/unable-to-display-plot-using-the-repl-gks-errors/12826/18
ENV["GKSwstype"] = "nul"
using Plots

using ArgParse

using Oceananigans
using Oceananigans.Units
using Oceananigans.BuoyancyModels: g_Earth

using LESbrary
using LESbrary.TurbulenceStatistics: TurbulentKineticEnergy, ViscousDissipation
using LESbrary.TurbulenceStatistics: first_through_second_order, turbulent_kinetic_energy_budget
using Oceanostics.TKEBudgetTerms: ZShearProduction

"Returns a dictionary of command line arguments."
function parse_command_line_arguments()
    settings = ArgParseSettings()

    @add_arg_table! settings begin
        "--Nh"
            help = "The number of grid points in x, y."
            default = 32
            arg_type = Int

        "--Nz"
            help = "The number of grid points in z."
            default = 32
            arg_type = Int

        "--Lh"
            help = "The physical extent of the domain in x, y."
            default = 128.0
            arg_type = Float64

        "--Lz"
            help = "The physical extent of the domain in z."
            default = 96.0
            arg_type = Float64

        "--stop-hours"
            help = "Number of hours to run the simulation for"
            default = 12.0
            arg_type = Float64

        "--advection-scheme"
            help = "Advection scheme: CenteredSecondOrder, WENO5, etc"
            default = :WENO5
            arg_type = Symbol

        "--pickup"
            help = "Whether or not to pick the simulation up from latest checkpoint"
            default = false
            arg_type = Bool
    end

    return parse_args(settings)
end

args = parse_command_line_arguments()

#####
##### Parameters
#####

Nh = args["Nh"]
Nz = args["Nz"]
Lh = args["Lh"]
Lz = args["Lz"]

f = 1e-4
Qᵘ = -3.72e-5 # m² s⁻²
Qᵇ = 2.307e-9 # m³ s⁻²
N² = 1.936e-5 # s⁻²

initial_mixed_layer_depth = 33 # m

inertial_period = 2π / f

stop_time = args["stop-hours"] * hours
snapshot_time_interval = 10minutes
averages_time_interval = 24hours
averages_time_window = inertial_period
averages_stride = 100
slice_depth = 8.0

amplitude = 0.8 # m

const wavenumber = 2π / 60 # m⁻¹
const Uˢ = amplitude^2 * wavenumber * sqrt(g_Earth * wavenumber) # m s⁻¹

uˢ(z) = Uˢ * exp(2wavenumber * z)
@inline ∂z_uˢ(z, t) = 2wavenumber * Uˢ * exp(2wavenumber * z)

#####
##### Grid
#####

grid = RectilinearGrid(CPU(), size=(Nh, Nh, Nz), extent=(Lh, Lh, Lz))

#####
##### Boundary conditions
#####

u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵘ))

b_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qᵇ),
                                bottom = GradientBoundaryCondition(N²))

#####
##### Sponge layer
#####

gaussian_mask = GaussianMask{:z}(center=-grid.Lz, width=grid.Lz/10)

u_sponge = v_sponge = w_sponge = Relaxation(rate=1/hour, mask=gaussian_mask)

b_sponge = Relaxation(rate = 1/hour,
                      target = LinearTarget{:z}(intercept=0.0, gradient=N²),
                      mask = gaussian_mask)

#####
##### Model setup
#####

advection = eval(args["advection-scheme"])()

model = NonhydrostaticModel(; advection, grid,
                            timestepper = :RungeKutta3,
                            tracers = :b,
                            buoyancy = BuoyancyTracer(),
                            coriolis = FPlane(f=f),
                            closure = AnisotropicMinimumDissipation(),
                            stokes_drift = UniformStokesDrift(∂z_uˢ=∂z_uˢ),
                            boundary_conditions = (u=u_bcs, b=b_bcs),
                            forcing = (u=u_sponge, v=v_sponge, w=w_sponge, b=b_sponge))

#####
##### Initial conditions: Stokes drift + stratification + noise
#####

Ξ(z) = randn() * exp(z / 4)

stratification(z) = z < - initial_mixed_layer_depth ? N² * z : - N² * initial_mixed_layer_depth

bᵢ(x, y, z) = stratification(z)     + 1e-2 * Ξ(z) * N² * model.grid.Lz
uᵢ(x, y, z) = uˢ(z) + sqrt(abs(Qᵘ)) * 1e-2 * Ξ(z)
wᵢ(x, y, z) =         sqrt(abs(Qᵘ)) * 1e-2 * Ξ(z)

set!(model, u=uᵢ, w=wᵢ, b=bᵢ)

#####
##### Simulation setup
#####

simulation = Simulation(model; Δt=1.0, stop_time)

wizard = TimeStepWizard(cfl=1.0, max_change=1.1, max_Δt=30.0)
simulation.callbacks[:wizard] = Callback(wizard, IterationInterval(10))

wall_clock = [time_ns()]

function print_progress(simulation)
    model = simulation.model

    umax = maximum(abs, model.velocities.u)
    vmax = maximum(abs, model.velocities.v)
    wmax = maximum(abs, model.velocities.w)

    ## Print a progress message
    msg = @sprintf("i: %04d, t: %s, Δt: %s, umax = (%.1e, %.1e, %.1e) ms⁻¹, wall time: %s\n",
                   iteration(simulation),
                   prettytime(simulation),
                   prettytime(simulation.Δt),
                   umax, vmax, wmax,
                   prettytime(1e-9 * (time_ns() - wall_clock[1]))
                  )

    wall_clock[1] = time_ns()

    @info msg

    return nothing
end

simulation.callbacks[:progress] = Callback(print_progress, IterationInterval(10))

#####
##### Output setup
#####

prefix = @sprintf("langmuir_turbulence_Nx%d_Nz%d", grid.Nx, grid.Nz)
data_directory = joinpath(@__DIR__, "..", "data", prefix) # save data in /data/prefix

# "Primitive" statistics

b = BuoyancyField(model)
p = model.pressures.pHY′ + model.pressures.pNHS
w_scratch = ZFaceField(model.grid)
c_scratch = CenterField(model.grid)

primitive_statistics = first_through_second_order(model, b=b, p=p, w_scratch=w_scratch, c_scratch=c_scratch)

U = primitive_statistics[:u]
V = primitive_statistics[:v]

# Turbulent kinetic energy budget terms

e = TurbulentKineticEnergy(model, U=U, V=V)
shear_production = ZShearProduction(model, U=U, V=V)
dissipation = ViscousDissipation(model)

tke_budget_statistics = turbulent_kinetic_energy_budget(model, b=b, p=p, U=U, V=V, e=e,
                                                        shear_production=shear_production, dissipation=dissipation)

statistics_to_output = merge(primitive_statistics, tke_budget_statistics)

fields_to_output = merge(model.velocities, model.tracers,
                         (νₑ = model.diffusivity_fields.νₑ,
                           e = Field(e),
                          sp = Field(shear_production),
                           ϵ = Field(dissipation)))

# Output configured for pickup

pickup = args["pickup"]
force = pickup ? false : true

k_xy_slice = searchsortedfirst(znodes(Face, grid), -slice_depth)

simulation.output_writers[:checkpointer] =
    Checkpointer(model, schedule = TimeInterval(26hour), prefix = prefix * "_checkpointer", dir = data_directory)

simulation.output_writers[:xz] =
    JLD2OutputWriter(model, fields_to_output,
                         schedule = TimeInterval(snapshot_time_interval),
                           prefix = prefix * "_xz",
                     field_slicer = FieldSlicer(j=1),
                              dir = data_directory,
                            force = force)

simulation.output_writers[:yz] =
    JLD2OutputWriter(model, fields_to_output,
                         schedule = TimeInterval(snapshot_time_interval),
                           prefix = prefix * "_yz",
                     field_slicer = FieldSlicer(i=1),
                              dir = data_directory,
                            force = force)

simulation.output_writers[:xy] =
    JLD2OutputWriter(model, fields_to_output,
                         schedule = TimeInterval(snapshot_time_interval),
                           prefix = prefix * "_xy",
                     field_slicer = FieldSlicer(k=k_xy_slice),
                              dir = data_directory,
                            force = force)

simulation.output_writers[:statistics] =
    JLD2OutputWriter(model, statistics_to_output,
                     schedule = TimeInterval(snapshot_time_interval),
                       prefix = prefix * "_statistics",
                          dir = data_directory,
                        force = force)

simulation.output_writers[:averaged_statistics] =
    JLD2OutputWriter(model, statistics_to_output,
                     schedule = AveragedTimeInterval(averages_time_interval,
                                                     window = averages_time_window,
                                                     stride = averages_stride),
                       prefix = prefix * "_averaged_statistics",
                          dir = data_directory,
                        force = force)

#####
##### Run
#####

run!(simulation)

#####
##### Animations and analysis
#####

xw, yw, zw = nodes(model.velocities.w)
xu, yu, zu = nodes(model.velocities.u)
xc, yc, zc = nodes((Center, Center, Center), grid)

xyfile = jldopen(joinpath(data_directory, prefix * "_xy.jld2"))
xzfile = jldopen(joinpath(data_directory, prefix * "_xz.jld2"))

iterations = parse.(Int, keys(xyfile["timeseries/t"]))

function divergent_levels(c, clim, nlevels=31)
    levels = range(-clim, stop=clim, length=nlevels)
    cmax = maximum(abs, c)
    clim < cmax && (levels = vcat([-cmax], levels, [cmax]))
    return (-clim, clim), levels
end

function sequential_levels(c, clims, nlevels=31)
    levels = collect(range(clims[1], stop=clims[2], length=nlevels))
    cmin = minimum(c)
    cmax = maximum(c)
    cmin < clims[1] && pushfirst!(levels, cmin)
    cmax > clims[2] && push!(levels, cmax)
    return clims, levels
end

@info "Making an animation from the saved data..."

anim = @animate for (i, iter) in enumerate(iterations)

    @info "Drawing frame $i from iteration $iter \n"

    ## Load 3D fields from file
    t = xyfile["timeseries/t/$iter"]
    wxy = xyfile["timeseries/w/$iter"][:, :, 1]
    uxz = xzfile["timeseries/u/$iter"][:, 1, :]
    exz = xzfile["timeseries/e/$iter"][:, 1, :]

    wlims, wlevels = divergent_levels(wxy, 0.02)
    ulims, ulevels = divergent_levels(uxz, 0.1)
    elims, elevels = sequential_levels(exz, (0, 1e-4))

    wxy_plot = contourf(xw, yw, clamp.(wxy, wlims[1], wlims[2])';
                              color = :balance,
                          linewidth = 0,
                        aspectratio = :equal,
                              clims = wlims,
                             levels = wlevels,
                              xlims = (0, grid.Lx),
                              ylims = (0, grid.Ly),
                             xlabel = "x (m)",
                             ylabel = "y (m)")

    uxz_plot = contourf(xu, zu, clamp.(uxz, ulims[1], ulims[2])';
                              color = :balance,
                          linewidth = 0,
                        aspectratio = :equal,
                              clims = ulims,
                             levels = ulevels,
                              xlims = (0, grid.Lx),
                              ylims = (-grid.Lz, 0),
                             xlabel = "x (m)",
                             ylabel = "z (m)")

    exz_plot = contourf(xc, zc, clamp.(exz, elims[1], elims[2])';
                              color = :thermal,
                          linewidth = 0,
                        aspectratio = :equal,
                              clims = elims,
                             levels = elevels,
                              xlims = (0, grid.Lx),
                              ylims = (-grid.Lz, 0),
                             xlabel = "x (m)",
                             ylabel = "z (m)")

    w_title = @sprintf("w(z=0, t=%s) (m s⁻¹)", prettytime(t))
    u_title = @sprintf("u(y=0, t=%s) (m s⁻¹)", prettytime(t))
    e_title = @sprintf("e(y=0, t=%s) (m s⁻¹)", prettytime(t))

    plot(wxy_plot, uxz_plot, exz_plot, layout=(3, 1), size=(1200, 1000),
         title = [w_title u_title e_title])

    if iter == iterations[end]
        close(xyfile)
        close(xzfile)
    end
end

mp4(anim, prefix * ".mp4", fps = 8)

##### Turbulent kinetic energy budget analysis

xu, yu, zu = nodes(model.velocities.u)

statistics_file = jldopen(joinpath(data_directory, prefix * "_averaged_statistics.jld2"))
statistics_iterations = parse.(Int, keys(statistics_file["timeseries/t"]))

last_iteration = statistics_iterations[end]

U = statistics_file["timeseries/u/$last_iteration"][1, 1, :]
V = statistics_file["timeseries/v/$last_iteration"][1, 1, :]
E = statistics_file["timeseries/e/$last_iteration"][1, 1, :]

# Budget terms
shear_production = statistics_file["timeseries/tke_shear_production/$last_iteration"][1, 1, :]
dissipation = statistics_file["timeseries/tke_dissipation/$last_iteration"][1, 1, :]
buoyancy_flux = statistics_file["timeseries/tke_buoyancy_flux/$last_iteration"][1, 1, :]
advective_flux = statistics_file["timeseries/tke_advective_flux/$last_iteration"][1, 1, :]
pressure_flux = statistics_file["timeseries/tke_pressure_flux/$last_iteration"][1, 1, :]

close(statistics_file)

Δz = grid.Δzᵃᵃᶜ
pressure_flux_divergence = @. (pressure_flux[2:end] - pressure_flux[1:end-1]) ./ Δz
advective_flux_divergence = @. (advective_flux[2:end] - advective_flux[1:end-1]) ./ Δz

tendency = @. - pressure_flux_divergence - advective_flux_divergence + shear_production + buoyancy_flux - dissipation

kwargs = (linewidth=2, alpha=0.8, ylims=(-48, 0))

velocities_plot = plot(U, zu; label="\$ U \$", legend=:bottomleft, xlabel="Velocities (m s⁻¹)", ylabel="z (m)", kwargs...)
plot!(velocities_plot, V, zu; label="\$ V \$", kwargs...)
plot!(velocities_plot, sqrt.(E), zu; label="\$ \\sqrt{E} \$", kwargs...)

budget_plot = plot(-dissipation, zu; label="dissipation", legend=:bottomleft, xlabel="Turbulent kinetic energy budget", ylabel="z (m)", kwargs...)
plot!(budget_plot, shear_production, zu; label="shear production", kwargs...)
plot!(budget_plot, buoyancy_flux, zu; label="buoyancy flux", kwargs...)
plot!(budget_plot, -pressure_flux_divergence, zu; label="pressure flux divergence", kwargs...)
plot!(budget_plot, -advective_flux_divergence, zu; label="advective flux divergence", kwargs...)
plot!(budget_plot, tendency, zu; label="tendency", linestyle=:dash, kwargs...)

tendency_plot = plot(tendency, zu; xlabel="TKE tendency", ylabel="z (m)", legend=nothing, linestyle=:dash, kwargs...)

statistics_plot = plot(velocities_plot, budget_plot, tendency_plot,
                       layout = (1, 3), size = (1000, 600))

savefig(statistics_plot, prefix * "_statistics.png")


# ellipse to multiple forward simulated shapes
# only final configurations observed

using BridgeLandmarks
using Random
using Distributions
using StaticArrays
using LinearAlgebra
using JLD2

workdir = @__DIR__
cd(workdir)

Random.seed!(9)

σobs = 0.01

n = 15
nshapes = 10

a0 = 2.0
c0 = 0.1
γ0 = 0.7

q0 = [PointF(2.0cos(t), sin(t))  for t in collect(0:(2pi/n):2pi)[2:end]]
xobsT = Vector{PointF}[]
Ptrue = MarslandShardlow(a0, c0, γ0, 0.0, n)

for k in 1:nshapes
    # Ptrue = MarslandShardlow(a0 * exp(0.5*randn()), c0 * exp(0.5*randn()),
    #                 γ0 *  exp(0.5*randn()), 0.0, n)
    x0 = State(q0, randn(PointF,n))
    T = 1.0; dt = 0.01; t = 0.0:dt:T
    Wf, Xf = landmarksforward(t, x0, Ptrue)
    push!(xobsT, [Xf.yy[end].q[i] for i in 1:n ] + σobs * randn(PointF,n))
end
xobs0 = []

JLD2.@save "data_exp2.jld2" xobs0 xobsT n nshapes q0


# parallel-random-block-gauss-seidel
function prbgs(A::SparseMatrixCSC, x0::Vector{Float64}, b::Vector{Float64}; block_fraction=0.2, Np=2)

    Nk = length(x0)
    chunksize = Int64(floor(block_fraction*Nk/Np))
    offset = div(chunksize, 2)
    slices = [collect(Iterators.partition(1:Nk, chunksize))...,
              collect(Iterators.partition(offset:Nk, chunksize))...]
    Ns = length(slices)

    x = x0[:]
    r = b - A*x
    bnorm = norm(b)

    err = norm(r) / bnorm
    while err > 1e-10

        picked = [slices[rand(1:end)] for _=1:Np]
        As = []
        for slc in picked
            append!(As, [sparse(view(A, :, slc))])
        end

        deltas = pmap(As) do M
            M \ r
        end

        for i=1:Np
            slc = picked[i]
            x[slc] = x[slc] + deltas[i]
        end

        r[:] = b - A*x
        err = norm(r) / bnorm
    end

    return x
end