function process_likelihood_model(ρ_list::Vector,Eₘ_list::Vector)
    # Generate the A matrix used to calculate likelihoods
    # The A matrix depends on the input states and measurement operators
    dimsmatch(ρ_list,Eₘ_list)
    sum(abs,data(sum(Eₘ_list))-I)<1E-15 ||
        throw(ArgumentError("Eₘ operators do not form a valid POVM!"))
    prep = (Eₘ,ρ) -> transpose(vec(Array(ρ⊗transpose(Eₘ))))
    return reduce(vcat,vec(prep.(Eₘ_list,permutedims(ρ_list))))
end

# below is the actual projected gradient descent algorithm from Knee, et al.
# Quantum process tomography via completely positive and trace-preserving
# projection. Phys. Rev. A 98, 062336 (2018).

function pgd_process_tomo(M::Matrix{T}, A::Matrix; tol=1E-10, cptp_tol=1E-8, info=false, cbfun=nothing) where T
    # Choi process matrix reconstruction by maximum likelihood projected gradient descent
    size(A,1)==length(M) || throw(ArgumentError("A matrix inconsistent with number of measurements!"))
    abs(sum(M)-1)<1/4 || throw(ArgumentError("measurement counts not normalized!"))
    # infer space dimensions from A matrix
    d = isqrt(isqrt(size(A,2)))
    # initial Choi matrix guess, the identity map
    C = Matrix{Complex{T}}(I/d,d^2,d^2)
    # objective and gradient functions setup
    f = C -> loglikelihood(M,C,A)
    ∇f = C -> loglikelihood_gradient(M,C,A)
    # metaparameters & initial cost calculation
    μ⁻¹ = 2d^2/T(3); γ = T(0.3)
    c₁ = T(1E6); c₂ = f(C)
    info && println("start cost = $c₂")
    # iterate through projected gradient descent steps, with backtracking
    h = CPTP_helpers(C)
    stp = 0
    while c₁ - c₂ > tol
        stp += 1
        c₁, ∇C = c₂, ∇f(C)
        D = project_CPTP(C .- μ⁻¹.*∇C, h, cptp_tol) - C
        α = one(T); Π = γ*real(D⋅∇C)
        while (c₂ = f(C .+ α.*D)) > c₁ + α*Π
            α = α/2 # backtrack
        end
        C .= C .+ α.*D
        if cbfun !== nothing # run callback function; e.g. to calc fidelity at each step
            cbfun(stp,c₂,C)
        end
    end
    info && println("final cost = $c₂, Δ = $(c₁-c₂), number of steps: $stp")
    return C
end

function loglikelihood(M::Matrix{T}, C::Matrix{Complex{T}}, A::Matrix) where {T<:Real}
    # Binomial statistics for the measurement count probability, up to some irrelevant constant
    logP = log.(max.(real.(A*vec(C)), T(1E-12)))
    return -(M⋅logP)
end

