text:
# # Navier--Stokes equations
#
# In this example, we numerically solve the incompressible Navier--Stokes
# equations
#
# ```math
# ∂_t \bm{v} + (\bm{v} ⋅ \bm{∇}) \bm{v} = -\frac{1}{ρ} \bm{∇} p + ν ∇^2 \bm{v},
# \quad \bm{∇} ⋅ \bm{v} = 0,
# ```
#
# where ``\bm{v}(\bm{x}, t)`` and ``p(\bm{x}, t)`` are respectively the velocity
# and pressure fields, ``ν`` is the fluid kinematic viscosity and ``ρ`` is the
# fluid density.
#
# We solve the above equations a 3D periodic domain using a standard Fourier
# pseudo-spectral method.

# ## First steps
#
# We start by loading the required packages, initialising MPI and setting the
# simulation parameters.

using MPI
using PencilFFTs

MPI.Init()
comm = MPI.COMM_WORLD
procid = MPI.Comm_rank(comm) + 1

## Simulation parameters
Ns = (64, 64, 64)  # = (Nx, Ny, Nz)
Ls = (2π, 2π, 2π)  # = (Lx, Ly, Lz)

## Collocation points ("global" = over all processes).
## We include the endpoint (length = N + 1) for convenience.
xs_global = map((N, L) -> range(0, L; length = N + 1), Ns, Ls)  # = (x, y, z)

# Let's check the number of MPI processes over which we're running our
# simulation:
MPI.Comm_size(comm)

# We can now create a partitioning of the domain based on the number of grid
# points (`Ns`) and on the number of MPI processes.
# There are different ways to do this.
# For simplicity, here we do it automatically following the [PencilArrays.jl
# docs](https://jipolanco.github.io/PencilArrays.jl/stable/Pencils/#pencil-high-level):

pen = Pencil(Ns, comm)

# The subdomain associated to the local MPI process can be obtained using
# [`range_local`](https://jipolanco.github.io/PencilArrays.jl/dev/Pencils/#PencilArrays.Pencils.range_local-Tuple{Pencil,%20LogicalOrder}):

range_local(pen)

# We now construct a distributed vector field that follows the decomposition
# configuration we just created:

v⃗₀ = (
    PencilArray{Float64}(undef, pen),  # vx
    PencilArray{Float64}(undef, pen),  # vy
    PencilArray{Float64}(undef, pen),  # vz
)
summary(v⃗₀[1])

# We still need to fill this array with interesting values that represent a
# physical velocity field.

# ## Initial condition

# Let's set the initial condition in physical space.
# In this example, we choose the [Taylor--Green
# vortex](https://en.wikipedia.org/wiki/Taylor%E2%80%93Green_vortex)
# configuration as an initial condition:
#
# ```math
# \begin{aligned}
# v_x(x, y, z) &=  u₀ \sin(k₀ x) \cos(k₀ y) \cos(k₀ z) \\
# v_y(x, y, z) &= -u₀ \cos(k₀ x) \sin(k₀ y) \cos(k₀ z) \\
# v_z(x, y, z) &= 0
# \end{aligned}
# ```
#
# where ``u₀`` and ``k₀`` are two parameters setting the amplitude and the
# period of the velocity field.
#
# To set the initial condition, each MPI process needs to know which portion of
# the physical grid it has been attributed.
# For this, PencilArrays.jl includes a
# [`localgrid`](https://jipolanco.github.io/PencilArrays.jl/dev/LocalGrids/#PencilArrays.LocalGrids.localgrid)
# helper function:

grid = localgrid(pen, xs_global)

# We can use this to initialise the velocity field:

u₀ = 1.0
k₀ = 2π / Ls[1]  # should be integer if L = 2π (to preserve periodicity)

@. v⃗₀[1] =  u₀ * sin(k₀ * grid.x) * cos(k₀ * grid.y) * cos(k₀ * grid.z)
@. v⃗₀[2] = -u₀ * cos(k₀ * grid.x) * sin(k₀ * grid.y) * cos(k₀ * grid.z)
@. v⃗₀[3] =  0
nothing # hide

# Let's plot a 2D slice of the velocity field managed by the local MPI process:

using GLMakie

## Compute the norm of a vector field represented by a tuple of arrays.
function vecnorm(v⃗::NTuple)
    vnorm = similar(v⃗[1])
    for n ∈ eachindex(v⃗[1])
        w = zero(eltype(vnorm))
        for v ∈ v⃗
            w += v[n]^2
        end
        vnorm[n] = sqrt(w)
    end
    vnorm
end

let fig = Figure(resolution = (700, 600))
    ax = Axis3(fig[1, 1]; aspect = :data, xlabel = "x", ylabel = "y", zlabel = "z")
    vnorm = vecnorm(v⃗₀)
    ct = contour!(
        ax, grid.x, grid.y, grid.z, vnorm;
        alpha = 0.2, levels = 4,
        colormap = :viridis, colorrange = (0.0, 1.0),
    )
    cb = Colorbar(fig[1, 2], ct; label = "Velocity magnitude")
    fig
end

# ## Velocity in Fourier space
#
# In the Fourier pseudo-spectral method, the periodic velocity field is
# discretised in space as a truncated Fourier series
#
# ```math
# \bm{v}(\bm{x}, t) =
# ∑_{\bm{k}} \hat{\bm{v}}_{\bm{k}}(t) \, e^{i \bm{k} ⋅ \bm{x}},
# ```
#
# where ``\bm{k} = (k_x, k_y, k_z)`` are the discrete wave numbers.
#
# The wave numbers can be obtained using the
# [`fftfreq`](https://juliamath.github.io/AbstractFFTs.jl/dev/api/#AbstractFFTs.fftfreq)
# function.
# Since we perform a real-to-complex transform along the first dimension, we use
# [`rfftfreq`](https://juliamath.github.io/AbstractFFTs.jl/dev/api/#AbstractFFTs.rfftfreq) instead for ``k_x``:

using AbstractFFTs: fftfreq, rfftfreq

ks_global = (
    rfftfreq(Ns[1], 2π * Ns[1] / Ls[1]),  # kx | real-to-complex
     fftfreq(Ns[2], 2π * Ns[2] / Ls[2]),  # ky | complex-to-complex
     fftfreq(Ns[3], 2π * Ns[3] / Ls[3]),  # kz | complex-to-complex
)

ks_global[1]'
#
ks_global[2]'
#
ks_global[3]'

# To transform the velocity field to Fourier space, we first create a
# real-to-complex FFT plan to be applied to one of the velocity components:

plan = PencilFFTPlan(v⃗₀[1], Transforms.RFFT())

# See [`PencilFFTPlan`](@ref) for details on creating plans and on optional
# keyword arguments.
#
# We can now apply this plan to the three velocity components to obtain the
# respective Fourier coefficients ``\hat{\bm{v}}_{\bm{k}}``:

v̂s = plan .* v⃗₀
summary(v̂s[1])

# Note that, in Fourier space, the domain decomposition is performed along the
# directions ``x`` and ``y``:

pencil(v̂s[1])

# This is because the 3D FFTs are performed one dimension at a time, with the
# ``x`` direction first and the ``z`` direction last.
# To efficiently perform an FFT along a given direction (taking advantage of
# serial FFT implementations like FFTW), all the data along that direction must
# be contained locally within a single MPI process.
# For that reason, data redistributions (or *transpositions*) among MPI
# processes are performed behind the scenes during each FFT computation.
# Such transpositions require important communications between MPI processes,
# and are usually the most time-consuming aspect of massively-parallel
# simulations using this kind of methods.
#
# To solve the Navier--Stokes equations in Fourier space, we will
# also need the respective wave numbers ``\bm{k}`` associated to the local MPI
# process.
# Similarly to the local grid points, these are obtained using the `localgrid`
# function:

grid_fourier = localgrid(v̂s[1], ks_global)

# As an example, let's first use this to compute and plot the vorticity
# associated to the initial condition.
# The vorticity is defined as the curl of the velocity,
# ``\bm{ω} = \bm{∇} × \bm{v}``.
# In Fourier space, this becomes ``\hat{\bm{ω}} = i \bm{k} × \hat{\bm{v}}``.

using StaticArrays: SVector
using LinearAlgebra: ×

function curl_fourier!(
        ω̂s::NTuple{N, <:PencilArray}, v̂s::NTuple{N, <:PencilArray}, grid_fourier,
    ) where {N}
    @inbounds for I ∈ eachindex(grid_fourier)
        ## We use StaticArrays for the cross product between small vectors.
        ik⃗ = im * SVector(grid_fourier[I])
        v⃗ = SVector(getindex.(v̂s, Ref(I)))  # = (v̂s[1][I], v̂s[2][I], ...)
        ω⃗ = ik⃗ × v⃗
        for n ∈ eachindex(ω⃗)
            ω̂s[n][I] = ω⃗[n]
        end
    end
    ω̂s
end

ω̂s = similar.(v̂s)
curl_fourier!(ω̂s, v̂s, grid_fourier);

# We finally transform back to physical space and plot the result:

ωs = plan .\ ω̂s

let fig = Figure(resolution = (700, 600))
    ax = Axis3(fig[1, 1]; aspect = :data, xlabel = "x", ylabel = "y", zlabel = "z")
    ω_norm = vecnorm(ωs)
    ct = contour!(
        ax, grid.x, grid.y, grid.z, ω_norm;
        alpha = 0.1, levels = 0.8:0.2:2.0,
        colormap = :viridis, colorrange = (0.8, 2.0),
    )
    cb = Colorbar(fig[1, 2], ct; label = "Vorticity magnitude")
    fig
end

# ## Computing the non-linear term
#
# One can show that, in Fourier space, the incompressible Navier--Stokes
# equations can be written as
#
# ```math
# ∂_t \hat{\bm{v}}_{\bm{k}} =
# - \mathcal{P}_{\bm{k}} \! \left[ \widehat{(\bm{v} ⋅ \bm{∇}) \bm{v}} \right]
# - ν |\bm{k}|^2 \hat{\bm{v}}_{\bm{k}}
# \quad \text{ with } \quad
# \mathcal{P}_{\bm{k}}(\hat{\bm{F}}_{\bm{k}}) = \left( I - \frac{\bm{k} ⊗
# \bm{k}}{|\bm{k}|^2} \right) \hat{\bm{F}}_{\bm{k}},
# ```
#
# where ``\mathcal{P}_{\bm{k}}`` is a projection operator allowing to preserve the
# incompressibility condition ``\bm{∇} ⋅ \bm{v} = 0``.
# This operator encodes the action of the pressure gradient term, which serves
# precisely to enforce incompressibility.
# Note that, because of this, the pressure gradient dissapears from the
# equations.
#
# Now that we have the wave numbers ``\bm{k}``, computing the linear viscous
# term in Fourier space is straighforward once we have the Fourier coefficients
# ``\hat{\bm{v}}_{\bm{k}}`` of the velocity field.
# What is slightly more challenging (and much more costly) is the computation of
# the non-linear term in Fourier space, ``\hat{\bm{F}}_{\bm{k}} =
# \left[ \widehat{(\bm{v} ⋅ \bm{∇}) \bm{v}} \right]_{\bm{k}}``.
# In the pseudo-spectral method, the quadratic nonlinearity is computed
# by collocation in physical space (i.e. this term is evaluated at grid points),
# while derivatives are computed in Fourier space.
# This requires transforming fields back and forth between both spaces.
#
# Below we implement a function that computes the non-linear term in Fourier
# space based on its convective form ``(\bm{v} ⋅ \bm{∇}) \bm{v} = \bm{∇} ⋅
# (\bm{v} ⊗ \bm{v})``.
# Note that this equivalence uses the incompressibility condition ``\bm{∇} ⋅ \bm{v} = 0``.

using LinearAlgebra: mul!, ldiv!  # for applying FFT plans in-place

## Compute non-linear term in Fourier space from velocity field in physical
## space. Optional keyword arguments may be passed to avoid allocations.
function ns_nonlinear!(
        F̂s, vs, plan, grid_fourier;
        vbuf = similar(vs[1]), v̂buf = similar(F̂s[1]),
    )
    ## Compute F_i = ∂_j (v_i v_j) for each i.
    ## In Fourier space: F̂_i = im * k_j * FFT(v_i * v_j)
    w, ŵ = vbuf, v̂buf
    @inbounds for (i, F̂i) ∈ enumerate(F̂s)
        F̂i .= 0
        vi = vs[i]
        for (j, vj) ∈ enumerate(vs)
            w .= vi .* vj     # w = v_i * v_j in physical space
            mul!(ŵ, plan, w)  # same in Fourier space
            ## Add derivative in Fourier space
            for I ∈ eachindex(grid_fourier)
                k⃗ = grid_fourier[I]  # = (kx, ky, kz)
                kj = k⃗[j]
                F̂i[I] += im * kj * ŵ[I]
            end
        end
    end
    F̂s
end

# As an example, let's use this function on our initial velocity field:

F̂s = similar.(v̂s)
ns_nonlinear!(F̂s, v⃗₀, plan, grid_fourier);

