text:
# # Compare ADCs
#
# The [apparent diffusion coefficient](@ref adc) (ADC) may be sufficient to describe the
# signal attenuation for small ``b``-values. SpinDoctor comes with multiple approaches for
# computing or estimating the ADC for a `ScalarGradient` ``\vec{g}(t) = f(t) g \vec{d}``:
#
# - Using the free diffusion coefficient ``\vec{d}^\mathsf{T} \mathbf{D} \vec{d}``, which
#   represents unrestricted diffusion in the absence of boundaries;
# - Computing the short diffusion time approximation for the ADC
# - Fitting the log-signal obtained by solving the BTPDE for different ``b``-values
# - Solving a homogenized model (HADC) assuming negligible permeability
# - Using the matrix formalism effective diffusion tensor
#
# In this example we will compare the different approaches.

# ## Building a biological model
#
# We start by loading SpinDoctor and a Makie plotting backend.

# LSP indexing solution                                                          #src
# https://github.com/julia-vscode/julia-vscode/issues/800#issuecomment-650085983 #src
if isdefined(@__MODULE__, :LanguageServer)                                       #src
    include("../src/SpinDoctor.jl")                                              #src
    using .SpinDoctor                                                            #src
end                                                                              #src

using SpinDoctor
using LinearAlgebra

if haskey(ENV, "GITHUB_ACTIONS")
    using CairoMakie
else
    using GLMakie
end

# Here we create a recipe for five stacked plates with isotropic diffusion tensors. They
# should allow for free diffusion in the horizontal direction, but a rather restricted
# vertical diffusion with the permeable membranes.

ncell = 5
setup = PlateSetup(;
    name = "Plates",
    width = 50.0,
    depth = 50.0,
    heights = fill(5.0, ncell),
    bend = 0.0,
    twist = 0.0,
    refinement = 10.0,
)
coeffs = coefficients(
    setup;
    D = [0.002 * I(3) for _ = 1:ncell],
    T₂ = fill(Inf, ncell),
    ρ = fill(1.0, ncell),
    κ = (; interfaces = fill(1e-4, ncell - 1), boundaries = fill(0.0, ncell)),
    γ = 2.67513e-4,
)

# We then proceed to build the geometry and finite element mesh.

mesh, = create_geometry(setup; recreate = true)
plot_mesh(mesh)

# The mesh looks good, so we may then proceed to assemble the biological model and the
# associated finite element matrices.

model = Model(; mesh, coeffs...)
matrices = assemble_matrices(model);

# We may also compute some useful quantities, including a scalar diffusion coefficient from
# the diffusion tensors.

volumes = get_cmpt_volumes(model.mesh)
D_avg = 1 / 3 * tr.(model.D)' * volumes / sum(volumes)
ncompartment = length(model.mesh.points)

# The gradient pulse sequence will be a PGSE with both vertical and horizontal components.
# This allows for both restricted vertical diffusion and almost unrestricted horizontal
# diffusion. The different approaches should hopefully confirm this behaviour.

dir = [1.0, 0.0, 1.0]
profile = PGSE(2500.0, 4000.0)
b = 1000
g = √(b / int_F²(profile)) / model.γ
gradient = ScalarGradient(dir, profile, g)

# ## Computing the ADC using different methods

# ### Short term approximation

# A simple model for the ADC is given by the [short term approximation](@ref STA).

adc_sta_cmpts = compute_adc_sta(model, gradient)
adc_sta = volumes'adc_sta_cmpts / sum(volumes)

# ### Fitting the ADC

# A more robust approach is to directly fit the BTPDE signal to a series of b-values. This
# is however more computationally expensive. We start by building a set of gradients.

bvalues = 0:400:4000
gvalues = map(b -> √(b / int_F²(profile)) / coeffs.γ, bvalues)
gradients = [ScalarGradient(gradient.dir, gradient.profile, g) for g ∈ gvalues]

# The [`solve_multigrad`](@ref) function computes the magnetization for each gradient. Since
# the gradients are interval-wise constant, we can use a specialized solver.

btpde = BTPDE(; model, matrices)
solver = IntervalConstantSolver(; θ = 0.5, timestep = 5.0)
ξ, = solve_multigrad(btpde, gradients, solver)

# The signals can be computed from the magnetization fields.

signals = [compute_signal(matrices.M, ξ) for ξ ∈ ξ]
signals_cmpts = [compute_signal.(matrices.M_cmpts, split_field(model.mesh, ξ)) for ξ ∈ ξ]

# Fitting the ADC is straightforward.

adc_fit = fit_adc(bvalues, signals)
adc_fit_cmpts =
    [fit_adc(bvalues, [s[icmpt] for s ∈ signals_cmpts]) for icmpt = 1:ncompartment]

# ### Homogenized ADC model

# The [HADC](@ref)-model uses homogenization and assumes negligible permeability between the
# compartments. This does require solving an ODE involving all the degrees of freedom.

hadc = HADC(; model, matrices)
adc_homogenized_cmpts = solve(hadc, gradient)
adc_homogenized = volumes'adc_homogenized_cmpts / sum(volumes)

# ### Matrix Formalism

# If a Laplace eigendecomposition is available, the HADC can be approximated with little
# additional computational expense.

laplace = Laplace(; model, matrices, neig_max = 400)
lap_eig = solve(laplace)

# An effective diffusion tensor can be computed.

D_mf = compute_mf_diffusion_tensor(model.mesh, matrices.M, lap_eig, gradient)

# We observe that ``D_{x x}`` and ``D_{y y}`` are very close to the intrinsic diffusion
# coefficient ``D = 0.002``, which confirms that diffusion in the horizontal direction is
# almost unrestricted. ``D_{z z}`` is significantly smaller, confirming the presence of
# membranes along the vertical direction.

# In particular, we may deduce the MF-ADC in our direction.

adc_mf = dir'D_mf * dir / dir'dir

# ## Comparing results

# ### Compartment ADCs

# Here we make a comparison between the compartment ADCs. For MF, only a global ADC was computed.

n = ncompartment
fig = Figure()
ax = Axis(fig[1, 1];
    xticks = (1:4, ["STA", "Fit BTPDE", "HADC", "MF"]),
    ylabel = "ADC / D",
    title = "Compartment ADCs",
)
barplot!(ax, fill(1, n), adc_sta_cmpts ./ D_avg; dodge = 1:n)
barplot!(ax, fill(2, n), adc_fit_cmpts ./ D_avg; dodge = 1:n)
barplot!(ax, fill(3, n), adc_homogenized_cmpts ./ D_avg; dodge = 1:n)
barplot!(ax, [4], [adc_mf / D_avg])
fig

# The STA and HADC are the same for all compartments, as they consider them separately and
# all the compartments have the same size. The fitted ADC is larger for the three inner
# compartments, as they all have permeable membranes both below and above, in contrast to
# the top and bottom compartments that have hard walls.

# ### Signal attenuation

# We may also inspect the resulting signal attenuations.

fig = Figure()
ax = Axis(fig[1,1]; xlabel = "b", yscale = log10, title = "Signal attenuation")
lines!(ax, [0, bvalues[end]], [1, exp(-adc_sta * bvalues[end])]; linestyle = :dash, label = "ADC STA")
lines!(ax, [0, bvalues[end]], [1, exp(-adc_fit * bvalues[end])]; linestyle = :dash, label = "ADC Fit")
lines!(ax, [0, bvalues[end]], [1, exp(-adc_homogenized * bvalues[end])]; linestyle = :dash, label = "HADC")
lines!(ax, [0, bvalues[end]], [1, exp(-adc_mf * bvalues[end])]; linestyle = :dash, label = "ADC MF")
lines!(ax, [0, bvalues[end]], [1, exp(-D_avg * bvalues[end])]; linestyle = :dash, label = "Free diffusion")
scatterlines!(ax, bvalues, abs.(signals) ./ abs(signals[1]); label = "BTPDE Signal")
axislegend(ax)
fig

# We observe the following:
#
# - The exact signal starts to visually deviate from the log-linear regime after ``b =
#   2000``. For higher ``b``-values, the ADC is no longer sufficient to describe the
#   attenuation, as higher order terms can no longer be neglected.
# - The fit-ADC signal coincides with the exact signal for the lowest ``b``-values. This
#   makes sense since that is how this ADC was obtained to begin with. This is considered to
#   be the "reference" ADC.
# - The free diffusion signal attenuates more thant the exact signal by many orders of
#   magnitude, which confirms the presence of restrictive membranes and boundaries in the
#   gradient direction.
# - The HADC signal attenuates less than the exact signal, as it assumes a more severe
#   restriction with impermeable membranes.
# - The MF-ADC signal coincides with the HADC signal, as the former is simply an MF
#   approximation of the latter.


text:
#-----------------------------
# initializing snapshots field
#-----------------------------
function init_snap(nt,nz,nr,nskip)
   #Initializing snapshots
   snapshots_vr,snapshots_vz,nsnap,itvec_snap=init_snapshots_v(nskip,nt,nz,nr)
   snapshots_trr,nsnap,itvec_snap=init_snapshots_t(nskip,nt,nz,nr)
   return nskip,snapshots_trr,snapshots_vr,snapshots_vz,nsnap,itvec_snap
end

function init_snapshots(nt,nz,nr,nskip)
    #---Snapshots (single field variable)
    itvec_snap=zeros(Float64,Int(ceil(nt/nskip))-1)
    tvec_snap=zeros(Float64,Int(ceil(nt/nskip))-1)

    cnt=1
    for itmp=1:Int(ceil(nt/nskip))-1
       cnt+=nskip
       itvec_snap[itmp]=cnt
    end

    if (maximum(itvec_snap)>nt)
       println("Error in snapshots settings!")
       return
    else
       nsnap=length(itvec_snap)
    end

    snapshots_field=zeros(nz,nr,nsnap)
    return snapshots_field,nsnap,itvec_snap
end

function init_snapshots_v(nskip,nt,nz,nr)
    #---Snapshots (velocity)
    #nskip=100
    itvec_snap=zeros(Float64,Int(ceil(nt/nskip))-1)
    tvec_snap=zeros(Float64,Int(ceil(nt/nskip))-1)

    cnt=1
    for itmp=1:Int(ceil(nt/nskip))-1
       cnt+=nskip
       itvec_snap[itmp]=cnt
    end

    if (maximum(itvec_snap)>nt)
       println("Error in snapshots settings!")
       return
    else
       nsnap=length(itvec_snap)
    end

    snapshots_vr=zeros(nz,nr,nsnap)
    snapshots_vz=zeros(nz,nr,nsnap)
    return snapshots_vr,snapshots_vz,nsnap,itvec_snap
end

function init_snapshots_t(nskip,nt,nz,nr)
    #---Snapshots (stress)
    #nskip=100
    itvec_snap=zeros(Float64,Int(ceil(nt/nskip))-1)
    tvec_snap=zeros(Float64,Int(ceil(nt/nskip))-1)

    cnt=1
    for itmp=1:Int(ceil(nt/nskip))-1
       cnt+=nskip
       itvec_snap[itmp]=cnt
    end

    if (maximum(itvec_snap)>nt)
       println("Error in snapshots settings!")
       return
    else
       nsnap=length(itvec_snap)
    end

    snapshots_trr=zeros(nz,nr,nsnap)
    return snapshots_trr,nsnap,itvec_snap
end

#----------------------------
# obtaining snapshots
#----------------------------

function get_snapshots!(snapshots_field,isnap,field,nr,nz)
@inbounds Threads.@threads for ir=1:nr
        for iz=1:nz
            snapshots_field[iz,ir,isnap]=field[iz,ir]
        end
    end
end

#--Snapshots (sparse version)
function get_snapshots_sp!(snapshots_field,isnap,field,nr,nz)

ir_sp=0
@inbounds for ir=1:10:nr
        ir_sp=ir_sp+1
        iz_sp=0
        for iz=1:10:nz
            iz_sp=iz_sp+1
            snapshots_field[iz_sp,ir_sp,isnap]=field[iz,ir]
        end
    end
end

ext: jl


text:
#=
            # P is the precision used for this value
type ArfFloat{P}  <: Real
  exponentOf2::Int # fmpz
  nwords_sign::UInt # mp_size_t
  significand1::UInt # significand_struct
  significand2::UInt
end
=#

precision{P}(x::ArfFloat{P}) = P
precision{P}(::Type{ArfFloat{P}}) = P
precision(::Type{ArfFloat}) = ArbFloatPrecision[1]
setprecision(::Type{ArfFloat}, x::Int) = setprecision(ArbFloat, x)

# a type specific hash function helps the type to 'just work'
const hash_arffloat_lo = (UInt === UInt64) ? 0x37e642589da3416a : 0x5d46a6b4
const hash_0_arffloat_lo = hash(zero(UInt), hash_arffloat_lo)
hash{P}(z::ArfFloat{P}, h::UInt) =
    hash(reinterpret(UInt,z.significand1)$z.exponentOf2,
         (h $ hash(z.significand2$(~reinterpret(UInt,P)), hash_arffloat_lo) $ hash_0_arffloat_lo))

# empty constructor
ArfFloat() = ArfFloat{precision(ArfFloat)}()

weakcopy{P}(x::ArfFloat{P}) = WeakRef(x)

function copy{T<:ArfFloat}(x::T)
    z = T()
    ccall(@libarb(arf_set), Void, (Ptr{T}, Ptr{T}), &z, &x)
    return z
end

function deepcopy{T<:ArfFloat}(x::T)
    z = T()
    ccall(@libarb(arf_set), Void, (Ptr{T}, Ptr{T}), &z, &x)
    return z
end

function zero{T<:ArfFloat}(::Type{T})
   z = T()
   return z
end
#zero{P}(x::ArfFloat{P}) = zero(ArfFloat{P})

function zero{T<:ArfFloat}(x::T)
    P = precision(T)
    return zero( ArfFloat{P} )
end

function one{T<:ArfFloat}(::Type{T})
    z = T()
    ccall(@libarb(arf_set_si), Void, (Ptr{ArfFloat}, Int), &z, one(Int))
    return z
end

function one{T<:ArfFloat}(x::T)
    P = precision(T)
    return one( ArfFloat{P} )
end

function isnan{T<:ArfFloat}(x::T)
    zero(Cint) != ccall(@libarb(arf_is_nan), Cint, (Ptr{T},), &x)
end
function isinf{T<:ArfFloat}(x::T)
    zero(Cint) != ccall(@libarb(arf_is_inf), Cint, (Ptr{T},), &x)
end
function isposinf{T<:ArfFloat}(x::T)
    zero(Cint) != ccall(@libarb(arf_is_pos_inf), Cint, (Ptr{T},), &x)
end
function isneginf{T<:ArfFloat}(x::T)
    zero(Cint) != ccall(@libarb(arf_is_neg_inf), Cint, (Ptr{T},), &x)
end
function iszero{T<:ArfFloat}(x::T)
    zero(Cint) != ccall(@libarb(arf_is_zero), Cint, (Ptr{T},), &x)
end
function isone{T<:ArfFloat}(x::T)
    zero(Cint) != ccall(@libarb(arf_is_one), Cint, (Ptr{T},), &x)
