text:
using Test
using LinearAlgebra
using UnPack
using MOONS.Grids
using MOONS.Fields

convergence_rate(err, Δh) =
    [log(err[i]/err[i-1])/log(Δh[i]/Δh[i-1]) for i in 2:length(Δh)]

@testset "Discrete operators - ∇(f_corn)=(∇f)_edge" begin
    FT = Float64
    a, b = FT(0.0), FT(1.0)
    nr = 2 .^ (3,4,5,6)
    Δh = zeros(length(nr))
    errx = zeros(length(nr))
    erry = zeros(length(nr))
    errz = zeros(length(nr))
    for (k,n) in enumerate(2 .^ (3,4,5,6))
        grid = Cube(a, b, n; warpfun=Roberts_both, args=(FT(1.01),))
        Δh[k] = grid.c[1].n.Δh[1]

        f = CellCorner(grid)
        ∇f = CellEdges(grid)
        ∇f_exact = CellEdges(grid)

        @inbounds for local_grid in GridField(grid,f)
            @unpack ijk,x,y,z = local_grid
            f[ijk] = sin(3*π*x)*sin(3*π*y)*sin(3*π*z)
        end
        @inbounds for local_grid in GridField(grid,∇f_exact[1])
            @unpack ijk,x,y,z = local_grid
            ∇f_exact[1][ijk] = 3*π*cos(3*π*x)*sin(3*π*y)*sin(3*π*z)
        end
        @inbounds for local_grid in GridField(grid,∇f_exact[2])
            @unpack ijk,x,y,z = local_grid
            ∇f_exact[2][ijk] = 3*π*sin(3*π*x)*cos(3*π*y)*sin(3*π*z)
        end
        @inbounds for local_grid in GridField(grid,∇f_exact[3])
            @unpack ijk,x,y,z = local_grid
            ∇f_exact[3][ijk] = 3*π*sin(3*π*x)*sin(3*π*y)*cos(3*π*z)
        end
        ∇!(∇f, f, grid)
        errx[k] = norm(∇f_exact[1] .- ∇f[1]) / length(∇f[1])
        erry[k] = norm(∇f_exact[2] .- ∇f[2]) / length(∇f[2])
        errz[k] = norm(∇f_exact[3] .- ∇f[3]) / length(∇f[3])
    end

    @test errx[1] ≈ 0.010636252411867275
    @test errx[2] ≈ 0.0008811926803278781
    @test errx[3] ≈ 9.917217369054536e-5
    @test errx[4] ≈ 9.925963845483006e-6
    @test erry[1] ≈ 0.010636252411867273
    @test erry[2] ≈ 0.0008811926803278777
    @test erry[3] ≈ 9.917217369054523e-5
    @test erry[4] ≈ 9.925963845482962e-6
    @test errz[1] ≈ 0.010636252411867275
    @test errz[2] ≈ 0.0008811926803278781
    @test errz[3] ≈ 9.917217369054519e-5
    @test errz[4] ≈ 9.925963845482955e-6

    convx = convergence_rate(errx, Δh)
    convy = convergence_rate(erry, Δh)
    convz = convergence_rate(errz, Δh)

    @test convx[1] ≈ 2.328658024470567
    @test convx[2] ≈ 2.5112742673767268
    @test convx[3] ≈ 2.9570095611906884
    @test convy[1] ≈ 2.328658024470567
    @test convy[2] ≈ 2.511274267376728
    @test convy[3] ≈ 2.9570095611906924
    @test convz[1] ≈ 2.328658024470567
    @test convz[2] ≈ 2.511274267376729
    @test convz[3] ≈ 2.9570095611906932
end

@testset "Discrete operators - ∇(f_cent)=(∇f)_face" begin
    FT = Float64
    a, b = FT(0.0), FT(1.0)
    nr = 2 .^ (3,4,5,6)
    Δh = zeros(length(nr))
    errx = zeros(length(nr))
    erry = zeros(length(nr))
    errz = zeros(length(nr))
    for (k,n) in enumerate(2 .^ (3,4,5,6))
        grid = Cube(a, b, n; warpfun=Roberts_both, args=(FT(1.01),))
        Δh[k] = grid.c[1].n.Δh[1]

        f = CellCenter(grid)
        ∇f = CellFaces(grid)
        ∇f_exact = CellFaces(grid)

        @inbounds for local_grid in GridField(grid,f)
            @unpack ijk,x,y,z = local_grid
            f[ijk] = sin(3*π*x)*sin(3*π*y)*sin(3*π*z)
        end
        @inbounds for local_grid in GridField(grid,∇f_exact[1])
            @unpack ijk,x,y,z = local_grid
            ∇f_exact[1][ijk] = 3*π*cos(3*π*x)*sin(3*π*y)*sin(3*π*z)
        end
        @inbounds for local_grid in GridField(grid,∇f_exact[2])
            @unpack ijk,x,y,z = local_grid
            ∇f_exact[2][ijk] = 3*π*sin(3*π*x)*cos(3*π*y)*sin(3*π*z)
        end
        @inbounds for local_grid in GridField(grid,∇f_exact[3])
            @unpack ijk,x,y,z = local_grid
            ∇f_exact[3][ijk] = 3*π*sin(3*π*x)*sin(3*π*y)*cos(3*π*z)
        end
        ∇!(∇f, f, grid)

        # Ghost points are not populated:
        ∇f[1][1,:,:] .= 0;   ∇f_exact[1][1,:,:] .= 0
        ∇f[1][end,:,:] .= 0; ∇f_exact[1][end,:,:] .= 0
        ∇f[2][:,1,:] .= 0;   ∇f_exact[2][:,1,:] .= 0
        ∇f[2][:,end,:] .= 0; ∇f_exact[2][:,end,:] .= 0
        ∇f[3][:,:,1] .= 0;   ∇f_exact[3][:,:,1] .= 0
        ∇f[3][:,:,end] .= 0; ∇f_exact[3][:,:,end] .= 0

        errx[k] = norm(∇f_exact[1] .- ∇f[1]) / length(∇f[1])
        erry[k] = norm(∇f_exact[2] .- ∇f[2]) / length(∇f[2])
        errz[k] = norm(∇f_exact[3] .- ∇f[3]) / length(∇f[3])
    end
    @test errx[1] ≈ 0.006504291238410591
    @test errx[2] ≈ 0.0015861146672118875
    @test errx[3] ≈ 0.0001769587725572569
    @test errx[4] ≈ 1.7467942413498958e-5
    @test erry[1] ≈ 0.006504291238410591
    @test erry[2] ≈ 0.0015861146672118875
    @test erry[3] ≈ 0.00017695877255725676
    @test erry[4] ≈ 1.7467942413498904e-5
    @test errz[1] ≈ 0.006504291238410591
    @test errz[2] ≈ 0.0015861146672118875
    @test errz[3] ≈ 0.0001769587725572567
    @test errz[4] ≈ 1.7467942413498917e-5

    convx = convergence_rate(errx, Δh)
    convy = convergence_rate(erry, Δh)
    convz = convergence_rate(errz, Δh)

    @test convx[1] ≈ 1.3193403930281424
    @test convx[2] ≈ 2.5212841871228213
    @test convx[3] ≈ 2.974797468220822
    @test convy[1] ≈ 1.3193403930281424
    @test convy[2] ≈ 2.521284187122822
    @test convy[3] ≈ 2.9747974682208245
    @test convz[1] ≈ 1.3193403930281424
    @test convz[2] ≈ 2.521284187122822
    @test convz[3] ≈ 2.9747974682208236
end

@testset "Discrete operators - ∇(f_face)=(∇f)_edge" begin
    FT = Float64
    a, b = FT(0.0), FT(1.0)
    nr = 2 .^ (3,4,5,6)
    Δh = zeros(length(nr))
    err = zeros(length(nr))
    for (k,n) in enumerate(2 .^ (3,4,5,6))
        grid = Cube(a, b, n; warpfun=Roberts_both, args=(FT(1.01),))
        Δh[k] = grid.c[1].n.Δh[1]

        f = CellFace(grid, 1)
        ∇f = CellEdge(grid, 2)
        ∇f_exact = CellEdge(grid, 2)

        @inbounds for local_grid in GridField(grid,f)
            @unpack ijk,x,y,z = local_grid
            f[ijk] = sin(3*π*x)*sin(3*π*y)*sin(3*π*z)
        end
        @inbounds for local_grid in GridField(grid,∇f_exact)
            @unpack ijk,x,y,z = local_grid
            ∇f_exact[ijk] = 3*π*sin(3*π*x)*sin(3*π*y)*cos(3*π*z)
        end
        ∇!(∇f, f, grid, 3)

        # Ghost points are not populated:
        ∇f[:,:,1] .= 0;   ∇f_exact[:,:,1] .= 0
        ∇f[:,:,end] .= 0; ∇f_exact[:,:,end] .= 0

        err[k] = norm(∇f_exact .- ∇f) / length(∇f)
    end
    @test err[1] ≈ 0.007431514201977017
    @test err[2] ≈ 0.0014919756305751361
    @test err[3] ≈ 0.00017156454627127846
    @test err[4] ≈ 1.7198655291695e-5

    conv = convergence_rate(err, Δh)

    @test conv[1] ≈ 1.5011397571937848
    @test conv[2] ≈ 2.486531977140251
    @test conv[3] ≈ 2.9549859273624386
end

function ∇²_simple!(∇²f::T, f::T, grid::Grid) where {T<:AbstractField}
    ∇f = dual_fields(grid, f)
    ∇²f_x = deepcopy(f)
    ∇²f_y = deepcopy(f)
    ∇²f_z = deepcopy(f)
    ∇!(∇f[1], f, grid)
    ∇!(∇f[2], f, grid)
    ∇!(∇f[3], f, grid)
    ∇!(∇²f_x, ∇f[1], grid)
    ∇!(∇²f_y, ∇f[2], grid)
    ∇!(∇²f_z, ∇f[3], grid)
    ∇²f .= ∇²f_x .+ ∇²f_y .+ ∇²f_z
    assign!(∇²f, 0, Ghost())
end

@testset "Discrete operators - ∇²f test against simple" begin
    FT = Float64
    a, b = FT(0.0), FT(1.0)
    n = 10
    grid = Cube(a, b, n; warpfun=Roberts_both, args=(FT(1.01),))
    f = CellCenter(grid)
    ∇²f = CellCenter(grid)
    ∇²f_simple = CellCenter(grid)
    @inbounds for local_grid in GridField(grid,f)
        @unpack ijk,x,y,z = local_grid
        f[ijk] = sin(3*π*x)*sin(3*π*y)*sin(3*π*z)
    end
    ∇²_simple!(∇²f_simple, f, grid)
    ∇²!(∇²f, f, grid)
    @test all(∇²f .≈ ∇²f_simple)

    f = CellCorner(grid)
    ∇²f = CellCorner(grid)
    ∇²f_simple = CellCorner(grid)
    @inbounds for local_grid in GridField(grid,f)
        @unpack ijk,x,y,z = local_grid
        f[ijk] = sin(3*π*x)*sin(3*π*y)*sin(3*π*z)
    end
    ∇²_simple!(∇²f_simple, f, grid)
    ∇²!(∇²f, f, grid)

    # We get hit by precision error a bit with ∇², so need to relax approx:
    @test all(isapprox.(∇²f, ∇²f_simple; atol=10^4*eps()))
end


text:
#   Unit "coherence" of the FourierAnalysis Package for julia language
#
#   MIT License
#   Copyright (c) 2019-2022,
#   Marco Congedo, CNRS, Grenobe, France:
#   https://sites.google.com/site/marcocongedo/home

# ? CONTENTS :
#   This unit implements Welch coherence estimates using FFTW

#############################################################################
# Generic constructors of Coherence objects:
#
# Enable construction giving only `y`, `sr`, `wl`, `DC`, `taper`
# and `non-linear` argument.
# `flabels` is generated automatically, `smoothing` is set to `noSmoother`
# and `tril` is set to false
Coherence(y, sr, wl, DC, taper, nonlinear) =
    Coherence(y, sr, wl, DC, taper, fdf(sr, wl; DC=DC), nonlinear, noSmoother, false)
#
# As above, but setting by default `DC` and `nonlinear` to false
Coherence(y, sr, wl, taper) =
    Coherence(y, sr, wl, false, taper, fdf(sr, wl; DC=false), false, noSmoother, false)
#############################################################################