# Strictly speaking, computing the non-linear term by collocation can lead to
# [aliasing
# errors](https://en.wikipedia.org/wiki/Aliasing#Sampling_sinusoidal_functions),
# as the quadratic term excites Fourier modes that fall beyond the range of
# resolved wave numbers.
# The typical solution is to apply Orzsag's 2/3 rule to zero-out the Fourier
# coefficients associated to the highest wave numbers.
# We define a function that applies this procedure below.

function dealias_twothirds!(ŵs::Tuple, grid_fourier, ks_global)
    ks_max = maximum.(abs, ks_global)  # maximum stored wave numbers (kx_max, ky_max, kz_max)
    ks_lim = (2 / 3) .* ks_max
    @inbounds for I ∈ eachindex(grid_fourier)
        k⃗ = grid_fourier[I]
        if any(abs.(k⃗) .> ks_lim)
            for ŵ ∈ ŵs
                ŵ[I] = 0
            end
        end
    end
    ŵs
end

## We can apply this on the previously computed non-linear term:
dealias_twothirds!(F̂s, grid_fourier, ks_global);

# Finally, we implement the projection associated to the incompressibility
# condition:

function project_divergence_free!(ûs, grid_fourier)
    @inbounds for I ∈ eachindex(grid_fourier)
        k⃗ = grid_fourier[I]
        k² = sum(abs2, k⃗)
        iszero(k²) && continue  # avoid division by zero
        û = getindex.(ûs, Ref(I))  # (ûs[1][I], ûs[2][I], ...)
        for i ∈ eachindex(û)
            ŵ = û[i]
            for j ∈ eachindex(û)
                ŵ -= k⃗[i] * k⃗[j] * û[j] / k²
            end
            ûs[i][I] = ŵ
        end
    end
    ûs
end

# We can verify the correctness of the projection operator by checking that the
# initial velocity field is not modified by it, since it is already
# incompressible:

v̂s_proj = project_divergence_free!(copy.(v̂s), grid_fourier)
v̂s_proj .≈ v̂s  # the last one may be false because v_z = 0 initially

# ## Putting it all together
#
# To perform the time integration of the Navier--Stokes equations, we will use
# the timestepping routines implemented in the DifferentialEquations.jl suite.
# For simplicity, we use here an explicit Runge--Kutta scheme.
# In this case, we just need to write a function that computes the right-hand
# side of the Navier--Stokes equations in Fourier space:

function ns_rhs!(
        dvs::NTuple{N, <:PencilArray}, vs::NTuple{N, <:PencilArray}, p, t,
    ) where {N}
    ## 1. Compute non-linear term and dealias it
    (; plan, cache, ks_global, grid_fourier) = p
    F̂s = cache.F̂s
    ns_nonlinear!(F̂s, vs, plan, grid_fourier; vbuf = dvs[1], v̂buf = cache.v̂s[1])
    dealias_twothirds!(F̂s, grid_fourier, ks_global)

    ## 2. Project onto divergence-free space
    project_divergence_free!(F̂s, grid_fourier)

    ## 3. Transform velocity to Fourier space
    v̂s = cache.v̂s
    map((v, v̂) -> mul!(v̂, plan, v), vs, v̂s)

    ## 4. Add viscous term (and multiply projected non-linear term by -1)
    ν = p.ν
    for n ∈ eachindex(v̂s)
        v̂ = v̂s[n]
        F̂ = F̂s[n]
        @inbounds for I ∈ eachindex(grid_fourier)
            k⃗ = grid_fourier[I]  # = (kx, ky, kz)
            k² = sum(abs2, k⃗)
            F̂[I] = -F̂[I] - ν * k² * v̂[I]
        end
    end

    ## 5. Transform RHS back to physical space
    map((dv, dv̂) -> ldiv!(dv, plan, dv̂), dvs, F̂s)

    nothing
end

# For the time-stepping, we load OrdinaryDiffEq.jl from the
# DifferentialEquations.jl suite and set-up the simulation.
# Since DifferentialEquations.jl can't directly deal with tuples of arrays, we
# convert the input data to the
# [`ArrayPartition`](https://github.com/SciML/RecursiveArrayTools.jl#arraypartition)
# type and write an interface function to make things work with our functions
# defined above.

using OrdinaryDiffEq
using RecursiveArrayTools: ArrayPartition

ns_rhs!(dv::ArrayPartition, v::ArrayPartition, args...) = ns_rhs!(dv.x, v.x, args...)
vs_init_ode = ArrayPartition(v⃗₀)
summary(vs_init_ode)

# We now define solver parameters and temporary variables, and initialise the
# problem:

params = (;
    ν = 5e-3,  # kinematic viscosity
    plan, grid_fourier, ks_global,
    cache = (
        v̂s = similar.(v̂s),
        F̂s = similar.(v̂s),
    )
)

tspan = (0.0, 10.0)
prob = ODEProblem(ns_rhs!, vs_init_ode, tspan, params)
integrator = init(prob, RK4(); dt = 1e-3, save_everystep = false);

# We finally solve the problem over time and plot the vorticity associated to
# the solution.
# It is also useful to look at the energy spectrum ``E(k)``, to see if the small
# scales are correctly resolved.
# To obtain a turbulent flow, the viscosity ``ν`` must be small enough to allow
# the transient appearance of an energy cascade towards the small scales (i.e.
# from small to large ``k``), while high enough to allow the small-scale motions
# to be correctly resolved.

function energy_spectrum!(Ek, ks, v̂s, grid_fourier)
    Nk = length(Ek)
    @assert Nk == length(ks)
    Ek .= 0
    for I ∈ eachindex(grid_fourier)
        k⃗ = grid_fourier[I]  # = (kx, ky, kz)
        knorm = sqrt(sum(abs2, k⃗))
        i = searchsortedfirst(ks, knorm)
        i > Nk && continue
        v⃗ = getindex.(v̂s, Ref(I))  # = (v̂s[1][I], v̂s[2][I], ...)
        factor = k⃗[1] == 0 ? 1 : 2  # account for Hermitian symmetry and r2c transform
        Ek[i] += factor * sum(abs2, v⃗) / 2
    end
    MPI.Allreduce!(Ek, +, get_comm(v̂s[1]))  # sum across all processes
    Ek
end

ks = rfftfreq(Ns[1], 2π * Ns[1] / Ls[1])
Ek = similar(ks)
v̂s = plan .* integrator.u.x
energy_spectrum!(Ek, ks, v̂s, grid_fourier)
Ek ./= scale_factor(plan)^2  # rescale energy

curl_fourier!(ω̂s, v̂s, grid_fourier)
ldiv!.(ωs, plan, ω̂s)
ω⃗_plot = Observable(ωs)
k_plot = @view ks[2:end]
E_plot = Observable(@view Ek[2:end])
t_plot = Observable(integrator.t)

fig = let
    fig = Figure(resolution = (1200, 600))
    ax = Axis3(
        fig[1, 1][1, 1]; title = @lift("t = $(round($t_plot, digits = 3))"),
        aspect = :data, xlabel = "x", ylabel = "y", zlabel = "z",
    )
    ω_mag = @lift vecnorm($ω⃗_plot)
    ω_mag_norm = @lift $ω_mag ./ maximum($ω_mag)
    ct = contour!(
        ax, grid.x, grid.y, grid.z, ω_mag_norm;
        alpha = 0.3, levels = 3,
        colormap = :viridis, colorrange = (0.0, 1.0),
    )
    cb = Colorbar(fig[1, 1][1, 2], ct; label = "Normalised vorticity magnitude")
    ax_sp = Axis(
        fig[1, 2];
        xlabel = "k", ylabel = "E(k)", xscale = log2, yscale = log10,
        title = "Kinetic energy spectrum",
    )
    ylims!(ax_sp, 1e-8, 1e0)
    scatterlines!(ax_sp, k_plot, E_plot)
    ks_slope = exp.(range(log(2.5), log(25.0), length = 3))
    E_fivethirds = @. 0.3 * ks_slope^(-5/3)
    @views lines!(ax_sp, ks_slope, E_fivethirds; color = :black, linestyle = :dot)
    text!(ax_sp, L"k^{-5/3}"; position = (ks_slope[2], E_fivethirds[2]), align = (:left, :bottom))
    fig
end

using Printf # hide
with_xvfb = ENV["DISPLAY"] == ":99" # hide
nstep = 0  # hide
const tmpdir = mktempdir()  # hide
filename_frame(procid, nstep) = joinpath(tmpdir, @sprintf("proc%d_%04d.png", procid, nstep)) # hide
record(fig, "vorticity_proc$procid.mp4"; framerate = 10) do io
    with_xvfb && recordframe!(io) # hide
    while integrator.t < 20
        dt = 0.001
        step!(integrator, dt)
        t_plot[] = integrator.t
        mul!.(v̂s, plan, integrator.u.x)  # current velocity in Fourier space
        curl_fourier!(ω̂s, v̂s, grid_fourier)
        ldiv!.(ω⃗_plot[], plan, ω̂s)
        ω⃗_plot[] = ω⃗_plot[]  # to force updating the plot
        energy_spectrum!(Ek, ks, v̂s, grid_fourier)
        Ek ./= scale_factor(plan)^2  # rescale energy
        E_plot[] = E_plot[]
        global nstep += 1  # hide
        with_xvfb ?  # hide
        save(filename_frame(procid, nstep), fig) :  # hide
        recordframe!(io)
    end
end;

if with_xvfb  # hide
    run(pipeline(`ffmpeg -y -r 10 -i $tmpdir/proc$(procid)_%04d.png -c:v libx264 -vf "fps=25,format=yuv420p" vorticity_proc$procid.mp4`; stdout = "ffmpeg.out", stderr = "ffmpeg.err"))   # hide
end              # hide
nothing          # hide

# ```@raw html
# <figure class="video_container">
#   <video controls="true" allowfullscreen="true">
#     <source src="../vorticity_proc1.mp4" type="video/mp4">
#   </video>
# </figure>
# ```


text:
function get_grid_points(obj::InterpolationGrid, arg0::jdouble, arg1::jdouble)
    return jcall(obj, "getGridPoints", Vector{jdouble}, (jdouble, jdouble), arg0, arg1)
end

ext: jl


text:
import NMRDataSetup

using FFTW
import PyPlot
import BSON

PyPlot.close("all")
fig_num = 1

PyPlot.matplotlib["rcParams"][:update](["font.size" => 22, "font.family" => "serif"])

### user inputs.
# save_dir = "/home/roy/MEGAsync/outputs/NMR/calibrate/"
# project_name = "phenylalanine-700"
#
# experiment_full_path = "/home/roy/Documents/repo/NMRData/combination/BMRB-700/L-Phenylalanine"
# solvent_ppm_guess = 4.7
# solvent_window_ppm = 0.1

# save_dir = "/home/roy/MEGAsync/outputs/NMR/calibrate/"
# project_name = "isoleucine-700"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/combination/BMRB-700/L-Isoleucine"
# solvent_ppm_guess = 4.7
# solvent_window_ppm = 0.1

save_dir = "/home/roy/MEGAsync/outputs/NMR/calibrate/"
project_name = "glucose-700"
experiment_full_path = "/home/roy/Documents/repo/NMRData/combination/BMRB-700/D-(+)-Glucose"
solvent_ppm_guess = 4.7
solvent_window_ppm = 0.1

# save_dir = "/home/roy/MEGAsync/outputs/NMR/calibrate/"
# project_name = "glucose-500-2mM"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/combination/BMRB-500-2mM/D-(+)-Glucose"
# solvent_ppm_guess = 4.7
# solvent_window_ppm = 0.1

# save_dir = "/home/roy/MEGAsync/outputs/NMR/calibrate/"
# project_name = "glucose-500-0.5mM"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/combination/BMRB-500-0.5mM/D-(+)-Glucose"
# solvent_ppm_guess = 4.7
# solvent_window_ppm = 0.1

# save_dir = "/home/roy/MEGAsync/outputs/NMR/calibrate/"
# project_name = "glucose-nrc-2018"
# experiment_full_path = "/home/roy/Documents/repo/NMRData/src/experiments/NRC/misc/glucose/Sep-25-2018"
# solvent_ppm_guess = 4.7
# solvent_window_ppm = 0.1

# save_dir = "/home/roy/MEGAsync/outputs/NMR/calibrate/"
# project_name = "Nam-Jan2022"
# solvent_ppm_guess = 4.7
# solvent_window_ppm = 0.1
# experiment_full_path = "/home/roy/MEGAsync/data/NMR/NRC/Nam_4_amino_acid_mixture_Jan_2022/NHK220113_4AA_glu_PBS_d2o/1"

### end inputs.

## load.
s_t, S, hz2ppmfunc, ppm2hzfunc, ν_0ppm, fs, SW, α_0ppm, β_0ppm, λ_0ppm, Ω_0ppm,
    results_0ppm, dic, α_solvent, β_solvent, λ_solvent, Ω_solvent,
    results_solvent = NMRDataSetup.loadspectrum(experiment_full_path;
    solvent_ppm = solvent_ppm_guess,
    solvent_window_ppm = solvent_window_ppm)