end
function isfinite{T<:ArfFloat}(x::T)
    zero(Cint) != ccall(@libarb(arf_is_finite), Cint, (Ptr{T},), &x)
end
function isnormal{T<:ArfFloat}(x::T)
    zero(Cint) != ccall(@libarb(arf_is_normal), Cint, (Ptr{T},), &x)
end
function isspecial{T<:ArfFloat}(x::T)
    zero(Cint) != ccall(@libarb(arf_is_special), Cint, (Ptr{T},), &x)
end

function notnan{T<:ArfFloat}(x::T)
    zero(Cint) == ccall(@libarb(arf_is_nan), Cint, (Ptr{T},), &x)
end
function notinf{T<:ArfFloat}(x::T)
    zero(Cint) == ccall(@libarb(arf_is_inf), Cint, (Ptr{T},), &x)
end
function notposinf{T<:ArfFloat}(x::T)
    zero(Cint) == ccall(@libarb(arf_is_pos_inf), Cint, (Ptr{T},), &x)
end
function notneginf{T<:ArfFloat}(x::T)
    zero(Cint) == ccall(@libarb(arf_is_neg_inf), Cint, (Ptr{T},), &x)
end
function notzero{T<:ArfFloat}(x::T)
    zero(Cint) == ccall(@libarb(arf_is_zero), Cint, (Ptr{T},), &x)
end
function notone{T<:ArfFloat}(x::T)
    zero(Cint) == ccall(@libarb(arf_is_one), Cint, (Ptr{T},), &x)
end
function notfinite{T<:ArfFloat}(x::T)
    zero(Cint) == ccall(@libarb(arf_is_finite), Cint, (Ptr{T},), &x)
end
function notnormal{T<:ArfFloat}(x::T)
    zero(Cint) == ccall(@libarb(arf_is_normal), Cint, (Ptr{T},), &x)
end
function notspecial{T<:ArfFloat}(x::T)
    zero(Cint) == ccall(@libarb(arf_is_special), Cint, (Ptr{T},), &x)
end

function convert{P}(::Type{BigFloat}, x::ArfFloat{P})
    z = zero(BigFloat)
    r = ccall(@libarb(arf_get_mpfr), Int, (Ptr{BigFloat}, Ptr{ArfFloat{P}}, Cint), &z, &x, 0)
    return z
end

function convert{P}(::Type{ArfFloat{P}}, x::BigFloat)
    z = ArfFloat{P}()
    ccall(@libarb(arf_set_mpfr), Void, (Ptr{ArfFloat{P}}, Ptr{BigFloat}), &z, &x)
    z
end
convert(::Type{ArfFloat}, x::BigFloat) = convert(ArfFloat{precision(ArfFloat)}, x)

convert{T<:ArfFloat}(::Type{T}, x::BigInt) = convert(T, convert(BigFloat, x))
convert{P}(::Type{ArfFloat{P}}, x::BigInt) = convert(ArfFloat{P}, convert(BigFloat, x))

function convert{P}(::Type{ArfFloat{P}}, x::Int64)
    z = ArfFloat{P}()
    ccall(@libarb(arf_set_si), Void, (Ptr{ArfFloat{P}}, Ptr{Int64}), &z, &x)
    z
end
convert(::Type{ArfFloat}, x::Int64) = convert(ArfFloat{precision(ArfFloat)}, x)

function convert{P}(::Type{ArfFloat{P}}, x::Float64)
    z = ArfFloat{P}()
    ccall(@libarb(arf_set_d), Void, (Ptr{ArfFloat{P}}, Ptr{Float64}), &z, &x)
    z
end
convert(::Type{ArfFloat}, x::Float64) = convert(ArfFloat{precision(ArfFloat)}, x)

midpoint{P}(x::ArfFloat{P}) = x

radius{P}(x::ArfFloat{P}) = zero(ArfFloat{P})

#=
#define ARF_RND_DOWN FMPR_RND_DOWN
#define ARF_RND_UP FMPR_RND_UP
#define ARF_RND_FLOOR FMPR_RND_FLOOR
#define ARF_RND_CEIL FMPR_RND_CEIL
#define ARF_RND_NEAR FMPR_RND_NEAR
=#

function round{T<:ArfFloat}(x::T, prec::Int64)
    P = precision(T)
    z = ArfFloat{P}()
    ccall(@libarb(arf_set_round), Int, (Ptr{T}, Ptr{T}, Int64, Int), &z, &x, prec, 2)
    return z
end

#=
function frexp{P}(x::ArfFloat{P})
   significand = ArfFloat{P}()
   exponentOf2 = zero(Int64)
   ccall(@libarb(arf_frexp), Void, (Ptr{ArfFloat{P}}, Int64, Ptr{ArfFloat{P}}), &significand, exponentOf2, &x)
   significand, exponentOf2
end
=#

function min{T<:ArfFloat}(x::T, y::T)
    c = ccall(@libarb(arf_cmp), Cint, (Ptr{T}, Ptr{T}), &x, &y)
    return (c < 0) ? x : y
end

function max{T<:ArfFloat}(x::T, y::T)
    c = ccall(@libarb(arf_cmp), Cint, (Ptr{T}, Ptr{T}), &x, &y)
    return (c > 0) ? x : y
end


text:
export path_info

"""
    path_info(tracker::CoreTracker, x, t₁ = 1.0, t₀ = 0.0)
    path_info(tracker::PathTracker, x)

Track a path using the given `tracker` and start value `x`. This returns a struct containing
detailed informations about the tracked path.

# Example
```julia
julia> @polyvar x y;
julia> f = [x^2 + y^2 + 1, x + y - 3];
julia> tracker, starts = coretracker_startsolutions(f);
julia> path_info(tracker, first(starts))
CTPathInfo:
 • # return code → success
 • # steps (✓/✗) → 11 ( 10 / 1 )
 • # factorizations → 31
 • # ldivs → 66
┌───┬───────┬────────┬───────┬─────────┬─────────┬──────┬─────────┬───────────┬──────┬─────────┐
│   │     s │     Δs │     ω │   |Δx₀| │     acc │    κ │       ψ │ limit_acc │  |x| │     |r| │
├───┼───────┼────────┼───────┼─────────┼─────────┼──────┼─────────┼───────────┼──────┼─────────┤
│ ✓ │     0 │  0.117 │     1 │       0 │ 3.4e-17 │  1.5 │ 2.2e-17 │   3.4e-17 │    1 │ 6.7e-17 │
│ ✓ │ 0.117 │  0.117 │ 0.866 │ 0.00063 │ 2.7e-17 │ 2.31 │ 1.7e-17 │   2.7e-17 │ 1.13 │ 5.6e-17 │
│ ✓ │ 0.235 │    0.2 │ 0.817 │ 0.00086 │ 3.7e-14 │ 2.31 │ 1.7e-17 │   2.7e-17 │ 1.29 │ 5.6e-17 │
│ ✗ │ 0.435 │  0.166 │ 0.437 │   0.021 │ 3.4e-09 │ 2.31 │ 1.7e-17 │   2.7e-17 │ 1.53 │ 5.6e-17 │
│ ✓ │ 0.435 │  0.105 │  1.08 │   0.041 │ 3.4e-09 │ 2.31 │ 1.7e-17 │   2.7e-17 │ 1.53 │ 5.6e-17 │
│ ✓ │ 0.539 │  0.105 │  1.08 │  0.0053 │   1e-15 │ 7.03 │ 6.7e-16 │     1e-15 │ 1.47 │ 2.2e-15 │
│ ✓ │ 0.644 │ 0.0704 │ 0.652 │   0.039 │ 3.4e-08 │ 7.03 │ 6.7e-16 │     1e-15 │ 1.24 │ 2.2e-15 │
│ ✓ │ 0.714 │ 0.0816 │ 0.408 │  0.0057 │ 1.7e-12 │ 7.03 │ 6.7e-16 │     1e-15 │ 1.33 │ 2.2e-15 │
│ ✓ │ 0.796 │ 0.0999 │ 0.357 │  0.0049 │ 3.3e-12 │ 7.03 │ 6.7e-16 │     1e-15 │ 1.59 │ 2.2e-15 │
│ ✓ │ 0.896 │  0.104 │ 0.683 │  0.0032 │ 4.1e-12 │ 7.03 │ 6.7e-16 │     1e-15 │ 1.94 │ 2.2e-15 │
│ ✓ │     1 │      0 │ 0.748 │ 0.00082 │ 1.5e-16 │ 3.86 │ 6.8e-17 │   1.5e-16 │ 2.24 │ 1.3e-15 │
└───┴───────┴────────┴───────┴─────────┴─────────┴──────┴─────────┴───────────┴──────┴─────────┘

julia> tracker = pathtracker(f);

julia> path_info(tracker, first(starts))
PTPathInfo:
 • # return code → success
 • # steps (✓/✗) → 19 ( 18 / 1 )
 • # factorizations → 56
 • # ldivs → 114
┌───┬───────┬───────┬─────────┬──────────┬──────────┬──────┬───────────┬─────────┬───────┬─────────┬──────┐
│   │     s │    Δs │    |ν̇| │    min_ν │    max_ν │    κ │ limit_acc │     acc │     ω │   |Δx₀| │  |x| │
├───┼───────┼───────┼─────────┼──────────┼──────────┼──────┼───────────┼─────────┼───────┼─────────┼──────┤
│ ✓ │     0 │ 0.126 │     NaN │      NaN │      NaN │  1.5 │   9.6e-18 │ 9.6e-18 │     1 │       0 │    1 │
│ ✓ │ 0.126 │ 0.126 │     NaN │      NaN │      NaN │  1.5 │   5.4e-17 │ 5.4e-17 │ 0.874 │ 0.00098 │ 1.13 │
│ ✓ │ 0.251 │ 0.209 │     NaN │      NaN │      NaN │  1.5 │   5.4e-17 │ 5.6e-14 │ 0.829 │ 0.00094 │ 1.29 │
│ ✓ │  0.46 │ 0.234 │     NaN │      NaN │      NaN │  1.5 │   5.4e-17 │   3e-11 │ 0.733 │  0.0049 │ 1.51 │
│ ✓ │ 0.694 │  0.23 │     NaN │      NaN │      NaN │  1.5 │   5.4e-17 │ 2.6e-10 │ 0.767 │   0.008 │ 1.54 │
│ ✓ │ 0.924 │ 0.158 │     NaN │      NaN │      NaN │  1.5 │   5.4e-17 │   6e-08 │  1.16 │   0.028 │ 1.31 │
│ ✓ │  1.08 │ 0.154 │     NaN │      NaN │      NaN │  1.5 │   5.4e-17 │ 6.3e-11 │ 0.598 │  0.0093 │ 1.21 │
│ ✓ │  1.24 │  0.25 │     NaN │      NaN │      NaN │  1.5 │   5.4e-17 │ 8.9e-15 │ 0.395 │  0.0015 │  1.3 │
│ ✓ │  1.48 │  0.38 │     NaN │      NaN │      NaN │  1.5 │   5.4e-17 │ 8.7e-14 │ 0.293 │  0.0022 │ 1.51 │
│ ✗ │  1.87 │ 0.135 │     NaN │      NaN │      NaN │  1.5 │   5.4e-17 │   8e-13 │ 0.686 │  0.0022 │ 1.78 │
│ ✓ │     2 │ 0.314 │    0.39 │    -0.26 │    -0.26 │  1.5 │   5.8e-17 │ 5.8e-17 │ 0.566 │ 8.8e-06 │ 1.85 │
│ ✓ │  2.31 │ 0.314 │    0.23 │    -0.17 │    -0.16 │ 3.64 │   1.2e-16 │ 1.2e-16 │ 0.491 │  0.0001 │ 1.97 │
│ ✓ │  2.63 │  1.24 │    0.15 │    -0.11 │     -0.1 │ 3.64 │   1.2e-16 │ 3.7e-10 │  0.84 │   3e-05 │ 2.05 │
│ ✓ │  3.86 │  1.71 │   0.028 │   -0.025 │   -0.023 │ 3.64 │   1.2e-16 │ 1.4e-12 │ 0.751 │   0.002 │ 2.19 │
│ ✓ │  5.58 │  2.84 │  0.0043 │  -0.0042 │  -0.0037 │ 3.71 │   7.9e-17 │ 7.9e-17 │ 0.275 │  0.0017 │ 2.23 │
│ ✓ │  8.41 │  5.13 │ 0.00024 │ -0.00024 │ -0.00021 │ 3.71 │   7.9e-17 │ 2.3e-14 │ 0.244 │  0.0012 │ 2.24 │
│ ✓ │  13.5 │  13.3 │ 1.4e-06 │ -1.4e-06 │ -1.3e-06 │ 3.71 │   7.9e-17 │ 1.1e-08 │ 0.245 │  0.0003 │ 2.24 │
│ ✓ │  26.8 │    55 │ 4.3e-10 │ -2.5e-12 │ -2.2e-12 │ 3.67 │   1.7e-16 │ 1.7e-16 │ 0.275 │ 9.8e-06 │ 2.24 │
│ ✓ │  81.8 │     0 │ 4.3e-10 │ -2.5e-12 │ -2.2e-12 │ 3.67 │   1.2e-16 │ 1.2e-16 │ 0.275 │ 1.8e-09 │ 2.24 │
└───┴───────┴───────┴─────────┴──────────┴──────────┴──────┴───────────┴─────────┴───────┴─────────┴──────┘
```
"""
function path_info end

struct CTPathInfo
    # per step info
    s::Vector{Float64}
    Δs::Vector{Float64}
    ω::Vector{Float64}
    Δx₀::Vector{Float64}
    accepted_rejected::Vector{Bool}
    norm_x::Vector{Float64}
    cond::Vector{Float64}
    accuracy::Vector{Float64}
    residual::Vector{Float64}
    limit_accuracy::Vector{Float64}
    eval_err::Vector{Float64}
    # total info
    return_code::CoreTrackerStatus.states
    n_factorizations::Int
    n_ldivs::Int
end

function path_info(tracker::CoreTracker, x₀, t₁ = 1.0, t₀ = 0.0)
    state = tracker.state

    s = Float64[]
    Δs = Float64[]
    ω = Float64[]
    Δx₀ = Float64[]
    accepted_rejected = Bool[]
    norm_x = Float64[]
    cond = Float64[]
    accuracy = Float64[]
    limit_accuracy = Float64[]
    residual = Float64[]
    eval_err = Float64[]

    init!(tracker, x₀, t₁, t₀)
    push!(s, state.s)
    push!(Δs, state.Δs)
    push!(ω, state.ω)
    push!(Δx₀, state.norm_Δx₀)
    push!(norm_x, maximum(abs, state.x))
    push!(cond, state.jacobian.cond[])
    push!(accuracy, state.accuracy)
    push!(limit_accuracy, state.limit_accuracy)
    push!(residual, maximum(tracker.corrector.abs_r))
    push!(eval_err, state.eval_err)

    first = true
    for _ in tracker
        push!(accepted_rejected, !state.last_step_failed)
        push!(s, state.s)
        push!(Δs, state.Δs)
        push!(ω, state.ω)
        push!(Δx₀, state.norm_Δx₀)
        push!(norm_x, maximum(abs, state.x))
        push!(cond, state.jacobian.cond[])
        push!(accuracy, state.accuracy)
        push!(limit_accuracy, state.limit_accuracy)
        push!(residual, maximum(tracker.corrector.abs_r))
        push!(eval_err, state.eval_err)
        first = false
    end
    push!(accepted_rejected, !state.last_step_failed)

    CTPathInfo(
        s,
        Δs,
        ω,
        Δx₀,
        accepted_rejected,
        norm_x,
        cond,
        accuracy,
        residual,
        limit_accuracy,
        eval_err,
        status(tracker),
        state.jacobian.factorizations[],
        state.jacobian.ldivs[],
    )