"""
```julia
(1)
function coherence(𝙎    :: CrossSpectra;
               allkinds :: Bool = false)

(2)
function coherence(𝓢    :: CrossSpectraVector;
               allkinds :: Bool = false,
               check    :: Bool = true)

(3)
function coherence(X  :: Matrix{T},
                   sr :: Int,
                   wl :: Int;
               tapering  :: Union{Taper, TaperKind} = harris4,
               planner   :: Planner                 = getplanner,
               slide     :: Int                     = 0,
               DC        :: Bool                    = false,
               nonlinear :: Bool                    = false,
               smoothing :: Smoother                = noSmoother,
               tril      :: Bool                    = false,
               allkinds  :: Bool                    = false,
               ⏩       :: Bool                    = true) where T<:Real

(4)
function coherence(𝐗 :: Vector{Matrix{T}},
              < same argument sr, ..., ⏩ of method (3) > where T<:Real

```

**alias**: coh

!!! note "export conflict"
    This identifier may conflict with the DSP package.
    Invoke it as `FourierAnalysis.coherence`.

(1)

Construct a [Coherence](@ref) object from a [CrossSpectra](@ref) object,
allowing coherence estimates using the Welch method.
All non-data fields are copied from the cross-spectra object, i.e.,
all fields but `y`, which holds the coherence matrices that are computed
by this function.

If `𝙎.tril` is true, only the lower triangular part of the coherence matrices
is computed and the matrices in `y` are LowerTriangular matrices,
otherwise the full matrices are computed and `y` will hold Hermitian
matrices (see [LinearAlgebra](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#Linear-Algebra-1)).

If optional keyword argument `allkinds` is true, all five
[kinds of coherence](@ref) are returned. In this case the output
is a 5-tuple of Coherence objects, in the order:
- *total* coherence,
- *real* coherence,
- *instantaneous* coherence
- *imaginary* coherence,
- *lagged* coherence.

If `allkinds` is false (default), only the *total* (classical) coherence
is returned as a single `Coherence` object.

(2)

Construct a [CoherenceVector](@ref) object from the input
[CrossSpectraVector](@ref) object `𝓢`, allowing coherence estimates
using the Welch method. This method calls method (1) for all
objects in `𝓢`.

The `allkinds` optional keyword parameter has the same meaning
as in method (1). In this method if the argument is passed as true,
the output is a 5-tuple of `CoherenceVector` objects.

If optional keyword argument `check` is true (default), it is checked that all
`CrossSpectra` objects in `𝓢` have the same value in the ``.sr`, `.wl`, `.DC`,
`.taper`, `.nonlinear` and `.smoothing` field. If this is not the case,
an error message is printed pointing to the first field that is not
identical in all objects and `Nothing` is returned.

(3)

Given a multivariate time series data matrix `X` of dimension ``t`` x ``n``,
where ``t`` is the number of samples (rows) and ``n`` the number of
time-series (columns), sampling rate `sr` and epoch length `wl`,
compute the squared coherence matrices of `X`, that is,
the coherence matrices at all Fourier discrete frequencies
obtained from the Welch (sliding window) average cross-spectra.

**optioal keyword arguments**:

`sr`, `wl`, `tapering`, `planner` and `slide` have the same meaning as for the
[`spectra`](@ref) function.

`DC`, `nonlinear`, `smoothing`, `tril` and `⏩` have the same meaning as
for the [`crossSpectra`](@ref) function, to which they apply
for estimating the cross-spectra.

The `allkinds` optional keyword parameter has the same meaning
as in method (1).

(4)

Construct a [CoherenceVector](@ref) object from a vector of
real multivariate data matrices. Compute the coherence matrices
from the cross-spectral matrices estimated using the Welch method
as per method (3) for all ``k`` data matrices in `𝐗`.

The ``k`` matrices in `𝐗` must have the same number of columns
(*i.e.*, the same number of time-series), but may have any number of (at least
`wl`) rows (samples).
All other arguments have the same meaning as in method (3),
with the following difference:

`⏩`: if true (default), the method is run in multi-threaded mode across the
``k`` data matrices if ``k`` is at least twice the number of threads Julia
is instructed to use, otherwise this method attempts to run each coherence
estimation in multi-threaded mode across series as per method (3).
See [Threads](@ref).

If a `Planner` is not explicitly passed as an argument,
the FFTW plan is computed once and applied for all coherence estimations.

**note for methods (1) and (3)**:

If `tril` is false (default), the output coherence object(s) is of type
`Array{Hermitian,1}`, which is the `ℍVector` type used in package
[PosDefManifold](https://github.com/Marco-Congedo/PosDefManifold.jl).
Since coherence estimates are symmetric positive definite,
they can be straightaway used as
argument to PosDefManifold's functions, e.g., for computing matrix moves on
[geodesics](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#Geodesic-equations-1),
matrix [distances](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#PosDefManifold.distance),
etc. and the the whole vector output to compute
matrix [means](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#Means-1),
[spectral embedding](https://marco-congedo.github.io/PosDefManifold.jl/dev/riemannianGeometry/#PosDefManifold.spectralEmbedding)
and more.

If `tril` is true, the output is of type `Array{LowerTriangular,1}`,
which is the `𝕃Vector` type used in PosDefManifold, that is, only the
lower triangle of the coherencee is computed in order to save time
and memory.

**note for methods (2) and (4)**:

If `tril` is false, the output is of type `Array{Array{Hermitian,1},1}`,
which is the `ℍVector₂` type used in
[PosDefManifold](https://github.com/Marco-Congedo/PosDefManifold.jl).

If `tril` is true, the output is of type `Array{Array{LowerTriangular,1},1}`,
which is the `𝕃Vector₂` type used in PosDefManifold.

**See**: [crossspectra.jl](@ref), [Spectra](@ref), [Coherence](@ref).

**Examples**:
```julia
## common code for methods (1)-(4)

using FourierAnalysis, LinearAlgebra

function generateSomeData(sr::Int, t::Int; noise::Real=1.)
    # four sinusoids of length t samples and sr sampling rate
    # peak amplitude: 0.7, 0.6, 0.5, 0.4
    # frequency:        5,   7,  13,  27
    # phase:            0, π/4, π/2,   π
    v1=sinusoidal(0.7, 5,  sr, t, 0)
    v2=sinusoidal(0.6, 7,  sr, t, π/4)
    v3=sinusoidal(0.5, 13, sr, t, π/2)
    v4=sinusoidal(0.4, 27, sr, t, π)
    return hcat(v1, v2, v3, v4) + (randn(t, 4)*noise)
end

sr, wl, t = 128, 512, 8192

# (1)

X=generateSomeData(sr, t) # multivariate data matrix 8192x4

# cross-spectra using default harris4 tapering window
S=crossSpectra(X, sr, wl)

# Only classical coherence
C=FourierAnalysis.coherence(S)

# All 5 kinds of coherence
Ctot, C2real, C3inst, C4imag, C5lag=FourierAnalysis.coherence(S, allkinds=true);
Ctot

# (2)

# generate 3 multivariate data matrices 8192x4
X=[generateSomeData(sr, t) for i=1:3]

# cross-spectra using default harris4 tapering window
S=crossSpectra(X, sr, wl)

# Only classical coherence
C=FourierAnalysis.coherence(S)

# All 5 kinds of coherence
Ctot, C2real, C3inst, C4imag, C5lag=FourierAnalysis.coherence(S, allkinds=true);
Ctot

# (3)

X=generateSomeData(sr, t) # multivariate data matrix 8192x4

# coherence using default harris4 tapering window
C=FourierAnalysis.coherence(X, sr, wl)

# check the coherence matrix at frequency 5Hz
C.y[f2b(5, sr, wl)]

# coherence using hann tapering window
C=FourierAnalysis.coherence(X, sr, wl; tapering=hann)

# using Slepian's multi-tapering
C=FourierAnalysis.coherence(X, sr, wl; tapering=slepians(sr, wl))

# compute non-linear coherence (phase-locking value)
C=FourierAnalysis.coherence(X, sr, wl; tapering=slepians(sr, wl), nonlinear=true)

# compute only the lower triangle of coherence matrices
C=FourierAnalysis.coherence(X, sr, wl; tapering=slepians(sr, wl), tril=true)

# compute all five kinds of coherence
Ctot, Creal, Cinst, Cimag, Clag=FourierAnalysis.coherence(X, sr, wl;
    tapering=slepians(sr, wl), tril=true, allkinds=true);
Ctot

# smooth a-posteriori the coherence
C2=smooth(blackmanSmoother, C)

# or compute coherence already smoothed
C=FourierAnalysis.coherence(X, sr, wl;
            tapering=slepians(sr, wl), tril=true, smoothing=blackmanSmoother)

# mean coherence matrix in 8Hz-12Hz range
M=mean(C, (8, 12)) # or also M=mean(C, (8.0, 12.0))

# extract all coherence matrices in 8Hz-12Hz range
E=extract(C, (8, 12))

# coherence matrices averaged in 2Hz band-pass regions
B=bands(C, 2)

# (4)

# generate 3 multivariate data matrices 8192x4
X=[generateSomeData(sr, t) for i=1:3]

# The examples here below use exactly the same syntax as method (3).
# However, since the input here is a vector of data matrices
# and not a single data matrix, the examples here below create a vector
# of the object created by the examples of method (3).
# For example:

# coherence using the default harris4 tapering window
# this creates a CoherenceVector object
C=FourierAnalysis.coherence(X, sr, wl)

# check the first Coherence object
C[1]

# check the coherence matrix at fr. 5Hz for the first Coherence object
C[1].y[f2b(5, sr, wl)]

# coherence using Hann tapering window
C=FourierAnalysis.coherence(X, sr, wl; tapering=hann)

# using Slepian's multi-tapering
C=coherence(X, sr, wl; tapering=slepians(sr, wl))

# compute non-linear coherence
C=FourierAnalysis.coherence(X, sr, wl; tapering=slepians(sr, wl), nonlinear=true)

# compute only the lower triangle of coherence matrices
C=FourierAnalysis.coherence(X, sr, wl; tapering=slepians(sr, wl), tril=true)

# compute all five kinds of coherence
Ctot, Creal, Cinst, Cimag, Clag=FourierAnalysis.coherence(X, sr, wl;
    tapering=slepians(sr, wl), tril=true, allkinds=true);
Ctot

# smooth a-posteriori all coherence objects in S
C2=smooth(blackmanSmoother, C)

# or compute them all already smoothed
C=FourierAnalysis.coherence(X, sr, wl; tapering=parzen, smoothing=blackmanSmoother)

# mean coherence matrix in 8Hz-12Hz range for all coherence objects
M=mean(C, (8, 12)) # or also M=mean(C, (8.0, 12.0))

# grand-average mean of the above across all Coherence objects
meanM=mean(mean(C, (8, 12)))

# extract all coherence matrices in 8Hz-12Hz range for all coherence objects
E=extract(C, (8, 12))

# grand average of coherence matrices in 8Hz-12Hz range for all coh. objects
meanE=mean(extract(C, (8, 12)))

# coherence matrices averaged in 2Hz band-pass regions for all coh. objects
B=bands(C, 2)

# Pre-compute a FFT planner and pass it as argument
# (this interesting if the function is to be called repeatedly).
plan=Planner(plan_exhaustive, 10.0, wl, eltype(X[1])) # wait 10s
C=FourierAnalysis.coherence(X, sr, wl; planner=plan)

# how faster is this?
using BenchmarkTools
@benchmark(FourierAnalysis.coherence(X, sr, wl))
@benchmark(FourierAnalysis.coherence(X, sr, wl; planner=plan))
...
...
```

"""
function coherence(𝙎    :: CrossSpectra;
               allkinds :: Bool = false)
    # f = discrete frequencies, n = time-series in X, T is the real type corresponding to the elements in 𝙎.y
    f, n, T = length(𝙎.y), size(𝙎.y[1], 1), real(eltype(𝙎.y[1]))
    M=Matrix; H=Hermitian; L=LowerTriangular
    args=(𝙎.sr, 𝙎.wl, 𝙎.DC, 𝙎.taper, 𝙎.flabels, 𝙎.nonlinear, 𝙎.smoothing, 𝙎.tril)

    if allkinds

        # function to compute the lower-triangular part of:
        # coherence [1],
        # real coherence [2],
        # imaginary coherence [3],
        # instantaneous coherence [4] &
        # lagged coherence [5]
        # for DFT frequency `𝚏`. If `realCoeff` the i>j elements are filled with 0;
        # this is used for imaginary estimations at freq. which DFT coeff. are real.
        function coh(𝚏, realCoeff)
            if 𝙎.nonlinear
                @inbounds for j=1:n-1, i=j+1:n
                    reGij = real(𝙎.y[𝚏][i, j])^2
                    imGij = imag(𝙎.y[𝚏][i, j])^2
                                𝘾[1][𝚏][i, j] = (reGij + imGij)
                                𝘾[2][𝚏][i, j] = reGij
                    realCoeff ? 𝘾[3][𝚏][i, j] = 0. : 𝘾[3][𝚏][i, j] = imGij
                                𝘾[4][𝚏][i, j] = reGij / (1. - imGij)
                    realCoeff ? 𝘾[5][𝚏][i, j] = 0. : 𝘾[5][𝚏][i, j] = imGij / (1. - reGij)
                end
            else # linear
                @inbounds for j=1:n-1, i=j+1:n
                    GiGj  = real(𝙎.y[𝚏][i, i]) * real(𝙎.y[𝚏][j, j])
                    reGij = real(𝙎.y[𝚏][i, j])^2
                    imGij = imag(𝙎.y[𝚏][i, j])^2
                                𝘾[1][𝚏][i, j] = (reGij + imGij) / GiGj
                                𝘾[2][𝚏][i, j] = reGij / GiGj
                    realCoeff ? 𝘾[3][𝚏][i, j] = 0. : 𝘾[3][𝚏][i, j] = imGij / GiGj
                                𝘾[4][𝚏][i, j] = reGij / (GiGj - imGij)
                    realCoeff ? 𝘾[5][𝚏][i, j] = 0. : 𝘾[5][𝚏][i, j] = imGij / (GiGj - reGij)
                end
            end
        end

        # Allocate memory and compute all coherence matrices.
        # If `DC`=true the first discrete frequency (CD level) is real.
        # The last discrete frequency is always real.
        𝘾 = Vector{Vector{L}}([ [L(ones(T, n, n)) for i=1:f] for l=1:5 ])
        if 𝙎.DC
            coh(1, true)                        # DC level
            for 𝚏=2:f-1 coh(𝚏, false) end       # Frequencies with complex DFT coeff.
            coh(f, true)                        # Nyquist Frequency
        else
            for 𝚏=1:f-1 coh(𝚏, false) end       # Frequencies with complex DFT coeff.
            coh(f, true)                        # Nyquist Frequency
        end

        # if `tril`=false return the full matrices, otherwise the lower-triangular matrices
        CT(l)=Coherence(𝘾[l], args...)
        CH(l)=Coherence(Vector{H}([H(M(𝘾[l][𝚏]), :L) for 𝚏=1:f]), args...)
        return 𝙎.tril ? ( CT(1), CT(2), CT(3), CT(4), CT(5) ) :
                        ( CH(1), CH(2), CH(3), CH(4), CH(5) )

    else # only classical coherence

        # Allocate memory and compute the lower-triangular part of coherence matrices.
        𝘾 = Vector{L}([L(ones(T, n, n)) for i=1:f])
        if 𝙎.nonlinear
            for 𝚏=1:f, j=1:n-1, i=j+1:n
                𝘾[𝚏][i, j] = abs2(𝙎.y[𝚏][i, j])
            end
        else # linear
            for 𝚏=1:f, j=1:n-1, i=j+1:n
                𝘾[𝚏][i, j] = abs2(𝙎.y[𝚏][i, j]) / (real(𝙎.y[𝚏][i, i]) * real(𝙎.y[𝚏][j, j]))
            end
        end

        # if `tril`=false return the full matrices, otherwise the lower-triangular matrices
        return 𝙎.tril ? Coherence(𝘾, args...) :
                        Coherence(Vector{H}([H(M(𝘾[𝚏]), :L) for 𝚏=1:f]), args...)

    end # if allkinds
end

function coherence(𝓢    :: CrossSpectraVector;
               allkinds :: Bool = false,
               check    :: Bool = true)
    check && !sameParams(𝓢, "coherence") && return Nothing
    j=length(𝓢)
    if allkinds
        # pre-allocate memory
        𝓒=Vector{CoherenceVector}([CoherenceVector(undef, j) for l=1:5])

        for k=1:j
            𝓒[1][k], 𝓒[2][k], 𝓒[3][k], 𝓒[4][k], 𝓒[5][k] = coherence(𝓢[k]; allkinds = true)
        end
        return 𝓒[1], 𝓒[2], 𝓒[3], 𝓒[4], 𝓒[5]

    else
        return [coherence(𝓢[k]; allkinds = false) for k=1:j]
    end
end

coherence(X  :: Matrix{T},
          sr :: Int,
          wl :: Int;
      tapering  :: Union{Taper, TaperKind} = harris4,
      planner   :: Planner                 = getplanner,
      slide     :: Int                     = 0,
      DC        :: Bool                    = false,
      nonlinear :: Bool                    = false,
      smoothing :: Smoother                = noSmoother,
      tril      :: Bool                    = false,
      allkinds  :: Bool                    = false,
      ⏩       :: Bool                    = true) where T<:Real =
    coherence(crossSpectra(X, sr, wl;
                           tapering  = tapering,
                           planner   = planner,
                           slide     = slide,
                           DC        = DC,
                           nonlinear = nonlinear,
                           smoothing = smoothing,
                           tril      = tril,
                           ⏩       = ⏩);
         allkinds = allkinds)