## store.
save_folder_path = joinpath(save_dir, project_name)
isdir(save_folder_path) || mkpath(save_folder_path); # make save folder if it doesn't exist.

save_path = joinpath(save_folder_path, "$(project_name).bson")
BSON.bson(save_path,
s_t = s_t,
fs = fs,
SW = SW,
ν_0ppm = ν_0ppm,
α_0ppm = α_0ppm,
β_0ppm = β_0ppm,
λ_0ppm = λ_0ppm)

## visualize.
hz2ppmfunc = uu->(uu - ν_0ppm)*SW/fs
ppm2hzfunc = pp->(ν_0ppm + pp*fs/SW)

offset_Hz = ν_0ppm - (ppm2hzfunc(0.3)-ppm2hzfunc(0.0))

N = length(s_t)
DFT_s = fft(s_t)
U_DFT, U_y, U_inds = NMRDataSetup.getwraparoundDFTfreqs(N, fs, offset_Hz)

S_U = DFT_s[U_inds]
P_y = hz2ppmfunc.(U_y)

q = uu->NMRDataSetup.evalcomplexLorentzian(uu, α_0ppm, β_0ppm, λ_0ppm, 2*π*ν_0ppm)*fs
q_U = q.(U_y)

PyPlot.figure(fig_num)
fig_num += 1

PyPlot.plot(P_y, real.(S_U), label = "data")
PyPlot.plot(P_y, real.(q_U), "--", label = "estimated 0ppm resonance component")

PyPlot.legend()
PyPlot.xlabel("ppm")
PyPlot.ylabel("real")
PyPlot.title("real part of data spectrum (shifted and scaled Discrete Fourier Transform)")

###normalized.

## normalize by max magnitude in the 0 ppm neighbourhood.
# inds = findall(xx->(-0.3<xx<0.3), P_y)
# Z = maximum(abs.(S_U[inds]))/9 # α at 0ppm for DSS is 9 because there are 9 protons.

# normalize by the 0 ppm estimate, if it had an α of 9, which corresponds to the 9 protons for DSS' 0ppm component.
c = NMRDataSetup.evalcomplexLorentzian(ν_0ppm, 9.0, β_0ppm, λ_0ppm, 2*π*ν_0ppm)
Z = q(ν_0ppm)/c
y = S_U ./ Z

PyPlot.figure(fig_num)
fig_num += 1

PyPlot.plot(P_y, real.(y))

PyPlot.legend()
PyPlot.xlabel("ppm")
PyPlot.ylabel("real")
PyPlot.title("real part of normalized data spectrum s.t. 0ppm peak has magnitude 1")


text:
module FEFields

using StaticArrays

"""
    FEField{N, T, IT}

Type of a finite element field. Parameterized with
- `N`: number of degrees of freedom per entity,
- `T`: type of the degree of freedom value,
- `IT`: type of the index (integer value). This describes the serial numbers
  of the degrees of freedom.
"""
mutable struct FEField{N, T, IT}
    dofnums::Vector{SVector{N, IT}}
    isdatum::Vector{SVector{N, Bool}}
    dofvals::Vector{SVector{N, T}}

    """
        FEField(::Val{N}, ::Type{T}, ::Type{IT}, ntrms) where {N, T, IT}

    Construct a field with the given number of terms `ntrms`.
    """
    function FEField(::Val{N}, ::Type{T}, ::Type{IT}, ntrms) where {N, T, IT}
        z = fill(zero(IT), N)
        dofnums = [SVector{N}(z) for i in 1:ntrms]
        z = fill(false, N)
        isdatum = [z for i in 1:ntrms]
        z = fill(zero(T), N)
        dofvals = [SVector{N}(z) for i in 1:ntrms]
        return new{N, T, IT}(dofnums, isdatum, dofvals)
    end
end

"""
    doftype(fef::FEField{N, T, IT}) where {N, T, IT}

Type of a degree of freedom value.
"""
doftype(fef::FEField{N, T, IT}) where {N, T, IT} = T

"""
    dofnumtype(fef::FEField{N, T, IT}) where {N, T, IT}

Type of the index (serial number) of the degree of freedom. Integer.
"""
dofnumtype(fef::FEField{N, T, IT}) where {N, T, IT} = IT
"""
    nterms(fef::FEField)

Number of terms in the field.
"""
nterms(fef::FEField) = length(fef.dofnums)

"""
    ndofsperterm(fef::FEField{N}) where {N}

Number of degrees of freedom per term of the field.
"""
ndofsperterm(fef::FEField{N}) where {N} = N

"""
    ndofs(fef::FEField)

Total number of degrees of freedom in the field.
"""
ndofs(fef::FEField) = nterms(fef) * ndofsperterm(fef)

"""
    dofnums(fef::FEField, n)

Provide degree of freedom numbers for given entity.
"""
dofnums(fef::FEField, n) = fef.dofnums[n]

"""
    dofvals(fef::FEField, n)

Provide degree of freedom values for given entity.
"""
dofvals(fef::FEField, n) = fef.dofvals[n]

"""
    setebc!(self::FEField, tid, comp, val::T) where {T}

Set the value of one particular degree of freedom to a given number.

- `tid`: which term,
- `comp`: which component of the term,
- `flag`: true or false.
"""
function setisdatum!(self::FEField, tid, comp, flag)
    ik = MVector(self.isdatum[tid])
    ik[comp] = flag
    self.isdatum[tid] = ik
    return  self
end

"""
    setdofval!(self::FEField, tid, comp, val::T) where {T}

Set the value of one particular degree of freedom to a given number.

- `tid`: which term,
- `comp`: which component of the term,
- `val`: value to which the degree of freedom should be set.
"""
function setdofval!(self::FEField, tid, comp, val::T) where {T}
    d = MVector(self.dofvals[tid])
    d[comp] = val
    self.dofvals[tid] = d
    return  self
end

"""
    setebc!(self::FEField, tid, comp, val::T) where {T}

Set the value of one particular degree of freedom to a given number.

- `tid`: which term,
- `comp`: which component of the term,
- `val`: value to which the degree of freedom should be set.
"""
function setebc!(self::FEField, tid, comp, val::T) where {T}
    setisdatum!(self, tid, comp, true)
    setdofval!(self, tid, comp, val)
    return  self
end

"""
    numberfreedofs!(f::FEField, firstnum = 1)

Number the unknowns in the field, starting from the one supplied on input.

Note: The data degrees of freedom have their numbers zeroed out.
"""
function numberfreedofs!(f::FEField, firstnum = 1)
    lnum = fnum = zero(dofnumtype(f)) + firstnum
    tnum = 0
    for i in 1:nterms(f)
        n = MVector(f.dofnums[i])
        for j in 1:ndofsperterm(f)
            if (!f.isdatum[i][j]) # unknown degree of freedom
                n[j] = lnum
                tnum = tnum + 1
                lnum = lnum + 1
            else
                n[j] = zero(dofnumtype(f))
            end
        end
        f.dofnums[i] = n
    end
    return  f
end

"""
    numberdatadofs!(f::FEField, firstnum = 1)

Number the data degrees of freedom in the field. Start from the
number supplied on input.

Note: The free degrees of freedom must be numbered first.
"""
function numberdatadofs!(f::FEField, firstnum = 1)
    num = zero(dofnumtype(f)) + firstnum
    for i in 1:nterms(f)
        n = MVector(f.dofnums[i])
        for j in 1:ndofsperterm(f)
            if (f.isdatum[i][j]) # known (data) degree of freedom
                n[j] = num
                num = num + 1
            end
        end
        f.dofnums[i] = n
    end
    return  f
end

"""
    freedofnums(f::FEField)

Collect information about unknown (free) degree of freedom numbers.

First number, last number, and the total number of degrees of freedom are
returned as a tuple.
"""
function freedofnums(f::FEField)
    tnum = lnum = zero(dofnumtype(f))
    fnum = typemax(dofnumtype(f))
    for i in 1:nterms(f)
        for j in 1:ndofsperterm(f)
            # unknown degree of freedom
            if (!f.isdatum[i][j])
                tnum = tnum + 1
                lnum = max(lnum, f.dofnums[i][j])
                fnum = min(fnum, f.dofnums[i][j])
            end
        end
    end
    return  (fnum, lnum, tnum)
end

"""
    highestfreedofnum(f::FEField)

Compute the highest serial number of a free degree of freedom in the field.
"""
highestfreedofnum(f::FEField) = freedofnums(f)[2]

"""
    datadofnums(f::FEField)

Collect information about known (data) degree of freedom numbers.

First number, last number, and the total number of degrees of freedom are
returned as a tuple.
"""
function datadofnums(f::FEField)
    tnum = lnum = zero(dofnumtype(f))
    fnum = typemax(dofnumtype(f))
    for i in 1:nterms(f)
        for j in 1:ndofsperterm(f)
            # known degree of freedom
            if (f.isdatum[i][j])
                tnum = tnum + 1
                lnum = max(lnum, f.dofnums[i][j])
                fnum = min(fnum, f.dofnums[i][j])
            end
        end
    end
    return  (fnum, lnum, tnum)
end

"""
    highestdatadofnum(f::FEField)

Compute the highest serial number of a datum degree of freedom in the field.
"""
highestdatadofnum(f::FEField) = datadofnums(f)[2]

"""
    gathersysvec!(vec, self::FEField)

Gather system vector contributions from the field.
"""
function gathersysvec!(vec, self::FEField)
    nt = nterms(self)
    ndpt = ndofsperterm(self)
    for i in 1:nt
        en = self.dofnums[i]
        for j in 1:ndpt
            vec[en[j]] = self.dofvals[i][j]
        end
    end
    return vec
end

"""
    scattersysvec!(self::FEField, v)

Scatter a system vector into the field.
"""
function scattersysvec!(self::FEField, v)
    vl = length(v)
    nt = nterms(self)
    ndpt = ndofsperterm(self)
    for i in 1:nt
        en = self.dofnums[i]
        for j in 1:ndpt
            if en[j] <= vl
                d = MVector(self.dofvals[i])
                d[j] = v[en[j]]
                self.dofvals[i] = d
            end
        end
    end
    return self
end

end


text:
# Convert a vectors of integers to a unique integer.
# The components of the vector are 1-based rather than 0-based.

# Example 1:
# julia> vect_to_int( [3,4,5,6],[10,10,10,10])
#  2345
# Example 2:
# julia> vect_to_int([3,2,1,4],[7,2,5,8])
#  203

function vect_to_int( ints::Vector{Int64}, max::Vector{Int64} )
  result = Int128(ints[1])-1
  for i = 2:length(ints)
    result *= max[i]
    result += ints[i]-1
  end
  result
end

# julia> int_to_vect( 2345, [10,10,10,10])
# 4-element Array{Int64,1}:
#  Int64[ 3, 4, 5, 6 ]
# julia> int_to_vect(203,[7,2,5,8])
# 4-element Array{Int64,1}:
#  Int64[ 3, 2, 1, 4]

function int_to_vect( int::Integer, max::Vector{Int64} )
  ints = zeros(Int64,length(max))
  for i = length(max):-1:1
    ints[i] = int % max[i] + 1
    int = div(int,max[i])
  end
  ints
end


text:
using Test, DistributionVectors
using Distributions, StatsBase, MissingStrategies, Missings
using LinearAlgebra

@testset "Normal" begin

@testset "two vars uncorrelated" begin
    # generate nSample values of two lognormal random variables
    d1 = Normal(2, 0.15)
    d2 = Normal(3, 0.25)
    dv = SimpleDistributionVector(d1, d2);
    #dsum = @inferred mean_normals(dv)
    dsum = @inferred sum(dv)
    dmean = @inferred mean(dv)
    @testset "no missings" begin
        @test params(dsum)[1] ≈ 5
        # checked with random numbers
        #boot_dvsum(dv)
        @test params(dsum)[2] ≈ √(abs2(0.15) + abs2(0.25))
        @test dmean.μ ≈ 2.5
        @test dmean.σ/dmean.μ ≈ dsum.σ/dsum.μ
        # squared mean uncertainty decreases by √n
        @test dmean.σ ≈ sqrt(mean([abs2(0.15),abs2(0.25)]))/√2
    end;
    @testset "with missings" begin
        dv = SimpleDistributionVector(d1, missing);
        @test isequal(collect(dv), [d1, missing])
        @test_throws Exception dsum2 = sum(dv) # error: forgot SkipMissing()
        dsum2 = @inferred sum(skipmissing(dv))
        @test dsum2 == d1
        dsum3 = @inferred sum(dv, SkipMissing())
        @test dsum3 == d1
        #@btime sum(skipmissing($dv)) # does not allocate
        dmean2 = @inferred mean(dv, SkipMissing())
        @test dmean2 == d1
    end;
    @testset "with gapfilling flag" begin
      dv = SimpleDistributionVector(d1, d2, d1);
      #dv = SimpleDistributionVector(d1, d1, d1);
      isgapfilled = [true, false, false]
      dsum = @inferred sum(dv, isgapfilled=isgapfilled)
      #@code_warntype sum(dv, isgapfilled)
      @test mean(dsum) == mean(sum(dv)) # mean takes all into account
      @test std(dsum) > std(sum(dv))
      # relative error like non-gapfilled
      dsumn = sum(dv[.!isgapfilled])
      @test std(dsum)/mean(dsum) == std(dsumn)/mean(dsumn)
      # mean function
      dmean = @inferred mean(dv; isgapfilled=isgapfilled)
      @test dmean.μ ≈ dsum.μ / 3
      # relative error like non-gapfilled
      dmeann = mean(dv[.!isgapfilled])
      @test std(dmean)/mean(dmean) ≈ std(dmeann)/mean(dmeann)
    end;
    @testset "with missings and gapfilling flag" begin
      dv = SimpleDistributionVector(d1, d2, d1, missing);
      isgapfilled = [true, true, false, false]
      dsum5 = @inferred sum(dv[1:3], isgapfilled=isgapfilled[1:3])
      dsum = @inferred sum(dv, SkipMissing(); isgapfilled=isgapfilled)
      #@code_warntype sum(dv, SkipMissing(); isgapfilled)
      @test dsum == dsum5
      # mean function
      dmean = @inferred mean(dv, SkipMissing(); isgapfilled=isgapfilled)
      dmeann = @inferred mean(dv[1:3], isgapfilled=isgapfilled[1:3])
      @test dmean == dmeann
    end;