end

path_table(info::CTPathInfo) = path_table(stdout, info)
function path_table(io::IO, info::CTPathInfo)
    header = ["", "s", "Δs", "ω", "|Δx₀|", "acc", "κ", "ψ", "limit_acc", "|x|", "|r|"]
    h1 = PrettyTables.Highlighter(
        f = (data, i, j) -> j == 1 && data[i, 1] == :✗,
        crayon = PrettyTables.crayon"red",
    )
    h2 = PrettyTables.Highlighter(
        f = (data, i, j) -> j == 1 && data[i, 1] == :✓,
        crayon = PrettyTables.crayon"green",
    )
    ✓✗ = map(v -> v ? :✓ : :✗, info.accepted_rejected)
    data = hcat(
        ✓✗,
        _sigdigits3.(info.s),
        _sigdigits3.(info.Δs),
        _sigdigits3.(info.ω),
        _sigdigits2.(info.Δx₀),
        _sigdigits2.(info.accuracy),
        _sigdigits3.(info.cond),
        _sigdigits2.(info.eval_err),
        _sigdigits2.(info.limit_accuracy),
        _sigdigits3.(info.norm_x),
        _sigdigits2.(info.residual),
    )
    PrettyTables.pretty_table(io, data, header, crop = :none, highlighters = (h1, h2))
end

function Base.show(io::IO, info::CTPathInfo)
    println(io, "CTPathInfo:")
    println(io, " • # return code → ", info.return_code)
    println(
        io,
        " • # steps (✓/✗) → ",
        length(info.s),
        " ( ",
        count(info.accepted_rejected),
        " / ",
        count(!, info.accepted_rejected),
        " )",
    )
    println(io, " • # factorizations → ", info.n_factorizations)
    println(io, " • # ldivs → ", info.n_ldivs)
    path_table(io, info)
end

struct PTPathInfo
    # per step info
    s::Vector{Float64}
    Δs::Vector{Float64}
    ω::Vector{Float64}
    Δx₀::Vector{Float64}
    accepted_rejected::Vector{Bool}
    norm_x::Vector{Float64}
    cond::Vector{Float64}
    accuracy::Vector{Float64}
    limit_accuracy::Vector{Float64}
    norm_ν̇::Vector{Float64}
    min_ν::Vector{Float64}
    max_ν::Vector{Float64}
    # total info
    return_code::PathTrackerStatus.states
    n_factorizations::Int
    n_ldivs::Int
end

function path_info(tracker::PathTracker, x₀, t₀ = nothing)
    ct_state = tracker.core_tracker.state

    s = Float64[]
    Δs = Float64[]
    ω = Float64[]
    Δx₀ = Float64[]
    accepted_rejected = Bool[]
    norm_x = Float64[]
    cond = Float64[]
    accuracy = Float64[]
    limit_accuracy = Float64[]
    norm_ν̇ = Float64[]
    min_ν = Float64[]
    max_ν = Float64[]

    if t₀ !== nothing
        init!(tracker, x₀, 0.0, t₀)
    else
        init!(tracker, x₀)
    end

    push!(s, real(current_t(ct_state)))
    push!(Δs, real(current_Δt(ct_state)))
    push!(ω, ct_state.ω)
    push!(Δx₀, ct_state.norm_Δx₀)
    push!(norm_x, maximum(abs, ct_state.x))
    push!(cond, ct_state.jacobian.cond[])
    push!(accuracy, ct_state.accuracy)
    push!(limit_accuracy, ct_state.limit_accuracy)
    push!(norm_ν̇, maximum(abs, tracker.state.valuation.ν̇))
    push!(max_ν, maximum(tracker.state.valuation.ν))
    push!(min_ν, minimum(tracker.state.valuation.ν))

    first = true
    for _ in tracker
        push!(accepted_rejected, !ct_state.last_step_failed)
        push!(s, real(current_t(ct_state)))
        push!(Δs, real(current_Δt(ct_state)))
        push!(ω, ct_state.ω)
        push!(Δx₀, ct_state.norm_Δx₀)
        push!(norm_x, maximum(abs, ct_state.x))
        push!(cond, ct_state.jacobian.cond[])
        push!(accuracy, ct_state.accuracy)
        push!(limit_accuracy, ct_state.limit_accuracy)
        push!(norm_ν̇, maximum(abs, tracker.state.valuation.ν̇))
        push!(max_ν, maximum(tracker.state.valuation.ν))
        push!(min_ν, minimum(tracker.state.valuation.ν))
        first = false
    end
    push!(accepted_rejected, !ct_state.last_step_failed)

    PTPathInfo(
        s,
        Δs,
        ω,
        Δx₀,
        accepted_rejected,
        norm_x,
        cond,
        accuracy,
        limit_accuracy,
        norm_ν̇,
        min_ν,
        max_ν,
        status(tracker),
        ct_state.jacobian.factorizations[],
        ct_state.jacobian.ldivs[],
    )
end

path_table(info::PTPathInfo) = path_table(stdout, info)
function path_table(io::IO, info::PTPathInfo)
    header = [
        "",
        "s",
        "Δs",
        "|ν̇|",
        "min_ν",
        "max_ν",
        "κ",
        "limit_acc",
        "acc",
        "ω",
        "|Δx₀|",
        "|x|",
    ]
    h1 = PrettyTables.Highlighter(
        f = (data, i, j) -> j == 1 && data[i, 1] == :✗,
        crayon = PrettyTables.crayon"red",
    )
    h2 = PrettyTables.Highlighter(
        f = (data, i, j) -> j == 1 && data[i, 1] == :✓,
        crayon = PrettyTables.crayon"green",
    )
    ✓✗ = map(v -> v ? :✓ : :✗, info.accepted_rejected)
    data = hcat(
        ✓✗,
        _sigdigits3.(info.s),
        _sigdigits3.(info.Δs),
        _sigdigits2.(info.norm_ν̇),
        _sigdigits2.(info.min_ν),
        _sigdigits2.(info.max_ν),
        _sigdigits3.(info.cond),
        _sigdigits2.(info.limit_accuracy),
        _sigdigits2.(info.accuracy),
        _sigdigits3.(info.ω),
        _sigdigits2.(info.Δx₀),
        _sigdigits3.(info.norm_x),
    )
    PrettyTables.pretty_table(io, data, header, crop = :none, highlighters = (h1, h2))
end

function Base.show(io::IO, info::PTPathInfo)
    println(io, "PTPathInfo:")
    println(io, " • # return code → ", info.return_code)
    println(
        io,
        " • # steps (✓/✗) → ",
        length(info.s),
        " ( ",
        count(info.accepted_rejected),
        " / ",
        count(!, info.accepted_rejected),
        " )",
    )
    println(io, " • # factorizations → ", info.n_factorizations)
    println(io, " • # ldivs → ", info.n_ldivs)
    path_table(io, info)
end

_sigdigits3(x) = Printf.@sprintf("%.3g", x)
_sigdigits2(x) = Printf.@sprintf("%.2g", x)


text:
"""
    AmultB2D!(A, B, Asize=size(A))

Multiply contents of A in place with contents of B. Both A and B should be 2D
arrays and be the same size.

Required Arguments:

- `A::Array{T,2}`: 2D dimmensional array, `A`, with element type `T`
- `B::Array{T,2}`: 2D dimmensional array, `B`, with element type `T`

Optional Arguments:

- `Asize::Tuple(Int,Int)`: size of array `A`, assuming that `B` is the same size
                           size as `A` `(default = size(A))`

Modifies in Place and Returns:

- `A::Array{T,2}`
"""
function AmultB2D!(A, B, Asize=size(A))
    @inbounds for i in 1:Asize[1]
        @inbounds for j in 1:Asize[2]
            A[i,j] = A[i,j] * B[i,j]
        end
    end
    return A
end

"""
    AmultB1D!(A::Vector, B::Vector, Asize=size(A)[1])

Multiply contents of A in place with contents of B. Both A and B should be 1D
arrays and be the same size.

Required Arguments:

- `A::Array{T,1}`: 1D dimmensional array, `A`, with element type `T`
- `B::Array{T,1}`: 1D dimmensional array, `B`, with element type `T`

Optional Arguments:

- `Asize::Int`: size of array `A`, assuming that `B` is the same size
                size as `A` `(default = size(A)[1])`

Modifies in Place and Returns:

- `A::Array{T,1}`
"""
function AmultB1D!(A::Vector, B::Vector, Asize=size(A)[1])
    @threads for i in 1:Asize
        @inbounds A[i] = A[i] * B[i]
    end
    return A
end

"""
    conjAmultB1D!(A, B, Asize=size(A))

Multiply contents of conj(A) in place with contents of B. Both A and B should
be 1D arrays and be the same size.

Required Arguments:

- `A::Array{Complex{T},1}`: 1D dimmensional array, `A`, with element type `T`
- `B::Array{Complex{T},1}`: 1D dimmensional array, `B`, with element type `T`

Optional Arguments:

- `Asize::Int`: size of array `A`, assuming that `B` is the same size
                size as `A` `(default = size(A)[1])`

Modifies in Place and Returns:

- `A::Array{Complex{T},1}`
"""
function conjAmultB1D!(A::Vector, B::Vector, Asize=size(A)[1])
    @threads for i in 1:Asize
        @inbounds A[i] = conj(A[i]) * B[i]
    end
    return A
end

"""
    conjA!(A, Asize=size(A))

Takes the conjugate of A in place. A should be a 1D array.

Required Arguments:

- `A::Array{Complex{T},1}`: 1D dimmensional array, `A`, with element type `T`

Optional Arguments:

- `Asize::Int`: size of array `A` `(default = size(A)[1])`

Modifies in Place and Returns:

- `A::Array{Complex{T},1}`
"""
function conjA!(A::Vector, Asize=size(A)[1])
    @threads for i in 1:Asize
        @inbounds A[i] = conj(A[i])
    end
    return A
end

"""
    fft_correlate(data, reference)

Calculate the cyclical FFT based correlation between the data and the reference
signal.

Required Arguments:

- `data::Array{T,N}`: original `N` dimmensional data signal of element
                      type `T`
- `reference::Array{T,N}`: `N` dimmensional reference signal of element
                           type `T`

Returns:

- `Array{Complex{Float64},N}`: `N` dimmensional array containing the correlation
                               result
"""
function fft_correlate(data, reference)
    return ifft(conj!(fft(reference)).*fft(data))
end

"""
    gnsstypes

Dictionary containing the qyuivalent strings for each type used in `GNSSTools`.
"""
const gnsstypes = Dict(Val{:l5q}() => "l5q",
                       Val{:l5i}() => "l5i",
                       Val{:l1ca}() => "l1ca",
                       Val{:fft}() => "fft",
                       Val{:carrier}() => "carrier",
					   Val{:sc4}() => "sc4",
					   Val{:sc8}() => "sc8",
					   Val{:sc16}() => "sc16",
                       Val{:sc32}() => "sc32",
					   Val{:sc64}() => "sc64")

"""
	calcinitcodephase(code_length::Int, f_code_d::Number, f_code_dd::Number,
					  f_s::Number, code_start_idx::Number)

Calculates the initial code phase of a given code where f_d and fd_rate are the
Doppler affected code frequency and code frequency rate, respectively. The
returned value is the fractional index location in the original code vector in
`CodeType.codes[i]`, where `i` is the current code being processed.

Required Arguments:

**NOTE:** Arguments can be either `Float64` or `Int`.

- `code_length::Int`: number of bits in code
- `f_code_d::Number`: Doppler adjusted code chipping rate in Hz
- `f_code_dd::Number`: Doppler rate adjusted code chipping rate rate in Hz
- `f_s::Number`: sampling rate of data in Hz
- `code_start_idx::Number`: index in `ReplicaSignal.data` where all layers of
                            code start

Returns:

- `Float64`: Fractional code index location
"""
function calcinitcodephase(code_length::Int, f_code_d::Number,
	                       f_code_dd::Number, f_s::Number,
						   code_start_idx::Number)
    t₀ = (code_start_idx-1)/f_s
    init_phase = -f_code_d*t₀ - 0.5*f_code_dd*t₀^2
    return (init_phase%code_length + code_length)%code_length
end

"""
	calccodeidx(init_chip::Number, f_code_d::Number, f_code_dd::Number,
				t::Number, code_length::Int)

Calculates the index in the codes for a given t.

Required Arguments:

- `init_chip::Number`: initial code phase at `t=0s` returned from
                       `calcinitcodephase`
- `f_code_d::Number`: Doppler adjusted code chipping rate in Hz
- `f_code_dd::Number`: Doppler rate adjusted code chipping rate rate in Hz
- `t::Number`: current time elapsed in seconds
- `code_length::Int`: number of bits in code

Returns:

- `Int`: the current code index at a given `t`
"""
function calccodeidx(init_chip::Number, f_code_d::Number, f_code_dd::Number,
	                 t::Number, code_length::Int)
    return Int(floor(init_chip+f_code_d*t+0.5*f_code_dd*t^2)%code_length)+1
end

"""
	calc_doppler_code_rate(f_code::Number, f_carrier::Number, f_d::Number,
						   fd_rate::Number)

Calculates the adjusted code chipping rate and chipping rate rate based off the
Doppler and Doppler rate.

Required Arguments:

- `f_code::Number`: chipping rate of code with no Doppler and Doppler rate in Hz
- `f_carrier::Number`: carrier frequency of signal in Hz
- `f_d::Number`: Doppler frequency in Hz
- `fd_rate::Number`: Doppler frequency rate in Hz

Returns:

- `(f_code_d::Float64, f_code_dd::Float64)::Tuple` where
	* `f_code_d`: Doppler adjusted code chipping rate in Hz
	* `f_code_dd`: Doppler rate adjusted code chipping rate rate in Hz
"""
function calc_doppler_code_rate(f_code::Number, f_carrier::Number, f_d::Number,
	                            fd_rate::Number)
	f_code_d = f_code*(1. + f_d/f_carrier)
	f_code_dd = f_code*fd_rate/f_carrier
	return (f_code_d, f_code_dd)
end

"""
    calctvector(N::Int, f_s::Number)

Generates an `N` long time vector with time spacing `Δt` or `1/f_s`.

Required Arguments:

- `N::Int`: length of the time vector
- `f_s::Number`: sampling rate of signal in Hz

Returns:

- `t::Array{Float64,1}`: time vector of length `N` and spacing of `1/f_s`
                         seconds
"""
function calctvector(N::Int, f_s::Number)
    # Generate time vector
    t = Array{Float64}(undef, N)
    @threads for i in 1:N
        @inbounds t[i] = (i-1)/f_s
    end
    return t