coherence(𝐗 :: Vector{Matrix{T}}, sr :: Int, wl :: Int;
          tapering  :: Union{Taper, TaperKind} = harris4,
          planner   :: Planner                 = getplanner,
          slide     :: Int                     = 0,
          DC        :: Bool                    = false,
          nonlinear :: Bool                    = false,
          smoothing :: Smoother                = noSmoother,
          tril      :: Bool                    = false,
          allkinds  :: Bool                    = false,
         ⏩       :: Bool                    = true) where T<:Real =
    coherence(crossSpectra(𝐗, sr, wl;
                           tapering  = tapering,
                           planner   = planner,
                           slide     = slide,
                           DC        = DC,
                           nonlinear = nonlinear,
                           smoothing = smoothing,
                           tril      = tril,
                           ⏩       = ⏩);
         allkinds = allkinds)

# ++++++++++++++++++++  Show override  +++++++++++++++++++ # (REPL output)
function Base.show(io::IO, ::MIME{Symbol("text/plain")}, C::Coherence)
    r=size(C.y[1], 1)
    l=length(C.flabels)
    ll=length(C.y)
    println(io, titleFont, "▦ ⋯ ▦ Coherence type; $l freq. x $(r)²-matrices")
    #println(io, "□  □    □      □        □           □", defaultFont)
    println(io, separatorFont, "⭒  ⭒    ⭒      ⭒        ⭒           ⭒", defaultFont)
    println(io, "sampling rate    (.sr): $(C.sr)")
    println(io, "epoch length     (.wl): $(C.wl)")
    println(io, "DC level         (.DC): $(C.DC)")
    println(io, "taper kind (.tapering): $(C.taper)")
    println(io, "freq. lab.  (.flabels): $(l)-", typeof(C.flabels))
    println(io, "data              (.y): $(l)-", typeof(C.y))
    # println(io, C.y)
    println(io, "non-linear(.nonlinear): ", string(C.nonlinear))
    println(io, "smoother  (.smoothing): ", string(C.smoothing))
    println(io, "triang. form   (.tril): ", string(C.tril))
    ll≠l && @warn "number of frequency labels does not match the Coherence data size" l ll
end

function Base.show(io::IO, ::MIME{Symbol("text/plain")}, 𝘾::CoherenceVector)
    println(io, titleFont, "▦ ⋯ ▦ CoherenceVector Type")
    println(io, titleFont, "  ⋯")
    println(io, separatorFont, "⭒  ⭒    ⭒      ⭒        ⭒           ⭒", defaultFont)
    println(io, "$(length(𝘾))-element Vector{Coherence}")
end

# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ #

ext: jl


text:
import StaticArrays: SVector, MVector
import DelayEmbeddings: Dataset

include("induced_invariant_measure.jl")

export InducedRectangularInvariantMeasure, inducedrectangularinvariantmeasure

"""
    struct InducedRectangularInvariantMeasure{T} <: AbstractRectangularInvariantMeasure where {T}
        points::AbstractArray{T, 2}
        ϵF::Union{Int, Float64, Vector{Int}, Vector{Float64}}
        ϵj::Union{Int, Float64, Vector{Int}, Vector{Float64}}
        bins_ϵF::AbstractArray{Float64, 2}
        bins_visited_ϵj::AbstractArray{Float64, 2}
        transferoperator_ϵj::TransferOperatorRectangularBinning
        measure_ϵj::InvariantDistribution
        measure_induced::InvariantDistribution
    end

An induced RectangularInvariantMeasure. Created from a set of points by
discretizing the state space into rectangular bins with edge lengths dictated
by the binning scheme `ϵ`.

The invariant measure is first computed at resolution `ϵj`, then induced at the
final resolution `ϵF` by considering the overlap of the bins between resolutions.

## Fields
- **`points`**: The points for which to estimate the invariant measure.

- **`ϵF`**: The binning scheme at the final resolution, expressed as absolute
edge lengths. The following `ϵ` are valid, and will all be converted into
    the `ϵ::Vector{Float64` because knowing edge lengths at both resolutions
    is necessary.
        1. `ϵ::Int` divides each axis into `ϵ` intervals of the same size.
        2. `ϵ::Float` divides each axis into intervals of size `ϵ`.
        3. `ϵ::Vector{Int}` divides the i-th axis into `ϵᵢ` intervals of the same size.
        4. `ϵ::Vector{Float64}` divides the i-th axis into intervals of size `ϵᵢ`.
- **`ϵF_absolute`**: `ϵF` converted to absolute edge lengths.

- **`ϵj`**: The binning scheme at the resolution from which we induce the measure.

- **`ϵj_absolute`**: `ϵj` converted to absolute edge lengths.

- **`bins_ϵF`**: The origins of all bins covering the hyperrectangle spanned
    by a box covering of `points` at resolution `ϵF`.

- **`bins_visited_ϵj`**: As for `bins_ϵF`, but at resolution `ϵj` and only
    for nonempty bins.

- **`transferoperator_ϵj`**: The transfer operator from from which the invariant
    measure is obtained at resolution `ϵj`.

- **`measure_ϵj`**: The invariant measure obtained from at resolution `ϵj`,
    which is obtained from `transferoperator_ϵj`.

- **`measure_induced`**: The invariant measure induced at resolution `ϵF`
    from resolution `ϵj`.
"""
struct InducedRectangularInvariantMeasure{T} <: AbstractRectangularInvariantMeasure where {T}
    points::AbstractArray{T, 2}
    ϵF::Union{Int, Float64, Vector{Int}, Vector{Float64}}
    ϵF_absolute::Vector{Float64}
    ϵj::Union{Int, Float64, Vector{Int}, Vector{Float64}}
    ϵj_absolute::Vector{Float64}
    bins_ϵF::AbstractArray{Float64, 2}
    bins_visited_ϵj::AbstractArray{Float64, 2}
    transferoperator_ϵj::TransferOperatorRectangularBinning
    measure_ϵj::InvariantDistribution
    measure_induced::InvariantDistribution
end

"""
    inducedrectangularinvariantmeasure(points, ϵF, ϵⱼ)

Compute the invariant measure the boxes of a hyperrectangular subdivision
of the hypercube spanned by `points` at resolution
`ϵF` induced by the invariant measure computed for the same points at resolution
`ϵⱼ`. Returns a `InducedRectangularInvariantMeasure` instance.
"""
function inducedrectangularinvariantmeasure(points::AbstractArray{T, 2}, ϵF, ϵⱼ) where {T}
    ϵF_absolute = minima_and_stepsizes(points, ϵF)[2]
    ϵⱼ_absolute = minima_and_stepsizes(points, ϵⱼ)[2]

    # All possible bins in the final partition given by ϵF
    allbins_ϵF = boxorigins(points, ϵF_absolute)

    # The unique bins visited by the orbit in the partition given by ϵⱼ
    visitedbins_ϵⱼ = unique(assign_coordinate_labels(points, ϵⱼ_absolute), dims = 2)

    # The transfer operator in the partition given by ϵⱼ
    TOϵⱼ = TransferOperatorEstimatorRectangularBinning(points, ϵⱼ_absolute)

    # The measure of the of the visited bins in the partition given by ϵⱼ
    μϵⱼ = invariantmeasure(TOϵⱼ)

    # The induced measure at resolution ϵF. This does not filter
    # zero entries, so be careful when comparing with the distribution
    # obtained using invariantmeasure(::TransferOperatorRectangularBinning)
    μϵF = μ_allbins_ϵF_induced_by_binningscheme_ϵⱼ(
        points,
        ϵF,
        allbins_ϵF,
        ϵⱼ,
        visitedbins_ϵⱼ,
        μϵⱼ.dist)

    InducedRectangularInvariantMeasure(
        points,
        ϵF,
        ϵF_absolute,
        ϵⱼ,
        ϵⱼ_absolute,
        allbins_ϵF,
        visitedbins_ϵⱼ,
        TOϵⱼ,
        μϵⱼ,
        InvariantDistribution(μϵF, findall(μϵF .> 0))
    )
end

function inducedrectangularinvariantmeasure(points::Vector{Vector{T}}, ϵF, ϵⱼ) where {T}
    inducedrectangularinvariantmeasure(hcat(points...,), ϵF, ϵⱼ)
end

function inducedrectangularinvariantmeasure(points::Vector{SVector{T}}, ϵF, ϵⱼ) where {T}
    inducedrectangularinvariantmeasure(Array(hcat(points...,)), ϵF, ϵⱼ)
end

function inducedrectangularinvariantmeasure(points::Vector{MVector{T}}, ϵF, ϵⱼ) where {T}
    inducedrectangularinvariantmeasure(Array(hcat(points...,)), ϵF, ϵⱼ)
end

function inducedrectangularinvariantmeasure(points::Dataset, ϵF, ϵⱼ) where {T}
    inducedrectangularinvariantmeasure(transpose(Matrix(points)), ϵF, ϵⱼ)
end

function summarise(invm::InducedRectangularInvariantMeasure)
    ϵF = invm.ϵF
    ϵⱼ = invm.ϵj
    ϵF_str = " partition resulting from binning scheme ϵF = $ϵF"
    ϵⱼ_str = " partition resulting from binning scheme ϵⱼ = $ϵⱼ"
    D = size(invm.points, 1)
    npts = size(invm.points, 2)

    measure_type = typeof(invm)
    measure_type_str = "$measure_type"
    pts_str = "$npts $D-dimensional points"
    measure_type_str*" from "*pts_str*" at"*ϵF_str*" induced by"*ϵⱼ_str

end

Base.show(io::IO, invm::InducedRectangularInvariantMeasure) = println(io, summarise(invm))


text:
function time_stepping_works_with_coriolis(arch, FT, Coriolis)
    grid = RegularCartesianGrid(FT, size=(1, 1, 1), extent=(1, 2, 3))
    c = Coriolis(FT, latitude=45)
    model = IncompressibleModel(grid=grid, architecture=arch, float_type=FT, coriolis=c)

    time_step!(model, 1, euler=true)

    return true # Test that no errors/crashes happen when time stepping.
end

function time_stepping_works_with_closure(arch, FT, Closure)
    # Use halos of size 2 to accomadate time stepping with AnisotropicBiharmonicDiffusivity.
    grid = RegularCartesianGrid(FT; size=(1, 1, 1), halo=(2, 2, 2), extent=(1, 2, 3))

    model = IncompressibleModel(grid=grid, architecture=arch, float_type=FT, closure=Closure(FT))
    time_step!(model, 1, euler=true)

    return true  # Test that no errors/crashes happen when time stepping.
end

function time_stepping_works_with_advection_scheme(arch, advection_scheme)
    # Use halo=(3, 3, 3) to accomodate WENO-5 advection scheme
    grid = RegularCartesianGrid(size=(1, 1, 1), halo=(3, 3, 3), extent=(1, 2, 3))
    model = IncompressibleModel(grid=grid, architecture=arch, advection=advection_scheme)
    time_step!(model, 1, euler=true)
    return true  # Test that no errors/crashes happen when time stepping.
end

function time_stepping_works_with_nothing_closure(arch, FT)
    grid = RegularCartesianGrid(FT; size=(1, 1, 1), extent=(1, 2, 3))
    model = IncompressibleModel(grid=grid, architecture=arch, float_type=FT, closure=nothing)
    time_step!(model, 1, euler=true)
    return true  # Test that no errors/crashes happen when time stepping.
end

function time_stepping_works_with_nonlinear_eos(arch, FT, EOS)
    grid = RegularCartesianGrid(FT; size=(1, 1, 1), extent=(1, 2, 3))

    eos = EOS()
    b = SeawaterBuoyancy(equation_of_state=eos)

    model = IncompressibleModel(architecture=arch, float_type=FT, grid=grid, buoyancy=b)
    time_step!(model, 1, euler=true)

    return true  # Test that no errors/crashes happen when time stepping.
end

function run_first_AB2_time_step_tests(arch, FT)
    add_ones(args...) = 1.0

    # Weird grid size to catch https://github.com/CliMA/Oceananigans.jl/issues/780
    grid = RegularCartesianGrid(FT, size=(13, 17, 19), extent=(1, 2, 3))

    model = IncompressibleModel(grid=grid, architecture=arch, float_type=FT, forcing=ModelForcing(T=add_ones))
    time_step!(model, 1, euler=true)

    # Test that GT = 1, T = 1 after 1 time step and that AB2 actually reduced to forward Euler.
    @test all(interior(model.timestepper.Gⁿ.u) .≈ 0)
    @test all(interior(model.timestepper.Gⁿ.v) .≈ 0)
    @test all(interior(model.timestepper.Gⁿ.w) .≈ 0)
    @test all(interior(model.timestepper.Gⁿ.T) .≈ 1.0)
    @test all(interior(model.timestepper.Gⁿ.S) .≈ 0)

    @test all(interior(model.velocities.u) .≈ 0)
    @test all(interior(model.velocities.v) .≈ 0)
    @test all(interior(model.velocities.w) .≈ 0)
    @test all(interior(model.tracers.T)    .≈ 1.0)
    @test all(interior(model.tracers.S)    .≈ 0)

    return nothing
end

"""
    This tests to make sure that the velocity field remains incompressible (or divergence-free) as the model is time
    stepped. It just initializes a cube shaped hot bubble perturbation in the center of the 3D domain to induce a
    velocity field.
"""
function incompressible_in_time(arch, FT, Nt, timestepper)
    Nx, Ny, Nz = 32, 32, 32
    Lx, Ly, Lz = 10, 10, 10

    grid = RegularCartesianGrid(FT, size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz))
    model = IncompressibleModel(grid=grid, architecture=arch, float_type=FT, timestepper=timestepper)

    grid = model.grid
    u, v, w = model.velocities

    div_U = CellField(FT, arch, grid, TracerBoundaryConditions(grid))

    # Just add a temperature perturbation so we get some velocity field.
    @. model.tracers.T.data[8:24, 8:24, 8:24] += 0.01

    for n in 1:Nt
        if timestepper === :QuasiAdamsBashforth2
            time_step!(model, 0.05, euler = n==1)
        elseif timestepper === :RungeKutta3
            time_step!(model, 0.05)
        end
    end

    event = launch!(arch, grid, :xyz, divergence!, grid, u.data, v.data, w.data, div_U.data, dependencies=Event(device(arch)))
    wait(device(arch), event)

    min_div = minimum(interior(div_U))
    max_div = maximum(interior(div_U))
    max_abs_div = maximum(abs, interior(div_U))
    sum_div = sum(interior(div_U))
    sum_abs_div = sum(abs, interior(div_U))
    @info "Velocity divergence after $Nt time steps [$(typeof(arch)), $FT, $timestepper]: " *
          "min=$min_div, max=$max_div, max_abs_div=$max_abs_div, sum=$sum_div, abs_sum=$sum_abs_div"

    # We are comparing with 0 so we use absolute tolerances. They are a bit larger than eps(Float64) and eps(Float32)
    # because we are summing over the absolute value of many machine epsilons. A better atol value may be
    # Nx*Ny*Nz*eps(FT) but it's much higher than the observed abs_sum_div.
    return isapprox(max_abs_div, 0, atol=5e-8)