end;

@testset "few correlated vars" begin
  mu = [110,100,80,120,160.0]
  sigma = [20.0,50,10,30,10]
  acf1 = @inferred AutoCorrelationFunction([1,0.4,0.1])
  n = length(mu)
  corrM = @inferred cormatrix_for_acf(n, acf1)
  dv = SimpleDistributionVector(Normal{eltype(mu)}, mu, sigma);
  mum = allowmissing(mu); mum[1] = missing
  dvm = SimpleDistributionVector(Normal{eltype(mu)}, mum, sigma)
  #
  @testset "matrix without missing" begin
    dsum = @inferred sum(dv, Symmetric(corrM))
    # checked with random numbers
    #boot_dvsums_acf(dv, acf1)
    @test mean(dsum) ≈ sum(mean.(dv))
    dsumuncorr = sum(dv)
    @test params(dsum)[2] > params(dsumuncorr)[2]
    # explicit sum over covariance Sigma
    Sigma = Diagonal(sigma) * corrM * Diagonal(sigma)
    @test params(dsum)[2] == sqrt(sum(Sigma))
    # mean function
    dmean = @inferred mean(dv, Symmetric(corrM))
    @test dmean.μ ≈ dsum.μ / length(dv)
    @test dmean.σ/dmean.μ ≈ dsum.σ/dsum.μ
  end;
  @testset "matrix with missing" begin
    @test_throws Exception dsumm = sum(dvm, Symmetric(corrM))
    #S = similar(mum);
    dsumm = @inferred sum(dvm, Symmetric(corrM), SkipMissing())
    params(dsumm) == params(sum(dvm[2:end], Symmetric(corrM[2:end,2:end])))
    # mean function
    dmean = @inferred mean(dvm, Symmetric(corrM), SkipMissing())
    @test dmean.μ ≈ dsumm.μ / 4
    @test dmean.σ/dmean.μ ≈ dsumm.σ/dsumm.μ
  end;
  @testset "with gapfilling flag" begin
    isgapfilled = fill(false, length(dv)); isgapfilled[4:end] .= true
    dsum4 = sum(dv, Symmetric(corrM))
    dsum = @inferred sum(dv, Symmetric(corrM), isgapfilled=isgapfilled)
    #@code_warntype sum(dv, isgapfilled)
    @test mean(dsum) == mean(dsum4)
    @test std(dsum) > std(dsum4)
    # test with explicit sum over Sigma and same relative error
    ifin = .!isgapfilled
    Sigma = Diagonal(sigma[ifin]) * corrM[ifin,ifin] * Diagonal(sigma[ifin])
    @test std(dsum)/mean(dsum) == sqrt(sum(Sigma))/sum(mu[ifin])
    # mean function
    dmean = @inferred mean(dv, Symmetric(corrM), SkipMissing(); isgapfilled=isgapfilled)
    @test dmean.μ ≈ dsum.μ / 5
    @test dmean.σ/dmean.μ ≈ dsum.σ/dsum.μ
  end;
  @testset "with missings and gapfilling flag" begin
    isgapfilled = fill(false, length(dvm)); isgapfilled[4:end] .= true
    dsum4 = sum(dvm, Symmetric(corrM), SkipMissing())
    dsum4b = @inferred sum(
      dvm, Symmetric(corrM), SkipMissing(); isgapfilled=isgapfilled)
    #@code_warntype sum(dvm, isgapfilled)
    @test mean(dsum4b) == mean(dsum4)
    @test std(dsum4b) > std(dsum4)
    # test with explicit sum over Sigma and same relative error
    ifin = .!(isgapfilled .| ismissing.(dvm))
    Sigma = Diagonal(sigma[ifin]) * corrM[ifin,ifin] * Diagonal(sigma[ifin])
    @test std(dsum4b)/mean(dsum4b) == sqrt(sum(Sigma))/sum(mu[ifin])
    #
    # acf variant
    dsum4c = @inferred sum(
      dvm, acf1, SkipMissing(); isgapfilled=isgapfilled)
    @test dsum4c == dsum4b
    # mean function
    dmean = @inferred mean(dvm, Symmetric(corrM), SkipMissing();
      isgapfilled=isgapfilled)
    @test dmean.μ ≈ dsum4b.μ / 4
    @test dmean.σ/dmean.μ ≈ dsum4b.σ/dsum4b.μ
    dmean_acf = @inferred mean(dvm, acf1, SkipMissing(); isgapfilled=isgapfilled)
    @test dmean_acf == dmean
  end;
end;

end;


text:
"""
Execution state that encapsulates synchronization primitives and resources
bound to a command submission on the GPU.

Resources bound to the related execution can be either freed or released (dereferenced)
once the execution completes.
The execution is assumed to be completed when the execution
state has been waited on (see [`wait`](@ref)), or when an inherited execution state has
been completed execution.

Execution state can be inherited across submissions, where resource dependencies
are transfered over to a new execution state bound to a command that **must** include
the inherited `ExecutionState`'s semaphore as a wait semaphore.
"""
struct ExecutionState
    queue::Queue
    semaphore::Optional{Vk.SemaphoreSubmitInfoKHR}
    fence::Optional{Vk.Fence}
    free_after_completion::Vector{Ref}
    release_after_completion::Vector{Ref}
end

function ExecutionState(queue; semaphore = nothing, fence = nothing, free_after_completion = [], release_after_completion = [])
    ExecutionState(queue, semaphore, fence, free_after_completion, release_after_completion)
end

function finalize!(exec::ExecutionState)
    for resource in exec.free_after_completion
        finalize(resource[])
    end
    empty!(exec.free_after_completion)
    empty!(exec.release_after_completion)
    true
end

function _wait(fence::Vk.Fence, timeout)
    ret = unwrap(Vk.wait_for_fences(Vk.device(fence), [fence], true, timeout))
    ret == Vk.SUCCESS
end

function _wait(fences, timeout)
    isempty(fences) && return true
    ret = unwrap(Vk.wait_for_fences(Vk.device(first(fence)), fences, true, timeout))
    ret == Vk.SUCCESS
end

function Base.wait(exec::ExecutionState, timeout = typemax(UInt32))
    fence = exec.fence::Vk.Fence
    _wait(fence, timeout) && finalize!(exec)
end

function Base.wait(execs::AbstractVector{ExecutionState}, timeout = typemax(UInt32))
    fences = map(Base.Fix2(getproperty, :fence), execs)::Vector{Fence}
    _wait(fences, timeout) && all(finalize!, execs)
end