end

"""
	calctvector(signal::GNSSSignal)

Generates an `N` long time vector with time spacing `Δt` or `1/f_s`.

Required Arguments:

- `signal::GNSSSignal`: the signal structure

Returns:

- `t::Array{Float64,1}`: time vector of length `N` and spacing of `1/f_s`
                         seconds
"""
function calctvector(signal::GNSSSignal)
	return calctvector(signal.sample_num, signal.f_s)
end

"""
	calc_t_at_i(i::Int, start_t::Number, f_s::Number)

Calculates the time, `t`, at index `i`.

Required Arguments:

- `i`: vector index integer spanning range 1:n
- `start_t`: the value of `t` at `i=1` in seconds
- `f_s`: the sampling rate of the data in Hz

Returns:

- `t`: time in seconds at `i`
"""
function calc_t_at_i(i::Int, start_t::Number, f_s::Number)
	return start_t + (i-1)/f_s
end

"""
    meshgrid(x::Vector, y::Vector)

Generate a meshgrid the way Numpy would in Python.

Required Arguments:

- `x::Vector{T}`, 1D array of element type `T`
- `y::Vector{T}`, 1D array of element type `T`

Returns:

- `(X, Y)::Tuple` where
	* `X::Array{eltype(x),2}`: 2D array of size `(length(y), length(x))`
	* `Y::Array{eltype(y),2}`: 2D array of size `(length(y), length(x))`
"""
function meshgrid(x::Vector, y::Vector)
    xsize = length(x)
    ysize = length(y)
    X = Array{eltype(x)}(undef, ysize, xsize)
    Y = Array{eltype(y)}(undef, ysize, xsize)
    for i in 1:ysize
        for j in 1:xsize
            X[i,j] = x[j]
            Y[i,j] = y[i]
        end
    end
    return (X, Y)
end

"""
	find_sequence(file_name, digit_nums, separators=missing)

Find a number sequence in a file name. `digit_nums` can be a number or array.
`separators` can also be either a number or array containing the characters such
as `_` and `.` that may be between the numbers. Returns the string containing
the sequence if found. Returns `missing` if not found.

Required Arguments:

- `file_name`: name of file
- `digit_nums`: can be an array or single number indicating number of digits in
                `file_name` to find
    * e.g. to find `1970-01-01`, user would pass `digit_nums` as `[4, 2, 2]`

Optional Arguments:

- `separators`: the characer used to separate the number sequence
	* e.g. to find `1970-01-01`, user would pass `separators` as `"-"`
	* `(default = missing)`

Returns:

- `String`: if sequence is found, a copy of the sequence in `file_name` is
            returned
    * returns `missing` if sequence is not found
"""
function find_sequence(file_name, digit_nums, separators=missing)
	total_digits = sum(digit_nums)
	section_num = length(digit_nums)
	current_section = 1
	sequence_found = false
	sequence_complete = false
	between_sections = false
    sequence_counter = 0
	sequence_start = 1
    sequence_stop = 1
	for i in 1:length(file_name)
        if isdigit(file_name[i])
            if sequence_found == false
                sequence_start = i
            end
            sequence_found = true
            sequence_counter += 1
			if between_sections
				current_section += 1
				between_sections = false
			end
        else
            if sequence_found && (sequence_counter == sum(digit_nums[1:current_section])) && occursin(file_name[i], separators)
				between_sections = true
            else
                sequence_found = false
                sequence_counter = 0
                sequence_idx_start = 1
                sequence_idx_stop = 1
				current_section = 1
				between_sections = false
            end
        end
		if (sequence_counter == total_digits) && (current_section == section_num)
			sequence_stop = i
			sequence_complete = true
			break
		end
    end
	if sequence_complete
		return file_name[sequence_start:sequence_stop]
	else
		return missing
	end
end

"""
    find_and_get_timestamp(file_name)

Find a sequency of 8 digits with `_` separating it from a sequence of 6 digits.
Return the timestamp tuple. Expected sequence to find is `YYYYMMDD_HHMMSS`.

Required Arguments:

- `file_name`: name of data file

Returns:

- `timestamp`: `Tuple` containing date and time
	* format is `(year, month, day, hour, minute, second)`
- `timestamp_JD`: Julian date of `timestamp`
"""
function find_and_get_timestamp(file_name)
	timestamp_string = find_sequence(file_name, [8, 6], "_")
    if ~ismissing(timestamp_string)
        year = parse(Int, timestamp_string[1:4])
    	month = parse(Int, timestamp_string[5:6])
    	day = parse(Int, timestamp_string[7:8])
    	hour = parse(Int, timestamp_string[10:11])
    	minute = parse(Int, timestamp_string[12:13])
    	second = parse(Int, timestamp_string[14:15])
    	timestamp = (year, month, day, hour, minute, second)
    	timestamp_JD = DatetoJD(timestamp...)
        return (timestamp, timestamp_JD)
    else
        error("Data timestamp not found in file name. Please supply it manually.")
    end
end

"""
	random_databits(chipping_rate, t_length; prns=missing)

Generates single or multiple random bit streams. Multiple bit streams are stored
in a dictionary wit the keys being the contents of `prns`.

Required Arguments:

- `chipping_rate`: databit chipping rate in Hz
- `t_length`: length of databit sequence in seconds

Optioonal Arguments:

- `prns`: vector containing PRN numbers `(default = missing)`

Returns:

- Vector of databits or dictionary of `length(prns)` databit sequences
"""
function random_databits(chipping_rate, t_length; prns=missing)
	if ismissing(prns)
		return rand(0:1, ceil(Int, chipping_rate*t_length))
	else
		databits = Dict{eltype(prns),Vector{Int}}()
		for prn in prns
			databits[prn] = rand(0:1, ceil(Int, chipping_rate*t_length))
		end
		return databits
	end
end

"""
	copy_dictionary(dict, key_list)

Copies a dictionary keys and values with keys defined as `b`.

Required Arguments:

- `dict`: dictionary to copy
- `key_list`: list of keys to copy

Return:

- dictionary containing partial to full copy of key_list
"""
function copy_dictionary(dict, key_list)
	dict_copy = Dict{eltype(keys(dict)),Vector{eltype(dict[key_list[1]])}}()
	for key in key_list
		dict_copy[key] = dict[key]
	end
	return dict_copy
end

"""
    get_signal_type(file_name)

Find the signal type of the data based off its file name. Only checks if signal
type is L1 or L5 signal. L2 is not supported.

Required Arguments:

- `file_name`: name of data file

Returns:

- `f_s`: sampling rate in Hz
- `f_if`: IF frequency in Hz
- `f_center`: receiver center frequency in Hz
- `sig_freq`: signal carrier frequency in Hz
- `sigtype`: either `Val(:l1ca)` or `Val(:l5q)`
"""
function get_signal_type(file_name)
    # Determine sampling and IF frequency and frequency center
	if occursin("g1b1", file_name)
		f_s, f_if, f_center, sig_freq, sigtype = g1b1()
	elseif occursin("g2r2", file_name)
		error("L2 band signals not supported. Use either L1 (g1b1) or L5 (g5) files instead.")
	elseif occursin("g5", file_name)
		f_s, f_if, f_center, sig_freq, sigtype = g5()
	else
		# Must check for information on center frequency and sampling rate
		# in file name. If not present, user must specify manually.
		data_freq_string =  find_sequence(file_name, [4,2,2], "._M")
		if ismissing(data_freq_string)
			data_freq_string =  find_sequence(file_name, [4,2,1], "._M")
		end
		if ismissing(data_freq_string)
			error("Cannot determine f_s, f_if, & f_center. Manually specify f_s and f_if.")
		else
			data_freq_string = replace(data_freq_string, "M"=>"")
			data_freq_string = split(data_freq_string, "_")
			f_center = parse(Float64, data_freq_string[1]) * 1e6  # Hz
			f_s = parse(Float64, data_freq_string[2]) * 1e6  # Hz
			# Determine the smallest difference between `f_center` and
			# all GPS carrier frequencies. The closest difference determines
			# what the signal type is.
			Δf = abs.(f_center .- [L1_freq, L2_freq, L5_freq])
			idx = argmin(Δf)
			if idx == 1
				sigtype = Val(:l1ca)
				sig_freq = L1_freq
			elseif idx == 3
				sigtype = Val(:l5q)
				sig_freq = L5_freq
			else
				error("L2 signals are not supported. Supported signals are L1 and L5.")
			end
			f_if = abs(sig_freq-f_center)
		end
	end
	return (f_s, f_if, f_center, sig_freq, sigtype)
end

"""
	make_subplot(fig, row, col, i; projection3d=false, aspect="auto")

Makes a subplot using the PyPlot module. Provides easier setup for 3D subplots.

Required Arguments:

- `fig`: PyPlot figure object
- `row::Int`:specifies the number of rows in the figure
- `col::Int`:specifies the number of columns in the figure
- `i::Int`: specifies the subplot placement in the figure `(range = (1,row*col))`

Optional Arguments:

- `projection3d`: set to `true` if generating a 3D plot `(default = `false`)`
- `aspect`: subplot aspect ratio; refer to PyPlot `(default = "auto")`
	* does not work if `projection3d` is set to `true`

Returns:

- `fig`: PyPlot figure object
- `ax`: PyPlot subplot axis object
"""
function make_subplot(fig, row, col, i; projection3d=false, aspect="auto")
	if projection3d
		# Must be called in local scope in order to make 3D subplots
		mplot3d = PyPlot.PyObject(PyPlot.axes3D)
		ax = fig.add_subplot(row, col, i, projection="3d")
	else
		ax = fig.add_subplot(row, col, i, aspect=aspect)
	end
	return (fig, ax)
end

"""
    plot_spectrum(x::Vector, f_s, log_freq::Bool=false)

Plot a spectrum of the time domain data `x` with frequency range determined
by the sampling rate, `f_s`.

Required Arguments:

- `x::Vector{T}`: 1D array with element types `T`
- `f_s::Number`: sampling rate of `x`

Optional Arguments:

- `log_freq::Bool`: flag for if plot shows linear frequency or `log10(freqs)`
                    where `freqs` is the array of frequencies spanning
					`(0, f_s/2) Hz` `(default = false)`

Returns `nothing`

Displays plot of frequencies against power of `fft(x)`
"""
function plot_spectrum(x::Vector, f_s::Number; log_freq::Bool=true,
	                   new_fig=true, log_X=false)
    N = length(x)
    t_length = N/f_s
	X = fft(x)[1:floor(Int, N/2)+1]
	X = (1/(f_s*N)) .* abs2.(X)
	X[2:end-1] .= 2 .* X[2:end-1]
	freqs = Array(range(0, f_s/2, length=length(X)))
	if new_fig
    	fig = figure()
		ax = fig.add_subplot(1, 1, 1)
	end
    if log_freq
		if new_fig
			ax.set_xscale("log")
			if log_X
				ax.set_yscale("log")
			end
		else
			xscale("log")
			if log_X
				yscale("log")
			end
		end
	end
	if log_X
		plot(freqs, X)
	else
		plot(freqs, 10 .* log10.(X))
	end
	xlabel("Frequency (Hz)")
	ylabel("Sₓ(f)")
end

"""
	get_distribution_bound(distribution, p, tail=:center)

Gets the bounds of a normalized distribution where the sum is equal to `p`.

Required Arguments:

- `distribution`: histogram object that contains the bin locations and their
                  weights
- `p`: value rangin from 0 to 1 that specifies the ratio of the distribution to
       include

Optional Arguments:

- `tail`: whether start from the center and obtain bounds above and below the
          center or whether to start from the right or left.
		  `(default = :center)`

Returns:

- `bounds`: 2-element vector that describes the bounds of the distribution which
            include the bins, whose weights sum to equal `p`
"""
function get_distribution_bounds(distribution, p, tail=:center)
	if tail == :center
		center_idx = ceil(Int, length(distribution.edges[1])/2)
		center_idx = findall(x->x==0., distribution.edges[1])[1]
		side_limits = ceil(Int, length(distribution.edges[1])/2)
		distrbution_sum = distribution.weights[center_idx]
		bounds = 0
		for i in 1:side_limits
			if distrbution_sum < p
				distrbution_sum += distribution.weights[center_idx-i]
				distrbution_sum += distribution.weights[center_idx+i]
				bounds += 1
			end
		end
		bounds = [center_idx-bounds, center_idx+bounds]
	elseif tail == :left
		distrbution_sum = 0.
		bounds = [1, 1]
		for i in 1:length(distribution.weights)
			if distrbution_sum < p
				distrbution_sum += distribution.weights[i]
				bounds[2] = i
			end
		end
	elseif tail == :right
		distrbution_sum = 0.
		bounds = [1, 1] .* length(distribution.weights)
		for i in length(distribution.weights):-1:1
			if distrbution_sum < p
				distrbution_sum += distribution.weights[i]
				bounds[1] = i
			end
		end
	else
		error("$(tail) invalid. Must be either `:center`, `:left`, or `:right`.")
	end
	return bounds
end

"""
	snr2cn0(SNR, T)

Convert SNR to C/N₀ for an SNR estimate determined
at an integration time of `T`.

Required Arguments:

- `SNR`: the signal-to-noise ratio in dB
- `T`: the integration time in seconds

Returns:

- `CN0`: carrier-to-noise ratio in dB⋅Hz
"""
function snr2cn0(SNR, T)
	return SNR + 10*log10(1/T)
end

"""
	phase_noise_variance(CN0, T)

Calculate the expected unfiltered carrier phase measurement variance, `σᵩ²`.

Required Arguments:

- `CN0`: the carrier-to-noise ratio of the signal in dB⋅Hz
- `T`: the integration time in seconds

Returns:

- `σᵩ²`: the expected unfiltered carrier phase measurement variance in rad²
"""
function phase_noise_variance(CN0, T)
	CN0_linear = 10^(CN0/10)
	return (1 + 1/(2*T*CN0_linear)) / (2*T*CN0_linear)
end

"""
	phase_noise_variance(B, CN0, T, h₀=0, h₋₁=0, h₋₂=0)

Calculate the expected filtered carrier phase measurement variance, `σᵩ²`.

Required Arguments:

- `B`: the bandwidth of the filter in Hz
- `CN0`: the carrier-to-noise ratio of the signal in dB⋅Hz
- `T`: the integration time in seconds
- `h₀`: oscillator h-parameter `(default = 0)`
- `h₋₁`: oscillator h-parameter `(default = 0)`
- `h₋₂`: oscillator h-parameter `(default = 0)`

Returns:

- `σᵩ²`: the expected filtered carrier phase measurement variance in rad²
"""
function phase_noise_variance(B, CN0, T, h₀=0, h₋₁=0, h₋₂=0)
    CN0_linear = 10^(CN0/10)
    σ²ₒ = B*(1 + 1/(2*T*CN0_linear))/(CN0_linear*(1-2*B*T))
    σ²ᵪ = 0.9048*(2π)^3*h₋₂/B^3 + 0.8706*(2π)^2*h₋₁/B^2 + 1.2566*2π*h₀/B
    σ²ᵩ = σ²ₒ + σ²ᵪ
    return σ²ᵩ