function loglikelihood_gradient(M::Matrix{T}, C::Matrix{Complex{T}}, A::Matrix) where {T<:Real}
    mMP = .-vec(M)./max.(real.(A*vec(C)), T(1E-12))
    return unvec(A'*mMP)
end

function project_CPTP(C::Matrix{Complex{T}}, h::Tuple, tol::Real=1E-8) where {T<:Real}
    # generate storage objects
    X₁ = copy(C); Y₁ = zero(X₁);
    X₂ = copy(Y₁); Y₂ = copy(Y₁)
    P = copy(Y₁); Q = copy(Y₁)
    ΔP = one(T); ΔQ = one(T)
    D, V, MdagM, 𝕀 = h
    # iterate through TP & CP projections
    while ΔP^2 + ΔQ^2 + 2*abs(P⋅X₂-P⋅X₁) + 2*abs(Q⋅Y₂-Q⋅Y₁) > tol
        project_TP!(Y₂, X₁+P, MdagM, 𝕀)
        ΔP = norm2_diff(X₁,Y₂)
        P .= X₁ .- Y₂ .+ P
        project_CP!(X₂, Y₂+Q, D, V)
        ΔQ = norm2_diff(Y₂,X₂)
        Q .= Y₂ .- X₂ .+ Q
        X₁, X₂ = X₂, X₁
        Y₁, Y₂ = Y₂, Y₁
    end
    return X₁
end

function project_CP!(X::Matrix{Complex{T}}, C::Matrix{Complex{T}}, D::Vector{T}, V::Matrix) where {T<:Real}
    # Project the process onto the completely positive subspace by making the
    # Choi matrix positive semidefinite
    # We do this by taking the eigendecomposition, setting any negative
    # eigenvalues to 0, and reconstructing the Choi matrix
    hermfact!(D,V,Hermitian(hermitianize!(C)))
    D .= max.(D,zero(T))
    mul!(C,Diagonal(D),V')
    mul!(X,V,C)
end

function project_TP!(Y::Matrix{Complex{T}}, C::Matrix{Complex{T}}, MdagM::Matrix, 𝕀::Matrix) where {T<:Real}
    # Project the process onto the trace-preserving subspace
    d⁻¹ = one(T)/isqrt(size(C,1))
    mul!(vec(Y), MdagM, vec(C))
    Y .= C .- d⁻¹.*(Y.-𝕀)
end

function CPTP_helpers(C::Matrix{Complex{T}}) where {T<:Real}
    D = Vector{T}(undef,size(C,1))
    V = Matrix{Complex{T}}(undef,size(C))
    MdagM, 𝕀 = TP_helper_matrices(C)
    return D, V, MdagM, 𝕀
end

function TP_helper_matrices(C::Matrix{Complex{T}}) where {T<:Real}
    d = isqrt(size(C,1))
    𝕀 = Matrix{T}(I,d,d); k = zeros(Complex{T},1,d)
    # this could be done more efficiently, but it doesn't matter
    M = sum(i->(k[i]=1; k[mod1(i-1,d)]=0; (𝕀 ⊗ k) ⊗ (𝕀 ⊗ k)), 1:d)
    return M'*M, 𝕀⊗𝕀
end


"""
    lagFn(x,n=1)

Create a matrix or vector of lagged values.

# Input
- `x::Array`: T Vector or Txk matrix
- `n::Int`:   scalar, order of lag. For instance, 2 puts x[t-2,:] on row t

# Output
- `z::Array`:  Txk matrix of lags

"""
function lagFn(x,n=1)

  (T,k) = (size(x,1),size(x,2))

  if n < 0                                    #leads
    z = [ x[1-n:T,:]; fill(NaN,(abs(n),k)) ]
  elseif n == 0                               #no lag or lead
    z = copy(x)
  elseif n > 0                                #lags
    z = [ fill(NaN,(n,k)); x[1:T-n,:] ]
  end

  isa(x,AbstractVector) && (z = z[:,1])               #z should be vector if x is

  return z

end


# ---
# title: 930. Binary Subarrays With Sum
# id: problem930
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Hash Table, Two Pointers
# link: <https://leetcode.com/problems/binary-subarrays-with-sum/description/>
# hidden: true
# ---
#
# In an array `A` of `0`s and `1`s, how many **non-empty** subarrays have sum
# `S`?
#
#
#
# **Example 1:**
#
#
#
#     Input: A = [1,0,1,0,1], S = 2
#     Output: 4
#     Explanation:
#     The 4 subarrays are bolded below:
#     [ **1,0,1** ,0,1]
#     [ **1,0,1,0** ,1]
#     [1, **0,1,0,1** ]
#     [1,0, **1,0,1** ]
#
#
#
#
# **Note:**
#
#   1. `A.length <= 30000`
#   2. `0 <= S <= A.length`
#   3. `A[i]` is either `0` or `1`.
#
#
## @lc code=start
using LeetCode

function num_subarrays_with_sum(nums::Vector{Int}, goal::Int)
    cnt = OffsetArray(fill(0, length(nums) + 1), -1)
    s, res = 0, 0
    for num in nums
        cnt[s] += 1
        s += num
        res += get(cnt, s - goal, 0)
    end
    res
end
## @lc code=end


module Day9

export get_inputs, get_solution1, get_solution2

## Input getting
read_input(f_name) = parse.(Int, readlines(joinpath(@__DIR__, f_name)))

function get_inputs()
    test_input1 = test_input2 = (read_input("test_input1.txt"), 5)
    test_output1 = 127
    test_output2 = 62
    data = (read_input("input.txt"), 25)
    return (; test_input1, test_output1, test_input2, test_output2, data)
end

## Solution functions
function is_valid(number, preamble)
    for i in eachindex(preamble), j in i:length(preamble)
        if preamble[i] + preamble[j] == number return true end
    end
    return false
end

# Part 1
function get_solution1((data, preamble_length))
    for i in eachindex(data)
        j = i + preamble_length
        preamble = @view data[i:j-1]
        val = data[j]
        if !is_valid(val, preamble) return val end
    end
    return -1
end

# Part 2
function get_solution2((data, preamble_length))
    sol1 = get_solution1((data, preamble_length))
    for i in eachindex(data)
        val = 0
        for j in i:length(data)
            val += data[j]
            if val==sol1
                this_range = @view data[i:j]
                return sum(extrema(this_range))
            elseif val > sol1
                break
            end
        end
    end
    return -1
end

end


### A Pluto.jl notebook ###
# v0.14.4

using Markdown
using InteractiveUtils

# ╔═╡ 51c330f6-7675-475e-b0eb-0fc8486fd094
begin
	using PlutoUI
	PlutoUI.TableOfContents(aside=true)
end

# ╔═╡ 5c6ed844-279b-4576-a2fb-1c1d11fb62ee
using Arrowhead, LinearAlgebra

# ╔═╡ c3caaf95-bef6-446d-b197-23ad1a7da450
md"""
# Updating the SVD

In many applications which are based on the SVD, arrival of new data requires SVD of the new matrix. Instead of computing from scratch, existing SVD can be updated.

__Prerequisites__

The reader should be familiar with concepts of singular values and singular vectors, related perturbation theory, and algorithms.

__Competences__

The reader should be able to recognise applications where SVD updating can be sucessfully applied and apply it.
"""

# ╔═╡ e070e205-b8fb-4608-93c1-f5c5beb075c3
md"""
## Facts

For more details see
[M. Gu and S. C. Eisenstat, A Stable and Fast Algorithm for Updating the Singular Value Decomposition](http://www.cs.yale.edu/publications/techreports/tr966.pdf) and [M. Brand, Fast low-rank modifications of the thin singular value decomposition](http://www.sciencedirect.com/science/article/pii/S0024379505003812) and the references therein.

1. Let $A\in\mathbb{R}^{m\times n}$ with $m\geq n$ and $\mathop{\mathrm{rank}}(A)=n$, and  let $A=U\Sigma V^T$ be its SVD. Let $a\in\mathbb{R}^{n}$ be a vector, and let $\tilde A=\begin{bmatrix} A \\ a^T\end{bmatrix}$. Then

$$
\begin{bmatrix} A \\ a^T\end{bmatrix} =\begin{bmatrix} U & \\ & 1 \end{bmatrix}
\begin{bmatrix} \Sigma \\ a^TV \end{bmatrix}  V^T.$$

Let $\begin{bmatrix} \Sigma \\ a^T V \end{bmatrix} = \bar U \bar \Sigma \bar V^T$ be the SVD of the half-arrowhead matrix. _This SVD can be computed in $O(n^2)$ operations._ Then

$$
\begin{bmatrix} A \\ a^T\end{bmatrix} =
\begin{bmatrix} U & \\ & 1 \end{bmatrix} \bar U \bar\Sigma \bar V^T V^T \equiv
\tilde U \bar \Sigma \tilde V^T$$

is the SVD of $\tilde A$.

2. Direct computation of $\tilde U$ and $\tilde V$ requires $O(mn^2)$ and $O(n^3)$ operations. However, these multiplications can be performed using Fast Multipole Method. This is not (yet) implemented in Julia and is "not for the timid" (quote by Steven G. Johnson).

3. If $m<n$ and $\mathop{\mathrm{rank}}(A)=n$, then

$$
\begin{bmatrix} A \\ a^T\end{bmatrix} =\begin{bmatrix} U & \\ & 1 \end{bmatrix}
\begin{bmatrix} \Sigma & 0 \\ a^T V & \beta\end{bmatrix} \begin{bmatrix} V^T \\ v^T \end{bmatrix},$$

where $\beta=\sqrt{\|a\|_2^2-\|V^T a\|_2^2}$ and $v=(I-VV^T)a$. Notice that $V^Tv=0$ by construction. Let $\begin{bmatrix} \Sigma & 0 \\ a^T V &  \beta\end{bmatrix} = \bar U \bar \Sigma \bar V^T$ be the SVD of the half-arrowhead matrix. Then

$$
\begin{bmatrix} A \\ a^T\end{bmatrix} =
\begin{bmatrix} U & \\ & 1 \end{bmatrix} \bar U \bar\Sigma \bar V^T \begin{bmatrix} V^T \\ v^T \end{bmatrix}
\equiv \tilde U \bar \Sigma \tilde V^T$$

is the SVD of $\tilde A$.

4. Adding a column $a$ to $A$ is equivalent to adding a row $a^T$ to $A^T$.

5. If $\mathop{\mathrm{rank}}(A)<\min\{m,n\}$ or if we are using SVD approximation of rank $r$, and if we want to keep the rank of the approximation (this is the common case in practice), then the formulas in Fact 1 hold approximately. More precisely, the updated rank $r$ approximation is __not__ what we would get by computing the approximation of rank $r$ of the updated matrix, but is sufficient in many applications.
"""

# ╔═╡ 0ff9f1d6-1ebf-4a42-b51a-5a67f6192fe7
md"""

## Examples

### Adding row to a tall matrix

If $m\geq n$, adding row does not increase the size of $\Sigma$.
"""

# ╔═╡ 1f7f9f7a-df70-40a3-aa06-b5621b729a0c
function SVDaddrow(svdA::SVD,a::Vector)
    # Create the transposed half-arrowhead
    m,r,n=size(svdA.U,1),length(svdA.S),size(svdA.V,1)
    T=typeof(a[1])
    b=svdA.Vt*a
    if m>=n || r<m
        M=HalfArrow(svdA.S,b)
    else
        β=√(norm(a)^2-norm(b)^2)
        M=HalfArrow(svdA.S,[b;β])
    end
    # From Arrowhead package
    svdM,info=svd(M)
    # Return the updated SVD
    if m>=n || r<m
        return SVD([svdA.U zeros(T,m); zeros(T,1,r) one(T)]*svdM.V,
            svdM.S, adjoint(svdA.V*svdM.U))
    else
        # Need one more row of svdA.V - v is an orthogonal projection
        v=a-svdA.V*b
        normalize!(v)
        return SVD([svdA.U zeros(T,m); zeros(T,1,r) one(T)]*svdM.V,
            svdM.S, adjoint([svdA.V v]*svdM.U))
    end
end

# ╔═╡ 445cacab-3bb2-42f6-8d34-73c77ceecd0a
methods(SVD)

# ╔═╡ 2d0c21e5-dc68-4434-94ca-f433355f0e5c
begin
	import Random
	Random.seed!(421)
	A=rand(10,6)
	a=rand(6)
end

# ╔═╡ c034e218-09d3-4403-a413-98932865c481
S=svd(A)

# ╔═╡ 485e8222-2bb7-490e-9a51-675f0baef274
# Residual
norm(A*S.V-S.U*Diagonal(S.S))

# ╔═╡ 3dd2734d-db44-4226-9163-e6a55dcc7667
Sa=SVDaddrow(S,a)

# ╔═╡ 2fb74892-bade-47fe-bcaa-ef5bfdb19445
begin
	Aa=[A;transpose(a)]
	size(Aa),size(Sa.U),size(S.V)
end

# ╔═╡ c074206d-ed7d-44fb-a0ee-22e5cdd6e0f8
[svdvals(Aa) Sa.S]

# ╔═╡ 89990d20-4752-4558-a328-a4ae6a428906
# Residual and orthogonality
norm(Aa*Sa.V-Sa.U*Diagonal(Sa.S)),
norm(Sa.U'*Sa.U-I), norm(Sa.Vt*Sa.V-I)

# ╔═╡ bd35d5e4-b454-4a02-a026-ded3bdf76cdf
md"""
### Adding row to a flat matrix
"""

# ╔═╡ 8f31ef4a-a5da-4837-9320-ec9578c5a277
begin
	# Now flat matrix
	Random.seed!(421)
	A₁=rand(6,10)
	a₁=rand(10)
	S₁=svd(A₁);
end

# ╔═╡ e2ff1d44-7f3c-4100-80d3-df97579bc56c
A₁

# ╔═╡ 2b055add-b46f-4a2e-8ae1-89e59f12c21c
Aa₁=[A₁;transpose(a₁)]

# ╔═╡ e68d3450-da63-4b3f-9244-632a4eba5aff
Sa₁=SVDaddrow(S₁,a₁);

# ╔═╡ fac291b0-cf08-46d6-a233-3e1e8675c7e0
size(Aa₁),size(Sa₁.U),size(S₁.V)

# ╔═╡ 9f2d9dcc-ee9c-49e5-a839-3dadfed371e8
# Residual and orthogonality
norm(Aa₁*Sa₁.V-Sa₁.U*Diagonal(Sa₁.S)),
norm(Sa₁.U'*Sa₁.U-I), norm(Sa₁.Vt*Sa₁.V-I)

# ╔═╡ 9727df4c-52f5-4527-84e6-dfc02820d583
md"""
### Adding columns

This can be viewed as adding rows to the transposed matrix, an elegant one-liner in Julia.
"""

# ╔═╡ 3b34587c-fc0c-4677-91a7-4efb8a9abdba
function SVDaddcol(svdA::SVD,a::Vector)
    X=SVDaddrow(SVD(svdA.V,svdA.S,adjoint(svdA.U)),a)
    SVD(X.V,X.S,adjoint(X.U))
end

# ╔═╡ d8807efc-0cff-4aee-8bc5-4a5403c092eb
begin
	# Tall matrix
	Random.seed!(897)
	A₂=rand(10,6)
	a₂=rand(10)
	S₂=svd(A₂)
	Sa₂=SVDaddcol(S₂,a₂)
end

# ╔═╡ 8ce644c3-f71a-4904-bf15-d3f602f44c05
Aa₂=[A₂ a₂];

# ╔═╡ ca5fd757-5ab8-4826-8636-59a29a8c5843
# Residual and orthogonality
norm(Aa₂*Sa₂.V-Sa₂.U*Diagonal(Sa₂.S)),
norm(Sa₂.U'*Sa₂.U-I), norm(Sa₂.Vt*Sa₂.V-I)

# ╔═╡ 5f5e6293-17b5-405d-96c6-d661151796dc
begin
	# Flat matrix
	Random.seed!(332)
	A₃=rand(6,10)
	a₃=rand(6)
	S₃=svd(A₃)
	Sa₃=SVDaddcol(S₃,a₃)
end

# ╔═╡ 3c5d82e1-81f4-4cff-9c2a-6cc9bfbefd58
Aa₃=[A₃ a₃]

# ╔═╡ a2f5ab12-5c09-4f2b-9665-d2d3d44abd17
# Residual and orthogonality
norm(Aa₃*Sa₃.V-Sa₃.U*Diagonal(Sa₃.S)),
norm(Sa₃.U'*Sa₃.U-I), norm(Sa₃.Vt*Sa₃.V-I)

# ╔═╡ 6563cdc2-8577-4fac-997d-413edc56e2ac
begin
	# Square matrix
	A₄=rand(10,10)
	a₄=rand(10)
	S₄=svd(A₄)
end

# ╔═╡ 0f656a69-2243-4273-8600-83a12b694497
begin
	Sa₄=SVDaddrow(S₄,a₄)
	Aa₄=[A₄;transpose(a₄)]
	norm(Aa₄*Sa₄.V-Sa₄.U*Diagonal(Sa₄.S)),
	norm(Sa₄.U'*Sa₄.U-I), norm(Sa₄.Vt*Sa₄.V-I)
end

# ╔═╡ 35b4c90e-830f-487e-b683-988952aa12af
begin
	Sa₅=SVDaddcol(S₄,a₄)
	Aa₅=[A₄ a₄]
	norm(Aa₅*Sa₅.V-Sa₅.U*Diagonal(Sa₅.S)),
	 norm(Sa₅.U'*Sa₅.U-I), norm(Sa₅.Vt*Sa₅.V-I)
end

# ╔═╡ 22b4dea8-79a0-438a-a3e4-d7c725d47051
md"""
### Updating a low rank approximation

"""

# ╔═╡ 0a5298eb-ed33-443b-b9e4-c0a32b408605
begin
	# Adding row to a tall matrix
	A₆=rand(10,6)
	S₆=svd(A₆)
	a₆=rand(6)
	# Rank of the approximation
	r=4
end

# ╔═╡ 99b2a284-4a24-492d-b677-4b08430a4399
# Low-rank approximation
Sr=SVD(S₆.U[:,1:r], S₆.S[1:r],adjoint(S₆.V[:,1:r]));

# ╔═╡ 3677f364-e828-4a42-84b0-e86f0e3fff3b
begin
	# Eckart, Young, Mirsky
	Ar=Sr.U*Diagonal(Sr.S)*Sr.Vt
	Δ=Ar-A₆
	opnorm(Δ),svdvals(A₆)[5]
end

# ╔═╡ 32880612-cdd1-4e16-8a05-025eb35110f9
Sa₆=SVDaddrow(Sr,a₆);

# ╔═╡ 65c83ee2-36ad-4388-b9b3-05830ab3ee9a
Aa₆=[A₆; transpose(a₆)];

# ╔═╡ 13deeb53-9670-4351-85b3-a366690a94a8
svdvals(Aa₆),svdvals([Ar;transpose(a₆)]),Sa₆.S

# ╔═╡ a02b7022-b775-4dbd-bb83-0d6f88a47312
begin
	# Adding row to a flat matrix
	A₇=rand(6,10)
	S₇=svd(A₇)
	a₇=rand(10)
	# Rank of the approximation
	r₇=4
end

# ╔═╡ 2f22e978-2fb4-4ec4-80af-519384e98f6c
begin
	Sr₇=SVD(S₇.U[:,1:r₇], S₇.S[1:r₇],adjoint(S₇.V[:,1:r₇]))
	Sa₇=SVDaddrow(Sr₇,a₇);
end

# ╔═╡ 2beba2fd-cb70-4da2-a9f3-d4c39815ff18
begin
	Ar₇=Sr₇.U*Diagonal(Sr₇.S)*Sr₇.Vt
	svdvals(Sa₇),svdvals([Ar₇;transpose(a₇)]),Sa₇.S
end

# ╔═╡ Cell order:
# ╟─51c330f6-7675-475e-b0eb-0fc8486fd094
# ╟─c3caaf95-bef6-446d-b197-23ad1a7da450
# ╟─e070e205-b8fb-4608-93c1-f5c5beb075c3
# ╟─0ff9f1d6-1ebf-4a42-b51a-5a67f6192fe7
# ╠═5c6ed844-279b-4576-a2fb-1c1d11fb62ee
# ╠═1f7f9f7a-df70-40a3-aa06-b5621b729a0c
# ╠═445cacab-3bb2-42f6-8d34-73c77ceecd0a
# ╠═2d0c21e5-dc68-4434-94ca-f433355f0e5c
# ╠═c034e218-09d3-4403-a413-98932865c481
# ╠═485e8222-2bb7-490e-9a51-675f0baef274
# ╠═3dd2734d-db44-4226-9163-e6a55dcc7667
# ╠═2fb74892-bade-47fe-bcaa-ef5bfdb19445
# ╠═c074206d-ed7d-44fb-a0ee-22e5cdd6e0f8
# ╠═89990d20-4752-4558-a328-a4ae6a428906
# ╟─bd35d5e4-b454-4a02-a026-ded3bdf76cdf
# ╠═8f31ef4a-a5da-4837-9320-ec9578c5a277
# ╠═e2ff1d44-7f3c-4100-80d3-df97579bc56c
# ╠═2b055add-b46f-4a2e-8ae1-89e59f12c21c
# ╠═e68d3450-da63-4b3f-9244-632a4eba5aff
# ╠═fac291b0-cf08-46d6-a233-3e1e8675c7e0
# ╠═9f2d9dcc-ee9c-49e5-a839-3dadfed371e8
# ╟─9727df4c-52f5-4527-84e6-dfc02820d583
# ╠═3b34587c-fc0c-4677-91a7-4efb8a9abdba
# ╠═d8807efc-0cff-4aee-8bc5-4a5403c092eb
# ╠═8ce644c3-f71a-4904-bf15-d3f602f44c05
# ╠═ca5fd757-5ab8-4826-8636-59a29a8c5843
# ╠═5f5e6293-17b5-405d-96c6-d661151796dc
# ╠═3c5d82e1-81f4-4cff-9c2a-6cc9bfbefd58
# ╠═a2f5ab12-5c09-4f2b-9665-d2d3d44abd17
# ╠═6563cdc2-8577-4fac-997d-413edc56e2ac
# ╠═0f656a69-2243-4273-8600-83a12b694497
# ╠═35b4c90e-830f-487e-b683-988952aa12af
# ╟─22b4dea8-79a0-438a-a3e4-d7c725d47051
# ╠═0a5298eb-ed33-443b-b9e4-c0a32b408605
# ╠═99b2a284-4a24-492d-b677-4b08430a4399
# ╠═3677f364-e828-4a42-84b0-e86f0e3fff3b
# ╠═32880612-cdd1-4e16-8a05-025eb35110f9
# ╠═65c83ee2-36ad-4388-b9b3-05830ab3ee9a
# ╠═13deeb53-9670-4351-85b3-a366690a94a8
# ╠═a02b7022-b775-4dbd-bb83-0d6f88a47312
# ╠═2f22e978-2fb4-4ec4-80af-519384e98f6c
# ╠═2beba2fd-cb70-4da2-a9f3-d4c39815ff18

src/Module B - Eigenvalue and Singular Value Decompositions/L8 Updating the SVD.jl


using Random, Printf, Colors

import Luxor
const lx = Luxor

function triangle(points, degree)
    lx.sethue(cols[degree])
    lx.poly(points, :fill)
end

function sierpinski(points, degree)
    triangle(points, degree)
    if degree > 1
        p1, p2, p3 = points
        sierpinski([p1, lx.midpoint(p1, p2),
                        lx.midpoint(p1, p3)], degree-1)
        sierpinski([p2, lx.midpoint(p1, p2),
                        lx.midpoint(p2, p3)], degree-1)
        sierpinski([p3, lx.midpoint(p3, p2),
                        lx.midpoint(p1, p3)], degree-1)
    end
end

function draw(n)
    lx.circle(lx.O, 100, :clip)
    points = lx.ngon(lx.O, 150, 3, -pi/2, vertices=true)
    sierpinski(points, n)
end

lx.Drawing(400, 250)
lx.background("white")
lx.origin()

depth = 8
cols = distinguishable_colors(depth) # from Colors.jl
draw(depth)

lx.finish()
lx.preview()

import Winston
const wn = Winston

t = collect(range(0,stop=4pi,length=1000));
f(x::Array) = 10x .* exp.(-0.3x) .* sin.(3x);
g(x::Array) = 0.1x.*(2pi .- x).*(4pi .- x);
h(x::Array) = 10.0 ./ (1 .+ x.*x);  # Note need spaces tp clarify 1./ and 1 ./

y1 = f(t)
y2 = g(t)
y3 = h(t)
wn.plot(t,y1,"b",t,y2,"r--",t,y3,"k;")

x = randn(10000)
wn.plothist(x, nbins=50)

p = wn.FramedPlot( title="Frequency vs Time", ylabel="\\Theta",xlabel="\\Tau");
wn.add(p,wn.FillBetween(t,y1,t,y2));
wn.add(p,wn.Curve(t,y1,color="red"));
wn.add(p,wn.Curve(t,y2,color="blue"));

wn.display(p)

wn.semilogx(t,y1)
wn.title("log(t) vs 10x * exp(-0.3x) * sin(3x)")

wn.semilogy(y2,y3)
wn.title("0.1*(2\\pi - x)*(4\\pi - x)  vs  log(1 /( 1 + x*x))")

p = wn.FramedPlot(aspect_ratio=1,xrange=(-10,110),yrange=(-10,110));
n = 21;
x = collect(range(0.0, length=n, stop=100.0));

# Create a set of random variates

yA =  10.0*randn(n) .+ 40.0;
yB =  x .+ 5.0*randn(n);

# Set labels and symbol styles

a = wn.Points(x, yA, kind="circle");
wn.setattr(a,label="'a' points");

b = wn.Points(x, yB);
wn.setattr(b,label="'b' points");
wn.style(b, kind="filled circle");

# Plot a line which 'fits' through the yB points
# and add a legend in the top LHS part of the graph

s = wn.Slope(1, (0,0), kind="dotted");
wn.setattr(s, label="slope");
lg = wn.Legend(.1, .9, Any[a,b,s] );
wn.add(p, s, a, b, lg);

wn.display(p)

wn.savefig(p, "MyWPlot.png")

import Cairo
const ca = Cairo

c  = ca.CairoRGBSurface(512, 128);
cr = ca.CairoContext(c);
ca.save(cr);

ca.set_source_rgb(cr, 0.8, 0.8, 0.8);
ca.rectangle(cr, 0.0, 0.0, 512.0, 128.0); # background
ca.fill(cr);
ca.restore(cr);
ca.save(cr);

x0=51.2;  y0=64.0;
x1=204.8; y1=115.4;
x2=307.2; y2=12.8;
x3=460.8; y3=64.0;

ca.move_to(cr, x0, y0);
ca.curve_to(cr, x1, y1, x2, y2, x3, y3);
ca.set_line_width(cr, 10.0);
ca.stroke(cr);
ca.restore(cr);

ca.move_to(cr, 12.0, 12.0);
ca.set_source_rgb(cr, 0, 0, 0);
ca.show_text(cr,"Cairo Curve")
ca.write_to_png(c,"cstroke.png");

run(`display cstroke.png`)

import PyPlot
const py = PyPlot

x = collect(range(0.0,stop=2pi,length=1000))
y = sin.(3*x + 4*cos.(2*x));
py.title("A sinusoidally modulated sinusoid");
py.plot(x, y, color="red", linewidth=2.0, linestyle="--");

y = collect(range(0,stop=3π,length=250))
py.surf(y, y, y .* sin.(y) .* cos.(y)' .* exp.(-0.4y))

py.xkcd()
x = collect(range(1, length=101, stop=10));
y = sin.(3x + cos.(5x))
p = py.plot(x,y)

import Gadfly
gd = Gadfly

dd  =  gd.plot(x =  rand(10),  y = rand(10));
gd.draw(SVG("random-pts.svg",  15cm, 12cm) , dd);

x = collect(1:100);
y1 = ones(100) - 2*rand(100);
y2 = randn(100);

gd.plot(
 gd.layer(x=x,y=y1,gd.Geom.line,gd.Theme(default_color=gd.colorant"red")),
 gd.layer(x=x,y=y2,gd.Geom.line,gd.Theme(default_color=gd.colorant"blue"))
)

using RDatasets, DataFrames;

mlmf = dataset("mlmRev","Gcsemv")
df = mlmf[completecases(mlmf), :]

names(df)

# gd.set_default_plot_size(20cm, 12cm);
gd.plot(df, x="Course", y="Written", color="Gender")

gd.plot((x,y) -> exp.(-x) .* cos.(y), -2, 0.4, 0, 4)

gd.plot((x,y) -> x .* exp.(-(x - floor.(x))).^2 .- y.^2, -8.0, 8, -2.0, 2.0)

import Compose
const cm =  Compose

function sierpinski(n)
  if n == 0
    cm.compose(cm.context(), cm.polygon([(1,1), (0,1), (1/2, 0)]));
  else
    t = sierpinski(n - 1);
    cm.compose(cm.context(), (cm.context( 1/4, 0, 1/2, 1/2), t),
                             (cm.context( 0, 1/2, 1/2, 1/2), t),
                             (cm.context( 1/2, 1/2, 1/2, 1/2), t));
  end
end

cx1 = compose(sierpinski(1), linewidth(0.2mm),fill(nothing), stroke("black"));
img = SVG("sierp1.svg", 10cm, 8.66cm); draw(img,cx1);

cx3 = compose(sierpinski(3), linewidth(0.2mm),fill(nothing), stroke("black"));
img = SVG("sierp3.svg", 10cm, 8.66cm); draw(img,cx3);

cx5 = compose(sierpinski(5), linewidth(0.2mm),fill(nothing), stroke("black"));
img = SVG("sierp5.svg", 10cm, 8.66cm); draw(img,cx5);

using PGFPlots

p = Axis( [ Plots.Linear(x -> sin.(3x) .* exp.(-0.3x), (0,8),
                         legendentry = L"$\sin(3x)*exp(-0.3x)$"),
            Plots.Linear(x -> sqrt.(x) ./ (1+x.^2), (0,8),
                         legendentry = L"$\sqrt{2x}/(1+x^2)$")
        ])

# This requires installation of the pdf2svg utility
save("linear-plots.svg", p);

fq  =  randn(10000);
p = Axis(Plots.Histogram(fq, bins=100), ymin=0)

save("histogram-plot.svg", p);


module NMF_EM

include("../Misc.jl")

using .Misc
using Distributions, SpecialFunctions

export standard_EM, dual_EM

function standard_EM(X::Array{ℜ,2}, K::Ƶ; μ::ℜ=nanmean(X), γ::ℜ=0.1, EPOCHS::Ƶ=1, ϵ::ℜ=1e-16) where {ℜ<:Real, Ƶ<:Int}
    I::Ƶ, J::Ƶ = size(X)

    a::ℜ, b::ℜ = I*J*γ, γ/μ
    α_K, α_JK, α_IK = fill(a/K,K), fill(a/(K*J),J,K), fill(a/(K*I),I,K)

    log_λ = log(rand(Gamma(a,1.0/b)))
    log_θ_K = log.(rand(Dirichlet(α_K)))
    log_θ_JK = reshape([log(θ_jk) for k=1:K for θ_jk=rand(Dirichlet(α_JK[:,k]))],J,K)
    log_θ_IK = reshape([log(θ_ik) for k=1:K for θ_ik=rand(Dirichlet(α_IK[:,k]))],I,K)

    S_K, S_JK, S_IK, S⁺::ℜ = zeros(ℜ,K), zeros(ℜ,J,K), zeros(ℜ,I,K), 0.0
    s, p, log_ρ, log_ρ₊ = Array{ℜ}(undef,K), Array{ℜ}(undef,K), Array{ℜ}(undef,K), 0.0

    X_full = similar(X)
    for eph=1:EPOCHS
        S_K .= α_K
        S_JK .= α_JK
        S_IK .= α_IK
        S⁺ = a

        for j=1:J, i=1:I #order of traversal is important
            log_ρ .= log_λ .+ log_θ_K .+ log_θ_JK[j,:] .+ log_θ_IK[i,:]
            log_ρ₊ = logsumexp(log_ρ)
            p .= exp.(log_ρ .- log_ρ₊)

            X_full[i,j] = isnan(X[i,j]) ? exp(log_ρ₊) : X[i,j]
            s .= X_full[i,j] .* p

            S_JK[j,:] .+= s
            S_IK[i,:] .+= s
            S_K[:] .+= s
            S⁺ += X_full[i,j]
        end

        log_λ = log(max(S⁺ - 1.0, ϵ)) - log(b+1.0)
        log_θ_K .= log.(max.(S_K .- 1.0, ϵ))
        log_θ_JK .= log.(max.(S_JK .- 1.0, ϵ))
        log_θ_IK .= log.(max.(S_IK .- 1.0, ϵ))

        log_θ_K .-= logsumexp(log_θ_K;dims=1)
        log_θ_JK .-= logsumexp(log_θ_JK;dims=1)
        log_θ_IK .-= logsumexp(log_θ_IK;dims=1)
    end
    return X_full, log_λ, (log_θ_K, log_θ_IK, log_θ_JK)
end

function dual_EM(X::Array{ℜ,2}, K::Ƶ; μ::ℜ=nanmean(X), γ::ℜ=0.1, EPOCHS::Ƶ=1, ϵ::ℜ=1e-16) where {ℜ<:Real, Ƶ<:Int}
    I::Ƶ, J::Ƶ = size(X)

    a::ℜ, b::ℜ = I*J*γ, γ/μ
    α_K, α_JK, α_IK = fill(a/K,K), fill(a/(K*J),J,K), fill(a/(K*I),I,K)

    log_λ = log(rand(Gamma(a,1.0/b)))
    log_θ_K = log.(rand(Dirichlet(α_K)))
    log_θ_JK = reshape([log(θ_jk) for k=1:K for θ_jk=rand(Dirichlet(α_JK[:,k]))],J,K)
    log_θ_IK = reshape([log(θ_ik) for k=1:K for θ_ik=rand(Dirichlet(α_IK[:,k]))],I,K)

    S_K, S_JK, S_IK, S⁺::ℜ = zeros(ℜ,K), zeros(ℜ,J,K), zeros(ℜ,I,K), 0.0
    s, p, log_ρ, log_ρ₊ = Array{ℜ}(undef,K), Array{ℜ}(undef,K), Array{ℜ}(undef,K), 0.0

    X_full = similar(X)
    for eph=1:EPOCHS
        S_K .= α_K
        S_JK .= α_JK
        S_IK .= α_IK
        S⁺ = a

        for j=1:J, i=1:I #order of traversal is important
            log_ρ .= log_λ .+ log_θ_K .+ log_θ_JK[j,:] .+ log_θ_IK[i,:]
            log_ρ₊ = logsumexp(log_ρ)
            p .= exp.(log_ρ .- log_ρ₊)

            X_full[i,j] = isnan(X[i,j]) ? floor(exp(log_ρ₊)) : X[i,j]
            s .= finucan(X_full[i,j],p)

            S_JK[j,:] .+= s
            S_IK[i,:] .+= s
            S_K[:] .+= s
            S⁺ += X_full[i,j]
        end

        log_λ = digamma(S⁺+a) - log(b+1.0)
        log_θ_K .= digamma.(S_K.+α_K) .- digamma(S⁺+a)
        log_θ_JK .= digamma.(S_JK.+α_JK) .- digamma.(S_K.+α_K)'
        log_θ_IK .= digamma.(S_IK.+α_IK) .- digamma.(S_K.+α_K)'
    end
    return X_full, S⁺, (S_K, S_IK, S_JK)
end

end


import Test: @test, @testset, @test_broken
import Luna: Maths, Nonlinear, PhysData
import FFTW

# test full / nothg / envelope kerr effect
Nt = collect(range(0, length=2^16))
t = @. (Nt - 2^16/2)*3.430944979182369e-16/4
E = exp.(-0.5.*(t./10e-15).^2).*cos.(2π*PhysData.c/800e-9.*t)
kerrfield = Nonlinear.Kerr_field(1.0)
outf = similar(E)
fill!(outf, 0.0)
kerrfield(outf, E, 1.0)
kerrfieldn = Nonlinear.Kerr_field_nothg(1.0, 2^16)
outn = similar(E)
fill!(outn, 0.0)
kerrfieldn(outn, E, 1.0)
kerrfieldenv = Nonlinear.Kerr_env(1.0)
Eenv = Maths.hilbert(E)
oute = similar(Eenv)
fill!(oute, 0.0)
kerrfieldenv(oute, Eenv, 1.0)

outfω = FFTW.rfft(outf)
outnω = FFTW.rfft(outn)
outeω = FFTW.rfft(real.(oute))
# we compare only low (non THG) frequencies
# note that these are not expected to be exact, because we have dropped not just the THG term
# but also cross terms between positive and negative frequencies
@test isapprox(abs.(outnω[1800:2400]), abs.(outfω[1800:2400]), rtol=1e-15)
@test isapprox(abs.(outeω[1800:2400]), abs.(outfω[1800:2400]), rtol=1e-15)


input = joinpath(@__DIR__, "input")
lines = readlines(input)

function cubes(lines)
    acc = []
    for line in lines
        mode = startswith(line, "on") ? true : false
        l, r = parse.(Int, match(r"x=([-]?\d+)..([-]?\d+)", line).captures)
        xi = min(l, r)
        xa = max(l, r)
        l, r = parse.(Int, match(r"y=([-]?\d+)..([-]?\d+)", line).captures)
        yi = min(l, r)
        ya = max(l, r)
        l, r = parse.(Int, match(r"z=([-]?\d+)..([-]?\d+)", line).captures)
        zi = min(l, r)
        za = max(l, r)

        cube = (mode, xi:xa, yi:ya, zi:za)

        push!(acc, cube)
    end
    return acc
end

function boundcheck(cube)
    (_, x, y, z) = cube
    return issubset(x, -50:50) && issubset(y, -50:50) && issubset(z, -50:50)
end

function flip(switches)
    acc = Set()
    for cube in switches
        (mode, x, y, z) = cube
        for switch in collect(Iterators.product(x, y, z))
            if mode
                push!(acc, switch)
            else
                delete!(acc, switch)
            end
        end
    end
    return acc
end

function expand(cubes)
    acc = []
    for cube in cubes
        (emode, ex, ey, ez) = cube
        nacc = []
        for (nmode, nx, ny, nz) in acc
            x = intersect(ex, nx)
            y = intersect(ey, ny)
            z = intersect(ez, nz)
            if length(x) > 0 && length(y) > 0 && length(z) > 0
                push!(nacc, (!nmode, x, y, z))
            end
        end
        if emode
            push!(nacc, cube)
        end
        append!(acc, nacc)
    end
    return acc
end

function value(cube)
    (mode, x, y, z) = cube
    n = length(x) * length(y) * length(z)
    return mode ? n : -1 * n
end

cs = cubes(lines)
cs = filter(boundcheck, cs)
p1 = length(flip(cs))

cs = cubes(lines)
cs = expand(cs)
p2 = sum(value, cs)

println("-----------------------------------------------------------------------")
println("reactor reboot -- part one :: $p1")
println("reactor reboot -- part two :: $p2")
println("-----------------------------------------------------------------------")

@assert(p1 == 567496)
@assert(p2 == 1355961721298916)


using DrWatson
@quickactivate "EPExperiments"

using EPExperiments
using AnnealedIS

using Distributions
using Random
using StatsFuns: logsumexp
using LinearAlgebra: dot
using Logging
using LoggingExtras
using StatsPlots

const LOG_FILENAME = "out.log"

function is_sampling(logjoint, qi, qb, num_samples)
    qi_samples = rand(qi, num_samples)
    qb_samples = rand(qb, num_samples)

    i_b_samples = map(t -> (i₀ = t[1], β = t[2]), zip(qi_samples, qb_samples))
    is_weights = logjoint.(i_b_samples) - (
        logpdf.(qi, qi_samples) .+ logpdf.(qb, qb_samples)
    )
    normalised_is_weights = is_weights .- logsumexp(is_weights)

    return normalised_is_weights, qi_samples, qb_samples
end

@with_kw struct ISGroundTruthConfig
    num_samples::Int = Int(1e6)
    qi::Distributions.Distribution = truncated(Normal(10, 10), 0, 1_000)
    qb::Distributions.Distribution = truncated(Normal(1.5, 1.5), 0, Inf)
    cost_fn_name::Symbol = :cost_fn_default
    data_filename::String = "neg_binomial_data.jld2"
    seed::Int = 1234
end

DrWatson.allignore(::ISGroundTruthConfig) = (:data_filename,)

function calculate_ground_truth(is_config)
    @unpack num_samples, qi, qb, cost_fn_name, data_filename, seed = is_config
    cost_fn = COST_FUNCTIONS[cost_fn_name]

    Random.seed!(seed)

    result_folder = mkpath(joinpath(
        datadir("is_ground_truth"),
        savename(is_config)
    ))
    logger = TeeLogger(
        ConsoleLogger(),
        FileLogger(joinpath(result_folder, LOG_FILENAME))
    )

    input_data = load(datadir("sims", is_config.data_filename))
    @unpack tmax, total_population, i_0_true, β_true, γ = input_data["params"]
    Y = input_data["Y"]

    tm = bayes_sir(Y, total_population, γ, cost_fn).gamma2
    vi = Turing.VarInfo(tm)
    joint_dens = AnnealedIS.make_log_joint_density(tm, vi)

    @time begin
    log_weights, qi_samples, qb_samples = is_sampling(
        joint_dens, qi, qb, num_samples
    )
    end

    costs = cost_fn.(qb_samples, γ)

    ess = 1 / exp(logsumexp(2 * log_weights))
    target_weights = log_weights .+ log.(costs)
    ess_target = exp(2 * logsumexp(target_weights) - logsumexp(2 * target_weights))

    exp_estimate = dot(exp.(log_weights), costs)
    with_logger(logger) do
        @info "Proposals:" qi qb
        @info "Number of samples: $num_samples"
        @info "ESS p(x|y): $ess"
        @info "ESS p(x|y)f(x): $ess_target"
        @info "Expectation estimate: $exp_estimate"
    end

    @tagsave(joinpath(result_folder, "results.jld2"), Dict(
        "log_weights" => log_weights,
        "qi_samples" => qi_samples,
        "qb_samples" => qb_samples,
        "is_config" => is_config,
        "expectation_estimate" => exp_estimate
    ))

    savefig(
        plot_is_samples(qb_samples, log_weights),
        joinpath(result_folder, "beta_samples.png")
    )
end

calculate_ground_truth(ISGroundTruthConfig(
    num_samples=Int(1e8),
    cost_fn_name=:cost_fn_shift20
))


"""
`module JAC.Math`
    ... a submodel of JAC that contains various (adopted) procedures from different libraries and sources.
"""
module Math

  using GSL
  using JAC, ..Radial

  """
  `Math.finiteDifferenceWeights(position, npoints; order=1)`
    ... computes the weights for a n-point finite-difference approximation. position=0 corresponds to the centered stencil,
        while positive and negative values refer to the non-symmetric stencils, respecitvely. The order specifies for which
        derivative the coefficients are computed, the default ist the first derivative.
  """
  function finiteDifferenceWeights(position, npoints; order=1)

      if !isodd(npoints)
          error("Number of points must be odd")
      end

      x = collect(-Int64((npoints-1)/2):Int64((npoints-1)/2))
      c = Array{Float64}(undef, order+1, npoints)

      c1 = 1
      c4 = x[1] - position
      c[1, 1] = 1

      for i = 2:npoints
          mn = min(i, order+1)
          c2 = 1
          c5 = c4
          c4 = x[i] - position
          for j = 1:i-1
              c3 = x[i] - x[j]
              c2 = c2 * c3
              if j == i - 1
                  c[2:mn, i] = c1 / c2 * ( collect(1:mn -1 ) .* c[1:mn-1, i-1] - c5 * c[2:mn, i-1])
                  c[1, i] = -c1 * c5 * c[1, i-1]/c2
              end
              c[2:mn, j] = 1 / c3 * (c4 * c[2:mn, j] - collect(1:mn-1) .* c[1:mn-1, j])
              c[1, j] = c4 * c[1, j] / c3
          end
          c1 = c2
      end

      return( c )
   end

  """
  `Math.derivative(input::Array{Float64}, i::Int64)`
    ... computes the derivative of a given input array at position i; a Float64 is returned. An equally spaced grid is
        assumed, the result must be multiplied by the inverse lattice spacing h manually. If an exponential grid is
        used, the transformation must be applied too.
  """
  function derivative(input::Array{Float64}, i::Int64) :: Float64

      n = JAC.Defaults.FINITE_DIFFERENCE_NPOINTS
      result ::Float64 = 0.

      if i <= n

        for j = 1:2*n + 1
          result += JAC.Defaults.weights[i, j] * input[j]
        end
        return result
      end

      if size(input, 1) - i < n

        for j = 1:2*n + 1
          result += JAC.Defaults.weights[2*n + 1 + i - size(input, 1), j] * input[size(input, 1) - 2*n - 1 + j]
        end

        return result
      end

      for j = i - n:i+n
        result += JAC.Defaults.weights[n + 1, j - i + n + 1] * input[j]
      end

      return result
  end

  """
  `Math.integrateFitTransform(f::Function, mtp::Int64, grid::Radial.Grid)`
    ... to integrate a (point-wise) given function over mtp grid points on an exponential grid. The transformation of the
        exponential grid is automatically performed such that no modifications to the integrand are necessary; a value::Float64
        is returned.
  """
  function integrateFitTransform(f::Function, mtp::Int64, grid::Radial.Grid)
      n      = size(JAC.Defaults.newtonCotesCoefficients, 1)
      newtonCotesCoefficientsOnGrid = copy(JAC.Defaults.newtonCotesCoefficients) * grid.h
      i0     = 2
      gamma  = 0.
      result = 0.

      while abs(f(i0)) == 0 || f(i0)/f(i0+1) < 0
          i0 += 1
      end
      gamma = log(f(i0)/f(i0+1)) / log(grid.r[i0]/grid.r[i0+1])
      result = 1/(gamma + 1) * grid.r[i0] * f(i0)

      result += f(i0) * grid.rp[i0] * newtonCotesCoefficientsOnGrid[n]
      for i = i0+1:mtp
          result += newtonCotesCoefficientsOnGrid[(i - i0) % (n - 1) + 1] * f(i) * grid.rp[i]
      end

      if (mtp - i0) % (n - 1) + 1 == 1
          result -= newtonCotesCoefficientsOnGrid[n] * f(mtp) * grid.rp[mtp]
      end

    return( result )
  end

  """
  `Math.integrateFit(f::Function, mtp::Int64, grid::Radial.Grid)`  ... to integrate a (point-wise) given function over mtp grid points
       on an exponential grid. This function assumes that the ingegrand was already transformed to the exponential grid by multiplying it with
       the derivative of the transformation function; a value::Float64 is returned.
  """
  function integrateFit(f::Function, mtp::Int64, grid::Radial.Grid)
      n      = size(JAC.Defaults.newtonCotesCoefficients, 1)
      newtonCotesCoefficientsOnGrid = copy(JAC.Defaults.newtonCotesCoefficients) * grid.h

      i0     = 2
      gamma  = 0.
      result = 0.

      while abs(f(i0)) == 0 || f(i0)/f(i0+1) < 0
          i0 += 1
      end
      gamma = log(f(i0)/f(i0+1) * grid.rp[i0+1] / grid.rp[i0]) / log(grid.r[i0]/grid.r[i0+1])
      result = 1/(gamma + 1) * grid.r[i0] * f(i0) / grid.rp[i0]

      result += f(i0) * newtonCotesCoefficientsOnGrid[n]
      for i = i0+1:mtp
          result += newtonCotesCoefficientsOnGrid[(i - i0) % (n - 1) + 1] * f(i)
      end

      if (mtp - i0) % (n - 1) + 1 == 1
          result -= newtonCotesCoefficientsOnGrid[n] * f(mtp)
      end

      return result
  end

  """
  `Math.integrateTransform(f::Function, n0::Int64, mtp::Int64, grid::Radial.Grid)`  ... to integrate a (point-wise) given function over
       mtp grid points ...; a value::Float64 is returned.
  """
  function integrateTransform(f::Function, n0::Int64, mtp::Int64, grid::Radial.Grid)
      newtonCotesCoefficientsOnGrid             = Array{Float64}([2 * 7, 32, 12, 32, 7]) * 2 / 45  # * grid.h
      ## newtonCotesCoefficientsOnGrid            = Array{Float64}([2 * 41, 216, 27, 272, 27, 216, 41]) * 1 / 140
      ## newtonCotesCoefficientsOnGrid             = Array{Float64}([2 * 989, 5888, -928, 10496, -4540, 10496, -928, 5888, 989]) * 4 / 14175
      newtonCotesCoefficientsOnGrid *= grid.h
      n      = size(newtonCotesCoefficientsOnGrid, 1)

      i0     = n0
      gamma  = 0.
      result = 0.

      result += f(i0) * grid.rp[i0] * newtonCotesCoefficientsOnGrid[n]
      for i = i0+1:mtp
          result += newtonCotesCoefficientsOnGrid[(i - i0) % (n - 1) + 1] * f(i) * grid.rp[i]
      end

      if (mtp - i0) % (n - 1) + 1 == 1
          ##x result -= newtonCotesCoefficientsOnGrid[n] * f(mtp)
          result -= newtonCotesCoefficientsOnGrid[n] * f(mtp) * grid.rp[mtp]
      end

      return( result )
  end

  """
  `Math.polylogExp(x::Float64, s::Int64)`  ... to evaluate the polylog(s, -exp(x)) function.
  """
  function polylogExp(x::Float64, s::Int64)
      return -sf_fermi_dirac_int(s-1, x)
  end

end # module


#####
##### helper functions
#####

"function wavespeed_1D(rho,rhou,E)
    one-dimensional wavespeed (for DG penalization terms)"
function wavespeed_1D(rho,rhou,E)
    p = pfun_nd(rho,(rhou,),E)
    cvel = @. sqrt(γ*p/rho)
    return @. abs(rhou/rho) + cvel
end

#####
##### one-dimensional fluxes
#####

"function euler_fluxes_1D(rhoL,uL,betaL,rhologL,betalogL,
                        rhoR,uR,betaR,rhologR,betalogR)"
function euler_fluxes_1D(rhoL,uL,betaL,rhologL,betalogL,
                        rhoR,uR,betaR,rhologR,betalogR)

    rholog = logmean.(rhoL,rhoR,rhologL,rhologR)
    betalog = logmean.(betaL,betaR,betalogL,betalogR)

    # arithmetic avgs
    rhoavg = (@. .5*(rhoL+rhoR))
    uavg   = (@. .5*(uL+uR))

    unorm = (@. uL*uR)
    pa    = (@. rhoavg/(betaL+betaR))
    f4aux = (@. rholog/(2*(γ-1)*betalog) + pa + .5*rholog*unorm)

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. f4aux*uavg)

    return (FxS1,FxS2,FxS3)
end

#####
##### shim functions for 1D
#####

# dispatch to n-dimensional constitutive routines, with optional entropy scaling
function primitive_to_conservative(rho,u,p)
   rho,rhou,E = primitive_to_conservative_nd(rho,tuple(u),p)
   return rho,first(rhou),E
end
function v_ufun(rho,rhou,E)
    v1,vU,vE = v_ufun_nd(rho,tuple(rhou),E)
    return scale_entropy_output(v1,first(vU),vE)
end
function u_vfun(v1,vU,vE)
    v1,vU,vE = scale_entropy_input(v1,vU,vE)
    rho,rhoU,E = u_vfun_nd(v1,tuple(vU),vE)
    return rho,first(rhoU),E
end
function conservative_to_primitive_beta(rho,rhou,E)
    rho,U,beta = conservative_to_primitive_beta_nd(rho,tuple(rhou),E)
    return rho,first(U),beta
end

Sfun(rho,rhou,E) = Sfun_nd(rho,tuple(rhou),E)
pfun(rho,rhou,E) = pfun_nd(rho,tuple(rhou),E)
betafun(rho,rhou,E) = betafun_nd(rho,tuple(rhou),E)

"function euler_fluxes_1D(rhoL,uL,betaL,rhoR,uR,betaR)"
function euler_fluxes(rhoL,uL,betaL,rhoR,uR,betaR)
    rhologL,betalogL,rhologR,betalogR = map(x->log.(x),(rhoL,betaL,rhoR,betaR))
    return euler_fluxes_1D(rhoL,uL,betaL,rhologL,betalogL,
                           rhoR,uR,betaR,rhologR,betalogR)
end
euler_fluxes(rhoL,uL,betaL,rhologL,betalogL,rhoR,uR,betaR,rhologR,betalogR) =
    euler_fluxes_1D(rhoL,uL,betaL,rhologL,betalogL,rhoR,uR,betaR,rhologR,betalogR)

    euler_fluxes(varL,varR) = euler_fluxes(varL...,varR...)
    euler_fluxes(varL,varR,logL,logR) = euler_fluxes(varL...,logL...,varR...,logR...)


@doc doc"""
    TranslationGroup{T<:Tuple,F} <: GroupManifold{Euclidean{T,F},AdditionOperation}

Translation group $\mathrm{T}(n)$ represented by translation arrays.

# Constructor
    TranslationGroup(n₁,...,nᵢ; field = 𝔽)

Generate the translation group on
$𝔽^{n₁,\dots,nᵢ}$ = `Euclidean(n₁,...,nᵢ; field = 𝔽)`, which is isomorphic to group itself.
"""
const TranslationGroup{T<:Tuple,F} = GroupManifold{Euclidean{T,F},AdditionOperation}

function TranslationGroup(n::Int...; field::AbstractNumbers = ℝ)
    return TranslationGroup{Tuple{n...},field}(
        Euclidean(n...; field = field),
        AdditionOperation(),
    )
end

function show(io::IO, ::TranslationGroup{N,F}) where {N,F}
    print(io, "TranslationGroup($(join(N.parameters, ", ")); field = $(F))")
end


using Test
using ThreeBodyDecay

@testset "Circular permutations" begin
    m1 = 0.938; m2 = 0.49367; m3 = 0.13957; m0 = 2.46867
    tbs = ThreeBodySystem(m1,m2,m3,m0=m0)
    #
    τ1 = (sum(lims1(tbs.ms))/2.0, 0.3, 0.3, 0.3, 0.3)
    τ3  = change_basis_3from1(τ1, tbs.ms)
    # @show [τ1...,        tbs.msq...,    tbs.s]
    τ2  = change_basis_2from3(τ3, tbs.ms)
    τ1p = change_basis_1from2(τ2, tbs.ms)
    @test sum(τ1p .≈ τ1) == 5

    #
    @test τ1[1] ≈ σ1of2(τ2[4],τ2[1],tbs.ms^2)
    @test τ3[1] ≈ σ3of1(τ1[4],τ1[1],tbs.ms^2)
    @test τ2[1] ≈ σ2of3(τ3[4],τ3[1],tbs.ms^2)
    #
    # implementation cosθ12
    σs = (τ1[1],τ2[1],τ3[1])
    @test τ3[4] ≈ cosθ12(σs,tbs.ms^2)
    @test τ1[4] ≈ cosθ23(σs,tbs.ms^2)
    @test τ2[4] ≈ cosθ31(σs,tbs.ms^2)
    #
end


# all the plots for my presentation in IOR2017, Stavanger
using Plots
# include("rel_perms.jl")

# plot the fractional flow functions for a water-wet and an oil-wet system
include("rel_perms_real.jl")
# Water-wet Mixed-wet Oil-wet
# Swi 0.2 0.2 0.2
# Sor 0.1 0.2 0.3
# Ew=Krw(Sor) 0.6 0.8 0.8
# Eo=Kro(Swi) 0.8 0.8 0.5
# Nw 3 2 3
# No 2 2 3
krw0_ww = 0.33
krw0_ow = 0.28
kro0_ww = 0.79
kro0_ow = 0.72
nw_ww = 2.4
nw_ow= 2.4
no_ww = 2.0
no_ow= 1.6
sor_ww=0.145
sor_ow=0.37
sor_mw=0.08
swc_ww=0.08
swc_ow=0.08
teta_ow=deg2rad(180)
teta_ww=deg2rad(0)
teta_mw=deg2rad(90)
teta_prime=teta_mw-π/2 # see my report
a=(sor_mw-0.5*(sor_ww-sor_ow)*cos(teta_mw)-0.5*(sor_ww+sor_ow))/(cos(teta_mw)^2-1)
b=0.5*(sor_ww-sor_ow)
c=0.5*(sor_ww+sor_ow)-a

SF=0.0 # 1.0 is water-wet, 0.0 is oil-wet
sw_plot_ow=linspace(swc_ow,1-sor_ow,100)
sw_plot_ww=linspace(swc_ww,1-sor_ww,100)

plot(size=(500,400), xtickfont = font(12, "Courier"), ytickfont=font(12, "Courier"),
  ylabel="krw, kro", xlabel="Water saturation", legendfont=font(12, "Courier"),
  guidefont=font(12, "Courier"))
plot!(sw_plot_ow, krw.(sw_plot_ow, krw0_ow, sor_ow, swc_ow, nw_ow),
  xlims=(0.0,1.0), linewidth=3, linestyle=:dash, linecolor=:red,
  label="krw, Oil-wet")
plot!(sw_plot_ow, kro.(sw_plot_ow, kro0_ow, sor_ow, swc_ow, no_ow),
  xlims=(0.0,1.0), linewidth=3, linestyle=:dash, linecolor=:red,
  label="kro, Oil-wet")
plot!(sw_plot_ww, krw.(sw_plot_ww, krw0_ww, sor_ww, swc_ww, nw_ww),
  xlims=(0.0,1.0), linewidth=1, linestyle=:solid, linecolor=:blue,
  label="krw, Water-wet")
plot!(sw_plot_ww, kro.(sw_plot_ww, kro0_ww, sor_ww, swc_ww, no_ww),
  xlims=(0.0,1.0), linewidth=1, linestyle=:solid, linecolor=:blue,
  label="kro, Water-wet")
savefig("rel_perms_fathi.png")


"""
    firstInteractionStep!(r_j, mu_sum, a)

Calculates accelerations that are defined in the first step of the ``H_{\\mathrm{Interaction}}`` Hamiltonian.

```math
\\mathcal{H}_{\\mathrm{Interaction}_{1}}(\\textbf{q}, \\textbf{p}) = \\sum_{i=1}^{N} \\frac{\\sum_{i=2}^{N} \\frac{Gm'_{i}m_{sum_{i}}}{\\mid \\textbf{r'}_{i}\\mid} \\\\
```

# Args

* `r_j`: NxD matrix where N is the number of bodies and D is the number of dimensions, it contains the positions of the bodies in jacobi coordinates
* `mu_sum`: array of N elements, where the standard gravitational parameter of the bodies from 1 to i is saved in the i-th element (µ_sum)
* `a`: NxD matrix in which the accelerations defined by the first step of the interaction between bodies are saved.
"""
function firstInteractionStep!(r_j, mu_sum, a)

    for i in 2:size(r_j)[1]
        a[i,:] = (mu_sum[i] * r_j[i,:]) / norm(r_j[i,:])^3
    end

end

"""
    secondInteractionStep!(r, m, m_sum, m_j, mu, dp, dp_sum, a)

Calculates accelerations that are defined in the second step of the ``\\mathcal{H}_{Interaction_{2}}`` Hamiltonian.

```math
\\mathcal{H}_{Interaction_{2}}(\\textbf{q}, \\textbf{p}) = - \\sum_{i=1}^{N} \\sum_{j=i+1}^{N} \\frac{Gm_{i}m_{j}}{\\mid \\textbf{q}_{i} - \\textbf{q}_{j} \\mid} \\\\
```

# Args

* `r_j`: NxD matrix where N is the number of bodies and D is the number of dimensions, it contains the positions of the bodies in jacobi coordinates
* `m`: array of N elements, where the sum of the masses of the bodies from 1 to i is saved in the i-th element.
* `m_sum`: array of N elements, where the sum of the masses of the bodies from 1 to i is saved in the i-th element.
```math
m_{sum_{i}} = \\sum_{j = 1}^{i} m_{j} \\\\
```
* `m_j`: array of N elements, each containing the jacobi mass (m').
```math
m'_{1} = m_{sum_{N}} \\\\
m'_{i} = m_{i} \\frac{m_{sum_{i-1}}}{m_{sum_{i}}} \\\\
\\mathrm{,where } \\hspace{0.5cm} m_{sum_{i}} = \\sum_{j = 1}^{i} m_{j}
```
* `mu`: array of N elements, each containing the standard gravitational parameter of each body.
```math
mu_{i} = G \\cdot m_{j} \\\\
```
* `dp`: auxiliary array of D elements that is defined previously in order to decrease execution time and unneccesary allocations.
* `dp_sum`: NxD matrix that is defined previously in order to decrease execution time and unneccesary allocations.
* `a`: NxD matrix in which the accelerations defined by the first step of the interaction between bodies are saved.
"""
function secondInteractionStep!(r, m, m_sum, m_j, mu, dp, dp_sum, a)

    for i in 1:size(r)[1]
        for j in 1:length(dp)
            dp[j] = 0.0
        end
        for j in 1:size(r)[1]
            if i != j
                r_coor = r[i,:] - r[j,:]
                dp -= (mu[j]*m[i]*r_coor / (norm(r_coor)^3)) # dp[i] / dt
            end
        end

        if i == 1
            dp_sum[1,:] = dp
        else
            dp_sum[i,:] = dp + dp_sum[i-1,:]
            a[i,:] = ((m_sum[i-1]*dp - m[i]*dp_sum[i-1,:]) / m_sum[i]) / m_j[i] # azelerazioak jacobira pasa
        end
    end

    a[1,:] = dp_sum[size(r)[1],:] / m_j[1]

end


# ===============================
# Written by AAE
# TU Delft, Spring 2014
# simulkade.com
# ===============================

function transientTerm(phi_old::CellValue, dt::Real, alfa::Real)
transientTerm(phi_old, dt, alfa.*ones(Float64,tuple(phi_old.domain.dims...)))
end

function transientTerm(phi_old::CellValue, dt::Real)
transientTerm(phi_old, dt, 1.0)
end

function transientTerm(phi_old::CellValue, dt::Real, alfa::CellValue)
d = phi_old.domain.dimension
if d==1 || d==1.5
  transientTerm1D(phi_old, dt, alfa.value[2:end-1])
elseif d==2 || d==2.5 || d==2.8
  transientTerm2D(phi_old, dt, alfa.value[2:end-1,2:end-1])
elseif d==3 || d==3.2
  transientTerm3D(phi_old, dt, alfa[2:end-1,2:end-1,2:end-1])
end
end

function transientTerm(phi_old::CellValue, dt::Real, alfa::Array{T}) where T<:Real
d = phi_old.domain.dimension
if d==1 || d==1.5
  transientTerm1D(phi_old, dt, alfa)
elseif d==2 || d==2.5 || d==2.8
  transientTerm2D(phi_old, dt, alfa)
elseif d==3 || d==3.2
  transientTerm3D(phi_old, dt, alfa)
end

end

function transientTerm1D(phi_old::CellValue,
		    dt::Real, alfa::Array{T}) where T<:Real
# returns the matrix and RHS for a d(phi)/dt term

# extract data from the mesh structure
Nx = phi_old.domain.dims[1]
G = [1:Nx+2;]

# rearrange the matrix of k and build the sparse matrix for internal cells
row_index = reshape(G[2:Nx+1],Nx) # main diagonal (only internal cells)
AP_diag = reshape(alfa/dt,Nx)
M = sparse(row_index, row_index, AP_diag, Nx+2, Nx+2)

# define the RHS Vector
RHS = zeros(Nx+2)

# assign the values of the RHS vector
RHS[row_index] = reshape(alfa.*phi_old.value[2:Nx+1]/dt,Nx)

(M, RHS)

end

function transientTerm2D(phi_old::CellValue,
		    dt::Real, alfa::Array{T}) where T<:Real
# returns the matrix and RHS for a d(phi)/dt term

# extract data from the mesh structure
Nx = phi_old.domain.dims[1]
Ny = phi_old.domain.dims[2]
G=reshape([1:(Nx+2)*(Ny+2);], Nx+2, Ny+2)

# rearrange the matrix of k and build the sparse matrix for internal cells
row_index = reshape(G[2:Nx+1,2:Ny+1],Nx*Ny) # main diagonal (only internal cells)
AP_diag = reshape(alfa/dt,Nx*Ny)
M = sparse(row_index, row_index, AP_diag, (Nx+2)*(Ny+2), (Nx+2)*(Ny+2))

# define the RHS Vector
RHS = zeros((Nx+2)*(Ny+2))

# assign the values of the RHS vector
RHS[row_index] = reshape(alfa.*phi_old.value[2:Nx+1,2:Ny+1]/dt,Nx*Ny)

(M, RHS)

end

function transientTerm3D(phi_old::CellValue,
		    dt::Real, alfa::Array{T}) where T<:Real
# returns the matrix and RHS for a d(phi)/dt term

# extract data from the mesh structure
Nx = phi_old.domain.dims[1]
Ny = phi_old.domain.dims[2]
Nz = phi_old.domain.dims[3]
G=reshape([1:(Nx+2)*(Ny+2)*(Nz+2);], Nx+2, Ny+2, Nz+2)

# rearrange the matrix of k and build the sparse matrix for internal cells
row_index = reshape(G[2:Nx+1,2:Ny+1,2:Nz+1],Nx*Ny*Nz) # main diagonal (only internal cells)
AP_diag = reshape(alfa/dt,Nx*Ny*Nz)
M = sparse(row_index, row_index, AP_diag, (Nx+2)*(Ny+2)*(Nz+2), (Nx+2)*(Ny+2)*(Nz+2))

# define the RHS Vector
RHS = zeros((Nx+2)*(Ny+2)*(Nz+2))

# assign the values of the RHS vector
RHS[row_index] = reshape(alfa.*phi_old.value[2:Nx+1,2:Ny+1,2:Nz+1]/dt,Nx*Ny*Nz)

(M, RHS)

end


"""
    log_transformer(x)

Internal function used to get the exponent of x.
Used to select the nearest unit index.
"""
function log_transformer(x)
    return floor(Int, log10(x))
end

"""
    business_round(number::Number, names::Bool)

Internal function to round a number (num) to the nearest IG style.
"""
function business_round(number::Number, names::Bool)
    if number < 1000
        return string(number)
    else
        return extract_identifying_unit(number, names)
    end
end

"""
    number_groups(s::AbstractString)

Internal function to extract the two groups of numbers from a '.' separated string.
"""
function number_groups(s::AbstractString)
    group_1 = match(r"(\d+)", s)[1]
    group_2 = match(r"(\d+$)", s)[1]

    if length(group_2) > 1
        group_2 = group_2[1:2]
    else
        group_2 = group_2[1]
    end
    return group_1, group_2
end

"""
    extract_identifying_unit(num::Number, names::Bool)

Internal function to extract the identifying unit from a number.
"""
function extract_identifying_unit(num::Number, names::Bool)
    log_pos = log_transformer(num)
    div_3 = log_pos ÷ 3

    if names
        unit =  unit_names[div_3]
    else
        unit = units[div_3]
    end

    indentifier = string(num / ( BigInt(10) ^ (3 * div_3) ))
    main_int, rem_float = number_groups(indentifier)

    if names
        return string(main_int, ".", rem_float, " ", unit)
    else
        return string(main_int, ".", rem_float, unit)
    end
end

"""
    round(::Type{IGRound}, x::Number; names::Bool=false)

Function to round a number x to the nearest IG style.

# Arguments
 * x: The number to be rounded.
 * names: If true, the unit will be names, if false, the unit will be the number.

# Example

```julia-repl
julia> round(IGRound, 100_000; names=true)
"100Thousand"
```
"""
function round(::Type{IGRound}, x::Number; names::Bool=false)
    init_round =  Base.round(BigInt, x)
    return business_round(init_round, names)
end


using Poirot, Test, Random
using Poirot: Singleton

Random.seed!(0)

@test infer(() -> 1) == Singleton(1)
@test infer(() -> "foo") == Singleton("foo")

@test infer(() -> rand()) == Uniform(0, 1)
@test infer(() -> randn()) == Normal(0, 1)
@test infer(() -> rand(Bool)) == Bernoulli(1//2)

coin() = rand(Bool)

@test infer(() -> coin() & coin()) == Bernoulli(0.25)
@test infer(() -> coin() | coin()) == Bernoulli(0.75)

@test infer() do
  a = coin()
  b = coin()
  observe(a | b)
  a & b
end == Bernoulli(1/3)

infer() do
  a = coin()
  b = coin()
  observe(a | b)
  observe(!b)
  a
end == Singleton(true)

@test infer() do
  x = randn()
  observe(x > 0)
  x
end isa Poirot.Empirical

d = infer() do
  x = 10*rand()
  x^2 > 50
end

@test d isa Bernoulli
@test mean(d) ≈ (10-sqrt(50))/10

d = infer() do
  x = 10*rand()
  observe(x > 5)
  x^2 > 50
end

@test d isa Bernoulli
@test mean(d) ≈ (10-sqrt(50))/5

d = infer() do
  identical = rand(Bernoulli(1/3))
  boy1 = rand(Bernoulli(1/2))
  # TODO first version gives a different answer
  # boy2 = identical ? boy1 : rand(Bernoulli(1/2))
  boy2 = rand(Bernoulli(identical ? boy1*1.0 : 1/2))
  observe(boy1 & boy2)
  identical
end

@test mean(d) == 0.5


const c = Center()
const f = Face()

using Oceananigans.AbstractOperations: GridMetricOperation
import Oceananigans.Grids: solid_node, solid_interface

"""
    `solid_node` returns true only if a location is completely immersed
    `solid_interface` returns true if a location is partially immersed

    as an example (in a 1D immersed grid) :

     Immersed      Fluid
    ----------- ...........
   |     ∘     |     ∘
   f     c     f     c
  i-1   i-1    i     i

     `solid_interface(f, c, c, i, 1, 1, grid) = true`
     `solid_node(f, c, c, i, 1, 1, grid)      = false`

     `solid_node(Center(), Center(), Center(), args...) == solid_interface(Center(), Center(), Center(), args...)` as
     `Center(), Center(), Center()` can be only either fully immersed or not at all

     `solid_interface` is used in `GridMetricOperation` to assess the grid metric value.
     `metric = ifelse(solid_interface, 0.0, metric)` eliminating values inside the immersed domain and values lying on the
      immersed boundary from the reduction

      `is_immersed_boundary` returns true only if the interface has a solid and a fluid side (the actual immersed boundary)
      which is true only when `solid_node = false` and `solid_interface = true` (as the case of the face at `i` above)
"""

@inline solid_node(i, j, k, ibg::IBG) = is_immersed(i, j, k, ibg.grid, ibg.immersed_boundary) | solid_node(i, j, k, ibg.grid)

# Defining all the metrics for Immersed Boundaries

for LX in (:ᶜ, :ᶠ), LY in (:ᶜ, :ᶠ), LZ in (:ᶜ, :ᶠ)
    for dir in (:x, :y, :z), operator in (:Δ, :A)

        metric = Symbol(operator, dir, LX, LY, LZ)
        @eval begin
            import Oceananigans.Operators: $metric
            @inline $metric(i, j, k, ibg::IBG) = $metric(i, j, k, ibg.grid)
        end
    end

    volume = Symbol(:V, LX, LY, LZ)
    @eval begin
        import Oceananigans.Operators: $volume
        @inline $volume(i, j, k, ibg::IBG) = $volume(i, j, k, ibg.grid)
    end
end

@inline Δzᵃᵃᶜ(i, j, k, ibg::IBG) = Δzᵃᵃᶜ(i, j, k, ibg.grid)
@inline Δzᵃᵃᶠ(i, j, k, ibg::IBG) = Δzᵃᵃᶠ(i, j, k, ibg.grid)


using EntropyStableEuler
using Test
using StaticArrays

@testset "Logmean tests" begin
    uL,uR = 1,2
    @test logmean(uL,uR) == logmean(uL,uR,log(uL),log(uR))
    @test logmean(uL,uR) == logmean(uR,uL)
    @test logmean(uL,uL) ≈ uL
end

function init_prim(d)
    if d==1
        rho,u,p = 1,.1,2
        Q = (rho,u,p)
    elseif d==2
        rho,u,v,p = 1,.1,.2,2
        Q = (rho,u,v,p)
    elseif d==3
        rho,u,v,w,p = 1,.1,.2,.3,2
        Q = (rho,u,v,w,p)
    end
    return Q
end

@testset "Tests for d = $d" for d in (1:3)

    @testset "Entropy variable tests" begin

        U = prim_to_cons(Euler{d}(),init_prim(d))
        V = v_ufun(Euler{d}(),U)

        h = 1e-7
        central_diff(f,x) = (f(x+h) - f(x-h))/(2*h)
        swapentry(x,y,i) = (x[begin:i-1]...,y,x[i+1:end]...)
        for j = 1:d+2
            @test abs(V[j] - central_diff(x->Sfun(Euler{d}(),swapentry(U,x,j)),U[j])) < h
        end

        UV = u_vfun(Euler{d}(),V)
        @test all(UV .≈ U)
    end

    @testset "Symmetry" begin
        # test symmetry
        UL = prim_to_cons(Euler{d}(),init_prim(d))
        UR = prim_to_cons(Euler{d}(),init_prim(d).*1.1)
        FLR = fS(Euler{d}(),UL,UR)
        FRL = fS(Euler{d}(),UR,UL)
        @test all(FLR .≈ FRL)
    end

    @testset "Consistency" begin
        Q = init_prim(d)
        U = prim_to_cons(Euler{d}(),Q)
        F = fS(Euler{d}(),U,U)

        p = pfun(Euler{d}(),U)
        if d==1
            rho,rhou,E = U
            u = rhou./rho
            Fexact = SVector{3}(rho*u, rho*u^2 + p, u*(E+p))
        elseif d==2
            rho,rhou,rhov,E = U
            u,v = rhou./rho, rhov./rho
            Fx = SVector{4}(rho*u, rho*u^2 + p, rho*u*v,     u*(E+p))
            Fy = SVector{4}(rho*v, rho*u*v,     rho*v^2 + p, v*(E+p))
            Fexact = (Fx,Fy)
        elseif d==3
            rho,rhou,rhov,rhow,E = U
            u,v,w = rhou./rho, rhov./rho, rhow./rho
            Fx = SVector{5}(rho*u, rho*u^2 + p, rho*u*v,     rho*u*w,     u*(E+p))
            Fy = SVector{5}(rho*v, rho*u*v,     rho*v^2 + p, rho*v*w,     v*(E+p))
            Fz = SVector{5}(rho*w, rho*u*w,     rho*v*w,     rho*w^2 + p, w*(E+p))
            Fexact = (Fx,Fy,Fz)
        end

        @test all(F .≈ Fexact)
    end

    @testset "Entropy conservation property" begin
        γ = Euler{d}().γ

        UL = prim_to_cons(Euler{d}(),init_prim(d))
        UR = prim_to_cons(Euler{d}(),init_prim(d).*1.1)
        VL = v_ufun(Euler{d}(),UL)
        VR = v_ufun(Euler{d}(),UR)

        ψ(U) = (γ-1).*U[2:d+1]
        F = fS(Euler{d}(),UL,UR)
        vTF(VL,VR,F) = sum((VL .- VR).*F)
        for j = 1:d
            if d==1
                @test vTF(VL,VR,F) ≈ ψ(UL)[j]-ψ(UR)[j]
            else
                @test vTF(VL,VR,F[j]) ≈ ψ(UL)[j]-ψ(UR)[j]
            end
        end
    end

    @testset "Type stability tests" begin
        Q = init_prim(d)
        U = prim_to_cons(Euler{d}(),Q)
        V = v_ufun(Euler{d}(),U)

        @inferred prim_to_cons(Euler{d}(),Q)
        @inferred cons_to_prim_beta(Euler{d}(),U)
        @inferred v_ufun(Euler{d}(),U)
        @inferred u_vfun(Euler{d}(),V)
        @inferred fS(Euler{d}(),U,U)
    end
end


"""
    AbstractSpacing

Spacing discretization scheme supertype
"""
abstract type AbstractSpacing end

"""
    Uniform()

Uniform discretization scheme.
"""
struct Uniform <: AbstractSpacing end

"""
    Sine()

Sine-spaced discretization scheme.  Using sine-spacing on the right half of a
wing effectively results in cosine spacing once symmetry is applied.
"""
struct Sine <: AbstractSpacing end

"""
    Cosine()

Cosine-spaced discretization scheme.  This is typically one of the most accurate
spacing schemes for spanwise spacing.
"""
struct Cosine <: AbstractSpacing end

"""
   linearinterp(eta, rstart, rend)

Linearly interpolate between `rstart` and `rend` where eta is the fraction
between 0 (rstart) and 1 (rend)
"""
linearinterp(eta, rstart, rend) = (1-eta)*rstart + eta*rend

"""
    spanwise_spacing(n, spacing::AbstractSpacing)

Distribute `n` panel endpoints and `n-1` panel midpoints on the interval between
0 and 1 according to the discretization strategy in `spacing`.
"""
spanwise_spacing

# uniform
function spanwise_spacing(n, ::Uniform)

    eta = range(0, 1.0, length=n)
    eta_mid = linearinterp(0.5, eta[1:end-1], eta[2:end])

    return eta, eta_mid
end

# sine
function spanwise_spacing(n, ::Sine)

    theta = range(0, pi/2, length=n)
    eta = sin.(theta)

    # note that control points are also placed with cosine spacing as this improves accuracy tremendously
    theta_mid = linearinterp(0.5, theta[1:end-1], theta[2:end])
    eta_mid = sin.(theta_mid)

    return eta, eta_mid
end

# cosine
function spanwise_spacing(n, ::Cosine)
    theta = range(0, pi, length=n)
    eta = (1.0 .- cos.(theta))/2.0

    # note that control points are also placed with cosine spacing as this improves accuracy tremendously
    theta_mid = linearinterp(0.5, theta[1:end-1], theta[2:end])
    eta_mid = (1.0 .- cos.(theta_mid))/2.0

    return eta, eta_mid
end

"""
    chordwise_spacing(n, spacing::AbstractSpacing)

Distribute `n` panel edge, `n-1` vortex, and `n-1` control point chordwise
locations on the interval between 0 and 1 according to the discretization
strategy in `spacing`.
"""
chordwise_spacing

# uniform
function chordwise_spacing(n, ::Uniform)

    eta = range(0, 1.0, length=n)

    eta_qtr = linearinterp(0.25, eta[1:end-1], eta[2:end])
    eta_thrqtr = linearinterp(0.75, eta[1:end-1], eta[2:end])

    return eta, eta_qtr, eta_thrqtr
end

# sine
function chordwise_spacing(n, ::Sine)

    theta = range(0, pi/2, length=n)
    eta = sin.(theta)

    eta_qtr = linearinterp(0.25, eta[1:end-1], eta[2:end])
    eta_thrqtr = linearinterp(0.75, eta[1:end-1], eta[2:end])

    return eta, eta_qtr, etq_thrqtr
end

# cosine
function chordwise_spacing(n, ::Cosine)

    theta = range(0, pi, length=n)
    eta = (1.0 .- cos.(theta))/2.0

    eta_qtr = linearinterp(0.25, eta[1:end-1], eta[2:end])
    eta_thrqtr = linearinterp(0.75, eta[1:end-1], eta[2:end])

    return eta, eta_qtr, eta_thrqtr
end

"""
    interpolate_grid(xyz, eta, interp; xdir=0, ydir=1)

Interpolates the grid `xyz` along direction `dir`

# Arguments
 - `xyz`: Grid of size (3, ni, nj)
 - `eta`: New (normalized) coordinates in direction `dir` (0 <= eta <= 1)
 - `interp`: Interpolation method of form `ypt = f(x,y,xpt)`
 - `xdir`: Independent variable direction, defaults to arc length
 - `ydir`: Dependent variable direction `xyz` (`i=1`, `j=2`)
"""
function interpolate_grid(xyz, eta, interp; xdir=0, ydir=1)

    ydim = ydir + 1

    ni = size(xyz, ydim)

    xyz = mapslices(xyz; dims=[1,ydim]) do xyz_i

        # get distance between each grid location
        if xdir == 0
            ds = [norm(xyz_i[:,k] - xyz_i[:,max(1,k-1)]) for k = 1:ni]
        else
            ds = [xyz_i[xdir,k] - xyz_i[xdir,max(1,k-1)] for k = 1:ni]
        end

        # create interpolation vector
        t = cumsum(ds)

        # normalize interpolation vector
        t /= t[end]

        # interpolate x, y, and z
        x = interp(t, xyz_i[1,:], eta)
        y = interp(t, xyz_i[2,:], eta)
        z = interp(t, xyz_i[3,:], eta)

        vcat(x',y',z')
    end

    return xyz
end

"""
    grid_to_surface_panels(xyz; mirror = false, fcore = (c, Δs) -> 1e-3)

Construct a set of panels with associated vortex rings given a potentially curved
lifting surface defined by a grid with dimensions (3, i, j) where `i` corresponds
to the chordwise direction (ordered from leading edge to trailing edge) and `j`
corresponds to the spanwise direction (ordered from left to right).  The leading
edge of each ring vortex will be placed at the 1/4 chord and the control point
will be placed at the 3/4 chord of each panel.

Return a grid with dimensions (3, i, j) containing the panel corners and a matrix
with dimensions (i, j) containing the generated panels.

# Keyword Arguments
- `mirror`:  mirror the geometry across the X-Z plane? defaults to `false`.
- `fcore`: function for setting the finite core size based on the chord length
       (in the x-direction) and/or the panel width (in the y/z directions).
       Defaults to `(c, Δs) -> 1e-3`
"""
function grid_to_surface_panels(xyz;
    mirror = false,
    fcore = (c, Δs) -> 1e-3)

    TF = eltype(xyz)

    nc = size(xyz, 2)-1 # number of chordwise panels
    ns = size(xyz, 3)-1 # number of spanwise panels
    N = (1+mirror)*nc*ns # total number of panels

    # check which side we're working with
    right_side = sum(xyz[2,:,:]) > 0

    # initialize output
    xyz_panels = copy(xyz)
    panels = Matrix{SurfacePanel{TF}}(undef, nc, (1+mirror)*ns)

    # populate each panel
    for j = 1:ns

        # get leading edge panel corners
        r1n = SVector(xyz[1,1,j], xyz[2,1,j], xyz[3,1,j]) # top left
        r2n = SVector(xyz[1,1,j+1], xyz[2,1,j+1], xyz[3,1,j+1]) # top right
        r3n = SVector(xyz[1,2,j], xyz[2,2,j], xyz[3,2,j]) # bottom left
        r4n = SVector(xyz[1,2,j+1], xyz[2,2,j+1], xyz[3,2,j+1]) # bottom right

        # also get chord length for setting finite core size
        cl = norm(xyz[:,end,j] - xyz[:,1,j])
        cr = norm(xyz[:,end,j+1] - xyz[:,1,j+1])
        c = (cl + cr)/2

        for i = 1:nc-1

            # xyz corners of current panel
            r1 = r1n # top left of panel
            r2 = r2n # top right of panel
            r3 = r3n # bottom left of panel
            r4 = r4n # bottom right of panel

            # xyz corners of next panel
            r1n = r3 # top left
            r2n = r4 # top right
            r3n = SVector(xyz[1,i+2,j], xyz[2,i+2,j], xyz[3,i+2,j]) # bottom left
            r4n = SVector(xyz[1,i+2,j+1], xyz[2,i+2,j+1], xyz[3,i+2,j+1]) # bottom right

            # top left corner of ring vortex
            rtl = linearinterp(0.25, r1, r3)

            # top right corner of ring vortex
            rtr = linearinterp(0.25, r2, r4)

            # center of top bound vortex
            rtc = (rtl + rtr)/2

            # bottom left corner of ring vortex
            rbl = linearinterp(0.25, r1n, r3n)

            # bottom right corner of ring vortex
            rbr = linearinterp(0.25, r2n, r4n)

            # center of bottom bound vortex
            rbc = (rbl + rbr)/2

            # control point
            rtop = linearinterp(0.5, r1, r2)
            rbot = linearinterp(0.5, r3, r4)
            rcp = linearinterp(0.75, rtop, rbot)

            # surface normal
            ncp = cross(rcp - rtr, rcp - rtl)
            ncp /= norm(ncp)

            # set finite core size
            Δs = sqrt((rtr[2]-rtl[2])^2 + (rtr[3]-rtl[3])^2)
            core_size = fcore(c, Δs)

            # get chord length of current panel
            chord = norm((r1 + r2)/2 - (r3 + r4)/2)

            ip = i
            jp = mirror*right_side*ns + j
            panels[ip, jp] = SurfacePanel{TF}(rtl, rtc, rtr, rbl, rbc, rbr, rcp,
                ncp, core_size, chord)
        end

        # xyz corners of current panel
        r1 = r1n # top left
        r2 = r2n # top right
        r3 = r3n # bottom left
        r4 = r4n # bottom right

        # top left corner of ring vortex
        rtl = linearinterp(0.25, r1, r3)

        # top right corner of ring vortex
        rtr = linearinterp(0.25, r2, r4)

        # center of top bound vortex
        rtc = (rtl + rtr)/2

        # bottom left corner of ring vortex
        rbl = r3

        # bottom right corner of ring vortex
        rbr = r4

        # center of bottom bound vortex
        rbc = (rbl + rbr)/2

        # control point
        rtop = linearinterp(0.5, r1, r2)
        rbot = linearinterp(0.5, r3, r4)
        rcp = linearinterp(0.75, rtop, rbot)

        # surface normal
        ncp = cross(rcp - rtr, rcp - rtl)
        ncp /= norm(ncp)

        # set finite core size
        Δs = sqrt((rtr[2]-rtl[2])^2 + (rtr[3]-rtl[3])^2)
        core_size = fcore(c, Δs)

        # get chord length of current panel
        chord = norm((r1 + r2)/2 - (r3 + r4)/2)

        ip = nc
        jp = mirror*right_side*ns + j
        panels[ip,jp] = SurfacePanel{TF}(rtl, rtc, rtr, rbl, rbc, rbr, rcp, ncp,
            core_size, chord)
    end

    # other side
    if mirror

        # first reflect xyz
        if right_side
            xyz_l = reverse(xyz_panels, dims=3)
            xyz_l[2,:,:] .= -xyz_l[2,:,:]
            xyz_r = xyz_panels
        else
            xyz_l = xyz_panels
            xyz_r = reverse(xyz_panels, dims=3)
            xyz_r[2,:,:] .= -xyz_r[2,:,:]
        end
        xyz_panels = cat(xyz_l, xyz_r[:,:,2:end], dims=3)

        # then populate remaining panels
        for j = 1:ns
            for i = 1:nc
                ip = i
                jp1 = right_side*ns + j
                jp2 = 2*ns - right_side*ns - j + 1
                panels[ip,jp2] = reflect(panels[ip,jp1])
            end
        end
    end

    return xyz_panels, panels
end

"""
    grid_to_surface_panels(xyz, ns, nc;
        mirror = false,
        fcore = (c, Δs) -> 1e-3,
        spacing_s = Cosine(),
        spacing_c = Uniform(),
        interp_s = (x, y, xpt) -> LinearInterpolation(x, y)(xpt),
        interp_c = (x, y, xpt) -> LinearInterpolation(x, y)(xpt))

Discretize a potentially curved lifting surface defined by a grid with dimensions
(3, i, j) where `i` corresponds to the chordwise direction (ordered from leading
edge to trailing edge) and `j` corresponds to the spanwise direction (ordered from
left to right) into `ns` spanwise and `nc` chordwise panels with associated vortex
rings according to the spanwise discretization scheme `spacing_s` and chordwise
discretization scheme `spacing_c`.  The bound vortex will be placed at the
1/4 chord and the control point will be placed at the 3/4 chord of each panel.

Return a grid with dimensions (3, i, j) containing the interpolated panel corners
and a matrix with dimensions (i, j) containing the generated panels.

# Arguments
 - `xyz`: grid of dimensions (3, i, j) where where `i` corresponds to the
    chordwise direction and `j` corresponds to the spanwise direction.
 - `ns`: number of spanwise panels
 - `nc`: number of chordwise panels

# Keyword Arguments
 - `mirror`:  mirror the geometry across the X-Z plane? defaults to `false`.
 - `fcore`: function for setting the finite core size based on the chord length
        (in the x-direction) and/or the panel width (in the y/z directions).
        Defaults to `(c, Δs) -> 1e-3`
 - `spacing_s`: spanwise discretization scheme, defaults to `Cosine()`
 - `spacing_c`: chordwise discretization scheme, defaults to `Uniform()`
 - `interp_s`: spanwise interpolation function, defaults to linear interpolation
 - `interp_c`: chordwise interpolation function, defaults to linear interpolation
"""
function grid_to_surface_panels(xyz, ns, nc;
    mirror = false,
    fcore = (c, Δs) -> 1e-3,
    spacing_s = Cosine(),
    spacing_c = Uniform(),
    interp_s = (x, y, xpt) -> LinearInterpolation(x, y)(xpt),
    interp_c = (x, y, xpt) -> LinearInterpolation(x, y)(xpt))

    TF = eltype(xyz)

    # get spanwise and chordwise spacing
    etas, etabar = spanwise_spacing(ns+1, spacing_s)
    etac, eta_qtr, eta_thrqtr = chordwise_spacing(nc+1, spacing_c)

    # add leading and trailing edge to grid
    eta_qtr = vcat(0.0, eta_qtr, 1.0)

    # interpolate grid first along `i` then along `j`
    xyz_edge = interpolate_grid(xyz, etac, interp_c, ydir=1)
    xyz_panels = interpolate_grid(xyz_edge, etas, interp_s, ydir=2)

    xyz_bound = interpolate_grid(xyz, eta_qtr, interp_c, ydir=1)
    xyz_corner = interpolate_grid(xyz_bound, etas, interp_s, ydir=2)
    xyz_center = interpolate_grid(xyz_bound, etabar, interp_s, ydir=2)

    xyz_cp = interpolate_grid(xyz, eta_thrqtr, interp_c, ydir=1)
    xyz_cp = interpolate_grid(xyz_cp, etabar, interp_s, ydir=2)

    # check which side we're working with
    right_side = sum(xyz[2,:,:]) > 0

    # initialize output
    panels = Matrix{SurfacePanel{TF}}(undef, nc, (1+mirror)*ns)

    # populate each panel
    for j = 1:ns

        # get chord for setting finite core size

        # also get chord length for setting finite core size
        c = norm(xyz_center[:,end,j] - xyz_center[:,1,j])

        for i = 1:nc
            # note that `i==1` corresponds to the leading edge for `xyz_corner` and `xyz_center`
            rtl = SVector(xyz_corner[1,i+1,j], xyz_corner[2,i+1,j], xyz_corner[3,i+1,j])
            rtc = SVector(xyz_center[1,i+1,j], xyz_center[2,i+1,j], xyz_center[3,i+1,j])
            rtr = SVector(xyz_corner[1,i+1,j+1], xyz_corner[2,i+1,j+1], xyz_corner[3,i+1,j+1])
            rbl = SVector(xyz_corner[1,i+2,j], xyz_corner[2,i+2,j], xyz_corner[3,i+2,j])
            rbc = SVector(xyz_center[1,i+2,j], xyz_center[2,i+2,j], xyz_center[3,i+2,j])
            rbr = SVector(xyz_corner[1,i+2,j+1], xyz_corner[2,i+2,j+1], xyz_corner[3,i+2,j+1])
            rcp = SVector(xyz_cp[1,i,j], xyz_cp[2,i,j], xyz_cp[3,i,j])
            ncp = cross(rcp - rtr, rcp - rtl)
            ncp /= norm(ncp)

            # set finite core size
            Δs = sqrt((rtr[2]-rtl[2])^2 + (rtr[3]-rtl[3])^2)
            core_size = fcore(c, Δs)

            # set chord length of current panel
            r1 = xyz_panels[:, i, j] # top left
            r2 = xyz_panels[:, i, j+1] # top right
            r3 = xyz_panels[:, i+1, j] # bottom left
            r4 = xyz_panels[:, i+1, j+1] # bottom right
            chord = norm((r1 + r2)/2 - (r3 + r4)/2)

            ip = i
            jp = mirror*right_side*ns + j
            panels[ip, jp] = SurfacePanel{TF}(rtl, rtc, rtr, rbl, rbc, rbr, rcp,
                ncp, core_size, chord)
        end
    end

    # other side
    if mirror
        # first reflect grid
        if right_side
            xyz_l = reverse(xyz_panels, dims=3)
            xyz_l[2,:,:] .= -xyz_l[2,:,:]
            xyz_r = xyz_panels
        else
            xyz_l = xyz_panels
            xyz_r = reverse(xyz_panels, dims=3)
            xyz_r[2,:,:] .= -xyz_r[2,:,:]
        end
        xyz_panels = cat(xyz_l, xyz_r[:,:,2:end], dims=3)

        # then populate remaining panels
        for j = 1:ns
            for i = 1:nc
                ip = i
                jp1 = right_side*ns + j
                jp2 = 2*ns - right_side*ns - j + 1
                panels[ip,jp2] = reflect(panels[ip,jp1])
            end
        end
    end

    return xyz_panels, panels
end

"""
    wing_to_surface_panels(xle, yle, zle, chord, theta, phi, ns, nc;
        fc = fill(x -> 0, length(xle)),
        mirror = false,
        fcore = (c, Δs) -> 1e-3,
        spacing_s = Cosine(),
        spacing_c = Uniform(),
        interp_s = (x, y, xpt) -> LinearInterpolation(x, y)(xpt))

Discretize a wing into `ns` spanwise and `nc` chordwise panels with associated
vortex rings according to the spanwise discretization scheme `spacing_s` and
chordwise discretization scheme `spacing_c`.

Return a grid with dimensions (3, i, j) containing the panel corners and a matrix
with dimensions (i, j) containing the generated panels.

# Arguments
 - `xle`: leading edge x-coordinate of each airfoil section
 - `yle`: leading edge y-coordinate of each airfoil section
 - `zle`: leading edge z-coordinate of each airfoil section
 - `chord`: chord length of each airfoil section
 - `theta`: twist of each airfoil section
 - `phi`: dihedral angle of each airfoil section, defined by a right hand rotation about the x-axis
 - `ns`: number of spanwise panels
 - `nc`: number of chordwise panels
 - `fc`: (optional) camber line function y=f(x) of each airfoil section
 - `mirror`:  mirror the geometry across the X-Z plane?, defaults to `false`
 - `fcore`: function for setting the finite core size based on the chord length
        (in the x-direction) and/or the panel width (in the y/z directions).
        Defaults to `(c, Δs) -> 1e-3`
 - `spacing_s`: spanwise discretization scheme, defaults to `Cosine()`
 - `spacing_c`: chordwise discretization scheme, defaults to `Uniform()`
 - `interp_s`: interpolation function between spanwise stations, defaults to linear interpolation
"""
function wing_to_surface_panels(xle, yle, zle, chord, theta, phi, ns, nc;
    fc = fill(x->0, length(xle)),
    mirror = false,
    fcore = (c, Δs) -> 1e-3,
    spacing_s = Cosine(),
    spacing_c = Uniform(),
    interp_s = (x, y, xpt) -> LinearInterpolation(x, y)(xpt))

    TF = promote_type(eltype(xle), eltype(yle), eltype(zle), eltype(chord), eltype(theta), eltype(phi))

    N = (1+mirror)*nc*ns

    # get spanwise and chordwise spacing
    etas, etabar = spanwise_spacing(ns+1, spacing_s)
    etac, eta_qtr, eta_thrqtr = chordwise_spacing(nc+1, spacing_c)

    # add leading and trailing edge to grid
    eta_qtr = vcat(0.0, eta_qtr, 1.0)

    # check input dimensions
    n = length(xle)
    @assert n == length(yle) == length(zle) == length(chord) == length(theta) == length(phi)

    # set bound vortex and control point chordwise locations
    xyz_edge = Array{TF, 3}(undef, 3, nc+1, n)
    xyz_bound = Array{TF, 3}(undef, 3, nc+2, n)
    xyz_cp = Array{TF, 3}(undef, 3, nc, n)
    for j = 1:n

        # rotation matrix for twist
        st, ct = sincos(theta[j])
        Rt = @SMatrix [ct 0 st; 0 1 0; -st 0 ct]

        # rotation matrix for dihedral
        sp, cp = sincos(phi[j])
        Rp = @SMatrix [1 0 0; 0 cp -sp; 0 sp cp]

        # location of leading edge
        rle = SVector(xle[j], yle[j], zle[j])

        # panel edge chordwise locations
        for i = 1:nc+1
            # bound vortex location
            xc = etac[i]
            zc = fc[j](xc)

            # location on airfoil
            r = SVector(xc, 0, zc)

            # scale by chord length
            r = chord[j] * r

            # apply twist
            r = Rt * r

            # apply dihedral
            r = Rp * r

            # add leading edge offset
            r = r + rle

            # store final location
            xyz_edge[:,i,j] = r
        end

        # bound vortex chordwise locations (and leading and trailing edges)
        for i = 1:nc+2
            # bound vortex location
            xc = eta_qtr[i]
            zc = fc[j](xc)

            # location on airfoil
            r = SVector(xc, 0, zc)

            # scale by chord length
            r = chord[j] * r

            # apply twist
            r = Rt * r

            # apply dihedral
            r = Rp * r

            # add leading edge offset
            r = r + rle

            # store final location
            xyz_bound[:,i,j] = r
        end

        # control point chordwise locations
        for i = 1:nc
            # bound vortex location
            xc = eta_thrqtr[i]
            zc = fc[j](xc)

            # location on airfoil
            r = SVector(xc, 0, zc)

            # scale by chord length
            r = chord[j] * r

            # apply twist
            r = Rt * r

            # apply dihedral
            r = Rp * r

            # add leading edge offset
            r = r + rle

            # store final location
            xyz_cp[:,i,j] = r
        end
    end

    # interpolate the grids to match the specified spanwise spacing
    xyz_panels = interpolate_grid(xyz_edge, etas, interp_s; ydir=2)
    xyz_corner = interpolate_grid(xyz_bound, etas, interp_s; ydir=2)
    xyz_center = interpolate_grid(xyz_bound, etabar, interp_s; ydir=2)
    xyz_cp = interpolate_grid(xyz_cp, etabar, interp_s; ydir=2)

    # initialize output
    panels = Matrix{SurfacePanel{TF}}(undef,  nc, (1+mirror)*ns)

    # check which side we're working with
    right_side = sum(yle) > 0

    # populate each panel
    for j = 1:ns

        # get chord for setting finite core size
        c = norm(xyz_center[:,end,j] - xyz_center[:,1,j])

        for i = 1:nc
            # note that `i==1` corresponds to the leading edge for `xyz_corner` and `xyz_center`
            rtl = SVector(xyz_corner[1,i+1,j], xyz_corner[2,i+1,j], xyz_corner[3,i+1,j])
            rtc = SVector(xyz_center[1,i+1,j], xyz_center[2,i+1,j], xyz_center[3,i+1,j])
            rtr = SVector(xyz_corner[1,i+1,j+1], xyz_corner[2,i+1,j+1], xyz_corner[3,i+1,j+1])
            rbl = SVector(xyz_corner[1,i+2,j], xyz_corner[2,i+2,j], xyz_corner[3,i+2,j])
            rbc = SVector(xyz_center[1,i+2,j], xyz_center[2,i+2,j], xyz_center[3,i+2,j])
            rbr = SVector(xyz_corner[1,i+2,j+1], xyz_corner[2,i+2,j+1], xyz_corner[3,i+2,j+1])
            rcp = SVector(xyz_cp[1,i,j], xyz_cp[2,i,j], xyz_cp[3,i,j])
            ncp = cross(rcp - rtr, rcp - rtl)
            ncp /= norm(ncp)

            # set finite core size
            Δs = sqrt((rtr[2]-rtl[2])^2 + (rtr[3]-rtl[3])^2)
            core_size = fcore(c, Δs)

            # set chord length of current panel
            r1 = xyz_panels[:,i,j] # top left
            r2 = xyz_panels[:,i,j+1] # top right
            r3 = xyz_panels[:,i+1, j] # bottom left
            r4 = xyz_panels[:,i+1, j+1] # bottom right
            chord = norm((r1 + r2)/2 - (r3 + r4)/2)

            ip = i
            jp = mirror*right_side*ns + j
            panels[ip, jp] = SurfacePanel{TF}(rtl, rtc, rtr, rbl, rbc, rbr, rcp,
                ncp, core_size, chord)
        end
    end

    # other side
    if mirror
        # first reflect grid
        if right_side
            xyz_l = reverse(xyz_panels, dims=3)
            xyz_l[2,:,:] .= -xyz_l[2,:,:]
            xyz_r = xyz_panels
        else
            xyz_l = xyz_panels
            xyz_r = reverse(xyz_panels, dims=3)
            xyz_r[2,:,:] .= -xyz_r[2,:,:]
        end
        xyz_panels = cat(xyz_l, xyz_r[:,:,2:end], dims=3)

        # then populate remaining panels
        for j = 1:ns
            for i = 1:nc
                ip = i
                jp1 = right_side*ns + j
                jp2 = 2*ns - right_side*ns - j + 1
                panels[ip,jp2] = reflect(panels[ip,jp1])
            end
        end
    end

    return xyz_panels, panels
end

"""
    update_surface_panels!(surface, grid; fcore = (c, Δs) -> 1e-3)

Updates the surface panels in `surface` to correspond to the grid coordinates in
`grid`.
"""
function update_surface_panels!(surface, grid; fcore = (c, Δs) -> 1e-3)

    TF = eltype(eltype(surface))

    nc, ns = size(surface) # number of chordwise and spanwise panels

    # populate each panel
    for j = 1:ns

        # get leading edge panel corners
        r1n = SVector(grid[1,1,j], grid[2,1,j], grid[3,1,j]) # top left
        r2n = SVector(grid[1,1,j+1], grid[2,1,j+1], grid[3,1,j+1]) # top right
        r3n = SVector(grid[1,2,j], grid[2,2,j], grid[3,2,j]) # bottom left
        r4n = SVector(grid[1,2,j+1], grid[2,2,j+1], grid[3,2,j+1]) # bottom right

        # also get chord length for setting finite core size
        cl = norm(grid[:,end,j] - grid[:,1,j])
        cr = norm(grid[:,end,j+1] - grid[:,1,j+1])
        c = (cl + cr)/2

        for i = 1:nc-1

            # grid corners of current panel
            r1 = r1n # top left of panel
            r2 = r2n # top right of panel
            r3 = r3n # bottom left of panel
            r4 = r4n # bottom right of panel

            # grid corners of next panel
            r1n = r3 # top left
            r2n = r4 # top right
            r3n = SVector(grid[1,i+2,j], grid[2,i+2,j], grid[3,i+2,j]) # bottom left
            r4n = SVector(grid[1,i+2,j+1], grid[2,i+2,j+1], grid[3,i+2,j+1]) # bottom right

            # top left corner of ring vortex
            rtl = linearinterp(0.25, r1, r3)

            # top right corner of ring vortex
            rtr = linearinterp(0.25, r2, r4)

            # center of top bound vortex
            rtc = (rtl + rtr)/2

            # bottom left corner of ring vortex
            rbl = linearinterp(0.25, r1n, r3n)

            # bottom right corner of ring vortex
            rbr = linearinterp(0.25, r2n, r4n)

            # center of bottom bound vortex
            rbc = (rbl + rbr)/2

            # control point
            rtop = linearinterp(0.5, r1, r2)
            rbot = linearinterp(0.5, r3, r4)
            rcp = linearinterp(0.75, rtop, rbot)

            # surface normal
            ncp = cross(rcp - rtr, rcp - rtl)
            ncp /= norm(ncp)

            # set finite core size
            Δs = sqrt((rtr[2]-rtl[2])^2 + (rtr[3]-rtl[3])^2)
            core_size = fcore(c, Δs)

            # get chord length of current panel
            chord = norm((r1 + r2)/2 - (r3 + r4)/2)

            surface[i,j] = SurfacePanel{TF}(rtl, rtc, rtr, rbl, rbc, rbr, rcp,
                ncp, core_size, chord)
        end

        # grid corners of current panel
        r1 = r1n # top left
        r2 = r2n # top right
        r3 = r3n # bottom left
        r4 = r4n # bottom right

        # top left corner of ring vortex
        rtl = linearinterp(0.25, r1, r3)

        # top right corner of ring vortex
        rtr = linearinterp(0.25, r2, r4)

        # center of top bound vortex
        rtc = (rtl + rtr)/2

        # bottom left corner of ring vortex
        rbl = r3

        # bottom right corner of ring vortex
        rbr = r4

        # center of bottom bound vortex
        rbc = (rbl + rbr)/2

        # control point
        rtop = linearinterp(0.5, r1, r2)
        rbot = linearinterp(0.5, r3, r4)
        rcp = linearinterp(0.75, rtop, rbot)

        # surface normal
        ncp = cross(rcp - rtr, rcp - rtl)
        ncp /= norm(ncp)

        # set finite core size
        Δs = sqrt((rtr[2]-rtl[2])^2 + (rtr[3]-rtl[3])^2)
        core_size = fcore(c, Δs)

        # get chord length of current panel
        chord = norm((r1 + r2)/2 - (r3 + r4)/2)

        surface[end, j] = SurfacePanel{TF}(rtl, rtc, rtr, rbl, rbc, rbr, rcp, ncp,
            core_size, chord)
    end

    return surface
end

"""
    lifting_line_geometry(grids, xc=0.25)
Construct a lifting line representation of the surfaces in `grids` at the
normalized chordwise location `xc`.  Return the lifting line coordinates
and chord lengths.
# Arguments
 - `grids`: Vector with length equal to the number of surfaces.  Each element of
    the vector is a grid with shape (3, nc, ns) which defines the discretization
    of a surface into panels. `nc` is the number of chordwise panels and `ns` is
    the number of spanwise panels.
 - `xc`: Normalized chordwise location of the lifting line from the leading edge.
    Defaults to the quarter chord
# Return Arguments:
 - `r`: Vector with length equal to the number of surfaces, with each element
    being a matrix with size (3, ns+1) which contains the x, y, and z coordinates
    of the resulting lifting line coordinates
 - `c`: Vector with length equal to the number of surfaces, with each element
    being a vector of length `ns+1` which contains the chord lengths at each
    lifting line coordinate.
"""
function lifting_line_geometry(grids, xc=0.25)
    TF = eltype(eltype(grids))
    nsurf = length(grids)
    r = Vector{Matrix{TF}}(undef, nsurf)
    c = Vector{Vector{TF}}(undef, nsurf)
    for isurf = 1:nsurf
        ns = size(grids[isurf], 3) - 1
        r[isurf] = Matrix{TF}(undef, 3, ns+1)
        c[isurf] = Vector{TF}(undef, ns+1)
    end
    return lifting_line_geometry!(r, c, grids, xc)
end

"""
    lifting_line_geometry!(r, c, grids, xc=0.25)
In-place version of [`lifting_line_geometry`](@ref)
"""
function lifting_line_geometry!(r, c, grids, xc=0.25)
    nsurf = length(grids)
    # iterate through each lifting surface
    for isurf = 1:nsurf
        # extract current grid
        grid = grids[isurf]
        # dimensions of this grid
        nc = size(grid, 2)
        ns = size(grid, 3)
        # loop through each spanwise section
        for j = 1:ns
            # get leading and trailing edges
            le = SVector(grid[1,1,j], grid[2,1,j], grid[3,1,j])
            te = SVector(grid[1,end,j], grid[2,end,j], grid[3,end,j])
            # get quarter-chord
            r[isurf][:,j] = linearinterp(xc, le, te)
            # get chord length
            c[isurf][j] = norm(le - te)
        end
    end
    return r, c
end

"""
    translate(grid, r)

Return a copy of the grid points in `grid` translated the distance specified by vector `r`
"""
translate(grid, r) = grid .+ r

"""
    translate!(grid, r)

Translate the grid points in `grid` the distance specified by vector `r`
"""
function translate!(grid, r)
    grid .+= r
    return grid
end

"""
    rotate(grid, R, r = [0,0,0])

Return a copy of the grid points in `grid` rotated about point `r` using the rotation matrix `R`
"""
rotate(grid, args...) = rotate!(copy(grid), args...)

"""
    rotate!(grid, R, r = [0,0,0])

Rotate the grid points in `grid` about point `r` using the rotation matrix `R`
"""
function rotate!(grid, R, r = (@SVector zeros(3)))
    nc, ns = size(grid, 2), size(grid, 3)

    grid .-= r

    for i = 1:nc, j = 1:ns
        grid[:,i,j] = R*grid[:,i,j]
    end

    grid .+= r

    return grid
end

"""
    trailing_edge_points(surface[s])

Return points on the trailing edge of each surface.
"""
trailing_edge_points(surfaces) = [bottom_left.(surfaces[end, :])..., bottom_right(surfaces[end, end])]

"""
    repeated_trailing_edge_points(surface[s])

Generates a dictionary of the form `Dict((isurf, i) => [(jsurf1, j1), (jsurf2, j2)...]` which
defines repeated trailing edge points.  Trailing edge point `i` on surface
`isurf` is repeated on surface `jsurf1` at point `j1`, `jsurf2` at point `j2`,
and so forth.
"""
repeated_trailing_edge_points

repeated_trailing_edge_points(surface::AbstractMatrix) = repeated_trailing_edge_points([surface])

function repeated_trailing_edge_points(surfaces::AbstractVector{<:AbstractMatrix})

    nsurf = length(surfaces)
    points = trailing_edge_points.(surfaces)
    repeated_points = Dict{NTuple{2, Int}, Vector{NTuple{2, Int}}}()

    # loop through all points
    for isurf = 1:nsurf, i = 1:length(points[isurf])
        # start with empty vector
        repeats = NTuple{2, Int}[]
        # add all repeated points to the vector
        for jsurf = 1:nsurf, j = 1:length(points[jsurf])
            if (isurf, i) != (jsurf, j) && isapprox(points[isurf][i], points[jsurf][j])
                push!(repeats, (jsurf, j))
            end
        end
        # only store vector if nonempty
        if !isempty(repeats)
            repeated_points[(isurf, i)] = repeats
        end
    end

    return repeated_points
end

"""
    flipy(r)

Flip sign of y-component of vector `r` (used for symmetry)
"""
@inline flipy(r) = SVector{3}(r[1], -r[2], r[3])

"""
    on_symmetry_plane(args...; tol=eps())

Test whether points in `args` are on the symmetry plane (y = 0)
"""
on_symmetry_plane

on_symmetry_plane(r; tol=eps()) = isapprox(r[2], 0.0, atol=tol)

@inline function on_symmetry_plane(r, args...; tol=eps())
    return on_symmetry_plane(r; tol=tol) && on_symmetry_plane(args...; tol=tol)
end

"""
    not_on_symmetry_plane(args...; tol=eps())

Test whether and of the points in `args` are not on the symmetry plane (y = 0)
"""
not_on_symmetry_plane(args...; tol=eps()) = !on_symmetry_plane(args...; tol=tol)


type Fourier2d <: SuperRes #1D Static model
    freqs :: Vector{Tuple{Int, Int}}
    filter :: Vector{Float64}
    x_max :: Float64
    z_max :: Float64
    approx_grid_x :: Vector{Float64}
    approx_grid_z :: Vector{Float64}
    bounds :: Tuple{Array{Float64}, Array{Float64}}
    function Fourier2d(x_max, z_max, filter, n_approx_x, n_approx_z)
        ind = find(filter)
        f_c_x = div(size(filter, 1) - 1, 2)
        f_c_z = div(size(filter, 2) - 1, 2)
        freqs_x = -f_c_x:f_c_x
        freqs_z = -f_c_z:f_c_z
        freqs2d = [(freq1, freq2) for freq1 in freqs_x, freq2 in freqs_z]
        freqs2d = freqs2d[ind]
        filter2d = filter[ind]
        grid_x = linspace(0, x_max, n_approx_x)
        grid_z = linspace(0, z_max, n_approx_z)
        return new(freqs2d,
                   filter2d,
                   x_max,
                   z_max,
                   grid_x,
                   grid_z,
                   ([0.0; 0.0], [x_max; z_max]))
    end
    function Fourier2d(x_max, filter, n_approx_x, n_approx_z)
        return Fourier2d(x_max, x_max, filter, n_approx_x, n_approx_z)
    end
end
function setBounds(model :: Fourier2d, bounds)
    model.bounds = bounds
    model.approx_grid_x = linspace(bounds[1][1], bounds[2][1], length(model.approx_grid_x))
    model.approx_grid_z = linspace(bounds[1][2], bounds[2][2], length(model.approx_grid_z))
end

function psi(model :: Fourier2d, theta :: Vector{Float64})
  # This function computes the direct problem for a single point theta
  psi_cpx = vec(Complex128[model.filter[k] *
    exp(-2im * pi * (model.freqs[k][1] / model.x_max * theta[1] + model.freqs[k][2] / model.z_max * theta[2] )) for k in eachindex(model.freqs)])
  return [real(psi_cpx); imag(psi_cpx)]
end

function dpsi(model :: Fourier2d, theta :: Vector{Float64})
  # This function computes the gradient of psi
  dx1 = vec(Complex128[-model.filter[k] *
    model.freqs[k][1] / model.x_max * 2im * pi *
    exp(-2im * pi * (model.freqs[k][1]  * theta[1]/model.x_max +  model.freqs[k][2] * theta[2]/model.z_max )) for k in eachindex(model.freqs)])
  dx2 = vec(Complex128[-model.filter[k] *
    model.freqs[k][2] / model.z_max * 2im * pi *
    exp(-2im * pi * (model.freqs[k][1] * theta[1]/model.x_max + model.freqs[k][2] * theta[2]/model.z_max )) for k in eachindex(model.freqs)])
  return [[real(dx1); imag(dx1)] [real(dx2); imag(dx2)]]
end

function getStartingPoint(model :: Fourier2d, v :: Vector{Float64})
  # This function gets an initial value for the location
  values = Float64[dot(v,psi(model, [x; y])) for x = model.approx_grid_x, y = model.approx_grid_z]
  ind = indmin(values);
  i,j = ind2sub(values, ind)
  return vec([model.approx_grid_x[i]; model.approx_grid_z[j]])
end

parameterBounds(model :: Fourier2d) =
    model.bounds

dim(model :: Fourier2d) = 2

type DynamicFourier2d <: DynamicSuperRes #1D Space-time model
    static :: Fourier2d
    times :: Vector{Float64}
    v_max :: Float64
    grid_v :: Vector{Float64}
    bounds :: Tuple{Array{Float64}, Array{Float64}}
    DynamicFourier2d(static, v_max, tau, K, num_v = 20) = new(static,
                                                              linspace(-tau*K, tau*K, 2*K+1),
                                                              v_max,
                                                              linspace(-v_max, v_max, num_v),
                                                              ([0; 0; -v_max;v_max], [static.x_max; static.z_max; v_max; v_max]))

end

function getStartingPoint(model :: DynamicFourier2d, v :: Vector{Float64})
    # This function gets an initial value for the location
    vals = Float64[dot(v,psi(model, [x; y; vel_x; vel_z])) for x = model.static.approx_grid_x, y = model.static.approx_grid_z, vel_x = model.grid_v, vel_z = model.grid_v]
    ind = indmin(vals);
    i,j,k,l = ind2sub(vals, ind)
    return [model.static.approx_grid_x[i]; model.static.approx_grid_z[j]; model.grid_v[k]; model.grid_v[l]]
end

parameterBounds(model :: DynamicFourier2d) =
    model.bounds


function rotate(	q0::Float64,
										q1::Float64,
										q2::Float64,
										q3::Float64,
										sigma_rotation::Float64)

	ux_rot::Float64 = randn()
	uy_rot::Float64 = randn()
	uz_rot::Float64 = randn()
	u_norm::Float64 = sqrt( ux_rot^2 + uy_rot^2 + uz_rot^2 )
	ux_rot /= u_norm
	uy_rot /= u_norm
	uz_rot /= u_norm

	theta_rot::Float64 = sigma_rotation * randn()

	q0_rot::Float64 = cos(0.5 * theta_rot)
	q1_rot::Float64 = ux_rot * sin(0.5 * theta_rot)
	q2_rot::Float64 = uy_rot * sin(0.5 * theta_rot)
	q3_rot::Float64 = uz_rot * sin(0.5 * theta_rot)

	(q0_star::Float64, q1_star::Float64, q2_star::Float64, q3_star::Float64) = quaternion_mult(q0_rot, q1_rot, q2_rot, q3_rot, q0, q1, q2, q3)

	# Normalize proposal quaternion to avoid numerical problems.
	q_norm::Float64 = sqrt( q0_star^2 + q1_star^2 + q2_star^2 + q3_star^2 )
	q0_star /= q_norm
	q1_star /= q_norm
	q2_star /= q_norm
	q3_star /= q_norm

	return (q0_star, q1_star, q2_star, q3_star)
end


using Distributions, BayesianNonparametricStatistics, PyPlot

sde = SDEWithConstantVariance(x->sinpi(2*x), 1.0, 0.0, 10000.0, 0.01)
X = rand(sde)
distribution = MvNormal([k^(-1.0) for k in 1:10])
Π = GaussianProcess([fourier(k) for k in 1:10], distribution)
M = SDEModel(1.0, 0.0)
postΠ = calculateposterior(Π, X, M)
clf()
x=0.0:0.01:1.0
for k in 1:10
    f = rand(postΠ)
    plot(x, f.(x))
end


using ReachabilityBenchmarks, MathematicalSystems, Reachability, Plots

include(@relpath "linearSwitching_model.jl")
include(@relpath "linearSwitching_specifications.jl")

S = linearSwitching_model()
X0, options = linearSwitching_specification()
options = Options(options)

# initial value problem
problem = InitialValueProblem(S, X0)

# reachability algorithm
algorithm_continuous = BFFPS19(:δ => 1e-4, :partition=>[1:2, 3:3, 4:4, 5:5])
algorithm_hybrid = DecomposedDiscretePost(:out_vars=>[1, 2], :clustering=>:none)

# compute flowpipe
options[:mode] = "reach"
solution = solve(problem, options, algorithm_continuous, algorithm_hybrid)

# project flowpipe to x₁ and x₂
solution.options[:plot_vars] = [1, 2]
solution_proj = project(solution)

# plot projection
plot(solution_proj)
savefig("linearSwitching")