end

"""
    tracer_conserved_in_channel(arch, FT, Nt)

Create a super-coarse eddying channel model with walls in the y and test that
temperature is conserved after `Nt` time steps.
"""
function tracer_conserved_in_channel(arch, FT, Nt)
    Nx, Ny, Nz = 16, 32, 16
    Lx, Ly, Lz = 160e3, 320e3, 1024

    α = (Lz/Nz)/(Lx/Nx) # Grid cell aspect ratio.
    νh, κh = 20.0, 20.0
    νz, κz = α*νh, α*κh

    topology = (Periodic, Bounded, Bounded)
    grid = RegularCartesianGrid(size=(Nx, Ny, Nz), extent=(Lx, Ly, Lz))
    model = IncompressibleModel(architecture = arch, float_type = FT, grid = grid,
                                closure = AnisotropicDiffusivity(νh=νh, νz=νz, κh=κh, κz=κz))

    Ty = 1e-4  # Meridional temperature gradient [K/m].
    Tz = 5e-3  # Vertical temperature gradient [K/m].

    # Initial temperature field [°C].
    T₀(x, y, z) = 10 + Ty*y + Tz*z + 0.0001*rand()
    set!(model, T=T₀)

    Tavg0 = mean(interior(model.tracers.T))

    for n in 1:Nt
        time_step!(model, 600, euler= n==1)
    end

    Tavg = mean(interior(model.tracers.T))
    @info "Tracer conservation after $Nt time steps [$(typeof(arch)), $FT]: " *
          "⟨T⟩-T₀=$(Tavg-Tavg0) °C"

    return isapprox(Tavg, Tavg0, atol=Nx*Ny*Nz*eps(FT))
end

Planes = (FPlane, NonTraditionalFPlane, BetaPlane, NonTraditionalBetaPlane)

Closures = (IsotropicDiffusivity, AnisotropicDiffusivity,
            AnisotropicBiharmonicDiffusivity, TwoDimensionalLeith,
            SmagorinskyLilly, BlasiusSmagorinsky,
            AnisotropicMinimumDissipation, RozemaAnisotropicMinimumDissipation)

advection_schemes = (CenteredSecondOrder(), UpwindBiasedThirdOrder(), CenteredFourthOrder(), WENO5())

timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3)

@testset "Time stepping" begin
    @info "Testing time stepping..."

    for arch in archs, FT in float_types
        @testset "Time stepping with DateTimes [$(typeof(arch)), $FT]" begin
            @info "  Testing time stepping with datetime clocks [$(typeof(arch)), $FT]"

            model = IncompressibleModel(grid = RegularCartesianGrid(size=(1, 1, 1), extent=(1, 1, 1)),
                                        clock = Clock(time=DateTime(2020)))

            time_step!(model, 7.883)
            @test model.clock.time == DateTime("2020-01-01T00:00:07.883")

            model = IncompressibleModel(grid = RegularCartesianGrid(size=(1, 1, 1), extent=(1, 1, 1)),
                                        clock = Clock(time=TimeDate(2020)))

            time_step!(model, 123e-9)  # 123 nanoseconds
            @test model.clock.time == TimeDate("2020-01-01T00:00:00.000000123")
        end
    end

    @testset "Coriolis" begin
        for arch in archs, FT in [Float64], Coriolis in Planes
            @info "  Testing that time stepping works [$(typeof(arch)), $FT, $Coriolis]..."
            @test time_stepping_works_with_coriolis(arch, FT, Coriolis)
        end
    end

    @testset "Advection schemes" begin
        for arch in archs, advection_scheme in advection_schemes
            @info "  Testing time stepping with advection schemes [$(typeof(arch)), $(typeof(advection_scheme))]"
            @test time_stepping_works_with_advection_scheme(arch, advection_scheme)
        end
    end

    @testset "Turbulence closures" begin
        for arch in archs, FT in [Float64]

            @info "  Testing that time stepping works [$(typeof(arch)), $FT, nothing]..."
            @test time_stepping_works_with_nothing_closure(arch, FT)

            for Closure in Closures
                @info "  Testing that time stepping works [$(typeof(arch)), $FT, $Closure]..."
                if Closure === TwoDimensionalLeith
                    # This test is extremely slow; skipping for now.
                    @test_skip time_stepping_works_with_closure(arch, FT, Closure)
                else
                    @test time_stepping_works_with_closure(arch, FT, Closure)
                end
            end
        end
    end

    @testset "Idealized nonlinear equation of state" begin
        for arch in archs, FT in [Float64]
            for eos_type in (SeawaterPolynomials.RoquetEquationOfState, SeawaterPolynomials.TEOS10EquationOfState)
                @info "  Testing that time stepping works with " *
                        "RoquetIdealizedNonlinearEquationOfState [$(typeof(arch)), $FT, $eos_type]"
                @test time_stepping_works_with_nonlinear_eos(arch, FT, eos_type)
            end
        end
    end

    @testset "2nd-order Adams-Bashforth" begin
        @info "  Testing 2nd-order Adams-Bashforth..."
        for arch in archs, FT in float_types
            run_first_AB2_time_step_tests(arch, FT)
        end
    end

    @testset "Incompressibility" begin
        @info "  Testing incompressibility..."
        for arch in archs, FT in float_types, Nt in [1, 10, 100], timestepper in timesteppers
            @test incompressible_in_time(arch, FT, Nt, timestepper)
        end
    end

    @testset "Tracer conservation in channel" begin
        @info "  Testing tracer conservation in channel..."
        for arch in archs, FT in float_types
            @test tracer_conserved_in_channel(arch, FT, 10)
        end
    end
end


text:
using Celeste: DeterministicVI, SensitiveFloats

using Base.Test
using Distributions
using DerivativeTestUtils

######################################
# Helper functions

function true_star_init()
    blob, ea, body = gen_sample_star_dataset(perturb=false)

    ea.vp[1][ids.a[:, 1]] = [ 1.0 - 1e-4, 1e-4 ]
    ea.vp[1][ids.r2] = 1e-4
    ea.vp[1][ids.r1] = log(sample_star_fluxes[3]) - 0.5 * ea.vp[1][ids.r2]
    #ea.vp[1][ids.r1] = sample_star_fluxes[3] ./ ea.vp[1][ids.r2]
    ea.vp[1][ids.c2] = 1e-4

    blob, ea, body
end

"""
Return a vector of (h, w) indices of tiles that contain this source.
"""
function find_source_tiles(s::Int, b::Int, ea::ElboArgs)
    [ind2sub(size(ea.tile_source_map[b]), ind) for ind in
        find([ s in sources for sources in ea.tile_source_map[b]])]
end

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

function test_set_hess()
    sf = zero_sensitive_float(CanonicalParams)
    set_hess!(sf, 2, 3, 5.0)
    @test_approx_eq sf.h[2, 3] 5.0
    @test_approx_eq sf.h[3, 2] 5.0

    set_hess!(sf, 4, 4, 6.0)
    @test_approx_eq sf.h[4, 4] 6.0
end

function test_bvn_cov()
        e_axis = .7
        e_angle = pi/5
        e_scale = 2.

        manual_11 = e_scale^2 * (1 + (e_axis^2 - 1) * (sin(e_angle))^2)
        util_11 = DeterministicVI.get_bvn_cov(e_axis, e_angle, e_scale)[1,1]
        @test_approx_eq util_11 manual_11

        manual_12 = e_scale^2 * (1 - e_axis^2) * (cos(e_angle)sin(e_angle))
        util_12 = DeterministicVI.get_bvn_cov(e_axis, e_angle, e_scale)[1,2]
        @test_approx_eq util_12 manual_12

        manual_22 = e_scale^2 * (1 + (e_axis^2 - 1) * (cos(e_angle))^2)
        util_22 = DeterministicVI.get_bvn_cov(e_axis, e_angle, e_scale)[2,2]
        @test_approx_eq util_22 manual_22
end

function test_tile_predicted_image()
    blob, ea, body = gen_sample_star_dataset(perturb=false)
    tile = ea.images[1].tiles[1, 1]
    tile_source_map = ea.tile_source_map[1][1, 1]
    pred_image =
        DeterministicVI.tile_predicted_image(tile, ea, tile_source_map; include_epsilon=true)

    # Regress the tile pixels onto the predicted image
    # TODO: Why isn't the regression closer to one?    Something in the sample data
    # generation?
    reg_coeff = dot(tile.pixels[:], pred_image[:]) / dot(pred_image[:], pred_image[:])
    residuals = pred_image * reg_coeff - tile.pixels
    residual_sd = sqrt(mean(residuals .^ 2))

    @test residual_sd / mean(tile.pixels) < 0.1
end

function test_derivative_flags()
    blob, ea, body = gen_two_body_dataset()
    keep_pixels = 10:11
    trim_tiles!(ea.images, keep_pixels)

    elbo = DeterministicVI.elbo(ea)

    elbo_noderiv = DeterministicVI.elbo(ea; calculate_derivs=false)
    @test_approx_eq elbo.v[1] elbo_noderiv.v
    @test_approx_eq elbo_noderiv.d zeros(size(elbo_noderiv.d))
    @test_approx_eq elbo_noderiv.h zeros(size(elbo_noderiv.h))

    elbo_nohess = DeterministicVI.elbo(ea; calculate_hessian=false)
    @test_approx_eq elbo.v[1] elbo_nohess.v
    @test_approx_eq elbo.d elbo_nohess.d
    @test_approx_eq elbo_noderiv.h zeros(size(elbo_noderiv.h))
end

function test_active_sources()
    # Test that the derivatives of the expected brightnesses partition in
    # active_sources.

    blob, ea, body = gen_two_body_dataset()
    keep_pixels = 10:11
    trim_tiles!(ea.images, keep_pixels)
    b = 1
    tile = ea.images[b].tiles[1,1]
    h, w = 10, 10

    ea.active_sources = [1, 2]
    elbo_lik_12 = DeterministicVI.elbo_likelihood(ea)

    ea.active_sources = [1]
    elbo_lik_1 = DeterministicVI.elbo_likelihood(ea)

    ea.active_sources = [2]
    elbo_lik_2 = DeterministicVI.elbo_likelihood(ea)

    @test_approx_eq elbo_lik_12.v[1] elbo_lik_1.v
    @test_approx_eq elbo_lik_12.v[1] elbo_lik_2.v

    @test_approx_eq elbo_lik_12.d[:, 1] elbo_lik_1.d[:, 1]
    @test_approx_eq elbo_lik_12.d[:, 2] elbo_lik_2.d[:, 1]

    P = length(CanonicalParams)
    @test_approx_eq elbo_lik_12.h[1:P, 1:P] elbo_lik_1.h
    @test_approx_eq elbo_lik_12.h[(1:P) + P, (1:P) + P] elbo_lik_2.h
end

function test_that_variance_is_low()
    # very peaked variational distribution---variance for F(m) should be low
    blob, ea, body = true_star_init();

    test_b = 3
    tile = ea.images[test_b].tiles[1,1];
    tile_source_map = ea.tile_source_map[test_b][1,1];

    h, w = 10, 12
    star_mcs, gal_mcs = DeterministicVI.load_bvn_mixtures(ea, tile.b);
    sbs = DeterministicVI.SourceBrightness{Float64}[
      DeterministicVI.SourceBrightness(ea.vp[s]) for s in 1:ea.S];

    elbo_vars = DeterministicVI.ElboIntermediateVariables(
      Float64, ea.S, length(ea.active_sources));

    clear!(elbo_vars.E_G);
    clear!(elbo_vars.var_G);
    DeterministicVI.get_expected_pixel_brightness!(
      elbo_vars, h, w, sbs, star_mcs, gal_mcs, tile, ea, tile_source_map);

    @test 0 < elbo_vars.var_G.v[1] < 1e-2 * elbo_vars.E_G.v[1]^2
end

function test_that_star_truth_is_most_likely()
    blob, ea, body = true_star_init();
    best = DeterministicVI.elbo_likelihood(ea);

    for bad_a in [.3, .5, .9]
        ea_a = deepcopy(ea)
        ea_a.vp[1][ids.a[:, 1]] = [ 1.0 - bad_a, bad_a ]
        bad_a_lik = DeterministicVI.elbo_likelihood(ea_a)
        @test best.v[1] > bad_a_lik.v[1]
    end

    for h2 in -2:2
        for w2 in -2:2
            if !(h2 == 0 && w2 == 0)
                ea_mu = deepcopy(ea)
                ea_mu.vp[1][ids.u] += [h2 * .5, w2 * .5]
                bad_mu = DeterministicVI.elbo_likelihood(ea_mu)
                @test best.v[1] > bad_mu.v[1]
            end
        end
    end

    for delta in [.7, .9, 1.1, 1.3]
        ea_r1 = deepcopy(ea)
        ea_r1.vp[1][ids.r1] += log(delta)
        bad_r1 = DeterministicVI.elbo_likelihood(ea_r1)
        @test best.v[1] > bad_r1.v[1]
    end

    for b in 1:4
        for delta in [-.3, .3]
            ea_c1 = deepcopy(ea)
            ea_c1.vp[1][ids.c1[b, 1]] += delta
            bad_c1 = DeterministicVI.elbo_likelihood(ea_c1)
            @test best.v[1] > bad_c1.v[1]
        end
    end

end