end

"""
	make3dplot(corr_result, fd_range, fd_center, Δfd)

Make a 3D plot of a correlation result 2D array.

Required Arguments:

- `corr_result`: the output of `courseacquisition` or `courseacquisition!`
- `fd_range`: a value that describes the ± Doppler range of the `corr_result`
              array (value is in Hz)
- `fd_center`: the Doppler at the middle index of the `corr_result` row axis
               (value is in Hz)
- `Δf`: the Doppler bin width in Hz
"""
function make3dplot(corr_result, fd_range, fd_center, Δf)
	fig = figure()
	ax = Axes3D(fig)
	fd_bins = Array(-fd_range+fd_center:Δf:fd_range+fd_center)
	x, y = meshgrid(Array(1:size(corr_result)[2]), fd_bins./1000.)
	surf(x, y, corr_result)
	max_val_idx = argmax(collect(Iterators.flatten(corr_result)))
	scatter3D(collect(Iterators.flatten(x))[max_val_idx],
			  collect(Iterators.flatten(y))[max_val_idx],
			  collect(Iterators.flatten(corr_result))[max_val_idx],
			  c="r")
	xlabel("Samples")
    ylabel("Doppler (kHz)")
end

"""
	calc_snr(A², σ²)

Calculates the SNR using the amplitude squared and the variance

Required Arguments:

- `A²`: amplitude squared
- `σ²`: noise variance

Returns:
- `SNR`: signal-to-noise ratio of signal in dB `{SNR = 10log10(A²/(2⋅σ²))}`
"""
function calc_snr(A², σ²)
	return 10log10(A²/(2*σ²))
end


text:
using Oceananigans
using Oceananigans.Units
using Oceananigans.Advection: VelocityStencil, VorticityStencil

using Printf
using GLMakie

include("bickley_utils.jl")

"""
    run_bickley_jet(output_time_interval = 2, stop_time = 200, arch = CPU(), Nh = 64, ν = 0,
                    momentum_advection = VectorInvariant())

Run the Bickley jet validation experiment until `stop_time` using `momentum_advection`
scheme or formulation, with horizontal resolution `Nh`, viscosity `ν`, on `arch`itecture.
"""
function run_bickley_jet(;
                         output_time_interval = 2,
                         stop_time = 200,
                         arch = CPU(),
                         Nh = 64,
                         free_surface = ImplicitFreeSurface(gravitational_acceleration=10.0),
                         momentum_advection = WENO5(),
                         tracer_advection = WENO5(),
                         experiment_name = string(nameof(typeof(momentum_advection))))

    grid = bickley_grid(; arch, Nh)
    model = HydrostaticFreeSurfaceModel(; grid, momentum_advection, tracer_advection,
                                        free_surface, tracers = :c, buoyancy=nothing)
    set_bickley_jet!(model)

    Δt = 0.2 * 2π / Nh
    wizard = TimeStepWizard(cfl=0.2, max_change=1.1, max_Δt=10.0)
    simulation = Simulation(model; Δt, stop_time)

    progress(sim) = @printf("Iter: %d, time: %.1f, Δt: %.1e, max|u|: %.3f, max|η|: %.3f\n",
                            iteration(sim), time(sim), sim.Δt,
                            maximum(abs, model.velocities.u),
                            maximum(abs, model.free_surface.η))

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

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

    # Output: primitive fields + computations
    u, v, w = model.velocities
    outputs = merge(model.velocities, model.tracers, (ζ=∂x(v) - ∂y(u), η=model.free_surface.η))

    @show output_name = "bickley_jet_Nh_$(Nh)_" * experiment_name

    simulation.output_writers[:fields] =
        JLD2OutputWriter(model, outputs,
                                schedule = TimeInterval(output_time_interval),
                                prefix = output_name,
                                overwrite_existing = true)

    @info "Running a simulation of an unstable Bickley jet with $(Nh)² degrees of freedom..."

    start_time = time_ns()

    run!(simulation)

    elapsed = 1e-9 * (time_ns() - start_time)
    @info "... the bickley jet simulation took " * prettytime(elapsed)

    return output_name
end

"""
    visualize_bickley_jet(experiment_name)

Visualize the Bickley jet data in `name * ".jld2"`.
"""
function visualize_bickley_jet(name)
    @info "Making a fun movie about an unstable Bickley jet..."

    filepath = name * ".jld2"

    ζt = FieldTimeSeries(filepath, "ζ")
    ct = FieldTimeSeries(filepath, "c")
    t = ζt.times
    Nt = length(t)

    fig = Figure(resolution=(1400, 800))
    slider = Slider(fig[2, 1:2], range=1:Nt, startvalue=1)
    n = slider.value

    ζtitle = @lift @sprintf("ζ at t = %.1f", t[$n])
    ctitle = @lift @sprintf("c at t = %.1f", t[$n])

    ax_ζ = Axis(fig[1, 1], title=ζtitle, aspect=1)
    ax_c = Axis(fig[1, 2], title=ctitle, aspect=1)

    ζ = @lift interior(ζt[$n], :, :, 1)
    c = @lift interior(ct[$n], :, :, 1)

    heatmap!(ax_ζ, ζ, colorrange=(-1, 1), colormap=:redblue)
    heatmap!(ax_c, c, colorrange=(-1, 1), colormap=:thermal)

    record(fig, name * ".mp4", 1:Nt, framerate=24) do nn
        @info "Drawing frame $nn of $Nt..."
        n[] = nn
    end
end

advection_schemes = [CenteredSecondOrder(), WENO5()]

#=
advection_schemes = [WENO5(vector_invariant=VelocityStencil()),
                     WENO5(vector_invariant=VorticityStencil()),
                     WENO5(),
                     VectorInvariant()]
=#

arch = CPU()
for Nh in [128]
    for momentum_advection in advection_schemes
        name = run_bickley_jet(; arch, momentum_advection, Nh)
        visualize_bickley_jet(name)
    end
end


text:
include("args.jl")

function sw_p_to_lambda_den(sw, p)
    sw = tf.reshape(sw, (1, m, n, 1))
    p = tf.reshape(p, (1, m, n, 1))
    sw = tf.image.resize_bilinear(sw, (nz, nx))
    p = tf.image.resize_bilinear(p, (nz, nx))
    sw = cast(sw, Float64)
    p = cast(p, Float64)
    sw = squeeze(sw)
    p = squeeze(p)
    # tran_lambda, tran_den = Gassman(sw)
    # tran_lambda, tran_den = RockLinear(sw) # test linear relationship
    tran_lambda, tran_den = Patchy(sw)
    # tran_lambda_pad =  tf.pad(tran_lambda, [nPml (nPml+nPad); nPml nPml], constant_values=3500.0^2*2200.0/3.0) /1e6
    # tran_den_pad = tf.pad(tran_den, [nPml (nPml+nPad); nPml nPml], constant_values=2200.0)
    return tran_lambda, tran_den
end

if !isdir("figures_summary")
  mkdir("figures_summary")
end

lambdasObs = Array{PyObject}(undef, n_survey-1)
densObs = Array{PyObject}(undef, n_survey-1)
for iSur = 2:n_survey
  lp = readdlm("./CO2/FWI_stage$(iSur)/loss.txt")
  Lp = Int64((lp[end,1]))
  lambdasObs[iSur-1] = readdlm("./CO2/FWI_stage$(iSur)/Lambda$Lp.txt")
  densObs[iSur-1] = readdlm("./CO2/FWI_stage$(iSur)/Den$Lp.txt")
end

m = 15
n = 30
h = 30.0
dz = 3.0 # meters
dx = 3.0
nz = Int64(round((m * h) / dz)) + 1
nx = Int64(round((n * h) / dx)) + 1
z_src = (collect(5:10:nz-5) .- 1 ) .* dz .+ dz/2.0
x_src = (5-1)ones(Int64, size(z_src)) .* dx .+ dx/2.0
z_rec = (collect(5:1:nz-5) .- 1) .* dz .+ dz/2.0
x_rec = (nx-5-1) .* ones(Int64, size(z_rec)) .*dx .+ dx/2.0
z_inj = (9-1)*h + h/2.0
x_inj = (3-1)*h + h/2.0
z_prod = (9-1)*h + h/2.0
x_prod = (28-1)*h + h/2.0

rc("axes", titlesize=20)
rc("axes", labelsize=18)
rc("xtick", labelsize=18)
rc("ytick", labelsize=18)
rc("legend", fontsize=20)
figure()
K = readdlm("CO2/flow_fit_results/K100.txt")
imshow(K, extent=[0,n*h,m*h,0]);
xlabel("Distance (m)")
ylabel("Depth (m)")
cb = colorbar()
clim([20, 120])
cb.set_label("Permeability (md)")
shot_inds = collect(1:length(z_src))
scatter(x_src[shot_inds], z_src[shot_inds], c="w", marker="*")
scatter(x_rec, z_rec, s=16.0, c="r", marker="v")
scatter(x_inj, z_inj, c="r", marker=">")
scatter(x_prod, z_prod, c="r", marker="<")
savefig("figures_summary/K_sep_fit.pdf", bbox_inches="tight",pad_inches = 0, dpi = 300);

rc("axes", titlesize=30)
rc("axes", labelsize=30)
rc("xtick", labelsize=28)
rc("ytick", labelsize=28)
rc("legend", fontsize=30)
fig1,axs = subplots(3,3, figsize=[30,15], sharex=true, sharey=true)
ims = Array{Any}(undef, 9)
for iPrj = 1:3
  for jPrj = 1:3
    ims[(iPrj-1)*3+jPrj] = axs[iPrj,jPrj].imshow(lambdasObs[(iPrj-1)*3+jPrj], extent=[0,n*h,m*h,0], vmin=5500, vmax=9000);
    axs[iPrj,jPrj].title.set_text("Snapshot $((iPrj-1)*3+jPrj+1)")
    if jPrj == 1 || jPrj == 1
      axs[iPrj,jPrj].set_ylabel("Depth (m)")
    end
    if iPrj == 3 || iPrj == 3
      axs[iPrj,jPrj].set_xlabel("Distance (m)")
    end
    # cb = fig1.colorbar(ims[(iPrj-1)*3+jPrj], ax=axs[iPrj,jPrj])
    # cb.set_label("λ")
    axs[iPrj,jPrj].scatter(x_inj, z_inj, c="r", marker=">", s=128)
    axs[iPrj,jPrj].scatter(x_prod, z_prod, c="r", marker="<", s=128)
  end
end
# fig1.subplots_adjust(wspace=0.02, hspace=0.042)
fig1.subplots_adjust(wspace=0.02, hspace=0.18)
cbar_ax = fig1.add_axes([0.91, 0.08, 0.01, 0.82])
cb1 = fig1.colorbar(ims[1], cax=cbar_ax)
cb1.set_label("λ (MPa)")
savefig("figures_summary/Lambda_FWI_sep_inv.pdf",bbox_inches="tight",pad_inches = 0);


text:
using Raylib
using Raylib: rayvector, RayVector2

function vector2rotate(v, angle)
    x = v[1] * cos(angle) - v[2] * sin(angle)
    y = v[1] * sin(angle) + v[2] * cos(angle)
    return rayvector(x, y)
end

function main()
    screenWidth = 800
    screenHeight = 450

    MAX_POINTS = 11

    Raylib.InitWindow(screenWidth, screenHeight, "raylib [textures] example - textured polygon")

    texcoords = [
        rayvector(0.75, 0.0),
        rayvector(0.25, 0.0),
        rayvector(0.0, 0.5),
        rayvector(0.0, 0.75),
        rayvector(0.25, 1.0),
        rayvector(0.375, 0.875),
        rayvector(0.625, 0.875),
        rayvector(0.75, 1.0),
        rayvector(1.0, 0.75),
        rayvector(1.0, 0.5),
        rayvector(0.75, 0.0),
    ]

    points = Vector{RayVector2}(undef, MAX_POINTS)
    for i = 1:MAX_POINTS
        points[i] = (texcoords[i] .- 0.5) .* 256
    end

    positions = copy(points)

    texture = Raylib.LoadTexture(joinpath(@__DIR__, "resources/cat.png"))

    angle = 0.0

    Raylib.SetTargetFPS(60)

    while !Raylib.WindowShouldClose()
        angle+=1
        map!(Base.Fix2(vector2rotate, deg2rad(angle)), positions, points)

        Raylib.BeginDrawing()
        Raylib.ClearBackground(Raylib.RAYWHITE)
        Raylib.DrawText("textured polygon", 20, 20, 20, Raylib.DARKGRAY)
        Raylib.DrawTexturePoly(
            texture, rayvector(Raylib.GetScreenWidth()/2, Raylib.GetScreenHeight()/2),
            pointer(positions), pointer(texcoords), MAX_POINTS, Raylib.WHITE
        )
        Raylib.EndDrawing()
    end

    Raylib.UnloadTexture(texture)
    Raylib.CloseWindow()

end


text:
@testset "basic attributes" begin
    a1 = SegmentedMean(1:4 |> collect |> f32)
    a2 = SegmentedMean(4)
    a3 = SegmentedMeanMax(4)

    @test length(a1) == 4
    @test size(a1) == (4,)
    @test length(a2) == 4
    @test size(a2) == (4,)
    @test length(a3) == 8
    @test size(a3) == (8,)
end

@testset "vcat" begin
    # We do not implement == for models because it messes with Flux implementation
    function test_equal(a1::T, a2::T) where T <: AggregationStack
        length(a1.fs) == length(a2.fs) && all(test_equal(f1, f2) for (f1, f2) in zip(a1.fs, a2.fs))
    end
    function test_equal(a1::T, a2::T) where T <: AbstractAggregation
        all(getfield(a1, f) == getfield(a2, f) for f in fieldnames(T))
    end

    # Aggregation Operators
    @test test_equal(
                     vcat(SegmentedMean(ones(5)), SegmentedMean(ones(5))),
                     SegmentedMean(ones(10))
                    )
    @test test_equal(
                     vcat(SegmentedMax(ones(5)), SegmentedMax(ones(5))),
                     SegmentedMax(ones(10))
                     )
    @test test_equal(
                     vcat(SegmentedSum(ones(5)), SegmentedSum(ones(5))),
                     SegmentedSum(ones(10))
                    )
    @test test_equal(
                     vcat(SegmentedLSE(ones(5), zeros(5)), SegmentedLSE(ones(5), zeros(5))),
                     SegmentedLSE(ones(10), zeros(10))
                    )
    @test test_equal(
                     vcat(SegmentedPNorm(ones(5), zeros(5), -ones(5)),
                          SegmentedPNorm(ones(5), zeros(5), -ones(5))),
                     SegmentedPNorm(ones(10), zeros(10), -ones(10))
                    )

    @test test_equal(vcat(SegmentedMean(5)), SegmentedMean(5))
    @test test_equal(vcat(SegmentedMean(5), SegmentedMax(5)), SegmentedMeanMax(5))
end