function submit(dispatch::QueueDispatch, queue_family_index, submit_infos;
    signal_fence = false,
    semaphore = nothing,
    free_after_completion = [],
    release_after_completion = [],
    inherit = nothing,
    check_inheritance = true,
)
    q = queue(dispatch, queue_family_index)

    fence = signal_fence ? Vk.Fence(Vk.device(q)) : nothing

    if inherit isa ExecutionState
        append!(free_after_completion, inherit.free_after_completion)
        append!(release_after_completion, inherit.release_after_completion)
        empty!(inherit.free_after_completion)
        empty!(inherit.release_after_completion)
        if check_inheritance && !any(Base.Fix1(in, inherit.semaphore), submit_infos.wait_semaphores)
            error("No wait semaphore has been registered that
                matches the inherited state.")
        end
    end

    for submit_info in submit_infos
        for command_buffer_info in submit_info.command_buffer_infos
            end_recording(command_buffer_info.command_buffer)
        end
    end

    unwrap(Vk.queue_submit_2_khr(q, submit_infos; fence = something(fence, C_NULL)))
    ExecutionState(q; semaphore, fence, free_after_completion, release_after_completion)
end

function submit(dispatch::QueueDispatch, queue_family_index, submit_info::Vk.SubmitInfo2KHR;
    signal_fence = false,
    semaphore = nothing,
    free_after_completion = [],
    release_after_completion = [],
    inherit = nothing,
    check_inheritance = true,
)
    if !isnothing(semaphore)
        semaphore in submit_info.signal_semaphore_infos || error("The provided semaphore was not included in the submission structure")
    end
    submit(dispatch, queue_family_index, [submit_info]; signal_fence, semaphore, free_after_completion, release_after_completion, inherit, check_inheritance)
end

function Base.show(io::IO, exec::ExecutionState)
    print(io, ExecutionState, "($(exec.queue)")
    if !isnothing(exec.fence)
        is_complete = _wait(exec.fence, 0)
        if is_complete
            print(io, ", completed execution")
        else
            print(io, ", in progress")
        end
    end
    print(io, ')')
end


text:
using ROOTFramework, Cxx

tfile = open(TFile, "out.root", "recreate")
ttree = create_ttree!(tfile, "data", "Data")

idx = Ref{Int32}(0)
s = Ref(0.0)
a = [0.0, 0.0]
cpp_a = icxx""" std::vector<double>(2); """

create_branch!(ttree, "idx", idx)
create_branch!(ttree, "s", s)
create_branch!(ttree, "a", cpp_a)

@time for entry in 1:1000000
    idx[] = entry
    s[] = 2.0 + 2.0 * randn()

    a[1] = 12.0 + 2.0 * randn()
    a[2] = 22.0 + 2.0 * randn()
    resize!(cpp_a, length(a))
    copy!(cpp_a, a)

    push!(ttree)
end

close(tfile)


text:
using Test
using SimpleTensorNetworks.ContractionOrder
using LightGraphs: SimpleEdge

@testset "edge info" begin
    egs = [SimpleEdge(1, 2), SimpleEdge(2, 3), SimpleEdge(3, 1), SimpleEdge(3, 4)]
    log2_shapes = [log2.((2,2)), log2.((2,3)), log2.((3,2,6)), log2.((6,))]
    neighbors = [[2,3], [1,3], [1,2,4], [3]]
    strategy = "max_reduce_tri"
    ei = edgeinfo(egs[1:3], log2_shapes, neighbors, strategy)
    @test ei.data[SimpleEdge(3, 1)] ≈ -24.0
    @test ei.data[SimpleEdge(1, 2)] ≈ -4
    @test ei.data[SimpleEdge(2, 3)] ≈ -18
    @test select(ei, strategy, neighbors, nothing) == SimpleEdge(3, 1)

    # add!, remove!
    ei2 = add!(copy(ei), [4], log2_shapes, neighbors, strategy)
    ei3 = edgeinfo(egs, log2_shapes, neighbors, strategy)
    for (k, v) in ei2.data
        @test ei3.data[k] ≈ v
    end

    ei4 = remove!(copy(ei2), [4])
    @test ei4.data == ei.data
end


text:
using Dates,JLD2,Statistics,Optim,Parameters,Distributions,DataFrames,CSV
import KenyaCoVSD
include("fitting_methods.jl");

## Load the collated fits
@load("data/N_kenya.jld2")
@load("data/linelist_data_with_pos_neg_20feb_to_27apr_c_age.jld2")
@load("data/serological_data_with_20feb_to_31dec_age.jld2")
@load("data/serological_data_with_20feb_to_10Mar2021_age.jld2")
@load("data/cleaned_linelist20210426_deaths_c__date_of_lab_confirmation.jld2")
@load("data/p_ID.jld2")
@load("data/smoothedPCR_detection_after_infection.jld2")

@load("modelfits/Nairobi_model.jld2")
#schools effect

fits_schools_effect = []

for filename in readdir("modelfits",join = true)
    modeldict = load(filename)
    model = modeldict[first(keys(modeldict))]
    println("Gathering data from $(model.areaname)")
    params = get(model.MCMC_results.chain,:schooleffect)
    push!(fits_schools_effect,(name = model.areaname,
                                mean = mean(params.schooleffect[:]),
                                lwr = quantile(params.schooleffect[:],0.025),
                                upr = quantile(params.schooleffect[:],0.975)))
end

mean_school_effects = [1 .- fit.mean for fit in fits_schools_effect]
idx_nai = findfirst([fit.name .== "Nairobi" for fit in fits_schools_effect])
idx_mom = findfirst([fit.name .== "Mombasa" for fit in fits_schools_effect])
1 - fits_schools_effect[idx_nai].mean
1 - fits_schools_effect[idx_nai].upr
1 - fits_schools_effect[idx_nai].lwr
1 - fits_schools_effect[idx_mom].mean
1 - fits_schools_effect[idx_mom].upr
1 - fits_schools_effect[idx_mom].lwr

quantile(mean_school_effects,[0.25,0.5,0.75])

## Ct min
fits_ct_min1 = []

for filename in readdir("modelfits",join = true)
    modeldict = load(filename)
    model = modeldict[first(keys(modeldict))]
    println("Gathering data from $(model.areaname)")
    params = get(model.MCMC_results.chain,:ct_min1)
    push!(fits_ct_min1,(name = model.areaname,
                                mean = mean(params.ct_min1[:]),
                                lwr = quantile(params.ct_min1[:],0.025),
                                upr = quantile(params.ct_min1[:],0.975)))
end

mean_ct_min1= [1 - fit.mean for fit in fits_ct_min1]
quantile(mean_ct_min1,[0.25,0.5,0.75])

idx_nai = findfirst([fit.name .== "Nairobi" for fit in fits_ct_min1])
idx_mom = findfirst([fit.name .== "Mombasa" for fit in fits_ct_min1])
1 - fits_ct_min1[idx_nai].mean
1 - fits_ct_min1[idx_nai].upr
1 - fits_ct_min1[idx_nai].lwr
1 - fits_ct_min1[idx_mom].mean
1 - fits_ct_min1[idx_mom].upr
1 - fits_ct_min1[idx_mom].lwr

## P_eff
fits_P_eff = []

for filename in readdir("modelfits",join = true)
    modeldict = load(filename)
    model = modeldict[first(keys(modeldict))]
    println("Gathering data from $(model.areaname)")
    params = get(model.MCMC_results.chain,:P_eff)
    push!(fits_P_eff,(name = model.areaname,
                                mean = mean(params.P_eff[:]),
                                lwr = quantile(params.P_eff[:],0.025),
                                upr = quantile(params.P_eff[:],0.975)))
end

mean_P_eff = [1 - fit.mean for fit in fits_P_eff]
quantile(mean_P_eff,[0.25,0.5,0.75])

pop_weighted_peff = [(1 - fit.mean)*sum(N_kenya[:,fit.name])/sum(N_kenya) for fit in fits_P_eff]
sum(pop_weighted_peff)

idx_nai = findfirst([fit.name .== "Nairobi" for fit in fits_P_eff])
idx_mom = findfirst([fit.name .== "Mombasa" for fit in fits_P_eff])
1 - fits_P_eff[idx_nai].mean
1 - fits_P_eff[idx_nai].upr
1 - fits_P_eff[idx_nai].lwr
1 - fits_P_eff[idx_mom].mean
1 - fits_P_eff[idx_mom].upr
1 - fits_P_eff[idx_mom].lwr

## Rel detection rate
fits_rel_detect = []

for filename in readdir("modelfits",join = true)
    modeldict = load(filename)
    model = modeldict[first(keys(modeldict))]
    println("Gathering data from $(model.areaname)")
    params = get(model.MCMC_results.chain,[:p_test₁,:p_test₂])
    prob_detect_1 = [dailytestrate_to_probofdetections(p_test,PCR_array) for p_test in params.p_test₁[:]]
    prob_detect_2 = [dailytestrate_to_probofdetections(p_test,PCR_array) for p_test in params.p_test₂[:]]

    push!(fits_rel_detect,(name = model.areaname,
                                mean = mean(prob_detect_2./prob_detect_1),
                                lwr = quantile(prob_detect_2./prob_detect_1,0.025),
                                upr = quantile(prob_detect_2./prob_detect_1,0.975)))
end

idxs_P_eff = mean_P_eff .> 0.3
mean_rel_detect = [fit.mean for fit in fits_rel_detect]
fits_rel_detect[idxs_P_eff]

quantile(mean_rel_detect,[0.25,0.5,0.75])
idx_nai = findfirst([fit.name .== "Nairobi" for fit in fits_rel_detect])
idx_mom = findfirst([fit.name .== "Mombasa" for fit in fits_rel_detect])
fits_rel_detect[idx_nai].mean
fits_rel_detect[idx_nai].upr
fits_rel_detect[idx_nai].lwr
fits_rel_detect[idx_mom].mean
fits_rel_detect[idx_mom].upr
fits_rel_detect[idx_mom].lwr

## p_choose1
fits_p_choose1 = []

for filename in readdir("modelfits",join = true)
    modeldict = load(filename)
    model = modeldict[first(keys(modeldict))]
    println("Gathering data from $(model.areaname)")
    params = get(model.MCMC_results.chain,:p_choose1)
    push!(fits_p_choose1,(name = model.areaname,
                                mean = mean(params.p_choose1[:]),
                                lwr = quantile(params.p_choose1[:],0.025),
                                upr = quantile(params.p_choose1[:],0.975)))
end

mean_p_choose1 = [fit.mean for fit in fits_p_choose1]
quantile(mean_p_choose1,[0.25,0.5,0.75])
## p_choose1 bias
fits_p_choose2_bias = []

for filename in readdir("modelfits",join = true)
    modeldict = load(filename)
    model = modeldict[first(keys(modeldict))]
    println("Gathering data from $(model.areaname)")
    params = get(model.MCMC_results.chain,[:p_choose1,:P_eff])
    push!(fits_p_choose2_bias,(name = model.areaname,
                                mean = mean((1 .- params.p_choose1[:])./(1 .- params.P_eff[:])),
                                lwr = quantile((1 .- params.p_choose1[:])./(1 .- params.P_eff[:]),0.025),
                                upr = quantile((1 .- params.p_choose1[:])./(1 .- params.P_eff[:]),0.975)))
end

mean_p_choose2_bias = [fit.mean for fit in fits_p_choose2_bias]
quantile(mean_p_choose2_bias,[0.25,0.5,0.75])

idx_nai = findfirst([fit.name .== "Nairobi" for fit in fits_p_choose2_bias])
idx_mom = findfirst([fit.name .== "Mombasa" for fit in fits_p_choose2_bias])
fits_P_eff[idx_nai].mean
fits_p_choose2_bias[idx_nai].upr
fits_p_choose2_bias[idx_nai].lwr
fits_p_choose2_bias[idx_mom].mean
fits_p_choose2_bias[idx_mom].upr
fits_p_choose2_bias[idx_mom].lwr

## extra_transmissibility
fits_extra_transmissibility = []

for filename in readdir("modelfits",join = true)
    modeldict = load(filename)
    model = modeldict[first(keys(modeldict))]
    println("Gathering data from $(model.areaname)")
    params = get(model.MCMC_results.chain,:extra_transmissibility)
    push!(fits_extra_transmissibility,(name = model.areaname,
                                mean = mean(params.extra_transmissibility[:]),
                                lwr = quantile(params.extra_transmissibility[:],0.025),
                                upr = quantile(params.extra_transmissibility[:],0.975)))
end

mean_extra_transmissibility = [fit.mean - 1 for fit in fits_extra_transmissibility]
quantile(mean_extra_transmissibility,[0.25,0.5,0.75])
idx_nai = findfirst([fit.name .== "Nairobi" for fit in fits_extra_transmissibility])
idx_mom = findfirst([fit.name .== "Mombasa" for fit in fits_extra_transmissibility])
fits_extra_transmissibility[idx_nai].mean
fits_extra_transmissibility[idx_nai].upr
fits_extra_transmissibility[idx_nai].lwr
fits_extra_transmissibility[idx_mom].mean
fits_extra_transmissibility[idx_mom].upr
fits_extra_transmissibility[idx_mom].lwr


text:
ENV["PYCALL_JL_RUNTIME_PYTHON"] = Sys.which("python")
using Distributed
# using 4 cores
if 4 - Distributed.nprocs() > 0
    Distributed.addprocs(4 - Distributed.nprocs());
end
@everywhere using SACBP
using SACBP
using LinearAlgebra
using Plots
pyplot()
using ProgressMeter
using Random
using JLD2
using FileIO

include(joinpath(@__DIR__, "../baseline_implementations/gradient_greedy.jl"))
include(joinpath(@__DIR__, "../baseline_implementations/mcts.jl"))
include(joinpath(@__DIR__, "../baseline_implementations/tlqg.jl"))

function get_initial_augmented_state(;Nt=20) # Number of Targets
    p_init = PhysPos(0.,[12.0,12.0]);
    b_init = [BelMvNormal(0.,zeros(2),300*Matrix(1.0I, 2, 2)) for ii = 1:Nt];
    s_init = AugState(p_init,b_init);
    return s_init
end

function get_target_motion_data(;Nt=20,    # Number of Targets,
                                 dtc=0.01, # Euler Approximation Interval [s]
                                 Q=0.1*Matrix(1.0I, 2, 2), # Process Noise for Each Target
                                 tf=200.0) # Simulation Horizon [s])
    try
        data_path = normpath(joinpath(@__DIR__, "..", "data",
                                      "multi_target_localization",
                                      "target_motion_data_localization.jld2"))
        data = FileIO.load(data_path)
        println("Loaded target motion data from $(data_path)")
        q_history = data["q_history"]
        @assert length(q_history) == Int64(round(tf/dtc,digits=5)) + 1
        @assert length(q_history[1]) == Nt
        @assert data["Q"] == Q "Q in loaded data does not match desired Q"
        return q_history
    catch e
        println("Data file does not exist. Generating taget motion data")
        rng = MersenneTwister(1234)  # RNG for generating Target Motion.

        # Target Motion Generation.
        q_init_g1 = [PhysPos(0.,rand(rng, 2).*10) for ii = 1:Int64(Nt/2)];
        q_init_g2 = [PhysPos(0.,rand(rng, 2).*10 + [20.,20.]) for ii = 1:Int64(Nt/2)];  # Initial target positions.
        q_init = vcat([q_init_g1,q_init_g2]...);
        q_history = [q_init];
        for t = 0.:dtc:tf-dtc
            push!(q_history,[trans(TransModel_Pos(),q,PosControl(t,zeros(2)),dtc,Q,rng) for q in q_history[end]])
        end
        data_path = normpath(joinpath(@__DIR__, "..", "data",
                                      "multi_target_localization",
                                      "target_motion_data_localization.jld2"))
        FileIO.save(data_path, "q_history", q_history, "Q", Q)
        println("Saved target motion data to $(data_path)")
        return q_history
    end
end

function visualizeSimEnv(s::AugState{<:PhysPos,<:VecBelMvNormal{T} where T<:Real},
                         q::Vector{<:PhysPos};plot_belief=true)
    plt = plot();
    plt = scatter!([qi.pos[1] for qi in q],[qi.pos[2] for qi in q],
                   aspect_ratio=1.0,
                   xlabel="x [m]", ylabel="y [m]", label="Targets",
                   color=:cyan,marker=(:star5,8.,1.0,stroke(0.5,:darkblue)),
                   size=(500,500))
    if plot_belief
        for bi in s.b
            plot_e_ellipse!(bi,0.99,plt);
        end
    end
    plt = scatter!((s.p.pos[1],s.p.pos[2]),
                  label="Robot", markershape = :circle, markersize = 7.0,
                  color=:red)
    xlims!((-10.,50.));
    ylims!((-10.,50.));
    return plt
end;

function visualize_results(s_history,q_history)
    plt = plot();
    Nt = length(q_history[1])

    qix_traj = [[q_history[t][ii].pos[1] for t = 1:length(q_history)] for ii = 1:Nt];
    qiy_traj = [[q_history[t][ii].pos[2] for t = 1:length(q_history)] for ii = 1:Nt];
    for ii = 1:Nt
        plt = plot!(qix_traj[ii],qiy_traj[ii],color=:paleturquoise,label="",
                    aspect_ratio=1.0);
        plot_e_ellipse!(s_history[end].b[ii],0.99,plt)
    end

    p_trajectory = hcat(map(s -> s.p.pos,s_history)...);
    plot!(p_trajectory[1,:],p_trajectory[2,:],color=:crimson,
          label="",size=(500,500))
    xlims!((-10.,50.));
    ylims!((-10.,50.));
    return plt,p_trajectory
end

function simulate(s_init::AugState{<:PhysPos,<:VecBelMvNormal},
                  q_history::Vector,
                  controlMode::String, rng_num::Int64,
                  simPosUKF::SimulatePosRangeLocalization2D;
                  numSamples=10, # for sacbp only
                  u_param_min::Vector{Float64}=[-2.0, -2.0],
                  u_param_max::Vector{Float64}=[2.0, 2.0],
                  Th::Float64=2.0, # planning horizon [s]
                  tf::Float64=200.0, # simulation horizon [s]
                  mcts_param::Union{Nothing, MCTSParams}=nothing,
                  tlqg::Union{Nothing, TLQG_Planner_PosRangeLocalization2D}=nothing,
                  animate=true, plotfig=true, verbose=true)
    rng = MersenneTwister(rng_num)
    @assert controlMode in ["greedy", "sacbp", "mcts", "tlqg", "tlqg_offline"];
    if animate
        # Animation Object.
        anim = Animation();
    end
    # Only for MCTS policies.
    if controlMode == "mcts"
        @assert !isnothing(mcts_param) "MCTS Parameters not given"
        amdp = AugMDP(simPosUKF, u_param_min, u_param_max, mcts_param.numActions)
        mctsSolverPlain = MCTS.DPWSolver(depth=Int64(Th/simPosUKF.dto),
                                         n_iterations=mcts_param.numSamples,
                                         exploration_constant=mcts_param.expConst,
                                         k_action=mcts_param.kAct,
                                         alpha_action=mcts_param.αAct,
                                         k_state=mcts_param.kState,
                                         alpha_state=mcts_param.αState,
                                         rng=rng);
        mctsPolicyPlain = solve(mctsSolverPlain,amdp);
    end

    ukfModel = simPosUKF.discreteTransModel;
    transPos = simPosUKF.continuousTransModel;
    dtc, dto, Q = simPosUKF.dtc, simPosUKF.dto, simPosUKF.Q;
    Nt = length(q_history[1])

    # Controller Compilation.
    U_default = [PosControl(round(dtc*t,digits=5),zeros(2)) for t = 0:Int64(Th/dtc)-1];
    if controlMode == "greedy"     # Greedy Heuristic.
        @time U_default = gradientControlUpdate(ukfModel,s_init,U_default,u_param_min,u_param_max,Q,dto)
    elseif controlMode == "sacbp" # SACBP.
        @time U_default = sacControlUpdate(simPosUKF,s_init,U_default,u_param_min,u_param_max,numSamples,rng,offline=true)[1]
    elseif controlMode == "mcts" # MCTS.
        @time U_default = mctsControlUpdate(mctsPolicyPlain,s_init,U_default,dto)
    elseif controlMode == "tlqg"
        @assert !isnothing(tlqg) "T-LQG Planner object not given"
        @time ~,U_default = tlqgControlUpdate(tlqg,s_init,U_default,verbose=verbose,online=true)
    elseif controlMode == "tlqg_offline"
        # offline == planner can take much longer time to compute NLP iterations (until convergence).
        # Only for collecting computation time statistics
        @assert !isnothing(tlqg) "T-LQG Planner object not given"
        @time ~,U_default = tlqgControlUpdate(tlqg,s_init,U_default,verbose=verbose,online=false)
    end
    println("Controller Compiled. Starting Simulation...")

    # Simulation.
    y_history = Vector{Float64}[];
    tcalc_true_history = Float64[];
    s_history = [s_init];
    u_history = PosControl{Float64}[];
    U_pool = copy(U_default);

    observeModel = ObserveModel_Range();
    numSteps = Int64(tf/dtc);
    nextSampleTime = dto;
    @showprogress 1 for t = 1:numSteps
        s = s_history[end];
        if round(t*dtc,digits=5) > nextSampleTime
        # Belief & Control Update.
            # Observation.
            q = q_history[t];
            yVec = Vector{Float64}(undef, Nt);
            for ii = 1:Nt
                yVec[ii] = observe(observeModel,s.p,q[ii].pos,covObservePos(s.p.pos,q[ii].pos),rng);
            end
            push!(y_history,yVec);
            # Belief Update.
            b_new = trans(ukfModel,map(bi -> BelMvNormal(s.p.t,bi.params),s.b),s.p,dto,yVec,Q);
            s_new = AugState(s.p,b_new);
            if verbose
                println((s_new.t,"Belief Updated."));
            end
            if animate
                plt = visualizeSimEnv(s_new,q);
                title!("Time: $(s.t) [s]")
                Plots.frame(anim);
            end
            # Control Update.
            if controlMode == "greedy";     # Greedy Heuristic.
                tcalc_true = @elapsed U_pool = gradientControlUpdate(ukfModel,s_new,U_pool,u_param_min,u_param_max,Q,dto)
                push!(tcalc_true_history,tcalc_true);
                if verbose
                    println((s.t,"Control Updated."));
                end
            elseif controlMode == "sacbp"; # SACBP.
                tcalc_true = @elapsed U_pool,act_time_init,act_time_final,tcalc_simulated = sacControlUpdate(simPosUKF,s_new,U_pool,u_param_min,u_param_max,numSamples,rng,offline=false);
                push!(tcalc_true_history,tcalc_true);
                if verbose
                    println((s.t,"Control Updated.","time_init=$(act_time_init)","time_final=$(act_time_final)"));
                end
            elseif controlMode == "mcts"; # MCTS.
                tcalc_true = @elapsed U_pool = mctsControlUpdate(mctsPolicyPlain,s_new,U_pool,dto);
                push!(tcalc_true_history,tcalc_true);
                if verbose
                    println((s.t,"Control Updated."));
                end
            elseif controlMode == "tlqg" # T-LQG
                tcalc_true, U_pool = tlqgControlUpdate(tlqg,s_new,U_pool,verbose=verbose,online=true)
                if !isnothing(tcalc_true)
                    push!(tcalc_true_history,tcalc_true);
                end
                if verbose
                    println((s.t,"Control Updated."));
                end
            elseif controlMode == "tlqg_offline" # T-LQG, offline
                tcalc_true, U_pool = tlqgControlUpdate(tlqg,s_new,U_pool,verbose=verbose,online=false)
                if !isnothing(tcalc_true)
                    push!(tcalc_true_history,tcalc_true);
                end
                if verbose
                    println((s.t,"Control Updated."));
                end
            end
            # Control Execution.
            u = popfirst!(U_pool);
            push!(u_history,u);
            p_new = trans(transPos,s.p,u,dtc);
            s_new = AugState(p_new,map(bi -> BelMvNormal(p_new.t,bi.params),b_new));
            push!(s_history,s_new);
            #println(s_new.p)

            # Append nominal control.
            u_nominal = PosControl(round(s.t+Th,digits=5),zeros(2));
            push!(U_pool,u_nominal);
            nextSampleTime = round(nextSampleTime+dto,digits=5);
        else
            #q = q_history[t];
            #if t%5 == 0
            #    plt = visualizeSimEnv(s,q);
            #    title!("Time: $(s.t) [s]")
            #    frame(anim);
            #end
            # No Control Update.
            # Control Execution.
            u = popfirst!(U_pool);
            push!(u_history,u);
            p_new = trans(transPos,s.p,u,dtc);
            s_new = AugState(p_new,map(bi -> BelMvNormal(p_new.t,bi.params),s.b));
            push!(s_history,s_new);
            #println(s_new.p);

            # Append nominal control.
            u_nominal = PosControl(round(s.t+Th,digits=5),zeros(2));
            push!(U_pool,u_nominal);
        end
    end
    s = s_history[end];
    # Observation.
    q = q_history[numSteps+1];
    yVec = Vector{Float64}(undef, Nt);
    for ii = 1:Nt
        yVec[ii] = observe(observeModel,s.p,q[ii].pos,covObservePos(s.p.pos,q[ii].pos),rng);
    end
    push!(y_history,yVec);
    # Belief Update.
    b_new = trans(ukfModel,map(bi -> BelMvNormal(s.p.t,bi.params),s.b),s.p,dto,yVec,Q);
    s_new = AugState(s.p,b_new);
    if verbose
        println((s_new.t,"Belief Updated."));
    end
    if animate
        plt = visualizeSimEnv(s_new,q);
        title!("Time: $(s.t) [s]")
        Plots.frame(anim);
    end
    push!(s_history,s_new);
    log_dir = normpath(joinpath(@__DIR__, "../data",
                                 "multi_target_localization"));
    if plotfig
        fig_path = joinpath(log_dir, "tracking_traj_$(controlMode)_tf_$(tf)_rng_$(rng_num).png")
        plt = visualize_results(s_history, q_history);
        if controlMode == "greedy"
            title = "Greedy"
        elseif controlMode == "sacbp"
            title = "SACBP (Ours)"
        elseif controlMode == "mcts"
            title = "MCTS-DPW"
        elseif controlMode == "tlqg"
            title = "T-LQG"
        elseif controlMode == "tlqg_offline"
            title = "T-LQG (Offline)"
        end
        plot!(title=title)
        savefig(fig_path)
    end
    if animate
        gif(anim,joinpath(log_dir, "$(controlMode)_tf_$(tf)_rng_$(rng_num).gif"),fps=50)
    end
    return y_history,tcalc_true_history,s_history,u_history,U_pool
end

function simulate_main(controlMode::String; rng_num::Int64,
                       animate=true, plotfig=true, verbose=true)
    # SAC Controller Parameters.
    dtc = 0.01;         # Euler Approximation Interval.
    dto = 0.2;          # Observation & Control Update Interval.
    Th = 2.0;           # Planning Horizon.
    dtexec = 0.8*dto;   # Control Execution Duration.
    Q = 0.1*Matrix(1.0I, 2, 2);     # Process Noise for Each Target.
    Cu = 0.1*Matrix(1.0I, 2, 2);    # Control Cost Coefficient.
    u_param_min = [-2.0,-2.0];  # Lower Limit for PosControl
    u_param_max = [2.0,2.0];    # Upper Limit for PosControl
    transPos = TransModel_Pos();
    ukfModel = UKFPosRange();
    simPosUKF = SimulatePosRangeLocalization2D(transPos,ukfModel,dtc,dto,dtexec,Q,Cu);
    numSamples = 10;     # Number of Monte Carlo Samples.

    # MCTS Parameters
    numActions = 5;
    # numSamples_mcts = 15;
    numSamples_mcts = 25;
    expConst = 10.0;
    kAct = 10.0;
    αAct = 0.3;
    kState = 3.0;
    αState = 0.1;
    mcts_param = MCTSParams(numActions, numSamples_mcts, expConst, kAct, αAct, kState, αState);

    # Simulation Parameters.
    Nt = 20;     # Number of Targets.
    tf = 200.;    # Simulation Time.

    # T-LQG Planner
    Th_tlqg = Th;
    kl_threshold = 1e6;
    tlqg = TLQG_Planner_PosRangeLocalization2D(Th_tlqg,dto,Nt,Q,Cu,u_param_min,u_param_max,kl_threshold)

    # Initial Augmented State
    s_init = get_initial_augmented_state(Nt=Nt);

    # Target Motion Generation
    q_history = get_target_motion_data(Nt=Nt, dtc=dtc, Q=Q,tf=tf);
    if plotfig
        frame = visualizeSimEnv(s_init, q_history[1], plot_belief=false)
        init_fig_path = normpath(joinpath(@__DIR__, "../data",
                                          "multi_target_localization",
                                          "tracking_initial.pdf"))
        savefig(init_fig_path);
    end

    y_history,tcalc_true_history,s_history,u_history,U_pool =
        simulate(s_init, q_history, controlMode, rng_num,
                 simPosUKF, numSamples=numSamples,
                 u_param_min=u_param_min, u_param_max=u_param_max,
                 Th=Th, tf=tf, mcts_param=mcts_param, tlqg=tlqg,
                 animate=animate, plotfig=plotfig, verbose=verbose);
    println("Average control computation time: $(round(mean(tcalc_true_history), digits=3)) [s]")
    return y_history,tcalc_true_history,s_history,u_history,U_pool
end

if abspath(PROGRAM_FILE) == @__FILE__
    println("1. Simulation with Greedy Controller")
    simulate_main("greedy", rng_num=169, animate=true, plotfig=true, verbose=false);
    println("2. Simulation with SACBP Controller")
    simulate_main("sacbp", rng_num=169, animate=true, plotfig=true, verbose=false);
    println("3. Simulation with MCTS-DPW Controller")
    simulate_main("mcts", rng_num=169, animate=true, plotfig=true, verbose=false);
    println("4. Simulation with T-LQG Controller")
    simulate_main("tlqg", rng_num=169, animate=true, plotfig=true, verbose=false);
end


text:
# Fetch Libcint

try
    @info "Fetching libcint library"
    if Sys.ARCH == :x86_64
        run(`git clone https://github.com/sunqm/qcint.git`)
        mv("qcint", "libcint")
    else
        run(`git clone https://github.com/sunqm/libcint.git`)
    end
    mkdir(joinpath(@__DIR__, "libcint/build"))
    cd(joinpath(@__DIR__, "libcint/build"))
catch
    @error("Failed to clone libcint")
    throw(SystemError)
end

try
    # Compile Libcint
    @info "Building Libcint"

    @info " => Running CMAKE..."
    run(`cmake ..`)
    @info " => Running MAKE..."
    run(`make`)
catch
    rm(joinpath(@__DIR__, "libcint"), recursive=true, force=true)
    @error("Libcint compilation failed!!")
    throw(SystemError)
end

# Fetch binary
@info " => Cleaning up..."
if Sys.islinux()
    run(`cp libcint.so ../../`)
elseif Sys.isapple()
    run(`cp libcint.dylib ../../`)
else
    @error "Could not resolve OS"
end

# Untar basis set library and Clean up
cd(@__DIR__)
run(`tar -zxf lib.gz`)
rm(joinpath(@__DIR__, "libcint"), recursive=true, force=true)


text:
################################################################################
## UNIT TESTING JUST ARITHMETIC operations

include("./unit/unit-addition.jl")
include("./unit/unit-subtraction.jl")
include("./unit/unit-multiplication.jl")
include("./unit/unit-division.jl")

include("./unit/unit-boundstesting.jl")

max_stars_repo_name: REX-Computing/unumjl

max_issues_repo_name: REX-Computing/unumjl

max_forks_repo_name: REX-Computing/unumjl


text:
using NiLang, Test

@testset "filter and mapfoldl" begin
    @i function f(z, y, x)
        i_filter!((@skip! x -> x < 0), y, x)
        i_mapfoldl((@skip! exp), (@skip! PlusEq(identity)), z, y)
    end
    @test f(0.0, Float64[], [-1, -0.5, 3])[1] == exp(-0.5) + exp(-1)
end


text:
# This file is a part of SimilaritySearch.jl

"""
    mutable struct RandomHintsCallback

Indicates that hints are a random sample of the dataset
"""
@with_kw mutable struct RandomHintsCallback <: Callback
    logbase::Float32 = 1.1
end

"""
    callback(opt::RandomHintsCallback, index)

SearchGraph's callback for selecting hints at random
"""
function callback(opt::RandomHintsCallback, index)
    n = length(index)
    m = ceil(Int, log(opt.logbase, length(index)))
    sample = unique(rand(1:n, m))
    empty!(index.hints)
    append!(index.hints, sample)
end

"""
    struct DisjointNeighborhoodHints

Indicates that hints are selected to have a disjoint neighborhood
"""
@with_kw struct DisjointNeighborhoodHints <: Callback
    logbase::Float32 = 1.3
    expansion::Int32 = 2
end

"""
    callback(opt::DisjointNeighborhoodHints, index)

SearchGraph's callback for selecting hints at random
"""
function callback(opt::DisjointNeighborhoodHints, index)
    n = length(index)
    m = ceil(Int, log(opt.logbase, length(index)))
    visited = Set{Int32}()
    empty!(index.hints)
    E = Pair{Int32,Int32}[]
    for i in 1:m
        p = rand(1:n)
        p in visited && continue
        push!(index.hints, p)
        push!(visited, p)
        # visit the neighborhood with some expansion factor
        push!(E, p => 0)
        while length(E) > 0
            parent, e = pop!(E)
            for child in keys(index.links[parent])
                if !(child in visited)
                    push!(visited, child)
                    e + 1 <= opt.expansion && push!(E, child => e + 1)
                end
            end
        end
    end

    # @info "disjoint hints $m --> $(length(index.hints))"
end


text:
# MIT License
#
# Copyright (c) 2020 Joshua E Gentry

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

@inline function ltoh_array(block::AbstractVector{UInt32})::SVector{16, UInt32}
    return SVector{16, UInt32}(
        ltoh(block[1]),  ltoh(block[2]),  ltoh(block[3]),  ltoh(block[4]),
        ltoh(block[5]),  ltoh(block[6]),  ltoh(block[7]),  ltoh(block[8]),
        ltoh(block[9]),  ltoh(block[10]), ltoh(block[11]), ltoh(block[12]),
        ltoh(block[13]), ltoh(block[14]), ltoh(block[15]), ltoh(block[16])
    )
end

@inline function ltoh_array(input::AbstractVector{UInt8}, offset::Int)::SVector{16, UInt32}
    block = reinterpret(UInt32, view(input, offset:(offset+BLOCK_LEN-1)))

    return ltoh_array(block)
end

mutable struct ChunkState
    chaining_value::SVector{8, UInt32}
    chunk_counter::UInt64
    block::Vector{UInt32}
    block_len::Int
    blocks_compressed::Int
    flags::UInt32
end

function ChunkState(
            key::SVector{8, UInt32},
            chunk_counter::UInt64,
            flags::UInt32
        )::ChunkState
    return ChunkState(key, chunk_counter, zeros(UInt32, BLOCK_LEN÷4), 0, 0, flags)
end

function reset(
            self::ChunkState,
            key::SVector{8, UInt32},
            chunk_counter::UInt64,
            flags::UInt32
        )::Nothing
    self.chaining_value    = key
    self.chunk_counter     = chunk_counter
    self.block_len         = 0
    self.blocks_compressed = 0
    self.flags             = flags
    fill!(self.block, UInt8(0))
    nothing
end

@inline function data(
            self::ChunkState
        )::UInt
    return BLOCK_LEN * self.blocks_compressed + self.block_len
end

@inline function start_flag(
            self::ChunkState
        )::UInt32
    if self.blocks_compressed == 0
        return CHUNK_START
    end

    return 0
end

@inline function update_buffer(
            self::ChunkState,
            input::AbstractVector{UInt8},
            offset::Int
        )::Int
    bytes  = min(BLOCK_LEN - self.block_len, length(input) - offset + 1)

    copyto!(reinterpret(UInt8, self.block), self.block_len+1, input, offset, bytes)

    self.block_len += bytes

    return offset + bytes
end

@inline function update_buffer_and_complete(
            self::ChunkState,
            input::AbstractVector{UInt8},
            offset::Int
        )::Tuple{Int, Output}
    if self.block_len < BLOCK_LEN
        offset = update_buffer(self, input, offset)
    end

    return (offset, Output(
        self.chaining_value,
        ltoh_array(self.block),
        self.chunk_counter,
        UInt32(BLOCK_LEN),
        self.flags | start_flag(self) | CHUNK_END
    ))
end

@inline function update_buffer_and_compress(
            self::ChunkState,
            input::AbstractVector{UInt8},
            offset::Int
        )::Int
    if self.block_len < BLOCK_LEN
        offset = update_buffer(self, input, offset)
    end

    self.chaining_value = first_8_words(compress(
        self.chaining_value,
        ltoh_array(self.block),
        self.chunk_counter,
        UInt32(BLOCK_LEN),
        self.flags | start_flag(self)
    ))
    self.blocks_compressed += 1
    fill!(self.block, 0)
    self.block_len = 0

    return offset
end

@inline function update_compress(
            self::ChunkState,
            input::AbstractVector{UInt8},
            offset::Int
        )::Int
    self.chaining_value = first_8_words(compress(
        self.chaining_value,
        ltoh_array(input, offset),
        self.chunk_counter,
        UInt32(BLOCK_LEN),
        self.flags | start_flag(self)
    ))
    self.blocks_compressed += 1

    return offset + BLOCK_LEN
end

@inline function update_complete(
            self::ChunkState,
            input::AbstractVector{UInt8},
            offset::Int
        )::Tuple{Int, Output}
    return (offset + BLOCK_LEN, Output(
        self.chaining_value,
        ltoh_array(input, offset),
        self.chunk_counter,
        UInt32(BLOCK_LEN),
        self.flags | start_flag(self) | CHUNK_END
    ))
end

function update(
            self::ChunkState,
            input::AbstractVector{UInt8},
            offset::Int
        )::Tuple{Int, Union{Output, Nothing}}

    # Not enough to process, copy the bytes to our buffer and return.
    if length(input) - offset < BLOCK_LEN - self.block_len
        return (update_buffer(self, input, offset), nothing)
    end

    # If we have data in our buffer and we're on the last block complete it.
    if self.block_len > 0 && self.blocks_compressed == 15
        return update_buffer_and_complete(self, input, offset)
    end

    # If we have data in our buffer, complete that.
    if self.block_len > 0
        offset = update_buffer_and_compress(self, input, offset)
    end

    # Process while we have MORE than a full block but NOT a full chunk.
    while self.blocks_compressed < 15 && length(input) - offset >= BLOCK_LEN
        offset = update_compress(self, input, offset)
    end

    # If we have MORE data than a full block, then we can complete the current chunk.
    if length(input) - offset >= BLOCK_LEN
        return update_complete(self, input, offset)
    end

    # If there is still data remaining save it.
    if offset <= length(input)
        offset = update_buffer(self, input, offset)
    end

    return (offset, nothing)
end

function complete(self::ChunkState)::Output
    return Output(
        self.chaining_value,
        ltoh_array(self.block),
        self.chunk_counter,
        self.block_len,
        self.flags | start_flag(self) | CHUNK_END
    )
end


text:
"""
    incident_planewave(mesh::Vector{SVector{3,Float64}}; k, expjω=false)

Evaluates the electric field of a x-polarized plane wave traveling along
the z-direction. This field impinging on a PEC sphere give rise to the scattered field
computed by [`mieseries`](@ref).
# Arguments
- `point::SVector{3,Float64}`: point in space where the electric field is evaluated.
- `k`: wavenumber.
- `expjω`: set to `true` (`false`) if `exp(jω)` (`exp(-jω)`) time-harmonic dependence is used.
"""
function incident_planewave(point::SVector{3,Float64}; k, expjω=false)
    _, _, z = point
    if expjω
        return SVector{3,ComplexF64}(exp(-im*k*z), 0, 0)
    else
        return SVector{3,ComplexF64}(exp(im*k*z), 0, 0)
    end
end

"""
    mieseries(mesh::Vector{SVector{3,Float64}}; k, a=1, n_terms=20, expjω=false)
    mieseries(point::SVector{3,Float64}; k, a=1, n_terms=20, expjω=false)

Evaluates the scattered field produced by an incident plane wave ([`incident_planewave`](@ref))
on a PEC sphere. This scattered field is known as Mie series.
Refer to: C. A. Balanis. Advanced Engineering Electromatnetics (2013), Chapter 11.8.
# Arguments
- `mesh::Vector{SVector{3,Float64}}`: list of points in space where the electric field is evaluated.
- `k`: wavenumber.
- `a`: radius of the PEC sphere.
- `n_terms`: the number of terms to compute in the Mie series.
- `expjω`: set to `true` (`false`) if `exp(jω)` (`exp(-jω)`) time-harmonic dependence is used.
- `farfield`: set to `false` to compute near field `E`. Set to `true` to compute far field `E₀`. The near and far field
are related by `E(x) ≈ E₀(x/|x|)e^{ik|x|}/|x|` for `|x| → ∞`, assuming a `exp(-jω)` time-harmonic dependence.
"""
function mieseries(mesh::Vector{SVector{3,Float64}}; k, a=1, n_terms=20, farfield=false, expjω=false)
    isreal(k) || @error "Mie series is not working for complex wavenumbers"
    # compute Ricatti-Hankel functions, second kind
    rh, rh_d, _ = riccatihankel2_and_derivatives(1:n_terms, k*a)
    # compute coefficients aₙ, bₙ, cₙ
    an = [(1.0im)^(-n)*(2n + 1)/n/(n + 1) for n in 1:n_terms]
    bn = -an.*real(rh_d)./rh_d
    cn = -an.*real(rh)./rh
    # compute scattered field
    n_points = length(mesh)
    Es = Vector{SVector{3,ComplexF64}}(undef,n_points)
    mieseries_func = farfield ? _mieseries_farfield : _mieseries
    for i in 1:n_points
        Es[i] = mieseries_func(mesh[i], bn, cn; k, n_terms)
    end
    # conjugate result if exp(-jω) time-harmonic dependence is used
    if !expjω
        Es = conj(Es)
    end
    return Es
end
function mieseries(point::SVector{3,Float64}; k, a=1, n_terms=20, farfield=false, expjω=false)
    Es_list = mieseries([point]; k, a, n_terms, farfield, expjω)
    return first(Es_list)
end

function _mieseries(point::SVector{3,Float64}, bn, cn; k, n_terms, _debug::Union{Val{false}, Val{true}}=Val(false))
    x, y, z = point
    # convert to spherical coordinates
    r = sqrt(x^2 + y^2 + z^2)
    kr = k*r
    cos_theta = z / r
    sin_theta = sqrt(1 - cos_theta^2)
    cos_phi = x / (r*sin_theta)
    sin_phi = y / (r*sin_theta)
    # compute Ricatti-Hankel functions, second kind
    rh, rh_d, rh_dd = riccatihankel2_and_derivatives(1:n_terms, kr)
    # compute associated Legendre functions, m=1, for l=1,..,n_term
    lf, lf_d = MieSeries.associated_legendre(cos_theta; l_max=n_terms, m=1)
    # compute scattered field in spherical coordinates
    # return all the terms of the series if !(_debug === Val(false))
    func = (_debug===Val(false)) ? sum : identity
    Es_r = -im*cos_phi*func(@. bn*(rh_dd + rh)*lf)
    Es_theta = cos_phi/kr*func(@. im*bn*rh_d*sin_theta*lf_d - cn*rh*lf/sin_theta)
    Es_phi = sin_phi/kr*func(@. im*bn*rh_d*lf/sin_theta - cn*rh*sin_theta*lf_d)
    # convert to cartesian coordinates
    Es_x = sin_theta*cos_phi*Es_r + cos_theta*cos_phi*Es_theta - sin_phi*Es_phi
    Es_y = sin_theta*sin_phi*Es_r + cos_theta*sin_phi*Es_theta + cos_phi*Es_phi
    Es_z = cos_theta*Es_r - sin_theta*Es_theta
    return Es_x, Es_y, Es_z
end

function _mieseries_farfield(point::SVector{3,Float64}, bn, cn; k, n_terms)
    x, y, z = point
    # convert to spherical coordinates
    r = sqrt(x^2 + y^2 + z^2)
    cos_theta = z / r
    sin_theta = sqrt(1 - cos_theta^2)
    cos_phi = x / (r*sin_theta)
    sin_phi = y / (r*sin_theta)
    # compute associated Legendre functions, m=1, for l=1,..,n_term
    lf, lf_d = MieSeries.associated_legendre(cos_theta; l_max=n_terms, m=1)
    im_n = im.^(1:n_terms)
    # compute scattered field in spherical coordinates
    Es_theta = im*cos_phi/k*sum(@. im_n*(bn*sin_theta*lf_d - cn*lf/sin_theta))
    Es_phi = im*sin_phi/k*sum(@. im_n*(bn*lf/sin_theta - cn*sin_theta*lf_d))
    # convert to cartesian coordinates
    Es_x = cos_theta*cos_phi*Es_theta - sin_phi*Es_phi
    Es_y = cos_theta*sin_phi*Es_theta + cos_phi*Es_phi
    Es_z = -sin_theta*Es_theta
    return Es_x, Es_y, Es_z
end

"""
    sphere_monostatic_rcs(;k, a, n_terms=20)

Computes the monostatic RCS of a PEC sphere.
# Arguments
- `k`: wavenumber.
- `a`: radius of the PEC sphere.
- `n_terms`: the number of terms to compute in the Mie series.
"""
function sphere_monostatic_rcs(;k, a, n_terms=20)
    λ = 2π/k
    cn = [(-1)^n*(2n+1) for n in 1:n_terms]
    rh, rh_d, _ = riccatihankel2_and_derivatives(1:n_terms, k*a)
    σ_mono = λ^2/(4π)*abs2(sum(@. cn/(rh_d*rh)))
    return σ_mono
end

"""
    mieseries_debug(point::SVector{3,Float64}; k, a=1, n_terms=20, expjω=false)

Same as ([`mieseries`](@ref)), but returns all the terms of the series, for each cartesian component.
"""
function mieseries_debug(point::SVector{3,Float64}; k, a=1, n_terms=20, expjω=false)
    isreal(k) || @error "Mie series is not working for complex wavenumbers"
    # compute Ricatti-Hankel functions, second kind
    rh, rh_d, _ = riccatihankel2_and_derivatives(1:n_terms, k*a)
    # compute coefficients aₙ, bₙ, cₙ
    an = [(1.0im)^(-n)*(2n + 1)/n/(n + 1) for n in 1:n_terms]
    bn = -an.*real(rh_d)./rh_d
    cn = -an.*real(rh)./rh
    # compute scattered field
    Es_x, Es_y, Es_z = _mieseries(point, bn, cn; k, n_terms, _debug=Val(true))
    # conjugate result if exp(-jω) time-harmonic dependence is used
    if !expjω
        Es_x, Es_y, Es_z = conj.((Es_x,Es_y,Es_z))
    end
    return Es_x, Es_y, Es_z
end

max_stars_repo_name: Riarrieta/MieSeries

max_issues_repo_name: Riarrieta/MieSeries

max_forks_repo_name: Riarrieta/MieSeries


text:
"""
    ancestors(node::Node,key,<keyword arguments>)

Get attribute values from the ancestors (basipetal).

# Arguments

## Mandatory arguments

- `node::Node`: The node to start at.
- `key`: The key, or attribute name. Make it a `Symbol` for faster computation time.

## Keyword Arguments

- `scale = nothing`: The scale to filter-in (i.e. to keep). Usually a Tuple-alike of integers.
- `symbol = nothing`: The symbol to filter-in. Usually a Tuple-alike of Strings.
- `link = nothing`: The link with the previous node to filter-in. Usually a Tuple-alike of Char.
- `all::Bool = true`: Return all filtered-in nodes (`true`), or stop at the first node that
is filtered out (`false`).
- `self = false`: is the value for the current node needed ?
- `filter_fun = nothing`: Any filtering function taking a node as input, e.g. [`isleaf`](@ref).
- `recursivity_level = -1`: The maximum number of recursions allowed (considering filters).
*E.g.* to get the parent only: `recursivity_level = 1`, for parent + grand-parent:
`recursivity_level = 2`. If a negative value is provided (the default), the function returns
all valid values from the node to the root.
- `ignore_nothing = false`: filter-out the nodes with `nothing` values for the given `key`
- `type::Union{Union,DataType}`: The type of the attribute. Makes the function run much
faster if provided (≈4x faster).

# Note

In most cases, the `type` argument should be given as a union of `Nothing` and the data type
of the attribute to manage missing or inexistant data, e.g. measurements made at one scale
only. See examples for more details.

# Examples

```julia
# Importing an example mtg from the package:
file = joinpath(dirname(dirname(pathof(MultiScaleTreeGraph))),"test","files","simple_plant.mtg")
mtg = read_mtg(file)

# Using a leaf node from the mtg:
leaf_node = mtg.children[2].children[3].children[4].children[5]

ancestors(leaf_node, :Length) # Short to write, but slower to execute

# Fast version, note that we pass a union of Nothing and Float64 because there are some nodes
# without a `Length` attribute:
ancestors(leaf_node, :Length, type = Union{Nothing,Float64})

# Filter by scale:
ancestors(leaf_node, :XX, scale = 1, type = Float64)
ancestors(leaf_node, :Length, scale = 3, type = Float64)

# Filter by symbol:
ancestors(leaf_node, :Length, symbol = "Internode")
ancestors(leaf_node, :Length, symbol = ("Axis","Internode"))
```
"""
function ancestors(
    node, key;
    scale = nothing,
    symbol = nothing,
    link = nothing,
    all::Bool = true, # like continue in the R package, but actually the opposite
    self = false,
    filter_fun = nothing,
    recursivity_level = -1,
    ignore_nothing = false,
    type::Union{Union,DataType} = Any)

    # Check the filters once, and then compute the ancestors recursively using `ancestors_`
    check_filters(node, scale = scale, symbol = symbol, link = link)

    # Change the filtering function if we also want to remove nodes with nothing values.
    filter_fun_ = filter_fun_nothing(filter_fun, ignore_nothing, key)

    val = Array{type,1}()
    # Put the recursivity level into an array so it is mutable in-place:

    if self
        if is_filtered(node, scale, symbol, link, filter_fun_)
            val_ = unsafe_getindex(node, key)
            push!(val, val_)
        elseif !all
            # We don't keep the value and we have to stop at the first filtered-out value
            return val
        end
    end

    ancestors_(node, key, scale, symbol, link, all, filter_fun, val, recursivity_level)
    return val
end

function ancestors_(node, key, scale, symbol, link, all, filter_fun, val, recursivity_level)

    if !isroot(node) && recursivity_level != 0
        parent = node.parent

        # Is there any filter happening for the current node? (FALSE if filtered out):
        keep = is_filtered(parent, scale, symbol, link, filter_fun)

        if keep
            val_ = unsafe_getindex(parent, key)
            push!(val, val_)
            # Only decrement the recursivity level when the current node is not filtered-out
            recursivity_level -= 1
        end

        # If we want to continue even if the current node is filtered-out
        if all || keep
            ancestors_(parent, key, scale, symbol, link, all, filter_fun, val, recursivity_level)
        end
    end
end


text:
# NOTE: get_unreserved must be defined in this order because otherwise
# the one with keyword argument "exclude" will be called either case

"Get unreserved excluding the active order in specified market (if affects)"
function get_unreserved(inv::AbstractInvestor, asset::AbstractAsset; exclude::AbstractMarket)
    position = get_position(inv, asset)
    total_reserved = get_reserved(inv, asset)
    active = get_active_order(inv, exclude)
    if isnothing(active) || (active.from != asset)
        # Active order does not affect here
        # This include:
        #   active is buy while asset is stock (buy needs only currency to execute)
        #   active is sell while asset is currency (sell needs only stock to execute)
        exclude_quantity = 0
    else
        # Active order is the same setindex
        # of the
        exclude_quantity =  active.quantity
    end
    reserved = total_reserved - exclude_quantity
    return position - reserved

end

function get_unreserved(inv::AbstractInvestor, asset::AbstractAsset)
    return get_position(inv, asset) - get_reserved(inv, asset)
end

function get_reserved(inv::AbstractInvestor, asset::AbstractAsset)
    reserved = get(inv.reserved, asset, 0)
    return reserved
end

function get_position(inv::AbstractInvestor, asset::AbstractAsset)
    return get(inv.positions, asset, 0)
end

function Base.getindex(inv::AbstractInvestor, asset::AbstractAsset)
    get(inv.positions, asset, 0)
end

function Base.setindex!(inv::AbstractInvestor, val::Int64, asset::AbstractAsset)
    inv.positions[asset] = val
end

"Get the current active order in the market (or nothing if none)"
function get_active_order(trader::AbstractInvestor, market::AbstractMarket)
    return get(trader.active_orders, market, nothing)
end


text:
# douglass.jl

# Todo implement for first version:
#
# @gen is probably suffering from type instability
# @egen is still really crap. It should expect a type (e.g. `egen Float64 :x = ...`)

# Nice to have, but not necessary for first version:
# assert <condition>
# collapse (<statistic>) <varlist> ... , by(<varlist>)
# append using <df>

# WARNING:
# Douglass does not sanitize your input. If you run @gen(df, :x, destory_world()),
# that's your responsibility. You've been warned.
#
# ANOTHER WARNING:
# There are some notable differences to how Stata behaves:
# - Variables in a DataFrame are denoted using symbols, e.g. ``:myvariable` as opposed to Stata's `myvariable`
# - Missing values are NOT considered greater than any real number. If you if-condition evaluates to a missing number (in julia),
#   such as for example the condition `missing > 5`, then we will treat that like `false`. In other words, only if the condition
#   evaluates explicitly to `true`, or is automatically converted to `true`, the condition is considered satisfied.
# - There is only one missing value, namely `missing`.
# - If `egen` does not have any nonmissing observations to work with, it returns `missing`, not zero (i.e. what the option `missing` does in Stata).
# - `egen` and `ereplace` operate on vectors of variables (in each group), whereas `gen` and `replace` operate on scalars. That means that in
#   the former, you can use functions that take vectors as arguments, but if you do element-wise operations you have to broadcast these operations
#   (e.g. `bysort groupvar: egen :z = mean(:x)`, or `bysort groupvar: egen :z = :x .+ :y`). In `gen` and `replace`, you can use indexing, e.g. using
#   `bysort groupvar: gen :z = :x[_n] - :x[_n-1]`.
#__precompile__()
module Douglass

    using Reexport
    @reexport using DataFramesMeta
    using DataFrames, DataFramesMeta, Tables
    using MacroTools
    using REPL

    import DataFramesMeta: @with, @where

    # types
    include("Command.jl")

    include("setup.jl")
    include("interface.jl")

    include("parse.jl")

    include("commands/drop.jl")
    include("commands/keep.jl")
    include("commands/sort.jl")
    include("commands/generate.jl")
    include("commands/replace.jl")
    include("commands/rename.jl")
    include("commands/egen.jl")
    include("commands/erep.jl")
    # include("commands/egen2.jl")
    include("commands/merge.jl")
    include("commands/reshape.jl")
    include("commands/duplicates.jl")

    include("repl.jl")

    include("helper.jl")

    global active_df

    export parse_varlist, @use, @d_str, set_active_df

end


text:
using Documenter # A package to manage documentation

# Convert the Jupyter notebooks to Markdown files
include("convert.jl")

# Create documentation
makedocs(
    # Specify a name for the site
    sitename = "GPKernels",

    # Specify the author
    authors = "Pavel Sobolev.",

    # Specify the pages on the left side
    pages = [
        "Home" => "index.md",
        "IDs" => "ids.md",
        "Notebooks" => [
            "$preamble" => pushfirst!(
                map(
                    s -> "$s" => "generated/$notebooks_folder_name/$preamble/$s/$s.md",
                    replace.(
                        basename.(
                            filter(notebook -> count(preamble, notebook) == 1, notebooks)
                        ),
                        ".ipynb" => ""
                    ),
                ),
                "Preamble" => "generated/$notebooks_folder_name/$preamble/$preamble.md",
            )
            for preamble in preambles
        ],
    ],

    # Specify a format
    format = Documenter.HTML(
        # A fallback for creating docs locally
        prettyurls = get(ENV, "CI", nothing) == "true"
    ),

    # Disable doctests
    doctest = false,

    # Fail if any error occurred
    strict = true,
)

# Deploy documentation
deploydocs(
    # Specify a repository
    repo = "github.com/paveloom-c/GPKernels.git",

    # Specify a development branch
    devbranch = "notebooks",
)

max_stars_repo_name: paveloom-c/GPKernels

max_issues_repo_name: paveloom-c/GPKernels

max_forks_repo_name: paveloom-c/GPKernels