function test_that_galaxy_truth_is_most_likely()
    blob, ea, body = gen_sample_galaxy_dataset(perturb=false);
    ea.vp[1][ids.a[:, 1]] = [ 0.01, .99 ]
    best = DeterministicVI.elbo_likelihood(ea);

    for bad_a in [.3, .5, .9]
        ea_a = deepcopy(ea);
        ea_a.vp[1][ids.a[:, 1]] = [ 1.0 - bad_a, bad_a ];
        bad_a =
          DeterministicVI.elbo_likelihood(ea_a; calculate_derivs=false);
        @test best.v[1] > bad_a.v[1];
    end

    for h2 in -2:2
        for w2 in -2:2
            if !(h2 == 0 && w2 == 0)
                ea_mu = deepcopy(ea)
                ea_mu.vp[1][ids.u] += [h2 * .5, w2 * .5]
                bad_mu = DeterministicVI.elbo_likelihood(
                  ea_mu; calculate_derivs=false)
                @test best.v[1] > bad_mu.v[1]
            end
        end
    end

    for bad_scale in [.8, 1.2]
        ea_r1 = deepcopy(ea)
        ea_r1.vp[1][ids.r1] += 2 * log(bad_scale)
        bad_r1 = DeterministicVI.elbo_likelihood(
          ea_r1; calculate_derivs=false)
        @test best.v[1] > bad_r1.v[1]
    end

    for n in [:e_axis, :e_angle, :e_scale]
        for bad_scale in [.8, 1.2]
            ea_bad = deepcopy(ea)
            ea_bad.vp[1][getfield(ids, n)] *= bad_scale
            bad_elbo = DeterministicVI.elbo_likelihood(
              ea_bad; calculate_derivs=false)
            @test best.v[1] > bad_elbo.v[1]
        end
    end

    for b in 1:4
        for delta in [-.3, .3]
            ea_c1 = deepcopy(ea)
            ea_c1.vp[1][ids.c1[b, 2]] += delta
            bad_c1 = DeterministicVI.elbo_likelihood(
              ea_c1; calculate_derivs=false)
            @test best.v[1] > bad_c1.v[1]
        end
    end
end

function test_coadd_cat_init_is_most_likely()  # on a real stamp
    stamp_id = "5.0073-0.0739_2kpsf"
    blob = SampleData.load_stamp_blob(datadir, stamp_id);

    cat_entries = SampleData.load_stamp_catalog(datadir, "s82-$stamp_id", blob);
    bright(ce) = sum(ce.star_fluxes) > 3 || sum(ce.gal_fluxes) > 3
    cat_entries = filter(bright, cat_entries)

    ce_pix_locs =
      [ [ WCS.world_to_pix(blob[b].wcs, ce.pos) for b=1:5 ]
        for ce in cat_entries ]

    function ce_inbounds(ce)
        pix_locs = [ WCS.world_to_pix(blob[b].wcs, ce.pos) for b=1:5 ]
        inbounds(pos) = pos[1] > -10. && pos[2] > -10 &&
                        pos[1] < 61 && pos[2] < 61
        reduce(&, [inbounds(pos) for pos in pix_locs])
    end
    cat_entries = filter(ce_inbounds, cat_entries)

    ea = make_elbo_args(blob, cat_entries)
    for s in 1:length(cat_entries)
        ea.vp[s][ids.a[2, 1]] = cat_entries[s].is_star ? 0.01 : 0.99
        ea.vp[s][ids.a[1, 1]] = 1.0 - ea.vp[s][ids.a[2, 1]]
    end
    best = DeterministicVI.elbo_likelihood(ea; calculate_derivs=false);

    # s is the brightest source.
    s = 1

    for bad_scale in [.7, 1.3]
        ea_r1 = deepcopy(ea)
        ea_r1.vp[s][ids.r1] += 2 * log(bad_scale)
        bad_r1 = DeterministicVI.elbo_likelihood(
          ea_r1; calculate_derivs=false)
        @test best.v[1] > bad_r1.v[1]
    end

    for n in [:e_axis, :e_angle, :e_scale]
        for bad_scale in [.6, 1.8]
            ea_bad = deepcopy(ea)
            ea_bad.vp[s][getfield(ids, n)] *= bad_scale
            bad_elbo = DeterministicVI.elbo_likelihood(
              ea_bad; calculate_derivs=false)
            @test best.v[1] > bad_elbo.v[1]
        end
    end

    for bad_a in [.3, .7]
        ea_a = deepcopy(ea)
        ea_a.vp[s][ids.a[:, 1]] = [ 1.0 - bad_a, bad_a ]

        bad_a = DeterministicVI.elbo_likelihood(
          ea_a; calculate_derivs=false)
        @test best.v[1] > bad_a.v[1]
    end

    for h2 in -2:2
        for w2 in -2:2
            if !(h2 == 0 && w2 == 0)
                ea_mu = deepcopy(ea)
                ea_mu.vp[s][ids.u] += [0.5h2, 0.5w2]
                bad_mu = DeterministicVI.elbo_likelihood(
                  ea_mu; calculate_derivs=false)
                @test best.v[1] > bad_mu.v[1]
            end
        end
    end

    for b in 1:4
        for delta in [-2., 2.]
            ea_c1 = deepcopy(ea)
            ea_c1.vp[s][ids.c1[b, :]] += delta
            bad_c1 = DeterministicVI.elbo_likelihood(
              ea_c1; calculate_derivs=false)
            info("$(best.v[1])  >  $(bad_c1.v[1])")
            @test best.v[1] > bad_c1.v[1]
        end
    end
end

function test_tiny_image_tiling()
  # Test that the tiling doesn't matter much for a body that is nearly a
  # point with a narrow psf.

  blob0 = SampleData.load_stamp_blob(datadir, "164.4311-39.0359_2kpsf");
  pc = PsfComponent(1./3, zeros(2), 1e-4 * eye(2));
  trivial_psf = [pc, pc, pc]
  pixels = ones(100, 1) * 12
  pixels[98:100, 1] = [1e3, 1e4, 1e5]
  img = Image(3, 1, pixels, 1, blob0[3].wcs, trivial_psf, 1, 1, 1,
              fill(3., size(pixels)), fill(4., size(pixels, 1)),
              blob0[3].raw_psf_comp);
  catalog = [sample_ce([100., 1], true),];
  catalog[1].star_fluxes = ones(5) * 1e5

  ea0 = make_elbo_args(
    [img], catalog, patch_radius_pix=Inf)

  elbo_lik = DeterministicVI.elbo_likelihood(ea0;
        calculate_derivs=false, calculate_hessian=false);

  tile_width = 2
  ea1 = make_elbo_args(
    [img], catalog, tile_width=tile_width, patch_radius_pix=10.);
  elbo_lik_tiles =
    DeterministicVI.elbo_likelihood(
      ea1, calculate_derivs=false, calculate_hessian=false);

  tile_width = 5
  ea2 = make_elbo_args(
      [img], catalog, tile_width=tile_width, patch_radius_pix=10.);
  elbo_lik_tiles2 =
    DeterministicVI.elbo_likelihood(
      ea2, calculate_derivs=false, calculate_hessian=false);

  @test_approx_eq elbo_lik_tiles.v[1] elbo_lik_tiles2.v[1]
  @test_approx_eq_eps elbo_lik.v[1] elbo_lik_tiles.v[1] 100.

end

function test_elbo_with_nan()
    blob, ea, body = gen_sample_star_dataset(perturb=false);

    for b in 1:5
        blob[b].pixels[1,1] = NaN
    end

    # Set tile width to 5 to test the code for tiles with no sources.
    ea = make_elbo_args(blob, body, tile_width=5);
    initial_elbo = DeterministicVI.elbo(ea; calculate_hessian=false);

    nan_elbo = DeterministicVI.elbo(ea);

    # We deleted a pixel, so there's reason to expect them to be different,
    # but importantly they're reasonably close and not NaN.
    @test_approx_eq_eps (nan_elbo.v[1] - initial_elbo.v[1]) / initial_elbo.v[1] 0. 1e-4
    deriv_rel_err = (nan_elbo.d - initial_elbo.d) ./ initial_elbo.d
    @test_approx_eq_eps deriv_rel_err fill(0., length(ea.vp[1])) 0.05
end

function test_trim_source_tiles()
  # Set a seed to avoid a flaky test.
  blob, ea, bodies = gen_n_body_dataset(3, seed=42);

  # With the above seed, this is near the middle of the image.
  s = 1
  ea.active_sources = [s]

  loc_ids = ids.u
  non_loc_ids = setdiff(1:length(ids), ids.u)
  for b=1:length(blob)
      ea_trimmed = deepcopy(ea);
      Infer.trim_source_tiles!(ea_trimmed; noise_fraction=0.0001);

      # Make sure pixels got NaN-ed out
      @test(
          sum([ sum(!Base.isnan(tile.pixels)) for
                tile in ea_trimmed.images[b].tiles]) <
          sum([ sum(!Base.isnan(tile.pixels)) for
                tile in ea.images[b].tiles]))

      # Make sure the elbo derivatives are nearly identical.
      elbo_trim = DeterministicVI.elbo(ea_trimmed; calculate_hessian=false);
      elbo_full = DeterministicVI.elbo(ea; calculate_hessian=false);
      @test_approx_eq_eps(
          elbo_full.d[loc_ids, 1] ./ elbo_trim.d[loc_ids, 1],
          fill(1.0, length(loc_ids)), 0.06)
      @test_approx_eq_eps(
          elbo_full.d[non_loc_ids, 1] ./ elbo_trim.d[non_loc_ids, 1],
          fill(1.0, length(non_loc_ids)), 4e-3)

      # Set the source to be very dim so that the number of non-NaN
      # pixels is determined only by min_radius_pix
      min_radius_pix = 6.0
      ea_trimmed = deepcopy(ea);
      ea_trimmed.vp[s][ids.r1] = log(0.1)
      ea_trimmed.vp[s][ids.r2] = 0.001
      Infer.trim_source_tiles!(ea_trimmed, noise_fraction=0.1,
                               min_radius_pix = min_radius_pix);

      s_tiles = find_source_tiles(1, b, ea)
      total_nonempty_pixels = 0.0
      for tile_index in s_tiles
          tile = ea_trimmed.images[b].tiles[tile_index...]
          total_nonempty_pixels += sum(!Base.isnan(tile.pixels))
      end
      @test_approx_eq_eps total_nonempty_pixels pi * min_radius_pix ^ 2 2.0

      # Make sure a very dim image is completely NaN-ed out with a dim image
      # and zero min radius.
      min_radius_pix = 0.0
      ea_trimmed = deepcopy(ea);
      ea_trimmed.vp[s][ids.r1] = log(0.1)
      ea_trimmed.vp[s][ids.r2] = 0.001
      Infer.trim_source_tiles!(
          ea_trimmed, noise_fraction=0.2, min_radius_pix = min_radius_pix);

      total_nonempty_pixels = 0.0
      for tile_index in s_tiles
          tile = ea_trimmed.images[b].tiles[tile_index...]
          total_nonempty_pixels += sum(!Base.isnan(tile.pixels))
      end
      @test total_nonempty_pixels == 0.0
  end
end

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

test_trim_source_tiles()
test_set_hess()
test_bvn_cov()
test_tile_predicted_image()
test_derivative_flags()
test_active_sources()
test_tiny_image_tiling()
test_that_variance_is_low()
test_that_star_truth_is_most_likely()
test_that_galaxy_truth_is_most_likely()
test_coadd_cat_init_is_most_likely()
test_elbo_with_nan()


text:
using FluxReconstruction, KitBase, Plots, OrdinaryDiffEq, LinearAlgebra
using ProgressMeter: @showprogress
using Base.Threads: @threads

cd(@__DIR__)
ps = TriFRPSpace("../assets/naca0012.msh", 2)
ps0 = deepcopy(ps)

for i in eachindex(ps.cellType)
    if ps.cellType[i] != 0 && norm(ps.cellCenter[i, :]) > 2
        ps.cellType[i] = 2
    end
end

γ = 5 / 3
mach = 0.5
angle = (3 / 180) * π
u0 = zeros(size(ps.cellid, 1), ps.np, 4)
for i in axes(u0, 1), j in axes(u0, 2)
    c = mach * (γ / 2)^0.5

    if true#ps.cellType[i] == 2
        prim = [1.0, c * cos(angle), c * sin(angle), 1.0]
    else
        prim = [1.0, 0.0, 0.0, 1.0]
    end

    u0[i, j, :] .= prim_conserve(prim, γ)
end

function dudt!(du, u, p, t)
    J, lf, cell_normal, fpn, ∂l, ϕ, γ = p
    ncell = size(u, 1)
    nsp = size(u, 2)
    deg = size(fpn, 3) - 1

    f = zeros(ncell, nsp, 4, 2)
    @inbounds @threads for i in axes(f, 1)
        for j in axes(f, 2)
            fg, gg = euler_flux(u[i, j, :], γ)
            for k in axes(f, 3)
                f[i, j, k, :] .= inv(J[i]) * [fg[k], gg[k]]
            end
        end
    end

    u_face = zeros(ncell, 3, deg + 1, 4)
    f_face = zeros(ncell, 3, deg + 1, 4, 2)
    @inbounds @threads for i = 1:ncell
        for j = 1:3, k = 1:deg+1, l = 1:4
            u_face[i, j, k, l] = sum(u[i, :, l] .* lf[j, k, :])
            f_face[i, j, k, l, 1] = sum(f[i, :, l, 1] .* lf[j, k, :])
            f_face[i, j, k, l, 2] = sum(f[i, :, l, 2] .* lf[j, k, :])
        end
    end

    n = [[0.0, -1.0], [1 / √2, 1 / √2], [-1.0, 0.0]]

    fn_face = zeros(ncell, 3, deg + 1, 4)
    @inbounds @threads for i = 1:ncell
        for j = 1:3, k = 1:deg+1, l = 1:4
            fn_face[i, j, k, l] = dot(f_face[i, j, k, l, :], n[j])
        end
    end

    fn_interaction = zeros(ncell, 3, deg + 1, 4)
    @inbounds for i = 1:ncell, j = 1:3, k = 1:deg+1
        uL = local_frame(u_face[i, j, k, :], cell_normal[i, j, 1], cell_normal[i, j, 2])

        ni, nj, nk = fpn[i, j, k]

        fwn_local = zeros(4)
        if ni > 0
            uR = local_frame(
                u_face[ni, nj, nk, :],
                cell_normal[i, j, 1],
                cell_normal[i, j, 2],
            )

            flux_hll!(fwn_local, uL, uR, γ, 1.0)
        elseif ps.cellType[i] == 1
            prim = conserve_prim(u_face[i, j, k, :], γ)
            pn = zeros(4)
            pn[2] = -prim[2]
            pn[3] = -prim[3]
            pn[4] = 2.0 - prim[4]
            tmp = (prim[4] - 1.0)
            pn[1] = (1 - tmp) / (1 + tmp) * prim[1]

            un = prim_conserve(pn, γ)
            uR = local_frame(un, cell_normal[i, j, 1], cell_normal[i, j, 2])

            flux_hll!(fwn_local, uL, uR, γ, 1.0)
        end
        fwn_global = global_frame(fwn_local, cell_normal[i, j, 1], cell_normal[i, j, 2])

        fwn_xy = zeros(4, 2)
        for idx = 1:4
            fwn_xy[idx, :] .= fwn_global[idx] .* cell_normal[i, j, :]
        end

        fws_rs = zeros(4, 2)
        for idx = 1:4
            fws_rs[idx, :] .= inv(J[i]) * fwn_xy[idx, :]
        end

        fwns = [sum(fws_rs[idx, :] .* n[j]) for idx = 1:4]
        #fwns = [sum(fws_rs[idx, :]) for idx in 1:4]

        fn_interaction[i, j, k, :] .= fwns
    end

    rhs1 = zeros(ncell, nsp, 4)
    @inbounds for i in axes(rhs1, 1)
        for j in axes(rhs1, 2), k = 1:4
            if ps.cellType[i] in [0, 1]
                rhs1[i, j, k] =
                    -sum(f[i, :, k, 1] .* ∂l[j, :, 1]) - sum(f[i, :, k, 2] .* ∂l[j, :, 2])
            end
        end
    end

    rhs2 = zero(rhs1)
    @inbounds for i = 1:ncell
        if ps.cellType[i] in [0, 1]
            for j = 1:nsp, k = 1:4
                rhs2[i, j, k] =
                    -sum((fn_interaction[i, :, :, k] .- fn_face[i, :, :, k]) .* ϕ[:, :, j])
            end
        end
    end

    @inbounds for i = 1:ncell
        if ps.cellType[i] in [0, 1]
            du[i, :, :] .= rhs1[i, :, :] .+ rhs2[i, :, :]
        end
    end

    return nothing