@testset "flattening" begin
    as = [
          AggregationStack(SegmentedMean(2), SegmentedMax(2)),
          AggregationStack(SegmentedMean(2), AggregationStack(SegmentedMax(2))),
          AggregationStack((SegmentedMean(2), SegmentedMax(2))),
          AggregationStack((AggregationStack(SegmentedMean(2)), SegmentedMax(2))),
          AggregationStack(SegmentedMeanMax(2)),
          AggregationStack(AggregationStack(SegmentedMeanMax(2))),
          AggregationStack(SegmentedMeanMax(2)),
          vcat(SegmentedMean(2), SegmentedMax(2)),
          SegmentedMeanMax(2)
         ]

    @test all(length.(as) .== 4)
    @test all(a -> all(f -> length(f) == 2, a.fs), as)
    @test all(isa.(getindex.(as, 1), SegmentedMean))
    @test all(isa.(getindex.(as, 2), SegmentedMax))
end

@testset "size checking" begin
    b = bags([1:2])
    for a in all_aggregations(Float64, 3).fs, x in [zeros(2, 2), zeros(4, 2)]
        @test_throws DimensionMismatch a(x, b)
        @test_throws DimensionMismatch BagCount(a)(x, b)
    end
end

@testset "basic aggregation functionality" begin
    W = [1, 1/2, 1/2, 1/8, 1/3, 13/24] |> f32
    X = reshape(1:12, 2, 6) |> f32
    bags = BAGS[1]
    @test SegmentedMean(2)(X, bags) ≈ [1.0 4.0 9.0; 2.0 5.0 10.0]
    @test SegmentedSum(2)(X, bags) ≈ length.(bags)' .* SegmentedMean(2)(X, bags)
    @test SegmentedMax(2)(X, bags) ≈ [1.0 5.0 11.0; 2.0 6.0 12.0]
    @test SegmentedMeanMax(2)(X, bags) ≈ vcat(SegmentedMean(2)(X, bags),
                                              SegmentedMax(2)(X, bags))
    @test SegmentedMean(2)(X, bags, W) ≈ [1.0 4.0 236/24; 2.0 5.0 260/24]
    bagnorms = [_bagnorm(W, b) for b in bags]
    @test SegmentedSum(2)(X, bags, W) ≈ bagnorms' .* SegmentedMean(2)(X, bags, W)
    @test SegmentedMax(2)(X, bags, W) ≈ SegmentedMax(2)(X, bags)
end

@testset "matrix weights" begin
    W = abs.(rand(Float32, 6))
    W_mat = vcat(W, 2*W)
    X = reshape(1:12, 2, 6) |> f32
    for bags in BAGS
        @test SegmentedSum(2)(X, bags, W_mat) ≈ SegmentedSum(2)(X, bags, W)
        @test SegmentedMean(2)(X, bags, W_mat) ≈ SegmentedMean(2)(X, bags, W)
        @test SegmentedMax(2)(X, bags, W_mat) ≈ SegmentedMax(2)(X, bags, W)
    end
end

@testset "pnorm functionality" begin
    dummy = randn(2)
    for t = 1:10
        a, b, c, d, ρ1, ρ2, c1, c2, w1, w2 = randn(10)
        p1, p2 = p_map(ρ1), p_map(ρ2)
        w1 = abs(w1)
        w2 = abs(w2)
        bags = ScatteredBags([[1,2]])
        agg = SegmentedPNorm(dummy, [ρ1, ρ2], [c1, c2])
        @test agg([a b; c d], bags) ≈ [
                                     (1/2*(abs(a-c1)^p1 + abs(b-c1)^p1))^(1/p1);
                                     (1/2*(abs(c-c2)^p2 + abs(d-c2)^p2))^(1/p2)
                                    ]
        @test agg([a b; c d], bags) ≈ [
                                     (1/2*(abs(a-c1)^p1 + abs(b-c1)^p1))^(1/p1);
                                     (1/2*(abs(c-c2)^p2 + abs(d-c2)^p2))^(1/p2)
                                    ]
        @test agg([a b; c d], bags, [w1, w2]) ≈ [
                                               (1/(w1+w2)*(w1*abs(a-c1)^p1 + w2*abs(b-c1)^p1))^(1/p1);
                                               (1/(w1+w2)*(w1*abs(c-c2)^p2 + w2*abs(d-c2)^p2))^(1/p2)
                                              ]
        for bags in BAGS
            X = randn(2, 6)
            agg = SegmentedPNorm(dummy, inv_p_map([1+1e-16, 1+1e-16]), [0.0, 0.0])
            @test agg(X, bags) ≈ SegmentedMean(dummy)(abs.(X), bags)
            agg = SegmentedPNorm(dummy, inv_p_map([2.0, 2.0]), [0.0, 0.0])
            @test agg(X, bags) ≈ hcat([sqrt.(sum(X[:, b] .^ 2, dims=2) ./ length(b)) for b in bags]...)
        end
    end
end

@testset "lse functionality" begin
    dummy = randn(2)
    for t = 1:10
        a, b, c, d, ρ1, ρ2 = randn(6)
        r1, r2 = r_map(ρ1), r_map(ρ2)
        @test SegmentedLSE(dummy, [ρ1, ρ2])([a b; c d], ScatteredBags([[1,2]])) ≈ [
                                                                               1/r1*log(1/2*(exp(a*r1)+exp(b*r1)));
                                                                               1/r2*log(1/2*(exp(c*r2)+exp(d*r2)))
                                                                              ]
        for bags in BAGS
            X = randn(2, 6)
            W = abs.(randn(6))
            r1, r2 = randn(2)
            # doesn't use weights
            @test SegmentedLSE(dummy, [ρ1, ρ2])(X, bags) == SegmentedLSE(dummy, [ρ1, ρ2])(X, bags, W)
            # the bigger value of r, the closer we are to the real maximum
            @test isapprox(SegmentedLSE(dummy, [100.0, 100.0])(X, bags), SegmentedMax(dummy)(X, bags), atol=0.1)
        end
    end
end

@testset "pnorm numerical stability" begin
    k, d = 10, 5
    dummy, c = randn(d), zeros(d)
    b = AlignedBags([1:k])
    ρ1 = inv_p_map(ones(d))
    ρ2 = randn(d)
    p2 = p_map(ρ2)
    Z = 1e5 .+ 1e3 .* randn(d, k)
    W = abs.(randn(k)) .+ 1e-2
    for X in [Z, -Z, randn(d, k)]
        @test SegmentedPNorm(dummy, ρ1, c)(X, b) ≈ sum(abs.(X); dims=2) ./ k
        @test SegmentedPNorm(dummy, ρ1, c)(X, b, W) ≈ sum(W' .* abs.(X); dims=2) / sum(W)
        @test SegmentedPNorm(dummy, ρ1, c)(X, b, W) ≈ sum(W' .* abs.(X); dims=2) / sum(W)
        for i in 1:k
            @test SegmentedPNorm(dummy, ρ1, c)(repeat(X[:, i], 1, k), b) ≈ abs.(X[:, i])
            @test SegmentedPNorm(dummy, ρ2, c)(repeat(X[:, i], 1, k), b) ≈ abs.(X[:, i])
            @test SegmentedPNorm(dummy, ρ1, c)(repeat(X[:, i], 1, k), b, W) ≈ abs.(X[:, i])
            @test SegmentedPNorm(dummy, ρ1, c)(repeat(X[:, i], 1, k), b, W) ≈ abs.(X[:, i])
            @test SegmentedPNorm(dummy, ρ2, c)(repeat(X[:, i], 1, k), b, W) ≈ abs.(X[:, i])
            @test SegmentedPNorm(dummy, ρ2, c)(repeat(X[:, i], 1, k), b, W) ≈ abs.(X[:, i])
        end
    end
end

@testset "lse numerical stability" begin
    k, d = 10, 5
    dummy = randn(d)
    b = AlignedBags([1:k])
    ρ1 = inv_r_map(1e15 .+ 1e5 .* randn(d))
    ρ2 = inv_r_map(abs.(1e5 .* randn(d)))
    Z = 1e5 .+ 1e3 .* randn(d, k)
    W = abs.(randn(k)) .+ 1e-2
    for X in [Z, -Z, randn(d, k)]
        @test SegmentedLSE(dummy, ρ1)(X, b) ≈ maximum(X; dims=2)
        # doesn't use weights
        @test SegmentedLSE(dummy, ρ1)(X, b, W) ≈ maximum(X; dims=2)

        # implementation immune to underflow not available yet
        @test_skip @test SegmentedLSE(dummy, ρ2)(X, b) ≈ sum(X; dims=2) ./ k
        # doesn't use weights
        @test_skip @test SegmentedLSE(dummy, ρ2)(X, b, W) ≈ sum(X; dims=2) ./ k

        for i in 1:k
            @test SegmentedLSE(dummy, randn(d))(repeat(X[:, i], 1, k), b) ≈ X[:, i]
            # doesn't use weights
            @test SegmentedLSE(dummy, randn(d))(repeat(X[:, i], 1, k), b, W) ≈ X[:, i]
        end
    end
end

@testset "bagcount" begin
    X = reshape(1:12, 2, 6) |> f32
    d = 2
    bags = BAGS[1]
    baglengths = log.(1 .+ [1.0 2.0 3.0]) |> f32

    function test_count(a)
        a1 = a(X, bags)
        a2 = BagCount(a)(X, bags)
        @test [a1; baglengths] == a2
    end

    test_count(SegmentedMean(2))
    test_count(SegmentedMax(2))
    test_count(SegmentedSumMeanMaxPNormLSE(2))
    test_count(nonparam_aggregations(Float32, 2))
    test_count(param_aggregations(Float32, 2))
    test_count(all_aggregations(Float32, 2))
end

@testset "type stability of the output" begin
    types = [Float16, Float32, Float64]
    for t1 in types, t2 in types, b in BAGS2
        x = rand(t2, 3, 10)
        w = abs.(randn(t2, size(x, 2))) .+ t2(0.1)
        w_mat = abs.(randn(t2, size(x))) .+ t2(0.1)
        agg = all_aggregations(t1, 3)
        for a in [agg, BagCount(agg)]
            @inferred a(x, b)
            @inferred a(x, b, w)
            @inferred a(x, b, w_mat)
            @test eltype(a(x, b)) ===
                    eltype(a(x, b, w)) ===
                    eltype(a(x, b, w_mat)) ===
                    promote_type(t1, t2)
        end
    end
end

@testset "r_map and p_map are stable" begin
    @test first(gradient(softplus, 10000)) ≈ σ(10000) ≈ 1.0
    @test first(gradient(softplus, -10000)) ≈ σ(-10000) ≈ 0
end

@testset "missing values" begin
    dummy = randn(2)
    ψ = randn(2)
    for bags in [AlignedBags([0:-1]), AlignedBags([0:-1, 0:-1, 0:-1])]
        @test SegmentedMean(ψ)(missing, bags) == repeat(ψ, 1, length(bags))
        @test SegmentedSum(ψ)(missing, bags) == repeat(ψ, 1, length(bags))
        @test SegmentedMax(ψ)(missing, bags) == repeat(ψ, 1, length(bags))
        @test SegmentedLSE(ψ, dummy)(missing, bags) == repeat(ψ, 1, length(bags))
        @test SegmentedPNorm(ψ, dummy, dummy)(missing, bags) == repeat(ψ, 1, length(bags))
    end

    # default values ψ are indeed filled in
    for bags in vcat(BAGS2)
        idcs = isempty.(bags.bags)
        l = maximum(maximum.(bags.bags[.!idcs]))
        X = randn(2, l)
        @test SegmentedMean(ψ)(X, bags)[:, idcs] == repeat(ψ, 1, sum(idcs))
        @test SegmentedSum(ψ)(X, bags)[:, idcs] == repeat(ψ, 1, sum(idcs))
        @test SegmentedMax(ψ)(X, bags)[:, idcs] == repeat(ψ, 1, sum(idcs))
        @test SegmentedLSE(ψ, dummy)(X, bags)[:, idcs] == repeat(ψ, 1, sum(idcs))
        @test SegmentedPNorm(ψ, dummy, dummy)(X, bags)[:, idcs] == repeat(ψ, 1, sum(idcs))
    end

    # missing array entries are propagated
    b = bags([1:3])
    X1 = [1 2 3; missing 5 missing; missing missing missing]
    X2 = fill(missing, 3, 3)
    for a in all_aggregations(Float32, 3).fs
        @test !ismissing(a(X1, b)[1])
        @test ismissing(a(X1, b)[2])
        @test ismissing(a(X1, b)[3])
        @test all(ismissing.(a(X2, b)))
        @test !ismissing(BagCount(a)(X1, b)[1])
        @test ismissing(BagCount(a)(X1, b)[2])
        @test ismissing(BagCount(a)(X1, b)[3])
        @test all(ismissing.(BagCount(a)(X2, b)[1:end-1,:]))
    end
end

@testset "aggregation grad check w.r.t. input" begin
    for bags in BAGS2
        d = rand(1:20)
        x = randn(d, 10)
        w = abs.(randn(size(x, 2))) .+ 0.1
        w_mat = abs.(randn(size(x))) .+ 0.1
        agg = all_aggregations(Float64, d)
        for a in [agg, BagCount(agg)]
            @test gradtest(x -> a(x, bags), x)
            @test gradtest(x -> a(x, bags, w), x)
            @test gradtest(x -> a(x, bags, w_mat), x)
        end
    end
end

@testset "aggregation gradcheck w.r.t weights" begin
    for bags in BAGS2
        d = rand(1:20)
        x = randn(d, 10)
        w = abs.(randn(size(x, 2))) .+ 0.1
        w_mat = abs.(randn(size(x))) .+ 0.1
        agg = all_aggregations(Float64, d)
        for a in tuple(agg.fs..., BagCount.(agg.fs)...)
            if a isa SegmentedPNorm || a isa BagCount{<:SegmentedPNorm}
                @test_throws NotImplementedException gradtest(w -> a(x, bags, w), w)
                @test_throws NotImplementedException gradtest(w -> a(x, bags, w), w_mat)
            else
                @test gradtest(w -> a(x, bags, w), w)
                @test gradtest(w -> a(x, bags, w), w_mat)
            end
        end
    end
end

@testset "aggregation grad check w.r.t. params" begin
    d = 5
    x = randn(d, 0)
    a1 = nonparam_aggregations(Float64, d)
    a2 = param_aggregations(Float64, d)
    for a in [a1, BagCount(a1), a2, BagCount(a2)]
        @test gradtest(() -> a(missing, ScatteredBags([Int[], Int[]])), Flux.params(a))
        @test gradtest(() -> a(missing, AlignedBags([0:-1]), nothing), Flux.params(a))
        @test gradtest(() -> a(x, ScatteredBags([Int[]])), Flux.params(a))
        @test gradtest(() -> a(x, AlignedBags([0:-1, 0:-1]), nothing), Flux.params(a))
    end

    for bags in BAGS2
        d = rand(1:20)
        x = randn(d, 10)
        a1 = nonparam_aggregations(Float64, d)
        a2 = param_aggregations(Float64, d)
        w = abs.(randn(size(x, 2))) .+ 0.1
        w_mat = abs.(randn(size(x))) .+ 0.1
        for a in [a1, BagCount(a1)]
            @test gradtest(() -> a(x, bags), Flux.params(a))
            @test gradtest(() -> a(x, bags, w), Flux.params(a))
            @test gradtest(() -> a(x, bags, w_mat), Flux.params(a))
        end
        for a in [a2, BagCount(a2)]
            @test gradtest(() -> a(x, bags), Flux.params(a); atol=1e-3)
            @test gradtest(() -> a(x, bags, w), Flux.params(a); atol=1e-3)
            @test gradtest(() -> a(x, bags, w_mat), Flux.params(a); atol=1e-3)
        end
    end
end


text:
# BAM Reader
# ==========

type BAMReader <: Bio.IO.AbstractParser
    filepath::String
    stream::BGZFStream
    header::SAMHeader
    start_offset::VirtualOffset
    refseqnames::Vector{String}
    refseqlens::Vector{Int}
    index::Nullable{BAI}
end

function Base.show(io::IO, reader::BAMReader)
    println(io, summary(reader), ":")
    println("  filepath: ", reader.filepath)
    println("  header: ", reader.header)
      print("  reference sequences:")
    for (i, (name, len)) in enumerate(zip(reader.refseqnames, reader.refseqlens))
        println(io)
        print("    [", lpad(i, 2), "]: ", name, " (length: ", len, ")")
    end
end

function header(reader::BAMReader, fillSQ::Bool=false)
    if fillSQ
        if haskey(reader.header, "SQ")
            # TODO: check consistency
        else
            records = Dict{String,Any}[]
            for (name, len) in zip(reader.refseqnames, reader.refseqlens)
                push!(records, Dict("SN" => name, "LN" => len))
            end
            reader.header["SQ"] = records
        end
    end
    return reader.header
end

function Base.eof(reader::BAMReader)
    return eof(reader.stream)
end

function Base.close(reader::BAMReader)
    close(reader.stream)
end

function Base.seek(reader::BAMReader, voffset::VirtualOffset)
    seek(reader.stream, voffset)
end

function Base.seekstart(reader::BAMReader)
    seek(reader.stream, reader.start_offset)
end

function Base.open(filename::AbstractString, ::Type{BAM})
    stream = BGZFStream(filename)

    # magic bytes
    B = read(stream, UInt8)
    A = read(stream, UInt8)
    M = read(stream, UInt8)
    x = read(stream, UInt8)
    if B != UInt8('B') || A != UInt8('A') || M != UInt8('M') || x != 0x01
        error("input was not a valid BAM file")
    end

    # SAM header
    textlen = read(stream, Int32)
    samheader = parse_samheader(read(stream, UInt8, textlen))

    # reference sequences
    refseqnames = String[]
    refseqlens = Int[]
    n_refs = read(stream, Int32)
    for _ in 1:n_refs
        namelen = read(stream, Int32)
        data = read(stream, UInt8, namelen)
        seqname = unsafe_string(pointer(data))
        seqlen = read(stream, Int32)
        push!(refseqnames, seqname)
        push!(refseqlens, seqlen)
    end

    reader = BAMReader(
        filename,
        stream,
        samheader,
        virtualoffset(stream),
        refseqnames,
        refseqlens,
        Nullable())

    index_filepath = reader.filepath * ".bai"
    if isfile(index_filepath)
        loadindex!(reader, index_filepath)
    end

    return reader
end

function Base.read!(reader::BAMReader, aln::BAMRecord)
    datasize = read(reader.stream, Int32) - BAM_FIXED_FIELDS_BYTES
    unsafe_read(reader.stream, pointer_from_objref(aln), BAM_FIXED_FIELDS_BYTES)
    if length(aln.data) < datasize
        resize!(aln.data, datasize)
    end
    unsafe_read(reader.stream, pointer(aln.data), datasize)
    aln.datasize = datasize
    aln.refseqnames = reader.refseqnames
    return aln
end

function Base.start(reader::BAMReader)
    return BAMRecord()
end

function Base.done(reader::BAMReader, rec)
    return eof(reader)
end

function Base.next(reader::BAMReader, rec)
    read!(reader, rec)
    return copy(rec), rec
end

function loadindex!(reader, filepath)
    open(filepath) do input
        reader.index = read(input, BAI)
    end
    return reader
end


text:
###
### Mer specific specializations of src/biosequence/predicates.jl
###

Base.cmp(x::T, y::T) where {T<:Kmer} = cmp(x.data, y.data)
Base.:(==)(x::T, y::T) where {T<:Kmer} = x.data == y.data
Base.isless(x::T, y::T) where {T<:Kmer} = isless(x.data, y.data)

# TODO: Ensure this is the right way to go.
# See https://github.com/BioJulia/BioSequences.jl/pull/121#discussion_r475234270
Base.hash(x::Kmer{A,K,N}, h::UInt) where {A,K,N} = hash(x.data, h ⊻ K)


text:
import MappedArrays

"""
    OneStateAutomaton

Automaton with one state and the `nt` events 1, ..., `nt`.
"""
struct OneStateAutomaton <: AbstractAutomaton
    nt::Int
end

"""
    OneStateTransition

Transition of `OneStateAutomaton` with label `σ`.
"""
struct OneStateTransition <: AbstractTransition
    σ::Int
end

states(A::OneStateAutomaton) = Base.OneTo(1)
nstates(A::OneStateAutomaton) = 1
transitiontype(A::OneStateAutomaton) = OneStateTransition
function transitions(A::OneStateAutomaton)
    # mappedarray avoid allocation
    MappedArrays.mappedarray(OneStateTransition, Base.OneTo(A.nt))
end
ntransitions(A::OneStateAutomaton) = A.nt
source(::OneStateAutomaton, t::OneStateTransition) = 1
event(::OneStateAutomaton, t::OneStateTransition) = t.σ
target(::OneStateAutomaton, t::OneStateTransition) = 1
in_transitions(A::OneStateAutomaton, s) = transitions(A)
out_transitions(A::OneStateAutomaton, s) = transitions(A)
function has_transition(A::OneStateAutomaton, q, r)
    @assert q == 1
    @assert r == 1
    return !iszero(A.nt)
end
function has_transition(A::OneStateAutomaton, t::OneStateTransition)
    return 1 <= t.σ <= A.nt
end

mutable struct OneStateStateProperty{T} <: StateProperty{T}
    automaton::OneStateAutomaton
    value::Union{Nothing, T}
end
state_property_type(::Type{OneStateAutomaton}, T::Type) = OneStateStateProperty{T}
function state_property(automaton::OneStateAutomaton, T::Type)
    return OneStateStateProperty{T}(automaton, nothing)
end
function Base.getindex(p::OneStateStateProperty,
                       s::Int)
    @assert isone(s)
    p.value
end
function Base.setindex!(p::OneStateStateProperty, value,
                        s::Int)
    @assert isone(s)
    p.value = value
end

struct OneStateTransitionProperty{T} <: TransitionProperty{T}
    automaton::OneStateAutomaton
    value::Vector{T}
end
transition_property_type(::Type{OneStateAutomaton}, T::Type) = OneStateTransitionProperty{T}
function transition_property(automaton::OneStateAutomaton, T::Type)
    return OneStateTransitionProperty(automaton, Vector{T}(undef, automaton.nt))
end
function Base.getindex(p::OneStateTransitionProperty,
                       t::OneStateTransition)
    p.value[t.σ]
end
function Base.setindex!(p::OneStateTransitionProperty, value,
                        t::OneStateTransition)
    p.value[t.σ] = value
end


text:
module NumberTheory
include("harmonic.jl")
include("functions.jl")
include("harmonic.jl")
include("mersenne.jl")
end

max_stars_repo_name: EllianCarlos/number-theory-but-in-julia

max_issues_repo_name: EllianCarlos/number-theory-but-in-julia

max_forks_repo_name: EllianCarlos/number-theory-but-in-julia


text:
# facilities for code generation

type ArgInfo{AN} end

arginfo(n::Int) = ArgInfo{n}()

immutable CodegenHelper
    args::Vector{Symbol}
    contiguous_aparams::Vector{Expr}
    dense_aparams::Vector{Expr}

    term::Function
    termtype::Expr
    inputlen::Expr
    inputsize::Expr
end

function codegen_helper(AN::Int)
    N = abs(AN)

    if AN == 0
        args = [:a]
        contiguous_aparams = [:(a::ContiguousNumericArray)]
        dense_aparams = [:(a::NumericArray)]
        term = idx->:(a[$idx])
        termtype = :(eltype(a))
    elseif AN == 1
        args = [:fun, :a]
        contiguous_aparams = [:(fun::Functor{1}), :(a::ContiguousNumericArray)]
        dense_aparams = [:(fun::Functor{1}), :(a::NumericArray)]
        term = idx->:(evaluate(fun, a[$idx]))
        termtype = :(result_type(fun, eltype(a)))
    else
        FN = AN
        if AN == -2
            FN = 1
        end
        aargs = [symbol('a' + (i-1)) for i = 1:N]
        args = [:fun, aargs...]
        contiguous_aparams = [:(fun::Functor{$FN}), [:($a::ContiguousArrOrNum) for a in aargs]...]
        dense_aparams = [:(fun::Functor{$FN}), [:($a::DenseArrOrNum) for a in aargs]...]

        if AN > 1
            term = idx->Expr(:call, :evaluate, :fun, [:(getvalue($a, $idx)) for a in aargs]...)
            termtype = Expr(:call, :result_type, :fun, [:(eltype($a)) for a in aargs]...)
        elseif AN == -2
            term = idx->:(evaluate(fun, getvalue(a, $idx) - getvalue(b, $idx)))
            termtype = :(result_type(fun, promote_type(eltype(a), eltype(b))))
        end
    end

    if N <= 1
        inputlen = :(length(a))
        inputsize = :(size(a))
    else
        inputlen = Expr(:call, :maplength, aargs...)
        inputsize = Expr(:call, :mapshape, aargs...)
    end

    CodegenHelper(args,
                  contiguous_aparams,
                  dense_aparams,
                  term,
                  termtype,
                  inputlen,
                  inputsize)
end

immutable CodegenHelperEx
    imparams1d::Vector{Expr}
    imparams2d::Vector{Expr}
    imargs1d::Vector
    imargs2d::Vector
    eviewargs::Vector

    getoffsets::Expr
    getstrides1::Expr
    getstrides2::Expr
    contcol::Expr
    nextcol::Expr
    pkerargs::Vector
    pkerargs1::Vector
    kerargs::Vector{Symbol}
    kerargs1::Vector{Symbol}
end

function codegen_helper_ex(AN::Int)
    if AN == 0
        imparams1d = [:(a::ContiguousArray), :(ia::Int), :(sa1::Int)]
        imparams2d = [:(a::ContiguousArray), :(ia::Int), :(sa1::Int), :(sa2::Int)]
        imargs1d = [:(parent(a)), :ia, :sa1]
        imargs2d = [:(parent(a)), :ia, :sa1, :sa2]
        eviewargs = [:(ellipview(a, i))]
        getoffsets = :(ia = offset(a) + 1)
        getstrides1 = :(sa1 = stride(a, 1)::Int)
        getstrides2 = :(sa1 = stride(a, 1)::Int; sa2 = stride(a, 2)::Int)
        contcol = :(sa1 == 1)
        nextcol = :(ia += sa2)
        pkargs = [:(parent(a)), :ia]
        pkargs1 = [:(parent(a)), :ia, :sa1]
        kargs = [:a, :ia]
        kargs1 = [:a, :ia, :sa1]
    elseif AN == 1
        imparams1d = [:(fun::Functor{1}), :(a::ContiguousArray), :(ia::Int), :(sa1::Int)]
        imparams2d = [:(fun::Functor{1}), :(a::ContiguousArray), :(ia::Int), :(sa1::Int), :(sa2::Int)]
        imargs1d = [:fun, :(parent(a)), :ia, :sa1]
        imargs2d = [:fun, :(parent(a)), :ia, :sa1, :sa2]
        eviewargs = [:fun, :(ellipview(a, i))]
        getoffsets = :(ia = offset(a) + 1)
        getstrides1 = :(sa1 = stride(a, 1)::Int)
        getstrides2 = :(sa1 = stride(a, 1)::Int; sa2 = stride(a, 2)::Int)
        contcol = :(sa1 == 1)
        nextcol = :(ia += sa2)
        pkargs = [:fun, :(parent(a)), :ia]
        pkargs1 = [:fun, :(parent(a)), :ia, :sa1]
        kargs = [:fun, :a, :ia]
        kargs1 = [:fun, :a, :ia, :sa1]
    elseif AN == 2 || AN == -2
        FN = AN == 2 ? 2 : 1
        imparams1d = [:(fun::Functor{$FN}), :(a::ContiguousArrOrNum), :(ia::Int), :(sa1::Int),
                                            :(b::ContiguousArrOrNum), :(ib::Int), :(sb1::Int)]
        imparams2d = [:(fun::Functor{$FN}), :(a::ContiguousArrOrNum), :(ia::Int), :(sa1::Int), :(sa2::Int),
                                            :(b::ContiguousArrOrNum), :(ib::Int), :(sb1::Int), :(sb2::Int)]
        imargs1d = [:fun, :(parent(a)), :ia, :sa1,
                          :(parent(b)), :ib, :sb1]
        imargs2d = [:fun, :(parent(a)), :ia, :sa1, :sa2,
                          :(parent(b)), :ib, :sb1, :sb2]
        eviewargs = [:fun, :(ellipview(a, i)), :(ellipview(b, i))]
        getoffsets = :(ia = offset(a) + 1;
                       ib = offset(b) + 1)
        getstrides1 = :(sa1 = stride(a, 1)::Int;
                        sb1 = stride(b, 1)::Int)
        getstrides2 = :(sa1 = stride(a, 1)::Int; sa2 = stride(a, 2)::Int;
                        sb1 = stride(b, 1)::Int; sb2 = stride(b, 2)::Int)
        contcol = :(sa1 == 1 && sb1 == 1)
        nextcol = :(ia += sa2; ib += sb2)
        pkargs = [:fun, :(parent(a)), :ia, :(parent(b)), :ib]
        pkargs1 = [:fun, :(parent(a)), :ia, :sa1, :(parent(b)), :ib, :sb1]
        kargs = [:fun, :a, :ia, :b, :ib]
        kargs1 = [:fun, :a, :ia, :sa1, :b, :ib, :sb1]
    elseif AN == 3
        imparams1d = [:(fun::Functor{3}), :(a::ContiguousArrOrNum), :(ia::Int), :(sa1::Int),
                                          :(b::ContiguousArrOrNum), :(ib::Int), :(sb1::Int),
                                          :(c::ContiguousArrOrNum), :(ic::Int), :(sc1::Int)]
        imparams2d = [:(fun::Functor{3}), :(a::ContiguousArrOrNum), :(ia::Int), :(sa1::Int), :(sa2::Int),
                                          :(b::ContiguousArrOrNum), :(ib::Int), :(sb1::Int), :(sb2::Int),
                                          :(c::ContiguousArrOrNum), :(ic::Int), :(sc1::Int), :(sc2::Int)]
        imargs1d = [:fun, :(parent(a)), :ia, :sa1,
                          :(parent(b)), :ib, :sb1,
                          :(parent(c)), :ic, :sc1]
        imargs2d = [:fun, :(parent(a)), :ia, :sa1, :sa2,
                          :(parent(b)), :ib, :sb1, :sb2,
                          :(parent(c)), :ic, :sc1, :sc2]
        eviewargs = [:fun, :(ellipview(a, i)), :(ellipview(b, i)), :(ellipview(c, i))]
        getoffsets = :(ia = offset(a) + 1;
                       ib = offset(b) + 1;
                       ic = offset(c) + 1)
        getstrides1 = :(sa1 = stride(a, 1)::Int;
                        sb1 = stride(b, 1)::Int;
                        sc1 = stride(c, 1)::Int)
        getstrides2 = :(sa1 = stride(a, 1)::Int; sa2 = stride(a, 2)::Int;
                        sb1 = stride(b, 1)::Int; sb2 = stride(b, 2)::Int;
                        sc1 = stride(c, 1)::Int; sc2 = stride(c, 2)::Int)
        contcol = :(sa1 == 1 && sb1 == 1 && sc1 == 1)
        nextcol = :(ia += sa2; ib += sb2; ic += sc2)
        pkargs = [:fun, :(parent(a)), :ia, :(parent(b)), :ib, :(parent(c)), :ic]
        pkargs1 = [:fun, :(parent(a)), :ia, :sa1, :(parent(b)), :ib, :sb1, :(parent(c)), :ic, :sc1]
        kargs = [:fun, :a, :ia, :b, :ib, :c, :ic]
        kargs1 = [:fun, :a, :ia, :sa1, :b, :ib, :sb1, :c, :ic, :sc1]
    else
        error("AN = $(AN) is unsupported.")
    end

    return CodegenHelperEx(imparams1d,
                           imparams2d,
                           imargs1d,
                           imargs2d,
                           eviewargs,
                           getoffsets,
                           getstrides1,
                           getstrides2,
                           contcol,
                           nextcol,
                           pkargs,
                           pkargs1,
                           kargs,
                           kargs1)
end


text:
module Monteiro2022Results
export julia_main

using JuMP, Ipopt, DataFrames, CSV, HCEstimator, JLD

include("./Scenario.jl")
import .Scenario

function load_sys_from_csv()::DataFrames.DataFrame
    return DataFrame(CSV.File("src/case33.csv"))
end

function build_substation()::HCEstimator.Substation
    return DistSystem.Substation(
        12660,              # Nominal voltage (V)
        1,                  # Bus
        1.02,                # Vˢᴮ: Voltage (p.u.)
        4.0,                # Pˢᴮ: Active power capacity (MW)
        2.5,                # Qˢᴮ: Reactive power capacity (MVAr)
        [0.003, 12, 240]    # Costs (Not used)
    )
end

function build_sys(data::DataFrames.DataFrame, add_der, hc_modes)::HCEstimator.System
    sub = build_substation()
    sys = DistSystem.factory_system(data, 0.93, 1.05, sub)
    sys = add_der(sys)
    sys.m_load = [0.6, 0.8, 1.0]
    sys.m_new_dg = hc_modes
    sys.time_curr = 1.0
    return sys
end

function set_options!(model)
    if solver_name(model) == "Ipopt"
        set_optimizer_attribute(model, "expect_infeasible_problem", "yes")
        set_optimizer_attribute(model, "timing_statistics", "yes")

        set_optimizer_attribute(model, "tol", 1e-5)
        set_optimizer_attribute(model, "constr_viol_tol", 1e-5)
        set_optimizer_attribute(model, "dual_inf_tol", 1e-5)
        set_optimizer_attribute(model, "compl_inf_tol", 1e-5)
        set_optimizer_attribute(model, "mumps_mem_percent", 1000)
        set_silent(model)
    end
    return model

end

function get_ders_operation!(sys, model, df, i, alpha)
    (Ω, bΩ, L, K, D, S) = Tools.Get.sets(sys)
    HC = objective_value(model)

    for d = D, l = L, k = K, s = S
        P = Tools.Get.power_active_DER(model, d, l, k, s)
        Q = Tools.Get.power_reactive_DER(model, d, l, k, s)
        push!(df, ("S$(i)", alpha, HC, P, Q, d, l, k, s))
    end
end

function julia_main(mode)::Cint
    println("################################## MODE: $(mode["name"]) ######################## ")
    println("Stated!")
    println("Loading data...")
    data = load_sys_from_csv()

    ders_scenario = [
        [
            Scenario.dg_dispached,
            Scenario.renewable,
            Scenario.ev,
            Scenario.ess
        ],
        [Scenario.dg_dispached, Scenario.renewable, Scenario.ev],
        [Scenario.dg_dispached, Scenario.renewable, Scenario.ess],
        [Scenario.dg_dispached, Scenario.renewable],
        [Scenario.dg_dispached],
        [Scenario.renewable, Scenario.ev],
        [Scenario.renewable, Scenario.ess],
        [Scenario.renewable],
        [],
    ]
    ders_scenario = reverse(ders_scenario)

    A = [0.0, 0.025, 0.05, 0.075, 0.1]

    df = DataFrame(
        Scenario=String[], alpha=Float64[], HC=Float64[],
        P=Float64[],
        Q=Float64[],
        der=Int[],
        l=Float64[],
        k=Float64[],
        s=Float64[],
    )

    for (i, ders) in enumerate(ders_scenario)
        # println("Scenario: ", i)
        for α in A
            # println("   with α = ", α)
            add_der(sys) = Scenario.scenario(α, ders)(sys)
            sys = build_sys(data, add_der, mode["hc"])
            # println("       with HC modes = ", sys.m_new_dg)
            model = Model(Ipopt.Optimizer)
            model = set_options!(model)
            println("       Building Model...")
            model = build_model(model, sys)
            println("       Solving...")
            @time optimize!(model)
            println("       STATUS: ", termination_status(model))
            if termination_status(model) == MOI.LOCALLY_SOLVED || termination_status(model) == MOI.OPTIMAL
                # println("       Hosting Capacity: ", round(objective_value(model), digits=6), " MVA")
                get_ders_operation!(sys, model, df, i, α)
                println(df[end, [:Scenario, :alpha, :HC]])
            else
                error("Model infeasible!")
            end

            println("   Finished!!")
        end
    end

    println("Salving results...")
    CSV.write("results/results_$(mode["name"]).csv", df)
    println("Exting...")
    return 0
end
end # module


text:
"""
Base struct for all DiffEqEnvironments.jl rewards.
"""
struct RewardFunction
    f
end

(rf::RewardFunction)(s,a,s_next) = rf.f(s, a, s_next)

"""
Default reward of form ``r(s,a,s')``.
"""
SASReward(r) = RewardFunction(r)

"""
Reward of type ``r(s,a)``.
"""
SAReward(r) = RewardFunction((s, a, _) -> r(s, a))

"""
Reward of type ``r(a,s')``.
"""
ASReward(r) = RewardFunction((_, a, s) -> r(a, s))

"""
Quadratic reward of type ``r(s,a)=-s^T\\mathbf{Q}s - s^T\\mathbf{Q}s``,
commonly used in optimal control / LQR.
"""
function QuadraticReward(Q, R)
    r(_, a, s) = - dot(s, Q, s) - dot(a, R, a)
    return RewardFunction(r)
end

"""
Reward ``r(s')=0``` if ``s'`` is terminal, ``r(s')=-1`` else
"""
function DecrementingReward(s_next, condition_fn)
    r(s, a, s_next) = condition_fn(s_next) ? 0 : -1
    return RewardFunction(r)
end


text:
"""
    inverse(ϕ, θ) -> Function

Extract a diffusivity function `D` from a solution to a semi-infinite one-dimensional nonlinear diffusion problem,
where the solution is given as a set of discrete points.

Interpolates the given solution with a PCHIP monotonic spline and uses the Bruce and Klute method to reconstruct `D`.

Due to the method used for interpolation, `D` will be continuous but will have discontinuous derivatives.

# Arguments
- `ϕ::AbstractVector`: values of the Boltzmann variable. See [`ϕ`](@ref).
- `θ::AbstractVector`: solution values at each point in `ϕ`.

# References
BRUCE, R. R.; KLUTE, A. The measurement of soil moisture diffusivity.
Soil Science Society of America Journal, 1956, vol. 20, no 4, p. 458-462.
"""
function inverse(ϕ::AbstractVector, θ::AbstractVector)
    @argcheck length(ϕ) ≥ 2
    @argcheck length(ϕ) == length(θ) DimensionMismatch

    indices = sortperm(θ)
    ϕ = ϕ[indices]
    θ = θ[indices]

    indices = unique(i -> θ[i], eachindex(θ))
    ϕ = ϕ[indices]
    θ = θ[indices]

    θi = θ[argmax(ϕ)]
    ϕ = Interpolator(θ, ϕ)

    let ϕ=ϕ, θi=θi
        function D(θ)
            dϕ_dθ = derivative(ϕ, θ)
            ∫ϕdθ = integrate(ϕ, θi, θ)
            return -(dϕ_dθ*∫ϕdθ)/2
        end
    end
end


text:
using LinearMaps, LinearAlgebra, Test

@testset "filled maps" begin
    M, N = 2, 3
    μ = rand()
    for λ in (true, false, 3, μ, μ + 2im)
        L = FillMap(λ, (M, N))
        @test L == FillMap(λ, M, N)
        @test occursin("$M×$N FillMap{$(typeof(λ))} with fill value: $λ", sprint((t, s) -> show(t, "text/plain", s), L))
        @test LinearMaps.MulStyle(L) === LinearMaps.FiveArg()
        A = fill(λ, (M, N))
        x = rand(typeof(λ) <: Real ? Float64 : ComplexF64, 3)
        X = rand(typeof(λ) <: Real ? Float64 : ComplexF64, 3, 4)
        w = similar(x, 2)
        W = similar(X, 2, 4)
        @test size(L) == (M, N)
        @test adjoint(L) == FillMap(adjoint(λ), (3,2))
        @test transpose(L) == FillMap(λ, (3,2))
        @test Matrix(L) == A
        @test L * x ≈ A * x
        @test mul!(w, L, x) ≈ A * x
        @test mul!(W, L, X) ≈ A * X
        for α in (true, false, 1, 0, randn()), β in (true, false, 1, 0, randn())
            @test mul!(copy(w), L, x, α, β) ≈ fill(λ * sum(x) * α, M) + w * β
            @test mul!(copy(W), L, X, α, β) ≈ λ * reduce(vcat, sum(X, dims=1) for _ in 1:2) * α + W * β
        end
    end
    @test issymmetric(FillMap(μ + 1im, (3, 3)))
    @test ishermitian(FillMap(μ + 0im, (3, 3)))
    @test isposdef(FillMap(μ, (1,1))) == isposdef(μ)
    @test !isposdef(FillMap(μ, (3,3)))
    α = rand()
    β = rand()
    @test FillMap(μ, (M, N)) + FillMap(α, (M, N)) == FillMap(μ + α, (M, N))
    @test FillMap(μ, (M, N)) - FillMap(α, (M, N)) == FillMap(μ - α, (M, N))
    @test α*FillMap(μ, (M, N)) == FillMap(α * μ, (M, N))
    @test FillMap(μ, (M, N))*α == FillMap(μ * α, (M, N))
    @test FillMap(μ, (M, N))*FillMap(μ, (N, M)) == FillMap(μ^2*N, (M, M))
    @test Matrix(FillMap(μ, (M, N))*FillMap(μ, (N, M))) ≈ fill(μ, (M, N))*fill(μ, (N, M))
end


text:
module Utils

using LightGraphs: LightGraphs

"""
    kron_alt(mats...)

Same as `kron`, but, at least for Pauli matrices, it is faster by up to a factor of
two. Also, `kron_alt(one_matrix)` returns `one_matrix` rather than throwing an error.
"""
function kron_alt(mats...)
    if length(mats) == 1
        # copy input. Fixes bug #51. This masks root cause of the bug, which has also been fixed.
        # But, this is the correct behavior in any case.
        return copy(only(mats))
    end
    if length(mats) < 6
        return kron(mats...)
    else
        return _kron(mats...)
    end
end

"""
    _kron(mats...)

This is at most 2x faster than `kron`. It is often slower, especially for n < 5 matrices.
`_kron` is called by `kron_alt` when the number of matrices in `mats` is greater than
six.
"""
function _kron(mats...)
    isempty(mats) && error("_kron requires at least one argument")
    n = length(mats)
    if n == 1
        return only(mats)
    elseif n <= 3
        return kron(mats...)
    end
    half_n = n ÷ 2
    krons = Vector{Any}(undef, half_n)
    for i in 1:half_n
        krons[i] = kron(mats[2*i-1], mats[2*i])
    end
    if ! iseven(n)
        krons[half_n] = kron(krons[half_n], last(mats))
    end
    return _kron(krons...)
end

"""
    isapprox_zero(x::Number)

Return `true` if `x` is approximately zero.
This function exists because we define methods for special cases, such as symbolic
libraries.
"""
function isapprox_zero(x::Number, atol=1e-16)
    return isapprox(x, zero(x), atol=atol)
end

## MIME{Symbol("text/input")} is meant to print objects in input
## form, that is Julia code that constructs the object
Base.show(m::MIME{Symbol("text/input")}, item) = show(stdout, m, item)
Base.show(io::IO, ::MIME{Symbol("text/input")}, item) = show(io, item)

"""
    pow_of_minus_one(n::Integer)

Returns `-1` to the power `n`.
"""
pow_of_minus_one(n::Integer) = iseven(n) ? 1 : -1

"""
    property_graph(vector, property_func)

Return a `LightGraphs.SimpleGraph` with `length(vector)` integer-valued vertices, and edges
between exactly those indices of `vector` for which `property_func(vector[i], vector[j])` is
`true`. No self-edges are included.
"""
function property_graph(vector, property_func)
    n_verts = length(vector)
    graph = LightGraphs.SimpleGraph(n_verts)
    @inbounds for i in 1:n_verts, j in i+1:n_verts
        if property_func(vector[i], vector[j])
            LightGraphs.add_edge!(graph, i, j)
        end
    end
    return graph
end

"""
    n_cliffords(n)

Returns the number of elements in the Clifford group on `n` qubits.
"""
n_cliffords(n) = 8 * prod(i -> 2*(4^i - 1) * 4^i, big(1):n)

"""
    triprod(iter)

Iterate over only the upper triangular part of the Cartesian product of `iter` with `iter`,
returning `Tuple`s of elements

# Examples
```jldocstring
julia> collect(triprod('a':'d'))
6-element Vector{Tuple{Char, Char}}:
 ('a', 'b')
 ('a', 'c')
 ('a', 'd')
 ('b', 'c')
 ('b', 'd')
 ('c', 'd')
```
"""
triprod(iter) = ((a, b) for (n, a) in enumerate(iter) for b in Iterators.drop(iter, n))

end # module Utils