end

tspan = (0.0, 0.1)
p = (ps.J, ps.lf, ps.cellNormals, ps.fpn, ps.∂l, ps.ϕ, γ)
prob = ODEProblem(dudt!, u0, tspan, p)
dt = 0.0005
itg = init(prob, Euler(), save_everystep = false, adaptive = false, dt = dt)

function output(ps, itg)
    prim = zero(itg.u)
    for i in axes(prim, 1), j in axes(prim, 2)
        prim[i, j, :] .= conserve_prim(itg.u[i, j, :], γ)
        prim[i, j, 4] = 1 / prim[i, j, 4]
    end
    write_vtk(ps.points, ps.cellid, prim[:, 2, :])
end

@showprogress for iter = 1:500
    step!(itg)

    if iter % 50 == 0
        output(ps, itg)
    end
end

output(ps, itg)


text:
include("06_two_site_state.jl")

println("
#######################################################
# Tutorial 7                                          #
#                                                     #
# 2-site operators                                    #
#######################################################
")

#
# Make a Hamiltonian (transverse field Ising)
#

H = ITensor(i1', i2', dag(i1), dag(i2))
# Could set individual elements...
H[2, 1, 2, 1] = -1.0
# etc.

# Easier to build from single-site operator basis
Id1 = op("Id", i1)
Id2 = op("Id", i2)
Z1 = op("Z", i1)
Z2 = op("Z", i2)
X1 = op("X", i1)
X2 = op("X", i2)

ZZ = Z1 * Z2
XI = X1 * Id2
IX = Id1 * X2

h = 0.5
H = -ZZ + h * (XI + IX)

# Compute some expectation values/matrix elements:
# |0⟩ = |Z+Z+⟩
ZpZp = Zp1 * Zp2

# ⟨H⟩ = ⟨0|H|0⟩ = ⟨Z+Z+|H|Z+Z+⟩
@show (dag(ZpZp)' * H * ZpZp)[]
@show inner(ZpZp', H, ZpZp)
@show inner(ZpZp, apply(H, ZpZp))

# Diagonalize
D, _ = eigen(H; ishermitian=true)
@show diag(D)


text:
using Test
using YaoBlocks.AD
using YaoBlocks, YaoArrayRegister
using Random

@testset "apply put" begin
    Random.seed!(5)
    for reg0 in [rand_state(3), rand_state(3, nbatch = 10)]
        # put block, diagonal
        @test test_apply_back(reg0, put(3, 1 => shift(0.5)), 0.5; δ = 1e-5)
        @test test_apply_back(reg0, control(3, (2, 3), 1 => Rz(0.5)), 0.5; δ = 1e-5)
        @test test_apply_back(reg0, control(3, 2, 1 => shift(0.5)), 0.5; δ = 1e-5)
        # dense matrix
        @test test_apply_back(reg0, put(3, 1 => cache(Rx(0.5))), 0.5; δ = 1e-5)
        @test test_apply_back(reg0, put(3, 1 => chain(Rz(0.5), Rx(0.8))), [0.5, 0.8]; δ = 1e-5)
        @test test_apply_back(reg0, control(3, (2, 3), 1 => Rx(0.5)), 0.5; δ = 1e-5)
        # sparse matrix csc
        @test test_apply_back(reg0, put(3, (1, 2) => rot(SWAP, 0.5)), 0.5; δ = 1e-5)
        @test test_apply_back(reg0, control(3, (3,), (1, 2) => rot(SWAP, 0.5)), 0.5; δ = 1e-5)

        # special cases: DiffBlock
        @test test_apply_back(reg0, put(3, 1 => Rz(0.5)), 0.5; δ = 1e-5)
        @test test_apply_back(reg0, put(3, 1 => Rx(0.5)), 0.5; δ = 1e-5)
        @test test_apply_back(rand_state(1), Rx(0.0), 0.5; δ = 1e-5)

        # TimeEvolution
        @test test_apply_back(reg0, TimeEvolution(put(3, 1 => X), 0.5), 0.5; δ = 1e-5)
    end
end

@testset "apply chain" begin
    Random.seed!(5)
    for reg0 in [rand_state(3), rand_state(3, nbatch = 10)]
        @test test_apply_back(
            reg0,
            chain(3, put(3, 1 => shift(0.0)), control(3, (2, 3), 1 => Rz(0.0))),
            [0.5, 0.5];
            δ = 1e-5,
        )
        c = chain(3, put(3, 1 => shift(0.0)), NoParams(control(3, (2, 3), 1 => Rz(0.0))))
        @test nparameters(c) == 1
        @test test_apply_back(reg0, c, 0.5; δ = 1e-5)
    end
end

@testset "apply dagger, scale" begin
    Random.seed!(5)
    for reg0 in [rand_state(3), rand_state(3, nbatch = 10)]
        @test test_apply_back(
            reg0,
            chain(put(3, 1 => Rx(0.0)), (3 + 2im) * control(3, (2, 3), 1 => Rz(0.0))),
            [0.5, 0.5];
            δ = 1e-5,
        )
        @test test_apply_back(reg0, Daggered(put(3, 1 => Rx(0.0))), 0.5; δ = 1e-5)
        @test test_apply_back(reg0, control(3, (2, 3), 1 => Daggered(Rz(0.0))), 0.5; δ = 1e-5)
        @test test_apply_back(
            reg0,
            chain(3, Daggered(put(3, 1 => Rx(0.0))), control(3, (2, 3), 1 => Daggered(Rz(0.0)))),
            [0.5, 0.5];
            δ = 1e-5,
        )
    end
end

@testset "apply subroutine" begin
    Random.seed!(5)
    for reg0 in [rand_state(3), rand_state(3, nbatch = 10)]
        @test test_apply_back(
            reg0,
            chain(3, put(3, 1 => Rx(0.0)), subroutine(3, control(2, 2, 1 => shift(0.0)), (3, 1))),
            [0.5, 0.5];
            δ = 1e-5,
        )
    end
end

@testset "apply kron repeated" begin
    Random.seed!(5)
    for reg0 in [rand_state(3), rand_state(3, nbatch = 10)]
        @test test_apply_back(
            reg0,
            chain(3, put(3, 1 => Rx(0.0)), kron(Rx(0.4), Y, Rz(0.5))),
            [0.5, 0.3, 0.8];
            δ = 1e-5,
        )
        @test test_apply_back(
            reg0,
            chain(repeat(3, Rx(0.5), 1:2), repeat(3, Ry(0.6), 2:3)),
            [0.5, 0.8];
            δ = 1e-5,
        )
        @test test_apply_back(
            reg0,
            chain(repeat(3, H, 1:2), repeat(3, Ry(0.6), 2:3)),
            [0.8];
            δ = 1e-5,
        )
    end
end

@testset "apply Add" begin
    # + is not reversible
    for reg0 in [rand_state(3), rand_state(3, nbatch = 10)]
        @test_broken test_apply_back(
            reg0,
            (3 * control(3, (2, 3), 1 => Rz(0.0)) + put(3, 2 => Rx(0.0))) / 5,
            [0.5, 0.5];
            δ = 1e-5,
            in = copy(reg0),
        )
    end
end


text:
# save channel to file
function save!(
    file  :: HDF5.File,
    label :: String,
    ch    :: Channel
    )     :: Nothing

    file[label * "/q1"]   = ch.q1
    file[label * "/q2_1"] = ch.q2_1
    file[label * "/q2_2"] = ch.q2_2
    file[label * "/q3"]   = ch.q3

    return nothing
end

# read channel from file
function read_channel(
    file  :: HDF5.File,
    label :: String,
    )     :: Channel

    # read kernels
    q1   = read(file, label * "/q1")
    q2_1 = read(file, label * "/q2_1")
    q2_2 = read(file, label * "/q2_2")
    q3   = read(file, label * "/q3")

    # build channel
    ch = Channel(q1, q2_1, q2_2, q3)

    return ch
end

# save self energy to file
function save_self!(
    file :: HDF5.File,
    Λ    :: Float64,
    m    :: Mesh,
    a    :: Action
    )    :: Nothing

    # save self energy mesh to file
    file["σ/$(Λ)"] = m.σ

    # save self energy to file
    file["Σ/$(Λ)"] = a.Σ

    return nothing
end

"""
    read_self(
        file :: HDF5.File,
        Λ    :: Float64,
        )    :: NTuple{2, Vector{Float64}}

Read self energy mesh and values from HDF5 file at cutoff Λ.
"""
function read_self(
    file :: HDF5.File,
    Λ    :: Float64,
    )    :: NTuple{2, Vector{Float64}}

    # filter out nearest available cutoff
    list    = keys(file["σ"])
    cutoffs = parse.(Float64, list)
    index   = argmin(abs.(cutoffs .- Λ))
    println("Λ was adjusted to $(cutoffs[index]).")

    # read self energy mesh
    σ = read(file, "σ/$(cutoffs[index])")

    # read self energy
    Σ = read(file, "Σ/$(cutoffs[index])")

    return σ, Σ
end

# save vertex to file
function save!(
    file  :: HDF5.File,
    label :: String,
    Γ     :: Vertex
    )     :: Nothing

    # save bare vertex
    file[label * "/bare"] = Γ.bare

    # save channels
    save!(file, label * "/ch_s", Γ.ch_s)
    save!(file, label * "/ch_t", Γ.ch_t)
    save!(file, label * "/ch_u", Γ.ch_u)

    return nothing
end

# read vertex from file
function read_vertex(
    file  :: HDF5.File,
    label :: String
    )     :: Vertex

    # read bare vertex
    bare = read(file, label * "/bare")

    # read channels
    ch_s = read_channel(file, label * "/ch_s")
    ch_t = read_channel(file, label * "/ch_t")
    ch_u = read_channel(file, label * "/ch_u")

    # build vertex
    Γ = Vertex(bare, ch_s, ch_t, ch_u)

    return Γ
end


text:
function test_twodturb_lambdipole(n, dt, dev::Device=CPU(); L=2π, Ue=1, Re=L/20, ν=0.0, nν=1, ti=L/Ue*0.01, nm=3)
  nt = round(Int, ti/dt)
  prob = TwoDTurb.Problem(nx=n, Lx=L, ν=ν, nν=nν, dt=dt, stepper="FilteredRK4", dev=dev)
  zeta₀ = lambdipole(Ue, Re, prob.grid)
  TwoDTurb.set_zeta!(prob, zeta₀)

  xzeta = zeros(nm) # centroid of abs(zeta)
  Ue_m = zeros(nm)  # measured dipole speed
  x, y = gridpoints(prob.grid)
  zeta = prob.vars.zeta

  for i = 1:nm # step forward
    stepforward!(prob, nt)
    TwoDTurb.updatevars!(prob)
    xzeta[i] = mean(@. abs(zeta)*x) / mean(abs.(zeta))
    if i > 1
      Ue_m[i] = (xzeta[i]-xzeta[i-1]) / ((nt-1)*dt)
    end
  end
  isapprox(Ue, mean(Ue_m[2:end]), rtol=rtol_lambdipole)
end

function test_twodturb_stochasticforcingbudgets(dev::Device=CPU(); n=256, L=2π, dt=0.005, ν=1e-7, nν=2, μ=1e-1, nμ=0, tf=0.1/μ)
  nt = round(Int, tf/dt)

  # Forcing parameters
  kf, dkf = 12.0, 2.0
  ε = 0.1

  gr  = TwoDGrid(dev, n, L)
  x, y = gridpoints(gr)

  Kr = ArrayType(dev)([ gr.kr[i] for i=1:gr.nkr, j=1:gr.nl])

  force2k = ArrayType(dev)(zero(gr.Krsq))
  @. force2k = exp(-(sqrt(gr.Krsq)-kf)^2/(2*dkf^2))
  @. force2k[gr.Krsq .< 2.0^2 ] = 0
  @. force2k[gr.Krsq .> 20.0^2 ] = 0
  @. force2k[Kr .< 2π/L] = 0
  ε0 = parsevalsum(force2k.*gr.invKrsq/2.0, gr)/(gr.Lx*gr.Ly)
  force2k .= ε/ε0 * force2k

  Random.seed!(1234)

  function calcF!(Fh, sol, t, cl, v, p, g)
    eta = ArrayType(dev)(exp.(2π*im*rand(Float64, size(sol)))/sqrt(cl.dt))
    eta[1, 1] = 0.0
    @. Fh = eta * sqrt(force2k)
    nothing
  end

  prob = TwoDTurb.Problem(nx=n, Lx=L, ν=ν, nν=nν, μ=μ, nμ=nμ, dt=dt,
   stepper="RK4", calcF=calcF!, stochastic=true, dev=dev)

  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid;

  TwoDTurb.set_zeta!(prob, 0*x)
  E = Diagnostic(TwoDTurb.energy,      prob, nsteps=nt)
  D = Diagnostic(TwoDTurb.dissipation, prob, nsteps=nt)
  R = Diagnostic(TwoDTurb.drag,        prob, nsteps=nt)
  W = Diagnostic(TwoDTurb.work,        prob, nsteps=nt)
  diags = [E, D, W, R]

  stepforward!(prob, diags, nt)
  TwoDTurb.updatevars!(prob)

  E, D, W, R = diags
  t = round(μ*cl.t, digits=2)

  i₀ = 1
  dEdt = (E[(i₀+1):E.i] - E[i₀:E.i-1])/cl.dt
  ii = (i₀):E.i-1
  ii2 = (i₀+1):E.i

  # dEdt = W - D - R?
  # If the Ito interpretation was used for the work
  # then we need to add the drift term
  # total = W[ii2]+ε - D[ii] - R[ii]      # Ito
  total = W[ii2] - D[ii] - R[ii]        # Stratonovich

  residual = dEdt - total
  isapprox(mean(abs.(residual)), 0, atol=1e-4)
end

function test_twodturb_deterministicforcingbudgets(dev::Device=CPU(); n=256, dt=0.01, L=2π, ν=1e-7, nν=2, μ=1e-1, nμ=0)
  n, L  = 256, 2π
  ν, nν = 1e-7, 2
  μ, nμ = 1e-1, 0
  dt, tf = 0.005, 0.1/μ
  nt = round(Int, tf/dt)

  gr  = TwoDGrid(dev, n, L)
  x, y = gridpoints(gr)

  # Forcing = 0.01cos(4x)cos(5y)cos(2t)
  f = @. 0.01cos(4x)*cos(5y)
  fh = rfft(f)
  function calcF!(Fh, sol, t, cl, v, p, g::AbstractGrid{T, A}) where {T, A}
    Fh = fh*cos(2t)
    nothing
  end

  prob = TwoDTurb.Problem(nx=n, Lx=L, ν=ν, nν=nν, μ=μ, nμ=nμ, dt=dt,
   stepper="RK4", calcF=calcF!, stochastic=false, dev=dev)

  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

  TwoDTurb.set_zeta!(prob, 0*x)

  E = Diagnostic(TwoDTurb.energy,      prob, nsteps=nt)
  D = Diagnostic(TwoDTurb.dissipation, prob, nsteps=nt)
  R = Diagnostic(TwoDTurb.drag,        prob, nsteps=nt)
  W = Diagnostic(TwoDTurb.work,        prob, nsteps=nt)
  diags = [E, D, W, R]

  # Step forward
  stepforward!(prob, diags, nt)
  TwoDTurb.updatevars!(prob)

  E, D, W, R = diags
  t = round(μ*cl.t, digits=2)

  i₀ = 1
  dEdt = (E[(i₀+1):E.i] - E[i₀:E.i-1])/cl.dt
  ii = (i₀):E.i-1
  ii2 = (i₀+1):E.i

  # dEdt = W - D - R?
  total = W[ii2] - D[ii] - R[ii]
  residual = dEdt - total

  isapprox(mean(abs.(residual)), 0, atol=1e-8)
end

"""
    testnonlinearterms(dt, stepper; kwargs...)

Tests the advection term in the twodturb module by timestepping a
test problem with timestep dt and timestepper identified by the string stepper.
The test problem is derived by picking a solution ζf (with associated
streamfunction ψf) for which the advection term J(ψf, ζf) is non-zero. Next, a
forcing Ff is derived according to Ff = ∂ζf/∂t + J(ψf, ζf) - νΔζf. One solution
to the vorticity equation forced by this Ff is then ζf. (This solution may not
be realized, at least at long times, if it is unstable.)
"""
function test_twodturb_advection(dt, stepper, dev::Device=CPU(); n=128, L=2π, ν=1e-2, nν=1, μ=0.0, nμ=0)
  n, L  = 128, 2π
  ν, nν = 1e-2, 1
  μ, nμ = 0.0, 0
  tf = 1.0
  nt = round(Int, tf/dt)

  gr   = TwoDGrid(dev, n, L)
  x, y = gridpoints(gr)

   psif = @.   sin(2x)*cos(2y) +  2sin(x)*cos(3y)
  zetaf = @. -8sin(2x)*cos(2y) - 20sin(x)*cos(3y)

  Ff = @. -(
    ν*( 64sin(2x)*cos(2y) + 200sin(x)*cos(3y) )
    + 8*( cos(x)*cos(3y)*sin(2x)*sin(2y) - 3cos(2x)*cos(2y)*sin(x)*sin(3y) )
  )

  Ffh = rfft(Ff)

  # Forcing
  function calcF!(Fh, sol, t, cl, v, p, g)
    Fh .= Ffh
    nothing
  end

  prob = TwoDTurb.Problem(nx=n, Lx=L, ν=ν, nν=nν, μ=μ, nμ=nμ, dt=dt, stepper=stepper, calcF=calcF!, stochastic=false, dev=dev)
  sol, cl, p, v, g = prob.sol, prob.clock, prob.params, prob.vars, prob.grid
  TwoDTurb.set_zeta!(prob, zetaf)

  stepforward!(prob, nt)
  TwoDTurb.updatevars!(prob)

  isapprox(prob.vars.zeta, zetaf, rtol=rtol_twodturb)
end

function test_twodturb_energyenstrophy(dev::Device=CPU())
  nx, Lx  = 128, 2π
  ny, Ly  = 128, 3π
  gr = TwoDGrid(dev, nx, Lx, ny, Ly)
  x, y = gridpoints(gr)

  k0, l0 = gr.k[2], gr.l[2] # fundamental wavenumbers
    psi0 = @. sin(2k0*x)*cos(2l0*y) + 2sin(k0*x)*cos(3l0*y)
   zeta0 = @. -((2k0)^2+(2l0)^2)*sin(2k0*x)*cos(2l0*y) - (k0^2+(3l0)^2)*2sin(k0*x)*cos(3l0*y)

  energy_calc = 29/9
  enstrophy_calc = 2701/162

  prob = TwoDTurb.Problem(nx=nx, Lx=Lx, ny=ny, Ly=Ly, stepper="ForwardEuler", dev=dev)

  sol, cl, v, p, g = prob.sol, prob.clock, prob.vars, prob.params, prob.grid;

  TwoDTurb.set_zeta!(prob, zeta0)
  TwoDTurb.updatevars!(prob)

  energyzeta0 = TwoDTurb.energy(prob)
  enstrophyzeta0 = TwoDTurb.enstrophy(prob)

  params = TwoDTurb.Params(p.ν, p.nν)

  (isapprox(energyzeta0, energy_calc, rtol=rtol_twodturb) &&
   isapprox(enstrophyzeta0, enstrophy_calc, rtol=rtol_twodturb) &&
   TwoDTurb.addforcing!(prob.timestepper.N, sol, cl.t, cl, v, p, g)==nothing && p == params)
end

function test_twodturb_problemtype(T=Float32, dev::Device=CPU())
  prob = TwoDTurb.Problem(T=T, dev=dev)

  (typeof(prob.sol)==Array{Complex{T},2} && typeof(prob.grid.Lx)==T && typeof(prob.grid.x)==Array{T,2} && typeof(prob.vars.u)==Array{T,2})
end


text:
# # ShallowWaters.jl - The equations
#
# The shallow water equations for the prognostic variables velocity $\mathbf{u} = (u,v)$ and sea surface elevation $\eta$ over the 2-dimensional domain $\Omega$ in $x,y$ are
#
# ```math
# \begin{align}
# \partial_t u &+ u\partial_xu + v\partial_yu - fv = -g\partial_x\eta + D_x(u,v,\eta) + F_x, \\
# \partial_t v &+ u\partial_xv + v\partial_yv + fu = -g\partial_y\eta + D_y(u,v,\eta) + F_y, \\
# \partial_t \eta &+ \partial_x(uh) + \partial_y(vh) = 0.
# \end{align}
# ```
# where the first two are the momentum equations for $u,v$ and the latter is the continuity equation. The layer thickness is $h = \eta + H$ with $H=H(x,y)$ being the bottom topography. The gravitational acceleration is $g$, the coriolis parameter $f = f(y)$ depends only (i.e. latitude) only and the beta-plane approximation is used. $(D_x,D_y) = \mathbf{D}$ are the dissipative terms
#
# ```math
# \mathbf{D} = -\frac{c_D}{h}\vert \mathbf{u} \vert \mathbf{u} - u \nabla^4\mathbf{u}
# ```

# which are a sum of a quadratic bottom drag with dimensionless coefficient $c_D$ and a biharmonic diffusion with viscosity coefficient $u$. $\mathbf{F} = (F_x,F_y)$ is the wind forcing which can depend on time and space, i.e. $F_x = F_x(x,y,t)$ and $F_y = F_y(x,y,t)$.

# ## The vector invariant formulation

# The Bernoulli potential $p$ is introduced as

# ```math
# p = \frac{1}{2}(u^2 + v^2) + gh
# ```

# The relative vorticity $\zeta = \partial_xv + \partial_yu$ lets us define the potential vorticity $q$ as

# ```math
# q = \frac{f + \zeta}{h}
# ```

# such that we can rewrite the shallow water equations as

# ```math
# \begin{align}
# \partial_t u &= qhv -g\partial_xp + D_x(u,v,\eta) + F_x, \\
# \partial_t v &= -qhu -g\partial_yp + D_y(u,v,\eta) + F_y, \\
# \partial_t \eta &= -\partial_x(uh) -\partial_y(vh).
# \end{align}
# ```

# ## Runge-Kutta time discretisation

# Let
# ```math
# R(u,v,\eta) = \begin{pmatrix}
#                 qhv-g\partial_xp+F_x
#                 -qhu-g\partial_yp+F_y
#                 -\partial_x(uh) -\partial_y(vh)
#                 \end{pmatrix}
# ```
#
# be the non-dissipative right-hand side, i.e. excluding the dissipative terms $\mathbf{D}$. Then we dicretise the time derivative with 4th order Runge-Kutta with $\mathbf{k}_n = (u_n,v_n,\eta_n)$ by
#
# ```math
# \begin{align}
# \mathbf{d}_1 &= R(\mathbf{k}_n) , \\
# \mathbf{d}_2 &= R(\mathbf{k}_n + \frac{1}{2}\Delta t \mathbf{d}_1), \\
# \mathbf{d}_3 &= R(\mathbf{k}_n + \frac{1}{2}\Delta t \mathbf{d}_2), \\
# \mathbf{d}_4 &= R(\mathbf{k}_n + \Delta t \mathbf{d}_3), \\
# u_{n+1}^*,v_{n+1}^*,\eta_{n+1} = \mathbf{k}_{n+1} &= \mathbf{k}_n + \frac{1}{6}\Delta t(\mathbf{d}_1 + 2\mathbf{d}_2 + 2\mathbf{d}_3 + \mathbf{d}_1),
# \end{align}
# ```

# and the dissipative terms are then added semi-implictly.

# ```math
# \begin{align}
# u_{n+1} = u_{n+1}^* + \Delta t D_x(u_{n+1}^*,v_{n+1}^*,\eta_{n+1}) , \\
# v_{n+1} = v_{n+1}^* + \Delta t D_y(u_{n+1}^*,v_{n+1}^*,\eta_{n+1}).
# \end{align}
# ```

# Consequently, the dissipative terms only have to be evaluated once per time step, which reduces the computational cost of the right=hand side drastically. This is motivated as the Courant number $C = \sqrt{gH_0}$ is restricted by gravity waves, which are caused by $\partial_t\mathbf{u} = -g\nabla\eta$ and $\partial_t\eta = -H_0\nabla \cdot \mathbf{u}$. The other terms have longer time scales, and it is therefore sufficient to solve those with larger time steps. With this scheme, the shallow water model runs stable at $C=1$."

# ## Strong stability preserving Runge-Kutta with semi-implicit continuity equation

# We split the right-hand side into the momentum equations and the continuity equation

# ```math
# R_m(u,v,\eta) = \begin{pmatrix}
#                 qhv-g\partial_xp+F_x
#                 -qhu-g\partial_yp+F_y
#                 \end{pmatrix}, \quad R_\eta(u,v,\eta) = -\partial_x(uh) -\partial_y(vh)
# ```

# The 4-stage strong stability preserving Runge-Kutta scheme, with a semi-implicit treatment of the continuity equation then reads as

# ```math
# \begin{align}
# \mathbf{u}_1 &= \mathbf{u}_n + \frac{1}{2} \Delta t R_m(u_n,v_n,\eta_n), \quad \text{then}
# \quad \eta_1 = \eta_n + \frac{1}{2} \Delta t R_\eta(u_1,v_1,\eta_n), \\
# \mathbf{u}_2 &= \mathbf{u}_1 + \frac{1}{2} \Delta t R_m(u_1,v_1,\eta_1), \quad \text{then}
# \quad \eta_2 = \eta_1 + \frac{1}{2} \Delta t R_\eta(u_2,v_2,\eta_1) , \\
# \mathbf{u}_3 &= \frac{2}{3}\mathbf{u}_n + \frac{1}{3}\mathbf{u}_2 + \frac{1}{6} \Delta t R_m(u_2,v_2,\eta_2), \quad \text{then}
# \quad \eta_3 = \frac{2}{3}\eta_n + \frac{1}{3}\eta_2 + \frac{1}{6} \Delta t R_\eta(u_3,v_3,\eta_2), \\
# \mathbf{u}_{n+1} &= \mathbf{u}_3 + \frac{1}{2} \Delta t R_m(u_3,v_3,\eta_3), \quad \text{then}
# \quad \eta_{n+1} = \eta_3 + \frac{1}{2} \Delta t R_\eta(u_{n+1},v_{n+1},\eta_3).
# \end{align}
# ```

# ## Splitting the continuity equation

# From
# ```math
# \partial_t = -\partial_x(uh) - \partial_y(vh)
# ```


text:
module TwoDimensionalDiffusion

using Printf
using Statistics

using Oceananigans
using Oceananigans.Fields: nodes

include("analysis.jl")

# 2D diffusing sinuosoid
c(x, y, t) = exp(-2t) * cos(x) * cos(y)

instantiate(t) = Tuple(ti() for ti in t)
Lx(topo) = Lx(instantiate(topo))
Ly(topo) = Ly(instantiate(topo))

Lx(::Tuple{Periodic, Y, Z}) where {Y, Z} = 2π
Lx(::Tuple{Bounded,  Y, Z}) where {Y, Z} = π

Ly(::Tuple{X, Periodic, Z}) where {X, Z} = 2π
Ly(::Tuple{X, Bounded,  Z}) where {X, Z} = π

#####
##### x, y
#####

const DATA_DIR = joinpath(@__DIR__, "..", "data")

function setup_simulation(; Nx, Δt, stop_iteration, architecture=CPU(), dir=DATA_DIR,
                          topo=(Periodic, Periodic, Bounded), output=false)

    grid = RectilinearGrid(architecture, size=(Nx, Nx, 1), x=(0, Lx(topo)), y=(0, Ly(topo)), z=(0, 1), topology=topo)

    model = NonhydrostaticModel(        grid = grid,
                                    coriolis = nothing,
                                    buoyancy = nothing,
                                     tracers = :c,
                                     closure = ScalarDiffusivity(κ=1))

    set!(model, c = (x, y, z) -> c(x, y, 0))

    simulation = Simulation(model, Δt=Δt, stop_iteration=stop_iteration, iteration_interval=stop_iteration)

    if output
        simulation.output_writers[:fields] =
            JLD2OutputWriter(model, model.tracers; dir = dir, overwrite_existing = true, field_slicer = nothing,
                             prefix = @sprintf("%s_%s_diffusion_Nx%d_Δt%.1e", "$(topo[1])", "$(topo[2])", Nx, Δt),
                             schedule = TimeInterval(stop_iteration * Δt / 10))
    end

    return simulation
end

function run_and_analyze(; setup...)
    simulation = setup_simulation(; setup...)
    @info "Running two dimensional diffusion simulation $(setup[:topo]) with Nx = Ny = $(setup[:Nx]), Δt = $(setup[:Δt])"
    @time run!(simulation)

    c_simulation = simulation.model.tracers.c

    x, y, z, t = (nodes(c_simulation, reshape=true)..., simulation.model.clock.time)

    c_analytical = c.(x, y, t)

    c_simulation = interior(c_simulation)

    @show size(c_analytical)
    @show size(c_simulation)

    return compute_error(c_simulation, c_analytical)
end

end # module


text:
using DelimitedFiles #, Plots
include("nbg_tts.jl")
include("ttvfast_tts.jl")

#ntransits = 500 # Number of transits for each planet
tmax = 4533. # Should give ~3000 transits of planet b
denom = 1000 # P of planet b / denom = time step

#println("Running NBG...")
#output_NBG_times(setup_sim(tmax,denom)...) # NbodyGradient transit times

# Run rebound and record times
println("Running Python...")
run(`python reb_tts.py`) # run from command line

# Run ttvfast and record times
cd("ttvfast/")
#run(`gcc -o ttvfast/run_TTVFast ttvfast/run_TTVFast.c ttvfast/TTVFast.c -lm -O3`)
run(`gcc -o run_TTVFast run_TTVFast.c TTVFast.c -lm -O3`)
#run(`ttvfast/run_TTVFast ttvfast/setup_trappist ttvfast/Times`)
println("Running TTVFast...")
run(`./run_TTVFast setup_trappist Times`)
cd("../")
#run(`cd ttvfast '&&' ./run_TTVFast setup_trappist Times '&&' cd ..`)
output_ttvfast_times()
println("Plotting...")
include("plot_times_new.jl")

max_forks_repo_name: langfzac/nbg-papers


text:
# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    SeqSim(var₁=>param₁, var₂=>param₂, ...)

A sequential simulation solver.

For each location in the simulation `path`, a maximum number
of neighbors `maxneighbors` is used to fit a distribution.
The neighbors are searched according to a `neighborhood`,
and in case there are none, use a `marginal` distribution.

## Parameters

* `estimator`    - CDF estimator
* `neighborhood` - Spatial neighborhood
* `maxneighbors` - Maximum number of neighbors
* `marginal`     - Marginal distribution
* `path`         - Simulation path
"""
@simsolver SeqSim begin
  @param estimator
  @param neighborhood
  @param minneighbors
  @param maxneighbors
  @param marginal
  @param path
end

function preprocess(problem::SimulationProblem, solver::SeqSim)
  # retrieve problem info
  pdomain = domain(problem)

  # result of preprocessing
  preproc = Dict{Symbol,NamedTuple}()

  for covars in covariables(problem, solver)
    for var in covars.names
      # get user parameters
      varparams = covars.params[(var,)]

      # determine maximum number of neighbors
      maxneighbors = varparams.maxneighbors

      # determine neighbor search method
      neigh     = varparams.neighborhood
      searcher  = NeighborhoodSearcher(pdomain, neigh)
      bsearcher = BoundedSearcher(searcher, maxneighbors)

      # save preprocessed input
      preproc[var] = (estimator=varparams.estimator,
                      minneighbors=varparams.minneighbors,
                      maxneighbors=varparams.maxneighbors,
                      marginal=varparams.marginal,
                      path=varparams.path,
                      bsearcher=bsearcher)
    end
  end

  preproc
end

function solvesingle(problem::SimulationProblem, covars::NamedTuple,
                     solver::SeqSim, preproc)
  # retrieve problem info
  pdata = data(problem)
  pdomain = domain(problem)
  N = ncoords(pdomain)
  T = coordtype(pdomain)

  varreals = map(covars.names) do var
    # unpack preprocessed parameters
    estimator, minneighbors, maxneighbors, marginal, path, bsearcher = preproc[var]

    # determine value type
    V = variables(problem)[var]

    # pre-allocate memory for result
    realization = Vector{V}(undef, nelms(pdomain))

    # pre-allocate memory for coordinates
    xₒ = MVector{N,T}(undef)

    # pre-allocate memory for neighbors coordinates
    neighbors = Vector{Int}(undef, maxneighbors)
    X = Matrix{T}(undef, N, maxneighbors)

    # keep track of simulated locations
    simulated = falses(nelms(pdomain))
    for (loc, datloc) in datamap(problem, var)
      realization[loc] = pdata[var][datloc]
      simulated[loc] = true
    end

    # simulation loop
    for location in traverse(pdomain, path)
      if !simulated[location]
        # coordinates of neighborhood center
        coordinates!(xₒ, pdomain, location)

        # find neighbors with previously simulated values
        nneigh = search!(neighbors, xₒ, bsearcher, mask=simulated)

        # choose between marginal and conditional distribution
        if nneigh < minneighbors
          # draw from marginal
          realization[location] = rand(marginal)
        else
          # final set of neighbors
          nview = view(neighbors, 1:nneigh)

          # update neighbors coordinates
          coordinates!(X, pdomain, nview)

          Xview = view(X,:,1:nneigh)
          zview = view(realization, nview)

          # fit estimator
          fitted = fit(estimator, Xview, zview)

          if status(fitted)
            # estimate mean and variance
            μ, σ² = predict(fitted, xₒ)

            # draw from conditional
            realization[location] = μ + √σ²*randn(V)
          else
            # draw from marginal
            realization[location] = rand(marginal)
          end
        end

        # mark location as simulated and continue
        simulated[location] = true
      end
    end

    var => realization
  end

  Dict(varreals)
end


text:
using CUDArt
using CUBLAS

import Base: A_mul_B!, eltype

using BandedMatrices

type CudaBandedMatrix{T<:Number}
    data::CudaMatrix{T}
    m::Int
    l::Int
    u::Int
end
CudaBandedMatrix{T<:Number}(BM::BandedMatrix{T}) =
    CudaBandedMatrix{T}(CudaArray(BM.data),BM.m,BM.l,BM.u)

A_mul_B!{T<:Number}(α::T, A::CudaBandedMatrix{T},
                    x::CudaVector{T},
                    β::T, y::CudaVector{T}) =
                        CUBLAS.gbmv!('N',A.m,A.l,A.u,α,A.data,x,β,y)

eltype{T}(M::CudaBandedMatrix{T}) = T

export CudaBandedMatrix


text:
struct State
    position::Position
    attitude::Attitude
    velocity::Array{T, 1} where T<:Number
    angular_velocity::Array{T, 1} where T<:Number
    acceleration::Array{T, 1} where T<:Number
    angular_acceleration::Array{T, 1} where T<:Number
end

# Position getters
get_position(state::State) = state.position
get_llh(state::State) = get_llh(state.position)
get_xyz_earth(state::State) = get_xyz_earth(state.position)
get_xyz_ecef(state::State) = get_xyz_ecef(state.position)
get_height(state::State) = get_height(state.position)
# Attitude getters
get_attitude(state::State) = state.attitude
get_euler_angles(state::State) = get_euler_angles(state.attitude)
get_quaternions(state::State) = get_quaternions(state.attitude)
# Velocity getters
get_body_velocity(state::State) = state.velocity
get_horizon_velocity(state::State) = body2hor(get_body_velocity(state)...,
                                              get_quaternions(state)...)
function get_flight_path_angle(state::State)
    vn, ve, vz = get_horizon_velocity(state)
    atan(-vz, norm([vn, ve]))
end

# Angular velocity getters
get_body_ang_velocity(state::State) = state.angular_velocity
get_turn_rate(state::State) = body2hor(get_body_ang_velocity(state)...,
                                       get_quaternions(state)...)[3]
# Acceleration getters
get_body_accel(state::State) = state.acceleration
get_horizon_accel(state::State) = body2hor(get_body_accel(state)...,
                                           get_quaternions(state)...)
# Angular acceleration getters
get_body_ang_accel(state::State) = state.angular_acceleration

function get_sixdof_euler_fixed_mass_state(state::State)
    [get_body_velocity(state)...,
     get_body_ang_velocity(state)...,
     get_euler_angles(state)...,
     get_xyz_earth(state)...]
end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

## dummy stub for https://github.com/JuliaBinaryWrappers/p7zip_jll.jl
baremodule p7zip_jll
using Base
Base.Experimental.@compiler_options compile=min optimize=0 infer=false

const PATH_list = String[]
const LIBPATH_list = String[]

export p7zip

# These get calculated in __init__()
PATH = Ref("")
LIBPATH = Ref("")
artifact_dir = ""
p7zip_path = ""
if Sys.iswindows()
    const p7zip_exe = "7z.exe"
else
    const p7zip_exe = "7z"
end

# These functions look a little strange, but they're mimicking the JLLWrappers signature
p7zip(f::Function; adjust_PATH::Bool = true, adjust_LIBPATH::Bool = true) = f(p7zip_path)
p7zip(; adjust_PATH::Bool = true, adjust_LIBPATH::Bool = true) = Cmd([p7zip_path])

function init_p7zip_path()
    # Prefer our own bundled p7zip, but if we don't have one, pick it up off of the PATH
    # If this is an in-tree build, `7z` will live in `bin`.  Otherwise, it'll be in `libexec`
    for bundled_p7zip_path in (joinpath(Sys.BINDIR, Base.LIBEXECDIR, p7zip_exe),
                               joinpath(Sys.BINDIR, p7zip_exe))
        if isfile(bundled_p7zip_path)
            global p7zip_path = abspath(bundled_p7zip_path)
            return
        end
    end
    global p7zip_path = Sys.which(p7zip_exe)
end

function __init__()
    global artifact_dir = dirname(Sys.BINDIR)
    global LIBPATH[] = joinpath(Sys.BINDIR, Base.LIBDIR, "julia")
    init_p7zip_path()
end

# JLLWrappers API compatibility shims.  Note that not all of these will really make sense.
# For instance, `find_artifact_dir()` won't actually be the artifact directory, because
# there isn't one.  It instead returns the overall Julia prefix.
is_available() = true
find_artifact_dir() = artifact_dir
dev_jll() = error("stdlib JLLs cannot be dev'ed")
best_wrapper = nothing

end  # module p7zip_jll

max_stars_repo_name: dmageeLANL/julia

max_issues_repo_name: dmageeLANL/julia

max_forks_repo_name: dmageeLANL/julia


text:
# TODO: Add HigherPrecision.jl tests here?
# TODO: Test Interior Point algorithm as well
@testset "Arbitrary Precision" begin
    prob = MVP.UnconstrainedProblems.examples["Rosenbrock"]

    f = MVP.objective(prob)
    g! = MVP.gradient(prob)
    h! = MVP.hessian(prob)

    @testset "BigFloat" begin
        x0 = big.(prob.initial_x)
        res = optimize(f, x0)
        debug_printing && @show res
        @test Optim.converged(res) == true
        @test Optim.minimum(res) < 1e-8
        @test Optim.minimizer(res) ≈ [1.0, 1.0] atol=1e-4 rtol=0

        res = optimize(f, g!, x0)
        debug_printing && @show res
        @test Optim.converged(res) == true
        @test Optim.minimum(res) < 1e-16
        @test Optim.minimizer(res) ≈ [1.0, 1.0] atol=1e-10 rtol=0

        res = optimize(f, g!, h!, x0)
        debug_printing && @show res
        @test Optim.converged(res) == true
        @test Optim.minimum(res) < 1e-16
        @test Optim.minimizer(res) ≈ [1.0, 1.0] atol=1e-10 rtol=0

        lower = big.([-Inf, -Inf])
        upper = big.([0.5, 1.5])
        res = optimize(f, g!, lower, upper, x0, Fminbox(), Optim.Options(outer_g_abstol=sqrt(eps(big(1.0))), g_abstol=sqrt(eps(big(1.0)))))
        debug_printing && @show res
        @test Optim.converged(res) == true
        @test Optim.minimum(res) ≈ 0.25 atol=1e-10 rtol=0
        @test Optim.minimizer(res) ≈ [0.5, 0.25] atol=1e-10 rtol=0

        res = optimize(f, lower, upper, x0, Fminbox(), Optim.Options(outer_g_abstol=sqrt(eps(big(1.0))), g_abstol=sqrt(eps(big(1.0)))))
        debug_printing && @show res
        @test Optim.converged(res) == true
        @test Optim.minimum(res) ≈ 0.25 atol=1e-10 rtol=0
        @test Optim.minimizer(res) ≈ [0.5, 0.25] atol=1e-10 rtol=0
    end
end

@testset "Float32 doesn't work in Fminbox" begin
    f(x::AbstractVector{T}) where {T} = (T(1.0) - x[1])^2 + T(100.0) * (x[2] - x[1]^2)^2

    function g!(storage::AbstractVector{T}, x::AbstractVector{T}) where {T}
        storage[1] = T(-2.0) * (T(1.0) - x[1]) - T(400.0) * (x[2] - x[1]^2) * x[1]
        storage[2] = T(200.0) * (x[2] - x[1]^2)
    end

    for T = (Float32, Float64, BigFloat)
        lower = T[1.25, -2.1]
        upper = T[Inf, Inf]
        initial_x = T[2.0, 2.0]

        od = OnceDifferentiable(f, g!, initial_x)

        results = optimize(od, lower, upper, initial_x, Fminbox(GradientDescent()))
    end
end


text:
using ArgParse

function ArgParse.parse_item(T::Type{<: Union{Int, AbstractRange}},
                             x::AbstractString)
  r = eval(Meta.parse(x))
  @assert r isa T
  return r
end

settings = ArgParseSettings()
@add_arg_table! settings begin
  "--blas_num_threads"
    help = "Number of BLAS threads."
    arg_type = Int
    default = 1
  "--omp_num_threads"
    help = "Number of OpenMP threads (for block sparse contractions)."
    arg_type = Union{Int, AbstractRange}
    default = 1
  "--maxdim"
    help = "Maximum DMRG bond dimension."
    arg_type = Int
    default = 10_000
  "--U"
    help = "Interaction strength of the Hubbard model."
    arg_type = Float64
    default = 4.0
  "--Nx"
    help = "Number of sites in the x-direction (along the cylinder)."
    arg_type = Int
    default = 8
  "--Ny"
    help = "Number of sites in the y-direction (around the cylinder)."
    arg_type = Int
    default = 4
end

args = parse_args(settings)

blas_num_threads = args["blas_num_threads"]
omp_num_threads = args["omp_num_threads"]
maxdim = args["maxdim"]
U = args["U"]
Nx = args["Nx"]
Ny = args["Ny"]


text:
function mark_cross!(r::Robot)
    while !isborder(r, Nord)
        move!(r, Nord)
    end
    while !isborder(r, Ost)
        move!(r, Ost)
    end
    while !isborder(r, Sud)
        move!(r, Sud)
        putmarket!(r)
    end
    while !isborder(r, West)
        move!(r, West)
        putmarker!(r)
    end
    while !isborder(r, Nord)
        move!(r, Nord)
        putmarker!(r)
    end
    while !isborder(r, West)
        move!(r, West)
        putmarker!(r)
    end
end