text:
@parallel function diffusion3D_step!(T2, T, Ci, lam, dt, dx, dy, dz)
    @inn(T2) = @inn(T) + dt*(lam*@inn(Ci)*(@d2_xi(T)/dx^2 + @d2_yi(T)/dy^2 + @d2_zi(T)/dz^2));
    return
end

function diffusion3D()
# Physics
lam        = 1.0;                                        # Thermal conductivity
cp_min     = 1.0;                                        # Minimal heat capacity
lx, ly, lz = 10.0, 10.0, 10.0;                           # Length of computational domain in dimension x, y and z

# Numerics
nx, ny, nz = 8, 8, 8;                              # Number of gridpoints in dimensions x, y and z
nt         = 3;                                          # Number of time steps
dx         = lx/(nx-1);                                  # Space step in x-dimension
dy         = ly/(ny-1);                                  # Space step in y-dimension
dz         = lz/(nz-1);                                  # Space step in z-dimension

# Array initializations
T   = @zeros(nx, ny, nz);
T2  = @zeros(nx, ny, nz);
Ci  = @zeros(nx, ny, nz);

# Initial conditions (heat capacity and temperature with two Gaussian anomalies each)
Ci .= 1.0./( cp_min .+ Data.Array([5*exp(-(((ix-1)*dx-lx/1.5))^2-(((iy-1)*dy-ly/2))^2-(((iz-1)*dz-lz/1.5))^2) +
                                   5*exp(-(((ix-1)*dx-lx/3.0))^2-(((iy-1)*dy-ly/2))^2-(((iz-1)*dz-lz/1.5))^2) for ix=1:size(T,1), iy=1:size(T,2), iz=1:size(T,3)]) )
T  .= Data.Array([100*exp(-(((ix-1)*dx-lx/2)/2)^2-(((iy-1)*dy-ly/2)/2)^2-(((iz-1)*dz-lz/3.0)/2)^2) +
                   50*exp(-(((ix-1)*dx-lx/2)/2)^2-(((iy-1)*dy-ly/2)/2)^2-(((iz-1)*dz-lz/1.5)/2)^2) for ix=1:size(T,1), iy=1:size(T,2), iz=1:size(T,3)])
T2 .= T;                                                 # Assign also T2 to get correct boundary conditions.

# Time loop
dt = min(dx^2,dy^2,dz^2)*cp_min/lam/8.1;                 # Time step for the 3D Heat diffusion
for it = 1:nt
    @parallel diffusion3D_step!(T2, T, Ci, lam, dt, dx, dy, dz);
    T, T2 = T2, T;
end

return true
end


text:
"""
    RectangularCluster <: NeutralLandscapeMaker

    RectangularCluster(; minimum=2, maximum=4)
    RectangularCluster(minimum, [maximum=4])

Fills the landscape with rectangles containing a random value. The size of each
rectangle/patch is between `minimum` and `maximum` (the two can be equal for a
fixed size rectangle).
"""
@kwdef struct RectangularCluster <: NeutralLandscapeMaker
    minimum::Integer = 2
    maximum::Integer = 4
    function RectangularCluster(minimum::T, maximum::T=4) where {T <: Integer}
        @assert 0 < minimum <= maximum
        new(minimum, maximum)
    end
end

function _landscape!(mat, alg::RectangularCluster)
    mat .= -1.0
    while minimum(mat) == -1.0
        width, height = rand(alg.minimum:alg.maximum, 2)
        row = rand(1:(size(mat,1)-(width-1)))
        col = rand(1:(size(mat,2)-(height-1)))
        mat[row:(row+(width-1)) , col:(col+(height-1))] .= rand()
    end
    return mat
end

ext: jl


text:
using ModelingToolkit, Test
@parameters α β δ
expr = (((1 / β - 1) + δ) / α) ^ (1 / (α - 1))
ref = [β, δ, α]
sol = ModelingToolkit.get_variables(expr)
@test all([sol[i].op.name == ref[i].op.name for i in 1:3])

@parameters γ
s = α => γ
expr = (((1 / β - 1) + δ) / α) ^ (1 / (α - 1))
ModelingToolkit.substitute(expr, s)
new = (((1 / β - 1) + δ) / γ) ^ (1 / (γ - 1))
@test isequal(expr, new)


text:
"""
    SimpleErrorModel <: ErrorModel

Abstract supertype for simple IID error models that generate errors based on the number of
qubits and a probability distribution.
"""
abstract type SimpleErrorModel <: ErrorModel end

"""
    generate(
        error_model::SimpleErrorModel, code, p::Real, [rng::AbstractRNG=GLOBAL_RNG]
    ) -> BitVector

Generate a new IID error based on [`Model.probability_distribution`](@ref). See also
[`Model.generate`](@ref).

!!! note

    The method [`Model.probability_distribution`](@ref) should be implemented for concrete
    subtypes of [`SimpleErrorModel`](@ref).

"""
function Model.generate(
    error_model::SimpleErrorModel, code, p::Real, rng::AbstractRNG=GLOBAL_RNG
)
    n_qubits = nkd(code)[1]
    weights = ProbabilityWeights(collect(probability_distribution(error_model, p)))
    error_pauli = String(sample(rng, ['I','X','Y','Z'], weights, n_qubits))
    return to_bsf(error_pauli)
end

@doc raw"""
    BitFlipErrorModel <: SimpleErrorModel

IID error model with probability vector: ``(p_I, p_X, p_Y p_Z) = (1-p, p, 0, 0)``,
where ``p`` is the probability of an error on a single-qubit.
"""
struct BitFlipErrorModel <: SimpleErrorModel end
Model.label(::BitFlipErrorModel) = "Bit-flip"
function Model.probability_distribution(::BitFlipErrorModel, p::Real)
    return (1 - p, p, zero(p), zero(p))
end

@doc raw"""
    BitPhaseFlipErrorModel <: SimpleErrorModel

IID error model with probability vector: ``(p_I, p_X, p_Y p_Z) = (1-p, 0, p, 0)``,
where ``p`` is the probability of an error on a single-qubit.
"""
struct BitPhaseFlipErrorModel <: SimpleErrorModel end
Model.label(::BitPhaseFlipErrorModel) = "Bit-phase-flip"
function Model.probability_distribution(::BitPhaseFlipErrorModel, p::Real)
    return (1 - p, zero(p), p, zero(0))
end

@doc raw"""
    DepolarizingErrorModel <: SimpleErrorModel

IID error model with probability vector: ``(p_I, p_X, p_Y p_Z) = (1-p, p/3, p/3, p/3)``,
where ``p`` is the probability of an error on a single-qubit.
"""
struct DepolarizingErrorModel <: SimpleErrorModel end
Model.label(::DepolarizingErrorModel) = "Depolarizing"
function Model.probability_distribution(::DepolarizingErrorModel, p::Real)
    px = py = pz = p / 3
    pi = 1 - sum((px, py, pz))
    return (pi, px, py, pz)
end

@doc raw"""
    PhaseFlipErrorModel <: SimpleErrorModel

IID error model with probability vector: ``(p_I, p_X, p_Y p_Z) = (1-p, 0, 0, p)``,
where ``p`` is the probability of an error on a single-qubit.
"""
struct PhaseFlipErrorModel <: SimpleErrorModel end
Model.label(::PhaseFlipErrorModel) = "Phase-flip"
function Model.probability_distribution(::PhaseFlipErrorModel, p::Real)
    return (1 - p, zero(p), zero(p), p)
end


text:
"""
    innereltype(x)

Recursively determine the 'innermost' type in by the collection `x` (which may be, for example,
a collection of a collection).
"""
function innereltype(x)
  T = eltype(x)
  T <: AbstractArray ? innereltype(T) : return T
end

"""
    cxtype(T)

Returns `T` when `T` is `Complex`, or `Complex{T}` when `T` is `Real`.
"""
cxtype(::Type{T}) where T<:Number = T
cxtype(::Type{T}) where T<:Real = Complex{T}

"""
    fltype(T)

Returns `T` when `T<:AbstractFloat` or `Tf` when `T<:Complex{Tf}`.
"""
fltype(::Type{T})          where T<:AbstractFloat = T
fltype(::Type{Complex{T}}) where T<:AbstractFloat = T
fltype(T::Tuple) = fltype(T[1])

"""
    superzeros(T, A)

Returns an array like `A`, but full of zeros. If `innereltype(A)` can be promoted to `T`, then
the innermost elements of the array will have type `T`.
"""
superzeros(T, A::AbstractArray) = T(0)*A
superzeros(A::AbstractArray) = superzeros(innereltype(A), A)
superzeros(T, dims::Tuple) = eltype(dims) <: Tuple ? [ superzeros(T, d) for d in dims ] : zeros(T, dims)
superzeros(dims::Tuple) = superzeros(Float64, dims) # default
superzeros(T::Tuple, dims::Tuple) = [ superzeros(T[i], dims[i]) for i=1:length(dims) ]

"""
    @superzeros T a b c d...
    @superzeros T dims b c d...

Generate arrays `b, c, d...` with the super-dimensions of `a` and innereltype `T`.
"""
macro superzeros(T, ad, vars...)
  expr = Expr(:block)
  append!(expr.args, [:( $(esc(var)) = superzeros($(esc(T)), $(esc(ad))); ) for var in vars])

  return expr
end
supersize(a) = Tuple([size(ai) for ai in a])
supersize(a::Array{T}) where T<:Number = size(a)

macro createarrays(T, dims, vars...)
  expr = Expr(:block)
  append!(expr.args, [:( $(esc(var)) = zeros($(esc(T)), $(esc(dims))); ) for var in vars])

  return expr
end

"""
    @zeros T dims a b c...

Create arrays of all zeros with element type `T`, size `dims`, and global names
`a`, `b`, `c` (for example). An arbitrary number of arrays may be created.
"""
macro zeros(T, dims, vars...)
  expr = Expr(:block)
  append!(expr.args, [:( $(esc(var)) = zeros($(esc(T)), $(esc(dims))); ) for var in vars])

  return expr
end

Base.zeros(::CPU, T, dims) = zeros(T, dims)

"""
    devzeros(dev, T, dims)

Returns an array like `A` of type `T`, but full of zeros.
"""
devzeros(dev, T, dims) = zeros(dev, T, dims)

"""
    @devzeros dev T dims a b c...

Create arrays of all zeros with element type `T`, size `dims`, and global names
`a`, `b`, `c` (for example) on device `dev`.
"""
macro devzeros(dev, T, dims, vars...)
  expr = Expr(:block)
  append!(expr.args, [:( $(esc(var)) = zeros($(esc(dev))(), $(esc(T)), $(esc(dims))); ) for var in vars])

  return expr
end

"""
    parsevalsum2(uh, grid)

Returns `Σ |uh|²` on the `grid`, which is equal to the domain integral of `u`. More specifically,
it returns
```math
\\sum_{𝐤} |û_{𝐤}|² L_x L_y = \\int u(𝐱)² \\, 𝖽x 𝖽y \\,,
```
where ``û_{𝐤} =`` `uh` ``/(`` `grid.nx` ``e^{- i 𝐤 ⋅ 𝐱₀})``, with ``𝐱₀`` the vector with components
the left-most position in each direction.
"""
function parsevalsum2(uh, grid::TwoDGrid)
  if size(uh, 1) == grid.nkr # uh is in conjugate symmetric form
    U = sum(abs2, uh[1, :])           # k=0 modes
    U += 2*sum(abs2, uh[2:end, :])    # sum k>0 modes twice
  else # count every mode once
    U = sum(abs2, uh)
  end

  norm = grid.Lx * grid.Ly / (grid.nx^2 * grid.ny^2) # normalization for dft

  return norm * U
end

function parsevalsum2(uh, grid::OneDGrid)
  if size(uh, 1) == grid.nkr                 # uh is conjugate symmetric
    U = sum(abs2, CUDA.@allowscalar uh[1])   # k=0 modes
    U += @views 2 * sum(abs2, uh[2:end])     # sum k>0 modes twice
  else # count every mode once
    U = sum(abs2, uh)
  end

  norm = grid.Lx / grid.nx^2 # normalization for dft

  return norm * U
end

"""
    parsevalsum(uh, grid)

Returns `real(Σ uh)` on the `grid`, i.e.
```math
ℜ [ \\sum_{𝐤} û_{𝐤} L_x L_y ] \\,,
```
where ``û_{𝐤} =`` `uh` ``/(`` `grid.nx` ``e^{- i 𝐤 ⋅ 𝐱₀})``, with ``𝐱₀`` the vector with components
the left-most position in each direction.
"""
function parsevalsum(uh, grid::TwoDGrid)
  if size(uh, 1) == grid.nkr    # uh is conjugate symmetric
    U = sum(uh[1, :])           # k=0 modes
    U += 2*sum(uh[2:end, :])    # sum k>0 modes twice
  else # count every mode once
    U = sum(uh)
  end

  norm = grid.Lx * grid.Ly / (grid.nx^2 * grid.ny^2) # normalization for dft
  return norm * real(U)
end

"""
    jacobianh(a, b, grid)

Returns the Fourier transform of the Jacobian of `a` and `b` on `grid`.
"""
function jacobianh(a, b, grid::TwoDGrid)
  if eltype(a) <: Real
    bh = rfft(b)
    bx = irfft(im * grid.kr .* bh, grid.nx)
    by = irfft(im * grid.l  .* bh, grid.nx)

    return im * grid.kr .* rfft(a .* by) .- im * grid.l .* rfft(a .* bx)
  else
    # J(a, b) = ∂(a * ∂b/∂y)/∂x - ∂(a * ∂b/∂x)/∂y
    bh = fft(b)
    bx = ifft(im * grid.k .* bh)
    by = ifft(im * grid.l .* bh)

    return im * grid.k .* fft(a .* by) .- im * grid.l .* fft(a .* bx)
  end
end

"""
    jacobian(a, b, grid)

Returns the Jacobian of `a` and `b` on `grid`.
"""
function jacobian(a, b, grid::TwoDGrid)
  if eltype(a) <: Real
    return irfft(jacobianh(a, b, grid), grid.nx)
  else
    return ifft(jacobianh(a, b, grid))
  end
end

"""
    radialspectrum(fh, grid; n=nothing, m=nothing, refinement=2)

Returns the radial spectrum of `fh`. `fh` lives on Cartesian wavenumber grid ``(k, l)``. To
compute the radial spectrum, we first interpolate ``f̂(k, l)`` onto a radial wavenumber grid
``(ρ, θ)``, where ``ρ² = k²+l²`` and ``θ = \\arctan(l/k)``. Note here that
``f̂ =`` `fh` ``/(`` `grid.nx` ``e^{- i 𝐤 ⋅ 𝐱₀})``,  with ``𝐱₀`` the vector with components the
left-most position in each direction. After interpolation, we integrate ``f̂``over angles ``θ``
to get `fρ`,

```math
  f̂_ρ = \\int f̂(ρ, θ) ρ 𝖽ρ 𝖽θ \\, .
```

The resolution `(n, m)` for the polar wavenumber grid is `n = refinement * maximum(nk, nl),
m = refinement * maximum(nk, nl)`, where `refinement = 2` by default. If `fh` is in conjugate
symmetric form then only the upper-half plane in ``θ`` is represented on the polar grid.
"""
function radialspectrum(fh, grid::TwoDGrid; n=nothing, m=nothing, refinement=2)

  n = n == nothing ? refinement * maximum([grid.nk, grid.nl]) : n
  m = m == nothing ? refinement * maximum([grid.nk, grid.nl]) : m

  # Calculate the shifted k and l
  lshift = range(-grid.nl/2+1, stop=grid.nl/2, length=grid.nl) * 2π/grid.Ly

  if size(fh)[1] == grid.nkr # conjugate symmetric form
    m = Int(m/2)                         # => half resolution in θ
    θ = range(-π/2, stop=π/2, length=m)  # θ-grid from k=0 to max(kr)
    fhshift = fftshift(fh, 2)            # shifted fh
    kshift = range(0, stop=grid.nkr-1, length=grid.nkr) * 2π/grid.Lx
  else # ordinary form
    θ = range(0, stop=2π, length=m)      # θ grid
    fhshift = fftshift(fh, [1, 2])       # shifted fh
    kshift = range(-grid.nk/2+1, stop=grid.nk/2, length=grid.nk) * 2π/grid.Lx
  end

  # Interpolator for fh
  itp = scale(interpolate(fhshift, BSpline(Linear())), kshift, lshift)

  # Get radial wavenumber vector
  ρmax = minimum([(grid.nk/2-1) * 2π/grid.Lx, (grid.nl/2-1) * 2π/grid.Ly])
  ρ = range(0, stop=ρmax, length=n)

  # Interpolate fh onto fine grid in (ρ, θ).
  fhρθ = zeros(eltype(fhshift), (n, m))

  for i₁=2:n, i₂=1:m # ignore zeroth mode; i₁≥2
    fhρθ[i₁, i₂] = itp(ρ[i₁] * cos(θ[i₂]), ρ[i₁] * sin(θ[i₂]))
  end

  # fhρ = ρ ∫ fh(ρ, θ) dθ  =>  Fh = ∫ fhρ dρ = ∫∫ fh dk dl
  dθ = θ[2] - θ[1]
  if size(fh)[1] == grid.nkr
    fhρ = 2ρ .* sum(fhρθ, dims=2) * dθ # multiply by 2 for conjugate symmetry
  else
    fhρ =  ρ .* sum(fhρθ, dims=2) * dθ
  end

  CUDA.@allowscalar fhρ[1] = fh[1, 1] # zeroth mode

  return ρ, fhρ
end

"""
    on_grid(func, grid)

Returns an array, of the ArrayType of the device `grid` lives on, that contains the values of
function `func` evaluated on the `grid`.
"""
on_grid(func, grid::OneDGrid{T, A}) where {T, A} = CUDA.@allowscalar A([func(grid.x[i]) for i=1:grid.nx])

on_grid(func, grid::TwoDGrid{T, A}) where {T, A} = CUDA.@allowscalar A([func(grid.x[i], grid.y[j]) for i=1:grid.nx, j=1:grid.ny])

on_grid(func, grid::ThreeDGrid{T, A}) where {T, A} = CUDA.@allowscalar A([func(grid.x[i], grid.y[j], grid.z[k]) for i=1:grid.nx, j=1:grid.ny, k=1:grid.nz])

"""
    ArrayType(::Device)
    ArrayType(::Device, T, dim)

Returns the proper array type according to the Device chosen, i.e., `Array` for CPU and
`CuArray` for GPU.
"""
ArrayType(::CPU) = Array
ArrayType(::CPU, T, dim) = Array{T, dim}


text:
"""
    WignerHsize(m′ₘₐₓ, ℓₘₐₓ=-2)

Total size of array of wedges of width m′ₘₐₓ up to ℓₘₐₓ

Parameters
----------
ℓₘₐₓ : int
m′ₘₐₓ : int, optional
    If nothing, it is assumed to be at least ℓ

See Also
--------
WignerHrange : Array of (ℓ, m', m) indices corresponding to this wedge
WignerHindex : Index inside these wedges

Notes
-----
Here, it is assumed that only data with m≥|m'| are stored, and only
corresponding values are passed.  We also assume |m|≤ℓ and |m'|≤ℓ.  Neither
of these are checked.  The wedge array that this function indexes is ordered as

    [
        H(ℓ, m′, m) for ℓ in range(ℓₘₐₓ+1)
        for m′ in range(-min(ℓ, m′ₘₐₓ), min(ℓ, m′ₘₐₓ)+1)
        for m in range(abs(m′), ℓ+1)
    ]

"""
function WignerHsize(m′ₘₐₓ, ℓₘₐₓ=-2)
    if ℓₘₐₓ == -2
        ℓₘₐₓ = m′ₘₐₓ
    elseif ℓₘₐₓ < 0
        return 0
    end
    if m′ₘₐₓ === nothing || m′ₘₐₓ >= ℓₘₐₓ
        return (ℓₘₐₓ+1) * (ℓₘₐₓ+2) * (2*ℓₘₐₓ+3) ÷ 6
    else
        return ((ℓₘₐₓ+1) * (ℓₘₐₓ+2) * (2*ℓₘₐₓ+3) - 2*(ℓₘₐₓ-m′ₘₐₓ)*(ℓₘₐₓ-m′ₘₐₓ+1)*(ℓₘₐₓ-m′ₘₐₓ+2)) ÷ 6
    end
end

"""
    WignerHrange(m′ₘₐₓ, ℓₘₐₓ=-1)

Create an array of (ℓ, m', m) indices as in H array

Parameters
----------
ℓₘₐₓ : int
m′ₘₐₓ : int, optional
    If nothing, it is assumed to be at least ℓ

See Also
--------
WignerHsize : Total size of wedge array
WignerHindex : Index inside these wedges

Notes
-----
Here, it is assumed that only data with m≥|m'| are stored, and only
corresponding values are passed.  We also assume |m|≤ℓ and |m'|≤ℓ.  Neither of
these are checked.  The wedge array that this function indexes is ordered as

    [
        H(ℓ, m′, m) for ℓ in range(ℓₘₐₓ+1)
        for m′ in range(-min(ℓ, m′ₘₐₓ), min(ℓ, m′ₘₐₓ)+1)
        for m in range(abs(m′), ℓ+1)
    ]

"""
function WignerHrange(m′ₘₐₓ, ℓₘₐₓ=-1)
    if ℓₘₐₓ < 0
        ℓₘₐₓ = m′ₘₐₓ
    end
    r = zeros(typeof(m′ₘₐₓ), (WignerHsize(m′ₘₐₓ, ℓₘₐₓ), 3))
    i = 1
    for ℓ in 0:ℓₘₐₓ
        for m′ in -min(ℓ, m′ₘₐₓ):min(ℓ, m′ₘₐₓ)
            for m in abs(m′):ℓ
                r[i, 1] = ℓ
                r[i, 2] = m′
                r[i, 3] = m
                i += 1
            end
        end
    end
    return r
end

function _WignerHindex(ℓ, m′, m, m′ₘₐₓ)
    """Helper function for `WignerHindex`"""
    m′ₘₐₓ = min(m′ₘₐₓ, ℓ)
    i = WignerHsize(m′ₘₐₓ, ℓ-1)  # total size of everything with smaller ℓ
    if m′<1
        i += (m′ₘₐₓ + m′) * (2*ℓ - m′ₘₐₓ + m′ + 1) ÷ 2  # size of wedge to the left of m'
    else
        i += (m′ₘₐₓ + 1) * (2*ℓ - m′ₘₐₓ + 2) ÷ 2  # size of entire left half of wedge
        i += (m′ - 1) * (2*ℓ - m′ + 2) ÷ 2  # size of right half of wedge to the left of m'
    end
    i += m - abs(m′)  # size of column in wedge between m and |m'|
    return i + 1
end

"""
    WignerHindex(ℓ, m′, m, m′ₘₐₓ)

Index to "wedge" arrays

Parameters
----------
ℓ : int
m′ : int
m : int
m′ₘₐₓ : int, optional
    If nothing, it is assumed to be at least ℓ

See Also
--------
WignerHsize : Total size of wedge array
WignerHrange : Array of (ℓ, m', m) indices corresponding to this wedge

Notes
-----
Here, it is assumed that only data with m≥|m'| are stored, and only
corresponding values are passed.  We also assume |m|≤ℓ and |m'|≤ℓ.  Neither
of these are checked.  The wedge array that this function indexes is ordered as

    [
        H(ℓ, m′, m) for ℓ in range(ℓₘₐₓ+1)
        for m′ in range(-min(ℓ, m′ₘₐₓ), min(ℓ, m′ₘₐₓ)+1)
        for m in range(abs(m′), ℓ+1)
    ]

"""
function WignerHindex(ℓ, m′, m, m′ₘₐₓ=nothing)
    if ℓ == 0
        return 1
    end
    m′max = ℓ
    if m′ₘₐₓ !== nothing
        m′max = min(m′ₘₐₓ, m′max)
    end
    if m < -m′
        if m < m′
            return _WignerHindex(ℓ, -m′, -m, m′max)
        else
            return _WignerHindex(ℓ, -m, -m′, m′max)
        end
    else
        if m < m′
            return _WignerHindex(ℓ, m, m′, m′max)
        else
            return _WignerHindex(ℓ, m′, m, m′max)
        end
    end
end

"""
    WignerDsize(ℓₘᵢₙ, m′ₘₐₓ, ℓₘₐₓ=-1)

Compute total size of Wigner 𝔇 matrix

Parameters
----------
ℓₘᵢₙ : int
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ
m′ₘₐₓ : int, optional
    Integer satisfying 0 <= m′ₘₐₓ.  Defaults to ℓₘₐₓ.
ℓₘₐₓ : int
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ

Returns
-------
i : int
    Total size of Wigner 𝔇 matrix arranged as described below

See Also
--------
WignerDrange : Array of (ℓ, m', m) indices corresponding to the 𝔇 matrix
WignerDindex : Index of a particular element of the 𝔇 matrix

Notes
-----
This assumes that the Wigner 𝔇 matrix is arranged as

    [
        𝔇(ℓ, m′, m)
        for ℓ in range(ℓₘᵢₙ, ℓₘₐₓ+1)
        for m′ in range(-min(ℓ, m′ₘₐₓ), min(ℓ, m′ₘₐₓ)+1)
        for m in range(-ℓ, ℓ+1)
    ]

"""
function WignerDsize(ℓₘᵢₙ, m′ₘₐₓ, ℓₘₐₓ=-1)
    # from sympy import symbols, summation, horner
    # from sympy.printing.pycode import pycode
    # ℓ,m′,m,ℓₘᵢₙ,ℓₘₐₓ,m′ₘₐₓ = symbols('ℓ,m′,m,ℓₘᵢₙ,ℓₘₐₓ,m′ₘₐₓ', integer=True)
    #
    # def nice(expr)
    #     return horner(expr.expand().simplify(), (m′ₘₐₓ, ℓₘᵢₙ, ℓₘₐₓ))
    #
    # # Assuming ℓₘᵢₙ <= ℓₘₐₓ <= m′ₘₐₓ:
    # WignerDsize_ℓmin_ℓmax_m′max = horner(
    #     (
    #         # sum over all m′=0 elements
    #         summation(summation(summation(1, (m, -ℓ, ℓ)), (m′, 0, 0)), (ℓ, ℓₘᵢₙ, ℓₘₐₓ))
    #         # sum over all ℓₘᵢₙ <= ℓ <= ℓₘₐₓ elements
    #         + summation(summation(summation(2, (m, -ℓ, ℓ)), (m′, 1, ℓ)), (ℓ, ℓₘᵢₙ, ℓₘₐₓ))
    #     ).expand().simplify(),
    #     (ℓₘᵢₙ, ℓₘₐₓ)
    # )
    # print(f"({pycode(nice(3*WignerDsize_ℓmin_ℓmax_m′max.subs(ℓₘₐₓ, ℓ-1)))}) ÷ 3")
    #
    # # Assuming ℓₘᵢₙ <= m′ₘₐₓ <= ℓₘₐₓ:
    # WignerDsize_ℓmin_m′max_ℓmax = horner(
    #     (
    #         # sum over all m′=0 elements
    #         summation(summation(summation(1, (m, -ℓ, ℓ)), (m′, 0, 0)), (ℓ, ℓₘᵢₙ, ℓₘₐₓ))
    #         # sum over all ℓ <= m′ₘₐₓ elements
    #         + summation(summation(summation(2, (m, -ℓ, ℓ)), (m′, 1, ℓ)), (ℓ, ℓₘᵢₙ, m′ₘₐₓ))
    #         # sum over all ℓ >= m′ₘₐₓ elements
    #         + summation(summation(summation(2, (m, -ℓ, ℓ)), (m′, 1, m′ₘₐₓ)), (ℓ, m′ₘₐₓ+1, ℓₘₐₓ))
    #     ).expand().simplify(),
    #     (m′ₘₐₓ, ℓₘᵢₙ, ℓₘₐₓ)
    # )
    # print(f"({pycode(nice(3*WignerDsize_ℓmin_m′max_ℓmax.subs(ℓₘₐₓ, ℓ-1)))}) ÷ 3")
    #
    # # Assuming m′ₘₐₓ <= ℓₘᵢₙ <= ℓₘₐₓ:
    # WignerDsize_m′max_ℓmin_ℓmax = horner(
    #     (
    #         # sum over all m′=0 elements
    #         summation(summation(summation(1, (m, -ℓ, ℓ)), (m′, 0, 0)), (ℓ, ℓₘᵢₙ, ℓₘₐₓ))
    #         # sum over all remaining |m′| <= m′ₘₐₓ elements
    #         + summation(summation(summation(2, (m, -ℓ, ℓ)), (m′, 1, m′ₘₐₓ)), (ℓ, ℓₘᵢₙ, ℓₘₐₓ))
    #     ).expand().simplify(),
    #     (m′ₘₐₓ, ℓₘᵢₙ, ℓₘₐₓ)
    # )
    # print(f"{pycode(nice(WignerDsize_m′max_ℓmin_ℓmax.subs(ℓₘₐₓ, ℓ-1)).factor())}")
    if ℓₘₐₓ < 0
        ℓₘₐₓ = m′ₘₐₓ
    end
    if m′ₘₐₓ >= ℓₘₐₓ
        return (
            ℓₘₐₓ * (ℓₘₐₓ * (4 * ℓₘₐₓ + 12) + 11)
            + ℓₘᵢₙ * (1 - 4 * ℓₘᵢₙ^2)
            + 3
        ) ÷ 3
    end
    if m′ₘₐₓ > ℓₘᵢₙ
        return (
            3 * ℓₘₐₓ * (ℓₘₐₓ + 2)
            + ℓₘᵢₙ * (1 - 4 * ℓₘᵢₙ^2)
            + m′ₘₐₓ * (
                3 * ℓₘₐₓ * (2 * ℓₘₐₓ + 4)
                + m′ₘₐₓ * (-2 * m′ₘₐₓ - 3) + 5
            )
            + 3
        ) ÷ 3
    else
        return (ℓₘₐₓ * (ℓₘₐₓ + 2) - ℓₘᵢₙ^2) * (1 + 2 * m′ₘₐₓ) + 2 * m′ₘₐₓ + 1
    end
end

"""
    WignerDrange(ℓₘᵢₙ, m′ₘₐₓ, ℓₘₐₓ=-1)

Create an array of (ℓ, m', m) indices as in 𝔇 array

Parameters
----------
ℓₘᵢₙ : int
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ
m′ₘₐₓ : int, optional
    Integer satisfying 0 <= m′ₘₐₓ.  Default is ℓₘₐₓ.
ℓₘₐₓ : int
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ

See Also
--------
WignerDsize : Total size of 𝔇 array
WignerDindex : Index inside these wedges

Notes
-----
This assumes that the Wigner 𝔇 matrix is arranged as

    [
        𝔇(ℓ, m′, m)
        for ℓ in range(ℓₘᵢₙ, ℓₘₐₓ+1)
        for m′ in range(-min(ℓ, m′ₘₐₓ), min(ℓ, m′ₘₐₓ)+1)
        for m in range(-ℓ, ℓ+1)
    ]

"""
function WignerDrange(ℓₘᵢₙ, m′ₘₐₓ, ℓₘₐₓ=-1)
    if ℓₘₐₓ < 0
        ℓₘₐₓ = m′ₘₐₓ
    end
    r = zeros(typeof(ℓₘᵢₙ), (WignerDsize(ℓₘᵢₙ, m′ₘₐₓ, ℓₘₐₓ), 3))
    i = 1
    for ℓ in ℓₘᵢₙ:ℓₘₐₓ
        for m′ in -min(ℓ, m′ₘₐₓ):min(ℓ, m′ₘₐₓ)
            for m in -ℓ:ℓ
                r[i, 1] = ℓ
                r[i, 2] = m′
                r[i, 3] = m
                i += 1
            end
        end
    end
    return r
end

"""
    WignerDindex(ℓ, m′, m, ℓₘᵢₙ=0, m′ₘₐₓ=-1)

Compute index into Wigner 𝔇 matrix

Parameters
----------
ℓ : int
    Integer satisfying ℓₘᵢₙ <= ℓ <= ℓₘₐₓ
m′ : int
    Integer satisfying -min(ℓₘₐₓ, m′ₘₐₓ) <= m′ <= min(ℓₘₐₓ, m′ₘₐₓ)
m : int
    Integer satisfying -ℓ <= m <= ℓ
ℓₘᵢₙ : int, optional
    Integer satisfying 0 <= ℓₘᵢₙ <= ℓₘₐₓ.  Defaults to 0.
m′ₘₐₓ : int, optional
    Integer satisfying 0 <= m′ₘₐₓ.  Defaults to ℓ.

Returns
-------
i : int
    Index into Wigner 𝔇 matrix arranged as described below

See Also
--------
WignerDsize : Total size of the 𝔇 matrix
WignerDrange : Array of (ℓ, m', m) indices corresponding to the 𝔇 matrix

Notes
-----
This assumes that the Wigner 𝔇 matrix is arranged as

    [
        𝔇(ℓ, m′, m)
        for ℓ in range(ℓₘᵢₙ, ℓₘₐₓ+1)
        for m′ in range(-min(ℓ, m′ₘₐₓ), min(ℓ, m′ₘₐₓ)+1)
        for m in range(-ℓ, ℓ+1)
    ]

"""
function WignerDindex(ℓ, m′, m, ℓₘᵢₙ=0, m′ₘₐₓ=-1)
    if m′ₘₐₓ < 0
        m′ₘₐₓ = ℓ
    end
    i = (m′ + min(m′ₘₐₓ, ℓ)) * (2 * ℓ + 1) + m + ℓ
    if ℓ > ℓₘᵢₙ
        i += WignerDsize(ℓₘᵢₙ, m′ₘₐₓ, ℓ-1)
    end
    return i + 1
end


text:
################################################################################
#
# Functional objective for FWI
#
################################################################################

export objFWI!

function objFWI!(F,
				 G::Union{Array{R, 1}, Nothing},
				 m::Array{R, 2},
				 h::NTuple{2, R},
				 freqs::Array{R, 1}, x_rcv::Array{R, 1}, z_rcv::R,
				 x_src::Array{R, 1}, z_src::R,
				 src_enc::Array{C, 2},
				 dat::Array{Array{C, 2}, 1};
				 nabc::Union{Nothing, NTuple{4, Int64}} = nothing,
				 gradprec_fun = g->g,
				 bounds::Union{NTuple{2, R}, NTuple{2, Array{R, 2}}, Nothing} = nothing)

    # Initializing objective value
    J = R(0)

    # Zero-ing gradient
    if G != nothing
        G .= R(0)
    end

	# Projecting m to box constrained set
	m = proj_bounds(m, bounds)

    # Updating obj value and gradient (loop over frequencies)
    n = size(m)
    Ns = size(src_enc, 2)
    Nf = length(freqs)
    for idxf = 1:Nf

		# Encoding data
		dat_enc = dat[idxf]*src_enc

		# Data normalization factor
		ηd = sqrt.(sum(abs.(dat_enc).^R(2), dims = 1))

         # Set absorbing boundary size (frequency-dependent)
        if nabc == nothing
            nabc = computeNabc(h, m, freqs[idxf])
        end

        # Setting source/receiver restriction/injection operators & source weigths
        Pr = restrOpMat(n, h, x_rcv, z_rcv, nabc)
        Ps = restrOpMat(n, h, x_src, z_src, nabc)
        Fsrc = Matrix(adjoint(Ps)*src_enc)

        # Solve for wavefield (forward)
        H = lu(HelmholtzLinOp(h, m, freqs[idxf], nabc))
        U = C.(H\Fsrc)

        # Receiver restriction
        res = dat_enc-Pr*U

		# Objective value update
        if F != nothing
            J += (sum(sum(abs.(res).^R(2), dims = 1)./ηd.^R(2), dims = 2)/(Nf*Ns))[1]
        end

		# Gradient update
        if G != nothing
            V = C.(adjoint(H)\(adjoint(Pr)*res)) # Solve for wavefield (adjoint)
            g = -R(2).*(R(2).*π.*freqs[idxf]).^R(2).*restrict(reshape(sum(real.(conj.(U).*V)./ηd.^R(2), dims = 2), sz_abc(n, nabc)), nabc)./(Nf*Ns)
            G .+= gradprec_fun(g)
        end

    end

    if F != nothing
        return J
    end

end


text:
# Communication with Maxima
export
    MaximaError,
    MaximaSyntaxError,
    mcall,
    mcallj,
    @mx,
    ==

import
    Base: ==

struct MaximaError <: Exception
    errstr::String
end
Base.showerror(io::IO, err::MaximaError) = print(io, err.errstr)

struct MaximaSyntaxError <: Exception
    errstr::String
end
Base.showerror(io::IO, err::MaximaSyntaxError) = print(io, err.errstr)

"""
    mcall(m::MExpr)

Evaluate a Maxima expression.

## Examples
```julia
julia> m\"integrate(sin(x), x)\"

                             integrate(sin(x), x)

julia> mcall(ans)

                                   - cos(x)

```
"""
function mcall(m::MExpr)
    write(ms, convert(String, m))
    output = read(ms)
    if occursin(maxerr, output)
        write(ms.input, "errormsg()\$")
        write(ms.input, "print(ascii(4))\$")
        message = read(ms)
        throw(MaximaError(message))
    elseif occursin(synerr, output)
        throw(MaximaSyntaxError(output))
    else
        r = String[]
        for k in split(output, ETX)
            push!(r, replace(k, r"[ \n]"=>""))
        end
        return normalize(MExpr(r))
    end
end
"""
  mcall(expr::T) where T

Evaluate a Julia expression or string using the Maxima interpretor.

Type `T` can be a String or a Julia expression

## Examples
```julia
julia> mcall(\"integrate(sin(y)^2, y)\")

julia> mcall(:(integrate(1/(1+x^2), x)))

```
"""
mcall(expr::T) where T = mcall(MExpr(expr))

# Convert a Julia statement saved in `js` to a MExpr
mcallj(js::AbstractString) = mcall(MExpr(Meta.parse(js)))

function ==(m::MExpr, n::MExpr)
    r, s = normalize(m).strv("$x : $s")ec, normalize(n).strvec
    length(r) != length(s) && (return false)

    return all(j=>!occursin("false", mcall("is($(r[j]) = $(s[j]))")), 1:length(r))
end

macro mx(x, args...)
    r = MExpr[]
    if x == :use
        @assert !isempty(args) "Incomplete arguments for @mx"
        for a in args[1:end]
            s = jexpr_to_mstr(Main.eval(a))
            push!(r, mcall("$a : $s"))
        end
        return (length(r) == 1 ? r[1] : r)
    elseif x == :julia_export
        # TODO: not implemented yet
        @assert !isempty(args) "Incomplete arguments for @mx"
        x = args[1]
        return mcall(MExpr(x))
    end
    isempty(args) && return mcall(MExpr(x))

    # each element can be a Maxima statement, or a Maxima equation label
    argv = append!(Any[x], args)
    arglen = length(argv)
    i = 1
    while i <= arglen
        if argv[i].head === :(=)
            if i+1 <= arglen && argv[i+1] isa Symbol
                # equation definition
                push!(r, mcall(MExpr(string(argv[i+1])*": "*jexpr_to_mstr(x, defeqn=true))))
                i += 2
            else
                # variable definition
                push!(r, mcall(jexpr_to_mstr(x)))
                i += 1
            end
        else
            # function call or others
            push!(r, mcall(jexpr_to_mstr(x)))
            i += 1
        end
    end
    return (length(r) == 1 ? r[1] : r)
end


text:
import Baysor

Baysor.run_cli(["--help"]);
Baysor.run_cli(["run", "--help"]);
Baysor.run_cli(["preview", "--help"]);

# To build, run create_sysimage(:Baysor; precompile_execution_file="$(dirname(pathof(Baysor)))/../bin/build.jl", replace_default=true)
# If you need to build app, comment @require part in __init__ function of ImageCore (~/.julia/packages/ImageCore/BTvmx/src/ImageCore.jl:167)
# and use create_app instead of create_sysimage

max_stars_repo_name: YinuoJin/Baysor

max_issues_repo_name: YinuoJin/Baysor

max_forks_repo_name: YinuoJin/Baysor


text:
#=
jinc.jl
=#

export jinc

using SpecialFunctions: besselj1

"""
    jinc(x)

Return `jinc(x) = J1(π*x)/(2x)`,
where `J1` is a Bessel function of the first kind.

The argument `x` must be unitless.

Return type is `promote_type(typeof(x), Float32)`.
"""
function jinc(x::X) where {X <: Real}
    T = promote_type(X, Float32)
    if (x == 0) return convert(T, π/4) end
    y = abs(x)
    convert(T, besselj1(π*y) / (2*y))
end


text:
# https://github.com/jheinen/GR.jl/issues/278#issuecomment-587090846
ENV["GKSwstype"] = "nul"
# some of the tutorials cannot be run on the GPU
ENV["CLIMATEMACHINE_SETTINGS_DISABLE_GPU"] = true
# avoid problems with Documenter/Literate when using `global_logger()`
ENV["CLIMATEMACHINE_SETTINGS_DISABLE_CUSTOM_LOGGER"] = true

using Distributed
using DocumenterCitations
bib = CitationBibliography(joinpath(@__DIR__, "bibliography.bib"))

@everywhere push!(LOAD_PATH, joinpath(@__DIR__, ".."))
@everywhere using ClimateMachine
@everywhere using Documenter, Literate

@everywhere const clima_dir = dirname(dirname(pathof(ClimateMachine)));

@everywhere GENERATED_DIR = joinpath(@__DIR__, "src", "generated") # generated files directory
rm(GENERATED_DIR, force = true, recursive = true)
mkpath(GENERATED_DIR)

include("list_of_getting_started_docs.jl")      # defines `getting_started_docs`

include("list_of_tutorials.jl")                 # defines `tutorials`

include("list_of_how_to_guides.jl")             # defines `how_to_guides`
include("list_of_apis.jl")                      # defines `apis`
include("list_of_theory_docs.jl")               # defines `theory_docs`
include("list_of_dev_docs.jl")                  # defines `dev_docs`

pages = Any[
    "Home" => "index.md",
    "Getting started" => getting_started_docs,
    "Tutorials" => tutorials,
    "How-to-guides" => how_to_guides,
    "APIs" => apis,
    "Contribution guide" => "Contributing.md",
    "Theory" => theory_docs,
    "Developer docs" => dev_docs,
    "References" => "References.md",
]

mathengine = MathJax(Dict(
    :TeX => Dict(
        :equationNumbers => Dict(:autoNumber => "AMS"),
        :Macros => Dict(),
    ),
))

format = Documenter.HTML(
    prettyurls = get(ENV, "CI", "") != "" ? true : false,
    mathengine = mathengine,
    collapselevel = 1,
    analytics = get(ENV, "CI", "") != "" ? "UA-191640394-1" : "",
)

makedocs(
    bib,
    sitename = "ClimateMachine",
    doctest = false,
    strict = true,
    format = format,
    checkdocs = :exports,
    clean = true,
    modules = [ClimateMachine],
    pages = pages,
)

include("clean_build_folder.jl")

deploydocs(
    repo = "github.com/CliMA/ClimateMachine.jl.git",
    target = "build",
    push_preview = true,
    forcepush = true,
)


text:
"""
Output structure
"""
mutable struct Solution{T1<:AbstractFace, T2<:Vector{Float64}, T3<:Float64}
    F::T1
    Fnew::T1
    c::T2
    cnew::T2
    feas::T3
    feasnew::T3
    r::T2
    function Solution(M::AbstractLinearOp, b::Vector{Float64}, A::AbstractAtomicSet)
        z = M'*b
        F = face(A, z)
        Fnew = deepcopy(F)
        c = zeros( size(M*F, 2))
        cnew = deepcopy(c)
        feas = norm(b)^2/2
        feasnew = deepcopy(feas)
        r = deepcopy(b)
        new{AbstractFace, Vector{Float64},Float64}(F, Fnew, c, cnew, feas, feasnew, r)
    end
end

function update!(sol::Solution)
    sol.F = deepcopy(sol.Fnew)
    sol.c = deepcopy(sol.cnew)
    sol.feas = sol.feasnew
end

function constructPrimal(sol::Solution)
    if typeof(sol.F) <: SumFace
        rs = [rank(face) for face in sol.F.faces]
        cs = BlockArray(sol.c, rs)
        xs = [sol.F.faces[i]*cs[Block(i)] for i in 1:length(sol.F.faces)]
        return xs
    else
        return sol.F * sol.c
    end
end


text:
"""
# References
[1] M. Islam, M. Okasha, and M. M. Idres, “Dynamics and control of quadcopter using linear model predictive control approach,” IOP Conf. Ser.: Mater. Sci. Eng., vol. 270, p. 012007, Dec. 2017, doi: 10.1088/1757-899X/270/1/012007.
# Variables
u ∈ R^4: square of rotor angular speed
"""
Base.@kwdef struct IslamQuadcopter <: Quadcopter
    J = Diagonal([7.5e-3, 7.5e-3, 1.3e-2])  # kg m^2
    l = 0.23  # m
    kf = 3.13e-5  # N s^2
    kM = 7.5e-7  # N m s^2
    m = 0.65  # kg
    g = 9.81  # m / s^2
    B = [kf  kf  kf  kf;
          0 -kf   0  kf;
         kf   0 -kf   0;
         kM -kM  kM -kM]
    # drags
    # kt = Diagonal(0.1*ones(3))  # Ns / m
    # kr = Diagonal(0.1*ones(3))  # N m s
    # input limits
    dim_input = 4
    u_min = zeros(dim_input)
    u_max = (m*g/kf) * ones(dim_input)
end

function saturate(env::IslamQuadcopter, u)
    @unpack u_min, u_max = env
    u_saturated = clamp.(u, u_min, u_max)
end

function input_to_force_moment(env::IslamQuadcopter, u)
    @unpack B = env
    ν = B * u
end


text:
#===============================
<-----------Lathe.jl----------->
~~~~~~ 0.1.8 "ButterBall" ~~~~~~
Programmed by Emmett Boudreau
    <emmett@emmettboudreau.com>
        <http://emmettboudreau.com>
MIT General Open Source License
    (V 3.0.0)
        Free for Modification and
        Redistribution
=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=|
         CONTRIBUTORS
=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=|
        ~ emmettgb
        ~ stefanches7
        ~ PallHaraldsson
/><><><><><><><><><><><><><><><><\
Thank you for your forks!
<-----------Lathe.jl----------->
38d8eb38-e7b1-11e9-0012-376b6c802672
#[deps]
DataFrames.jl
Random.jl
================================#
@doc """
# Lathe.jl 0.1.8 Butterball
## Easily ML
38d8eb38-e7b1-11e9-0012-376b6c802672

### [contributors]
- emmettgb
- stefanches7
- PallHaraldsson
### [deps]
- DataFrames.jl
- Random.jl
- Distributions.jl
### Modules
**Use ?(Module) for more information!**
- stats
- preprocess
- models
       """
module Lathe
#
# <------- DEPS ----->
using DataFrames
using Random
# <------- DEPS ----->
#
# <------- PARTS ----->
include("Stats.jl")
include("Models.jl")
include("Preprocess.jl")
# <------- PARTS ----->
#
end


text:
abstract type AbstractGrid{T, N} <: AbstractArray{T, N} end

"""
    T: tick type
    N: N dimensional
    S: System (Cartesian, Cylindrical...)
"""
struct Grid{T, N, S} <: AbstractGrid{T, N}
    axes::NTuple{N, DiscreteAxis{T}}
end

const CartesianGrid{T, N} = Grid{T, N, :cartesian}
const CartesianGrid1D{T} = CartesianGrid{T, 1}
const CartesianGrid2D{T} = CartesianGrid{T, 2}
const CartesianGrid3D{T} = CartesianGrid{T, 3}
const RadialGrid{T} = Grid{T, 1, :Radial}
const PolarGrid{T} = Grid{T, 2, :Polar}
const CylindricalGrid{T} = Grid{T, 3, :cylindrical}
const SphericalGrid{T} = Grid{T, 3, :Spherical}

@inline size(g::Grid{T, N, S}) where {T, N, S} = size.(g.axes, 1)
@inline length(g::Grid{T, N, S}) where {T, N, S} = prod(size(g))
@inline getindex(g::Grid{T, N, S}, I::Vararg{Int, N}) where {T, N, S} = broadcast(getindex, g.axes, I)
@inline getindex(g::Grid{T, N, S}, i::Int) where {T, N, S} = getproperty(g, :axes)[i]
@inline getindex(g::Grid{T, N, S}, s::Symbol) where {T, N, S} = getindex(g, Val{s}())

@inline getindex(g::CylindricalGrid{T}, ::Val{:r}) where {T} = g.axes[1]
@inline getindex(g::CylindricalGrid{T}, ::Val{:φ}) where {T} = g.axes[2]
@inline getindex(g::CylindricalGrid{T}, ::Val{:z}) where {T} = g.axes[3]
@inline getindex(g::CartesianGrid{T}, ::Val{:x}) where {T} = g.axes[1]
@inline getindex(g::CartesianGrid{T}, ::Val{:y}) where {T} = g.axes[2]
@inline getindex(g::CartesianGrid{T}, ::Val{:z}) where {T} = g.axes[3]

function sizeof(g::Grid{T, N, S}) where {T, N, S}
    return sum( sizeof.(g.axes) )
end

function print(io::IO, g::Grid{T, N, S}) where {T, N, S}
    for ax in g.axes
        print(io, ax, " | ")
    end
end
function println(io::IO, g::Grid{T, N, S}) where {T, N, S}
    for (i, ax) in enumerate(g.axes)
        println(io, "Axis $(i): ", ax)
    end
end
show(io::IO, g::Grid{T, N, S}) where {T, N, S} = println(io, g)
show(io::IO, ::MIME"text/plain", g::Grid{T, N, S}) where {T, N, S} = show(io, g)

function check_grid(grid::CylindricalGrid{T})::Nothing where {T}
    nr::Int, nφ::Int, nz::Int = size(grid)
    @assert iseven(nz) "GridError: Field simulation algorithm in cylindrical coordinates need an even number of grid points in z. This is not the case. #z-ticks = $(nz)."
    @assert (iseven(nφ) || (nφ == 1)) "GridError: Field simulation algorithm in cylindrical coordinates need an even number of grid points in φ or just one point (2D). This is not the case. #φ-ticks = $(nφ)."
    return nothing
end

function check_grid(grid::CartesianGrid3D{T})::Nothing where {T}
    nx::Int, ny::Int, nz::Int = size(grid)
    @assert iseven(nx) "GridError: Field simulation algorithm in cartesian coordinates need an even number of grid points in x. This is not the case. #x-ticks = $(nx)."
    return nothing
end

include("RefineGrid.jl")

@recipe function f(grid::Grid{T, N, S}) where {T, N, S}
    layout --> N
    st --> :stephist
    legend --> false

    for iax in 1:N
        @series begin
            subplot := iax
            nbins --> div(length(grid[iax]), 2)
            xlabel --> "Grid point density - Axis $(iax)"
            grid[iax]
        end
    end
end

function get_coordinate_type(grid::Grid{T, N, S}) where {T, N, S}
    return S
end
function get_number_of_dimensions(grid::Grid{T, N, S}) where {T, N, S}
    return N
end
function eltype(grid::Grid{T, N, S})::DataType where {T, N, S}
    return T
end

function get_boundary_types(grid::Grid{T, N, S}) where {T, N, S}
   return get_boundary_types.(grid.axes)
end

function Grid(nt::NamedTuple)
    if nt.coordtype == "Cylindrical"
        axr::DiscreteAxis = DiscreteAxis(nt.axes.r, unit=u"m")
        axφ::DiscreteAxis = DiscreteAxis(nt.axes.phi, unit=u"rad")
        axz::DiscreteAxis = DiscreteAxis(nt.axes.z, unit=u"m")
        T = typeof(axr.ticks[1])
        return Grid{T, 3, :cylindrical}( (axr, axφ, axz) )
    else
        error("`coordtype` = $(nt.coordtype) is not valid.")
    end
end

Base.convert(T::Type{Grid}, x::NamedTuple) = T(x)

function NamedTuple(grid::Grid{T, 3, :cylindrical}) where {T}
    axr::DiscreteAxis{T} = grid[:r]
    axφ::DiscreteAxis{T} = grid[:φ]
    axz::DiscreteAxis{T} = grid[:z]
    return (
        coordtype = "Cylindrical",
        ndims = 3,
        axes = (
            r = NamedTuple(axr, unit=u"m"),
            phi = NamedTuple(axφ, unit=u"rad"),
            z = NamedTuple(axz, unit=u"m"),
        )
    )
end

Base.convert(T::Type{NamedTuple}, x::Grid) = T(x)


text:
const IMGUI_DEMO_MENU_OPEN_RESENT_STATES = Menu(
    "Open Recent",
    [
        MenuItem("fish_hat.c"),
        MenuItem("fish_hat.inl"),
        MenuItem("fish_hat.h"),
        Menu(
            "More..",
            [
                MenuItem("Hello"),
                MenuItem("Sailor"),
                Menu("Recurse..", [], false)
            ],
            true,
        )
    ],
    true,
)

struct ItemSeparator <: MenuItems.AbstractGenericMenuItem end
(x::ItemSeparator)() = CImGui.Separator()

mutable struct MenuOptions <: MenuItems.AbstractGenericMenuItem
    enabled::Bool
    f::Cfloat
    n::Cint
    b::Bool
end
function (x::MenuOptions)()
    if CImGui.BeginMenu("Options")
        CImGui.CSyntax.@c CImGui.MenuItem("Enabled", "", &x.enabled)

        CImGui.BeginChild("child", CImGui.ImVec2(0, 60), true)
            foreach(i->CImGui.Text("Scrolling Text $i"), 0:9)
        CImGui.EndChild()

        CImGui.CSyntax.@c CImGui.SliderFloat("Value", &x.f, 0.0, 1.0)
        CImGui.CSyntax.@c CImGui.InputFloat("Input", &x.f, 0.1)
        CImGui.CSyntax.@c CImGui.Combo("Combo", &x.n, "Yes\0No\0Maybe\0\0")
        CImGui.CSyntax.@c CImGui.Checkbox("Check", &x.b)

        CImGui.EndMenu()
    end
end

struct MenuColors <: MenuItems.AbstractGenericMenuItem end
function (x::MenuColors)()
    if CImGui.BeginMenu("Colors")
        sz = CImGui.GetTextLineHeight()
        for i = 0:CImGui.ImGuiCol_COUNT-1
            name = CImGui.GetStyleColorName(i)
            p = CImGui.GetCursorScreenPos()
            CImGui.AddRectFilled(CImGui.GetWindowDrawList(), p, (p.x+sz,p.y+sz), CImGui.GetColorU32(i))
            CImGui.Dummy(sz, sz)
            CImGui.SameLine()
            CImGui.MenuItem(name)
        end
        CImGui.EndMenu()
    end
end

const IMGUI_DEMO_MENU_FILE_STATES = Menu(
    "File",
    [
        MenuItem("(dummy menu)", "", false),
        MenuItem("New"),
        MenuItem("Open", "Ctrl+O"),
        IMGUI_DEMO_MENU_OPEN_RESENT_STATES,
        MenuItem("Save", "Ctrl+S"),
        MenuItem("Save As.."),
        ItemSeparator(),
        MenuOptions(true, 0.5, 0, true),
        MenuColors(),
        ToggleMenuItem("Checked", "", true),
        MenuItem("Quit", "Alt+F4"),
    ],
    true,
)

const IMGUI_DEMO_MENU_EDIT_STATES = Menu(
    "Edit",
    [
        MenuItem("Undo", "CTRL+Z"),
        MenuItem("Redo", "CTRL+Y", false),
        ItemSeparator(),
        MenuItem("Cut", "CTRL+X"),
        MenuItem("Copy", "CTRL+C"),
        MenuItem("Paste", "CTRL+V"),
    ],
    true,
)

const IMGUI_DEMO_MENU_STATES = [
    IMGUI_DEMO_MENU_FILE_STATES,
    IMGUI_DEMO_MENU_EDIT_STATES,
]

const IMGUI_DEMO_MAIN_MENUBAR_STATES = MainMenuBar(
    "Example: MainMenuBar",
    IMGUI_DEMO_MENU_STATES,
    false,
)


text:
using OpenSpiel
using StatsBase
using Test
using CxxWrap
using Random

function evaluate_bots(
    state::Union{Ptr{Nothing}, CxxPtr{<:State}, CxxWrap.StdLib.UniquePtrAllocated{State}},
    bots::Vector{<:Bot},
    seed::Int)

    rng = MersenneTwister(seed)

    for bot in bots
        restart_at(bot, state)
    end

    while !is_terminal(state)
        if is_chance_node(state)
            outcomes_with_probs = chance_outcomes(state)
            actions, probs = zip(outcomes_with_probs...)
            action = actions[sample(rng, weights(collect(probs)))]
            apply_action(state, action)
        elseif is_simultaneous_node(state)
            chosen_actions = [
                legal_actions(state, pid)[pid+1] ? step(bot, state) : INVALID_ACTION
                for (pid, bot) in enumerate(bots)
            ]  # in julia, index starts with 1
            apply_action(state, chosen_actions)
        else
            apply_action(state, step(bots[current_player(state) + 1], state))
        end
    end
    returns(state)
end

@testset "OpenSpiel.jl" begin
    include("games_api.jl")
    include("games_simulation.jl")
    include("bots.jl")
    include("cfr.jl")
    include("trajector.jl")
end


text:
using PyCall
using PyFOOOF
using Test

# test that this works
fooof = PyFOOOF.fooof

# this example is adapted from the FOOOF tutorial
# https://fooof-tools.github.io/fooof/auto_tutorials/plot_07-TroubleShooting.html

FOOOF = fooof.FOOOF
# test that this works
FOOOFGroup = PyFOOOF.FOOOFGroup

param_sampler = fooof.sim.params.param_sampler
gen_power_spectrum = fooof.sim.gen.gen_power_spectrum
gen_group_power_spectra = fooof.sim.gen.gen_group_power_spectra
# XXX this fails for fooof.sim.utils.set_random_seed
# it seems to be a pseudo-bug in upstream: the python module's init
# doesn't import ("include") everything from that dir, but CPython
# seems to be able to cope:
# https://github.com/fooof-tools/fooof/blob/f7b4ed8d074a3eaa03707c2c3c09413e6ff74192/fooof/sim/__init__.py#L6
set_random_seed =  pyimport("fooof.sim.utils").set_random_seed

f_range = [1, 50]
ap_params = [20, 2]
# XXX PyCall's aggressive conversion to NumPy arrays creates problems here
# so we force it be treated as raw Python code (py-prefix)
# and NOT be converted to a native Julia array (o-suffix)
gauss_params = py"[[10, 1.0, 2.5], [20, 0.8, 2], [32, 0.6, 1]]"
nlv = 0.1
set_random_seed(21)

freqs, spectrum = gen_power_spectrum(f_range, ap_params, gauss_params, nlv)

fm = FOOOF(; peak_width_limits=py"[1, 8]"o, max_n_peaks=6, min_peak_height=0.4)
fm.fit(freqs, spectrum)

ground_truth = float.(py"$(gauss_params)")
actual_estimates = fm.gaussian_params_
expected_estimates = [9.82  0.91  2.43; 20.03  0.82  1.94; 31.85  0.57  1.11]
@test actual_estimates ≈ expected_estimates atol=0.01


text:
#!/usr/bin/env julia
#=
# ToDo
# * Text Processing, remove stopwords and punctuations
# * Tokenize words and sentences
# * Word frequency
# * Senctence score
# * Find the largest score
# * Join sentences
#
# include("lib/sampletext.jl")     ## providing SampleDocument
# tokenize(SampleDocument)        ## tokenize words
# split_sentences(SampleDocument) ## tokenize sentences

julia> include("sample-code/document-summarization.jl")
Main.TextSummarizer

julia> varinfo(TextSummarizer)
name                 size summary
–––––––––––––– –––––––––– –––––––––––––––––––––––––––––––––––––
TextSummarizer 55.644 KiB Module
summarize         0 bytes typeof(Main.TextSummarizer.summarize)
=#

"""
include("lib/document-summarization.jl")

usage: TextSummarizer.summarize(SampleDocument)
Returns summary of SampleDocument.
"""
module TextSummarizer
  export summarize

  using WordTokenizers

  include("lib/stopwords.jl")

  function word_frequencies(docx::AbstractString)
    word_freq = Dict()
    for word in tokenize(docx)
      if word in default_english_stopwords ## from stopwords.jl
        continue
      end
      if ! in(word, keys(word_freq))
        word_freq[word] = 1
      else
        word_freq[word] += 1
      end
    end
    word_freq
  end

  function relative_word_frequencies(docx::AbstractString)
    word_freq = word_frequencies(docx)
    max_freq  = maximum(values(word_freq))
    for word in keys(word_freq)
      word_freq[word] = (word_freq[word]/max_freq)
    end
    word_freq
  end

  function sentence_scores(docx::AbstractString, word_freq::Dict)
    _sentence_scores = Dict()
    words = keys(word_freq)
    for sentence in split_sentences(docx)
      for word in tokenize(lowercase(sentence))
        if ! in(word, words)
          continue
        end
        if length(split(sentence," ")) > 30
          continue
        end
        if in(sentence, keys(_sentence_scores))
          _sentence_scores[sentence] += word_freq[word]
        else
          _sentence_scores[sentence] = word_freq[word]
        end
      end
    end
    _sentence_scores
  end

  """
  summarize(docx::AbstractString)

  usage: summarize(SampleDocument)
  Return a summary of SampleDocument using extractive summary technique.
  """
  function summarize(docx::AbstractString)
    _word_freq = relative_word_frequencies(docx)
    _sentence_scores = sentence_scores(docx, _word_freq)
    ## finding N-largest using sort or datastructures heap nlargest
    summary_sentences = sort(
                            collect(
                                    zip(values(_sentence_scores), keys(_sentence_scores))
                                    ),
                            rev=true)
    summarized_list = [ w[2] for w in summary_sentences ]
    join(summarized_list, " ")
  end

end


text:
module Inference

using GZip, JLD2, FileIO
using LinearAlgebra, Statistics, StatsBase

include("io.jl")
using .DataIO

include("mixtures.jl")
using .Mixtures

# ------------------------------------------------------------------------
# globals

Maybe{T} = Union{T, Missing}
rank(x) = invperm(sortperm(x))

# ------------------------------------------------------------------------
# helper functions

function cumulative(data)
    d = sort(data)
    function F(x)
        i = searchsortedfirst(d, x)
        return (i-1)/length(d)
    end

    return F
end

const columns(genes) = Dict(g=>i for (i,g) ∈ enumerate(genes))
const badgenes = Set{String}(["CG14427", "cenG1A", "CG13333", "CG31670", "CG8965", "HLHm5", "Traf1"])

function cohortdatabase(stage::Int)
    cohort   = load("$root/drosophila/bdntp/database.jld2", "cohort")
    keepgene = [g ∉ badgenes for g in cohort[stage].gene]

    left = cohort[stage].point[:,2] .≥ 0
    gene = cohort[stage].gene[keepgene]
    return (
        expression = (
            real = cohort[stage].data[left,findall(keepgene)],
            data = hcat((rank(col)/length(col) for col in eachcol(cohort[stage].data[left,findall(keepgene)]))...),
            gene = columns(gene),
        ),
        position = cohort[stage].point[left,:],
        gene = gene
    )
end

"""
    virtualembryo(;directory="/home/nolln/mnt/data/drosophila/dvex")

Load the Berkeley Drosophila Transcriptional Network Project database.
`directory` should be path to folder containing two folders:
  1. bdtnp.txt.gz    : gene expression over point cloud of virtual cells
  2. geometry.txt.gz : spatial position (x,y,z) of point cloud of virtual cells.
"""
function virtualembryo(;directory="/home/nolln/mnt/data/drosophila/dvex")
    expression, _, genes = GZip.open("$directory/bdtnp.txt.gz") do io
        read_matrix(io; named_cols=true)
    end

    positions, _, _, = GZip.open("$directory/geometry_reduced.txt.gz") do io
        read_matrix(io; named_cols=true)
    end

    return (
        expression = (
            real = expression,
            data = hcat(fitmixture.(eachcol(expression))...),
            gene = columns(genes),
        ),
        position  = positions
    )
end

function scrna()
    expression, genes, _ = GZip.open("$root/drosophila/dvex/dge_normalized.txt.gz") do io
        read_matrix(io; named_cols=true, named_rows=true)
    end

    return (
        data = expression',
        gene = columns(genes),
    )
end

function match(x, y)
    index = Array{String}(undef,length(x))
    for (g,i) in x
        index[i] = g
    end
    return [k ∈ keys(y) ? y[k] : nothing for k in index]
end

# ------------------------------------------------------------------------
# main functions

"""
    cost(ref, qry; α=1, β=1, γ=0, ω=nothing)

Return the cost matrix ``J_{i\alpha}`` associated to matching cells in `qry` to cells in `ref`.
The cost matrix is computed by a heuristic distance between quantiles.
Deprecated.
"""
function cost(ref, qry; α=1, β=1, γ=0, ω=nothing)
    ϕ = match(ref.gene, qry.gene)

    Σ = zeros(size(ref.data,1), size(qry.data,1))
    for i in 1:size(ref.data,2)
        isnothing(ϕ[i]) && continue

        r  = ref.data[:,i]
        q  = qry.data[:,ϕ[i]]
        ω₀ = isnothing(ω) ? 1 : ω[i]

        f = sum(q .== 0) / length(q)
        χ = quantile(r, f)

        F₀ = cumulative(r[r.≤χ])
        F₊ = cumulative(r[r.>χ])
        F₌ = cumulative(q[q.>0])

        for j in 1:size(ref.data,1)
            for k in 1:size(qry.data,1)
                if r[j] > χ && q[k] > 0
                    Σ[j,k] += -ω₀*(2*F₊(r[j])-1)*(2*F₌(q[k])-1)
                elseif r[j] > χ && q[k] == 0
                    Σ[j,k] += +ω₀*(α*F₊(r[j])+γ)
                elseif r[j] <= χ && q[k] > 0
                    Σ[j,k] += +ω₀*(α*F₌(q[k])+γ)
                else
                    Σ[j,k] += +ω₀*β*F₀(r[j])
                end
            end
        end
    end

    return Matrix(Σ), ϕ
end

"""
    cost_simple(ref, qry)

Return the cost matrix ``J_{i\\alpha}`` associated to matching cells in `qry` to cells in `ref`.
The cost matrix is computed by hamming distance between cells via transforming quantiles to continuous spin variables.
Deprecated.
"""
function cost_simple(ref, qry)
    ϕ = match(ref.gene, qry.gene)
    Σ = zeros(size(ref.data,1), size(qry.data,1))

    σ(x) = x
    for i in 1:size(ref.data,2)
        isnothing(ϕ[i]) && continue

        r  = ref.real[:,i]
        q  = qry.data[:,ϕ[i]]

        R = 2*σ.((rank(r)./length(r))) .- 1
        Q = 2*σ.((rank(q)./length(q))) .- 1

        Σ -= (R*Q')

        #=
        for j in 1:size(ref.data,1)
            for k in 1:size(qry.data,1)
                Σ[j,k] += -*(2*σ(R(r[j]).^4)-1)*(2*σ(Q(q[k]).^4)-1)
            end
        end
        =#
    end

    return Matrix(Σ), ϕ
end

function cost_scan(ref, qry, ν, ω)
    ϕ = match(ref.gene, qry.gene)
    Σ = zeros(size(ref.data,1), size(qry.data,1))

    # XXX: try out the other option???
    # k    = 1
    # σ(x) = 1/(1+((1-x)/x)^k)
    σ(x) = x

    for i in 1:size(ref.data,2)
        isnothing(ϕ[i]) && continue

        r = ref.real[:,i]
        q = qry.data[:,ϕ[i]]

        # χ = σ.(rank(q)./(length(q)).^ν[i])

        R = (2 .* rank(r)./length(r)) .- 1
        Q = (2 .* σ.((rank(q)./length(q)).^ν[i])) .- 1
        # Q = (2 .* χ) .- 1

        Σ -= ω[i]*(R*Q')
    end

    return Matrix(Σ), ϕ
end

"""
    transform(src, dst, ν)

Transform distribution `src` to distribution `dst` by minimizing the Wasserstein metric.
This is equivalent to ``x \\to F^{-1}_{dst}\\left(F_{src}\\left(x\\right)\\right)`` where ``F`` denotes the cumulative density function.
"""
function transform(src, dst, ν)
    ref = sort(dst)
    pos = collect(1:length(dst))/length(dst)

    σ(x) = 1/(1+((1-x)/x)^ν)
    qry  = σ.(rank(src) / length(src))

    return [
        let
            i = searchsorted(pos, q)
            if first(i) == last(i)
                ref[first(i)]
            elseif last(i) == 0
                ref[1]
            else
                @assert first(i) > last(i)
                δy = ref[first(i)] - ref[last(i)]
                δx = pos[first(i)] - pos[last(i)]
                δq = q - pos[last(i)]

                ref[last(i)] + (δy/δx)*δq
            end
        end for q in qry
    ]
end

"""
    cost_transform(ref, qry; ω=nothing, ν=nothing)

Return the cost matrix ``J_{i\\alpha}`` associated to matching cells in `qry` to cells in `ref`.
The cost matrix is computed by:
  1. Transforming the `qry` distribution to the `ref` distribution.
  2. Looking at the SSE across transformed genes.
Use this unless you know what you are doing.
"""
function cost_transform(ref, qry; ω=nothing, ν=nothing)
    ϕ = match(ref.gene, qry.gene)
    Σ = zeros(size(ref.data,1), size(qry.data,1))

    ω = isnothing(ω) ? ones(size(ref.real,2)) : ω
    ω = ω / sum(ω)

    ν = isnothing(ν) ? ones(size(ref.real,2)) : ν

    for i in 1:size(ref.data,2)
        isnothing(ϕ[i]) && continue

        r = ref.real[:,i]
        q = transform(qry.data[:,ϕ[i]], r, ν[i])

        Σ += ω[i]*(reshape(r, length(r), 1) .- reshape(q, 1, length(q))).^2
    end

    return Matrix(Σ), ϕ
end

"""
    sinkhorn(M::Array{Float64,2};
                  a::Maybe{Array{Float64}} = missing,
                  b::Maybe{Array{Float64}} = missing,
                  maxᵢ::Integer            = 1000,
                  τ::Real                  = 1e-5,
                  verbose::Bool            = false
    )

Rescale matrix `M` to have row & column marginals `a` and `b` respectively.
Will terminate either when constraints are held to within tolerance `τ` or the number of iterations exceed `maxᵢ`.
"""
function sinkhorn(M::Array{Float64,2};
                  a::Maybe{Array{Float64}} = missing,
                  b::Maybe{Array{Float64}} = missing,
                  maxᵢ::Integer            = 1000,
                  τ::Real                  = 1e-5,
                  verbose::Bool            = false)
    c = 1 ./sum(M, dims=1)
    r = 1 ./(M*c')

    if ismissing(a)
        a = ones(size(M,1), 1) ./ size(M,1)
    end
    if ismissing(b)
        b = ones(size(M,2), 1) ./ size(M,2)
    end

    if length(a) != size(M,1)
        throw(error("invalid size for row prior"))
    end
    if length(b) != size(M,2)
        throw(error("invalid size for column prior"))
    end

    i = 0
    rdel, cdel = Inf, Inf
    while i < maxᵢ && (rdel > τ || cdel > τ)
        i += 1

        cinv = M'*r
        cdel = maximum(abs.(cinv.*c .- b))
        c    = b./cinv

        rinv = M*c
        rdel = maximum(abs.(rinv.*r .- a))
        r    = a./rinv

        if verbose
            println("Iteration $i. Row = $rdel, Col = $cdel")
        end

    end

    if verbose
        println("Terminating at iteration $i. Row = $rdel, Col = $cdel")
    end

    return M.*(r*c')
end

function inversion()
    ref, pointcloud = virtualembryo()
    qry = scrna()

    Σ, _ = cost(ref, qry; α=1.0, β=2.6, γ=0.65) # TODO: expose parameters?

    return (
        invert = (β) -> sinkhorn(exp.(-(1 .+ β*Σ))),
        cost = Σ,
        pointcloud = pointcloud,
    )
end

"""
    inversion(counts, genes; ν=nothing, ω=nothing, refdb=nothing)

Infer the original position of scRNAseq data `counts` where genes, given by `genes` are arranged along rows.
The sampling probability over space is computed by regularized optimal transport by comparing to the Berkeley Drosophila Transcription Network Project database.
The cost matrix is determined by summing over the 1D Wasserstein metric over all genes within the BDTNP databse.
Returns the inversion as a function of inverse temperature.
"""
function inversion(counts, genes; ν=nothing, ω=nothing, refdb=nothing)
    ref, pointcloud = refdb === nothing ? virtualembryo() : refdb
    qry = (
        data = counts',
        gene = columns(genes),
    )

    Σ, ϕ =
        if isnothing(ν) || isnothing(ω)
            # cost_simple(ref, qry)
            cost_transform(ref, qry)
        else
            cost_scan(ref, qry, ν, ω)
        end

    ψ = sinkhorn(exp.(-(1 .+ 1.0*Σ)))
    names = collect(keys(ref.gene))
    indx  = collect(values(ref.gene))

    return (
        invert     = (β) -> sinkhorn(exp.(-(1 .+ β*Σ))),
        pointcloud = pointcloud,
        database   = (
            data=ref.data',
            gene=ref.gene,
        ),
        match      = match,
        index      = ϕ,
        cost       = Σ,
    )
end

# basic statistics
mean(x)  = sum(x) / length(x)
cov(x,y) = mean(x.*y) .- mean(x)*mean(y)
var(x)   = cov(x,x)
std(x)   = sqrt(abs(var(x)))
cor(x,y) = cov(x,y) / (std(x) * std(y))

function make_objective(ref, qry)
    function objective(Θ)
        # β, ν, ω = #0.5, Θ[1:84], Θ[85:end] #ones(84)
        # Σ, ϕ    = cost_scan(ref, qry, ν, ω)
        β = 250
        ν, ω = Θ[1:84], Θ[85:end]
        Σ, ϕ = cost_transform(ref, qry; ω=ω, ν=ν)

        ψ  = sinkhorn(exp.(-(1 .+ β*Σ)))
        ψ *= minimum(size(ψ))

        ι   = findall(.!isnothing.(ϕ))
        db  = ref.real[:,ι]
        est = ψ*qry.data[:,ϕ[ι]]

        return 1-mean(cor(db[:,i], est[:,i]) for i in 1:size(db,2))
    end

    return objective
end

using BlackBoxOptim

function scan_params(qry)
    db = cohortdatabase(6)
    f  = make_objective(db.expression,qry)

    return bboptimize(f,
                  SearchRange=[(0.1, 10.0) for _ ∈ 1:79],
                  MaxFuncEvals=5000,
                  Method=:generating_set_search,
                  TraceMode=:compact
    )
end

function scan_params(count, genes)
    qry = (
        data = count',
        gene = columns(genes),
    )
    ref, _ = virtualembryo()

    f = make_objective(ref,qry)

    return bboptimize(f,
                  SearchRange=[[(0.01, 10.0) for _ ∈ 1:84]; [(0.1, 2.0) for _ ∈ 1:84]],
                  # SearchRange=[(0.01, 10.0) for _ ∈ 1:(2*84)],
                  MaxFuncEvals=5000,
                  # Method=:adaptive_de_rand_1_bin_radiuslimited,
                  Method=:generating_set_search,
                  TraceMode=:compact
    ) #, Method=:dxnes, NThreads=Threads.nthreads(), )
end

end

max_stars_repo_name: nnoll/seqspace

max_issues_repo_name: nnoll/seqspace

max_forks_repo_name: nnoll/seqspace


text:
"""
Fast union-find algorithm for finding disjoint islands of entities.
Uses path compression. union! and find! are not threadsafe.
"""
module IslandFind

export Archipelago, union!, find!, find, islands, find_islands!, find_islands, group_sets, isorphan, reorder

struct Node{T}
  parent::Int
  rank::Int
  data::T
end

struct Archipelago
  parent::Vector{Int}
  rank::Vector{Int}
  Archipelago(n) = new(collect(1:n), fill(1, n))
end

maxentity(isls::Archipelago) = length(isl.parent)

function resize(isls::Archipelago, n)
  currn = length(isls.parent)
  for i in (currn+1):n
    push!(isls.parent, i)
    push!(isls.rank, 1)
  end
end

"""
    union!(isls, (id_a,id_b))

connect entities a and b in disjoint set data structure.
"""
function union!(isls::Archipelago, (a,b)::Tuple{T, T} ) where {T<:Integer}
  nab = max(a,b)
  if nab > length(isls.parent)
    resize(isls, nab)
  end

  roota = find!(isls, a)
  rootb = find!(isls, b)

  if roota != rootb
    ranka = isls.rank[roota]
    rankb = isls.rank[rootb]
    if ranka == rankb
      # doesn't matter who becomes parent
      isls.rank[roota] += 1
      isls.parent[rootb] = roota
    elseif ranka < rankb
      # add to higher rank parent
      isls.parent[roota] = rootb
    else
      isls.parent[rootb] = roota
    end
  end
  roota
end

function union!(arp::Archipelago, c::V) where {T<:Integer, V<:AbstractVector{T}}
  for edge in zip(c[1:end-1], c[2:end])
    union!(arp, edge)
  end
  arp
end

"""
    isorphan(isls, id_a)

true if the entity represented by id_a is not connected to any other entity
"""
isorphan(isls::Archipelago, id_a) = isls.parent[id_a] == id_a && isls.rank[id_a] == 1

"""
    find!(isls, id_a)

the representative id for the disjoint set that id_a belogs to
path compression optimization can mutate isls
"""
function find!(isls::Archipelago, id::T) where {T<:Integer}
  node = id
  while true
    root = isls.parent[node]
    root == node && break
    node = root
  end
  isls.parent[id] = node
end

"""
    find(isls, id_a)

the representative id for the disjoint set that id_a belogs to
"""
function find(isls::Archipelago, id::T) where {T<:Integer}
  node = id
  while true
    root = isls.parent[node]
    root == node && break
    node = root
  end
  node
end

"""
    find_islands(connected, group_by_entity=true)

find unique disjoint set ids for each entity.

returns different data if group_by_entity is true or false.
returns list of island ids for each entity if group_by_entity is set false
returns disjoint sets of connected entities if group_by_entity is set true

if there are entites not represented in the connections but with ids < maximum id in the connections, they will be in singleton islands.

example:

  julia> find_islands([(1,3),(3,4)], true)
  2-element Array{Array{Int64,1},1}:
   [1, 3, 4]
   [2]

   julia> find_islands([(1,3),(3,4)], false)
   4-element Array{Int64,1}:
    1
    2
    1
    1

"""
function find_islands(connected::C,
                      group_by_entity=true) where {E<:Union{AbstractVector, Tuple, Set},
                                                   C<:AbstractVector{E}}
  arp = Archipelago(first(first(connected)))
  for c in connected
    for edge in zip(c[1:end-1], c[2:end])
      union!(arp, edge)
   end
  end

  archi = map(i->find!(arp, i), 1:length(arp.parent))
  if group_by_entity
    group_sets(archi)
  else
    arp
  end
end

function find_islands!(arp::Archipelago, group_by_entity=true)
  archi = map(i->find!(arp, i), 1:length(arp.parent))
  if group_by_entity
    group_sets(archi)
  else
    arp
  end
end

""" group index sets together where values at indices are equal """
function group_sets(island_id::V) where {V<:Vector{Int}}
  isls = map(i->Vector{Int}(), 1:length(island_id))

  for (i, isl) in enumerate(island_id)
    push!(isls[isl], i)
  end
  #filter(!isempty, isls)
  filter(c->!(isequal(1)∘length)(c) && !isempty(c), isls)
end

# single airity
for fn in (:find!, :find, :union!, :isorphan)
  @eval $fn(archi::Archipelago) = h->$fn(archi, h)
end

end


text:
# # Votemaps
#
# ## Aim
# In this tutorial, your will learn how to create Votemaps that compare different tomographic models and look for similarities between different models.
#
# ## Steps
#
# #### 1. Load data
# We assume that you have all tomographic models already available as *.JLD2 files. In our example we use the data uploaded to
# [https://seafile.rlp.net/d/22b0fb85550240758552/](https://seafile.rlp.net/d/22b0fb85550240758552/)
#
# Specifically, we will use the tomographic models of Paffrath, Zhao and Koulakov, as we have them all available in processed form
# Download the corresponding *.jld2 files to the same directory
using JLD2, GeophysicalModelGenerator

Pwave_Zhao = load("Zhao_Pwave.jld2","Data_set_Zhao2016_Vp")
PSwave_Koulakov = load("Koulakov_Europe.jld2","DataKoulakov_Alps")
Pwave_Paffrath = load("Paffrath_Pwave.jld2","Data_set_Paffrath2021_Vp")

# The 3 data sets all contain tomographic models for roughly the alpine area, but as you can see, they all have a different resolution and the fields are sometimes called different as well:
Pwave_Paffrath
GeoData
  size  : (162, 130, 42)
  lon   ϵ [ -13.3019 : 35.3019]
  lat   ϵ [ 30.7638 : 61.2362]
  depth ϵ [ -606.0385 km : 31.0385 km]
  fields: (:dVp_Percentage, :Vp, :Resolution)

PSwave_Koulakov
  GeoData
  size  : (108, 81, 35)
  lon   ϵ [ 4.0 : 20.049999999999997]
  lat   ϵ [ 37.035928143712574 : 49.01197604790419]
  depth ϵ [ -700.0 km : -10.0 km]
  fields: (:dVp_percentage, :dVs_percentage)

# #### 2. Creating a Votemap
# The idea of `Votemaps` is rather simple:
# - assume that a certain perturbation describes a feature (say, P wave anomalies >3% are interpreted to be slabs in the model of Paffrath)
# - Everything that fulfills this criteria gets the number 1, everything else 0.
# - Do the same with other tomographic models (using the same criteria or a different one).
# - Make sure that all tomographic models are interpolated to the same grid.
# - Sum up the 3 models.
# - The resulting 3D map will have the number `3` at locations where all 3 models see a high-velocity anomaly (say a slab), implying that this feature is consistent between all models. Features that have a number 1 or 2 are only seen in a single or in 2/3 models.
#
# The result of this gives a feeling which features are consistent between the 3 models. It is ofcourse still subjective, as you still have to choose a criteria and as we are assuming in this that the 3 tomographic models are comparable (which they are not as they are produced using different amounts of datam, etc. etc.)
#
# So how do we create Votemaps?
# Doing this is rather simple:
Data_VoteMap = VoteMap( [Pwave_Paffrath,       PSwave_Koulakov,    Pwave_Zhao],
                        ["dVp_Percentage>3.0","dVp_percentage>2.0","dVp_Percentage>2.0"], dims=(100,100,100))

# This will look at the common `lon`,`lat`,`depth` ranges between all 3 models, interpret each of the models to a common grid of size `(100,100,100)` and apply each of the criteria specified
# The resulting `GeoData` struct looks like:
GeoData
  size  : (100, 100, 100)
  lon   ϵ [ 4.0 : 18.0]
  lat   ϵ [ 38.0 : 49.01197604790419]
  depth ϵ [ -606.0385 km : -10.0 km]
  fields: (:VoteMap,)

# And from this, we can generate profiles, visualize 3D features in Paraview etc. etc.
Write_Paraview(Data_VoteMap, "VoteMap_Alps")

# In paraview, this gives
# ![Tutorial_VoteMap](../assets/img/Tutorial_VoteMap.png)

# You can ofcourse argue that newer tomographic models include more data & should therefore count more
# A simple way to take that into account is to list the model twice:
Data_VoteMap = VoteMap( [Pwave_Paffrath,       Pwave_Paffrath,       PSwave_Koulakov,    Pwave_Zhao],
                        ["dVp_Percentage>3.0","dVp_Percentage>3.0", "dVp_percentage>2.0","dVp_Percentage>2.0"],
                        dims=(100,100,100))

# Similarly, you could only look at a single model (even when that is perhaps easier done directly in paraview, where you can simply select a different isocontour value)


text:
using BinaryProvider # requires BinaryProvider 0.3.0 or later

# Parse some basic command-line arguments
const verbose = "--verbose" in ARGS
const prefix = Prefix(get([a for a in ARGS if a != "--verbose"], 1, joinpath(@__DIR__, "usr")))
products = [
    LibraryProduct(prefix, ["libcairo"], :libcairo),
]

# Download binaries from hosted location
bin_prefix = "https://github.com/JuliaBinaryWrappers/Cairo_jll.jl/releases/download/Cairo-v1.16.0+1"

# Listing of files generated by BinaryBuilder:
download_info = Dict(
    Linux(:aarch64, libc=:glibc) => ("$bin_prefix/Cairo.v1.16.0.aarch64-linux-gnu.tar.gz", "222d35c12fc2d40fb7ccc501ad50d27ccc08803c5b636cc5d83a8c8fe56c6ceb"),
    Linux(:aarch64, libc=:musl) => ("$bin_prefix/Cairo.v1.16.0.aarch64-linux-musl.tar.gz", "fb49455b7d677e01e565babd7f1f7eb062acb5991aaf9c0034a5e8faea63bf47"),
    Linux(:armv7l, libc=:glibc, call_abi=:eabihf) => ("$bin_prefix/Cairo.v1.16.0.arm-linux-gnueabihf.tar.gz", "69340376016f0525745fc483e3dc12962d0a1576c44f1243b64f7a56aa78bdb6"),
    Linux(:armv7l, libc=:musl, call_abi=:eabihf) => ("$bin_prefix/Cairo.v1.16.0.arm-linux-musleabihf.tar.gz", "8a9ab11c22dc7eefc58728c26857baed97004b5c0c5f4e8c4fada49fa5294cf4"),
    Linux(:i686, libc=:glibc) => ("$bin_prefix/Cairo.v1.16.0.i686-linux-gnu.tar.gz", "73bf7f7f1c2cbac14f74e16f201b963a8158aa69e797977a3cd3e5633a8281ee"),
    Linux(:i686, libc=:musl) => ("$bin_prefix/Cairo.v1.16.0.i686-linux-musl.tar.gz", "347db2e09c6e5add3202f44510036ab8f15d153dcbfb6f492b2083b2fb8d72b1"),
    Windows(:i686) => ("$bin_prefix/Cairo.v1.16.0.i686-w64-mingw32.tar.gz", "5993e7d41106670203614f887b684d22d456b36af7345a782a34c2146613bb64"),
    Linux(:powerpc64le, libc=:glibc) => ("$bin_prefix/Cairo.v1.16.0.powerpc64le-linux-gnu.tar.gz", "ca629c1b3f48f3659cca5ce953d2c9a8223da1722c122778cd61c013e7b8337e"),
    MacOS(:x86_64) => ("$bin_prefix/Cairo.v1.16.0.x86_64-apple-darwin14.tar.gz", "fbb3c311d23ccb800a1b18355fbf52328e90769e45550df98bd78342ef650f6c"),
    Linux(:x86_64, libc=:glibc) => ("$bin_prefix/Cairo.v1.16.0.x86_64-linux-gnu.tar.gz", "06bc332f8efd661d7806d9a02fd40813be1dd34d568df52a4e7e38e852ecfb6e"),
    Linux(:x86_64, libc=:musl) => ("$bin_prefix/Cairo.v1.16.0.x86_64-linux-musl.tar.gz", "ee07ba9ce973b47429ca32384550cee97b5ff285099f06fe56939365342f836b"),
    FreeBSD(:x86_64) => ("$bin_prefix/Cairo.v1.16.0.x86_64-unknown-freebsd11.1.tar.gz", "5eab2fca884e26edc741b7be2fee542f657ac6342c59f6a630e57fa05b42e60b"),
    Windows(:x86_64) => ("$bin_prefix/Cairo.v1.16.0.x86_64-w64-mingw32.tar.gz", "d85155e3031b0d81da5320a9944821dc47e1b7f299e3c528883a44d776bb533c"),
)

# Install unsatisfied or updated dependencies:
unsatisfied = any(!satisfied(p; verbose=verbose) for p in products)
dl_info = choose_download(download_info, platform_key_abi())
if dl_info === nothing && unsatisfied
    # If we don't have a compatible .tar.gz to download, complain.
    # Alternatively, you could attempt to install from a separate provider,
    # build from source or something even more ambitious here.
    error("Your platform (\"$(Sys.MACHINE)\", parsed as \"$(triplet(platform_key_abi()))\") is not supported by this package!")
end

# If we have a download, and we are unsatisfied (or the version we're
# trying to install is not itself installed) then load it up!
if unsatisfied || !isinstalled(dl_info...; prefix=prefix)
    # Download and install binaries
    install(dl_info...; prefix=prefix, force=true, verbose=verbose)
end

# Write out a deps.jl file that will contain mappings for our products
write_deps_file(joinpath(@__DIR__, "deps.jl"), products, verbose=verbose)


text:
@testset "Cubic" begin
    for (constructor, copier) in ((interpolate, identity), (interpolate!, copy))
        isinplace = constructor == interpolate!
        f0(x) = sin((x-3)*2pi/9 - 1)
        f1(x) = 1.0 + 0.1*x + 0.01*x^2 + 0.001*x^3

        xmax = 10
        A0 = Float64[f0(x) for x in 1:xmax]
        A1 = Float64[f1(x) for x in 1:xmax]

        f2(x, y) = sin(x/10)*cos(y/6)
        xmax2, ymax2 = 30, 10
        A2 = Float64[f2(x, y) for x in 1:xmax2, y in 1:ymax2]

        for BC in (Line, Flat, Free, Periodic), GT in (OnGrid, OnCell)
            for (A, f) in ((A0, f0), (A1, f1))
                itp1 = @inferred(constructor(copier(A), BSpline(Cubic(BC(GT())))))
                ax1 = axes(itp1)[1]
                @test Interpolations.lbounds(itp1) == (GT == OnGrid ? (first(ax1),) : (first(ax1) - 0.5,))
                @test Interpolations.ubounds(itp1) == (GT == OnGrid ? (last(ax1),) : (last(ax1) + 0.5,))
                @test_throws ArgumentError parent(itp1)
                check_axes(itp1, A, isinplace)
                check_inbounds_values(itp1, A)
                check_oob(itp1)
                can_eval_near_boundaries(itp1)

                # test that inner region is close to data
                for x in 3.1:.2:8.1
                    @test f(x) ≈ itp1(x) atol=abs(0.1 * f(x))
                end
            end

            itp2 = @inferred(constructor(copier(A2), BSpline(Cubic(BC(GT())))))
            @test_throws ArgumentError parent(itp2)
            check_axes(itp2, A2, isinplace)
            check_inbounds_values(itp2, A2)
            check_oob(itp2)
            can_eval_near_boundaries(itp2)

            for x in 3.1:.2:xmax2-3, y in 3.1:2:ymax2-3
                @test f2(x,y) ≈ itp2(x,y) atol=abs(0.1 * f2(x,y))
            end
        end
        # Issue 419: Incorrect results with BSpline(Cubic(Periodic(OnCell()))))
        let BC = Periodic, GT = OnCell
            for (A, f) in ((A0, f0), (A1, f1))
                itp1 = @inferred(constructor(copier(A), BSpline(Cubic(BC(GT())))))
                @test itp1(Interpolations.lbounds(itp1)...) ≈ itp1(Interpolations.ubounds(itp1)...)
            end
        end
    end

    ix = 1:15
    k = length(ix) - 1
    f(x) = cos((x-1)*2pi/k)
    g(x) = -2pi/k * sin((x-1)*2pi/k)

    A = map(f, ix)

    for (constructor, copier) in ((interpolate, identity), (interpolate!, copy))

        for BC in (Line, Flat, Free, Periodic), GT in (OnGrid,OnCell)

            itp = constructor(copier(A), BSpline(Cubic(BC(GT()))))
            # test that inner region is close to data
            for x in range(ix[5], stop=ix[end-4], length=100)
                @test g(x) ≈ Interpolations.gradient1(itp,x) atol=cbrt(cbrt(eps(g(x))))
            end
        end
        # Issue 419: Incorrect results with BSpline(Cubic(Periodic(OnCell()))))
        let BC = Periodic, GT = OnCell
            itp = @inferred(constructor(A[1:end-1], BSpline(Cubic(BC(GT())))))
            @test itp(Interpolations.lbounds(itp)...) ≈ itp(Interpolations.ubounds(itp)...)
        end
    end
    itp_flat_g = interpolate(A, BSpline(Cubic(Flat(OnGrid()))))
    @test Interpolations.gradient(itp_flat_g,1)[1] ≈ 0 atol=eps()
    @test Interpolations.gradient(itp_flat_g,ix[end])[1] ≈ 0 atol=eps()

    itp_flat_c = interpolate(A, BSpline(Cubic(Flat(OnCell()))))
    @test Interpolations.gradient(itp_flat_c,0.5)[1] ≈ 0 atol=eps()
    @test Interpolations.gradient(itp_flat_c,ix[end] + 0.5)[1] ≈ 0 atol=eps()

    # Can construct from coefficients
    itp = interpolate(rand(5, 9), BSpline(Cubic(Flat(OnGrid()))))
    @test itp == Interpolations.BSplineInterpolation(itp.coefs, itp.it, itp.parentaxes)
    @test_throws ArgumentError Interpolations.BSplineInterpolation(itp.coefs, itp.it, (2:3, 2:10))
end


text:
using Econometrics, Statistics, Distributions, StatsPlots, DelimitedFiles
function main()
# weekly close price of NSYE, data provided with GRETL
data = readdlm("nysewk.txt")
# compute weekly percentage growth
y = 100.0 * log.(data[2:end] ./ data[1:end-1])
p1 = plot(y, leg=false, show=true)
p2 = histogram(y, normed=true)
plot!(p2, Distributions.Normal(mean(y),std(y)), show=true, leg=false)
p3 = npdensity(y)
plot(p1,p2,p3, layout = (3,1))
gui()
savefig("nyse.svg")
return
end
main()

max_stars_repo_name: nilshg/Econometrics

max_issues_repo_name: nilshg/Econometrics

max_forks_repo_name: nilshg/Econometrics


text:
struct UpSample
  ratio
end

@treelike UpSample

function (up::UpSample)(input)
  ratio = (up.ratio[1], up.ratio[2], 1, 1)
  (h, w, c, n) = size(input)
  y = similar(input, (ratio[1], 1, ratio[2], 1, 1, 1))
  fill!(y, 1)
  z = reshape(input, (1, h, 1, w, c, n))  .* y
  return reshape(z, size(input) .* ratio)
end

function crop(original, desired)
    desired_h, desired_w = size(desired)
    original_h, original_w = size(original)

    delta_h = convert(Int8, (original_h - desired_h) ÷ 2)
    delta_w = convert(Int8, (original_w - desired_w) ÷ 2)

    return original[
        1+delta_h:delta_h+desired_h,
        1+delta_w:delta_w+desired_w,
        :, :]
end

depth_4 = Chain(
    MaxPool((2,2)),
    Conv((3,3), 512=>1024, relu, pad=(0,0), stride=(1,1)),
    Conv((3,3), 1024=>1024, relu, pad=(0,0), stride=(1,1)),
    UpSample((2,2)),
    Conv((1,1), 1024=>512, pad=(0,0), stride=(1,1)),
)

depth_3 = Chain(
    MaxPool((2,2)),
    Conv((3,3), 256=>512, relu, pad=(0,0), stride=(1,1)),
    Conv((3,3), 512=>512, relu, pad=(0,0), stride=(1,1)),

    SkipConnection(depth_4, (mx, x) -> cat(mx, crop(x, mx), dims=3)),
    Conv((3,3), 1024=>512, relu, pad=(0,0), stride=(1,1)),

    Conv((3,3), 512=>512, relu, pad=(0,0), stride=(1,1)),
    UpSample((2,2)),
    Conv((1,1), 512=>256, pad=(0,0), stride=(1,1)),
)

depth_2 = Chain(
    MaxPool((2,2)),
    Conv((3,3), 128=>256, relu, pad=(0,0), stride=(1,1)),
    Conv((3,3), 256=>256, relu, pad=(0,0), stride=(1,1)),

    SkipConnection(depth_3, (mx, x) -> cat(mx, crop(x, mx), dims=3)),
    Conv((3,3), 512=>256, relu, pad=(0,0), stride=(1,1)),

    Conv((3,3), 256=>256, relu, pad=(0,0), stride=(1,1)),
    UpSample((2,2)),
    Conv((1,1), 256=>128, pad=(0,0), stride=(1,1)),
)

depth_1 = Chain(
    MaxPool((2,2)),
    Conv((3,3), 64=>128, relu, pad=(0,0), stride=(1,1)),
    Conv((3,3), 128=>128, relu, pad=(0,0), stride=(1,1)),

    SkipConnection(depth_2, (mx, x) -> cat(mx, crop(x, mx), dims=3)),
    Conv((3,3), 256=>128, relu, pad=(0,0), stride=(1,1)),

    Conv((3,3), 128=>128, relu, pad=(0,0), stride=(1,1)),
    UpSample((2,2)),
    Conv((1,1), 128=>64, pad=(0,0), stride=(1,1)),
)

unet = Chain(
    Conv((3,3), 1=>64, relu, pad=(0,0), stride=(1,1)),
    Conv((3,3), 64=>64, relu, pad=(0,0), stride=(1,1)),

    SkipConnection(depth_1, (mx, x) -> cat(mx, crop(x, mx), dims=3)),

    Conv((3,3), 128=>64, relu, pad=(0,0), stride=(1,1)),
    Conv((3,3), 64=>64, relu, pad=(0,0), stride=(1,1)),

    Conv((1,1), 64=>1, pad=(0,0), stride=(1,1)),

    # Finally, normalise and softmax to get nice probabilities
    BatchNorm(1),
    x -> sigmoid.(x)
)


text:
using Test
using NQCDynamics

atoms = Atoms([:H])
model = Harmonic()

sim = Simulation{Langevin}(atoms, model; temperature=1, γ=1)

z = DynamicsVariables(sim, randn(1,1), randn(1,1))

solution = run_trajectory(z, (0.0, 500.0), sim; dt=1)


text:
using DataStructures
using LazySets

mutable struct LEAFNODE
    min_control::Float64
    max_control::Float64
    prob::Float64
end

mutable struct KDNODE
    dim::Int32
    split::Float64
    left::Union{KDNODE, LEAFNODE}
    right::Union{KDNODE, LEAFNODE}
    min_control::Float64
    max_control::Float64
end

mutable struct KDTREE
    lbs::Vector{Float64}
    ubs::Vector{Float64}
    root_node::KDNODE
end

function leafnode(;min_control = -10.0, max_control = 10.0, prob = 0.0)
    return LEAFNODE(min_control, max_control, prob)
end

function kdnode(;dim = 1, split = 0.0, left = leafnode(), right = leafnode(),
    min_control = -10.0, max_control = 10.0)
    return KDNODE(dim, split, left, right, min_control, max_control)
end

function get_leaf(root_node, point)
    curr_node = root_node
    while typeof(curr_node) != LEAFNODE
        val = point[curr_node.dim]
        split = curr_node.split
        curr_node = val < split ? curr_node.left : curr_node.right
    end
    return curr_node
end

function get_bounds(tree)
    lbs = []
    ubs = []

    lb_s = Stack{Vector{Float64}}()
    ub_s = Stack{Vector{Float64}}()
    s = Stack{Union{LEAFNODE, KDNODE}}()

    push!(lb_s, tree.lbs)
    push!(ub_s, tree.ubs)
    push!(s, tree.root_node)

    while !isempty(s)
        curr = pop!(s)
        curr_lbs = pop!(lb_s)
        curr_ubs = pop!(ub_s)

        if typeof(curr) == LEAFNODE
            push!(lbs, curr_lbs)
            push!(ubs, curr_ubs)
        else
            # Traverse tree and keep track of bounds
            dim = curr.dim
            split = curr.split
            # Go left, upper bounds will change
            left_ubs = copy(curr_ubs)
            left_ubs[dim] = split

            push!(lb_s, curr_lbs)
            push!(ub_s, left_ubs)
            push!(s, curr.left)

            # Go right, lower bounds will change
            right_lbs = copy(curr_lbs)
            right_lbs[dim] = split

            push!(lb_s, right_lbs)
            push!(ub_s, curr_ubs)
            push!(s, curr.right)
        end
    end

    return lbs, ubs
end

function get_leaf_and_bounds(tree, point)
    curr_node = tree.root_node
    lbs = copy(tree.lbs)
    ubs = copy(tree.ubs)
    while typeof(curr_node) != LEAFNODE
        dim = curr_node.dim
        val = point[dim]
        split = curr_node.split
        if val < split
        	curr_node = curr_node.left
        	ubs[dim] = split
        else
        	curr_node = curr_node.right
        	lbs[dim] = split
        end
    end
    return curr_node, lbs, ubs
end

function get_bounds_and_probs(tree)
    lbs = []
    ubs = []
    probs = []

    lb_s = Stack{Vector{Float64}}()
    ub_s = Stack{Vector{Float64}}()
    s = Stack{Union{LEAFNODE, KDNODE}}()

    push!(lb_s, tree.lbs)
    push!(ub_s, tree.ubs)
    push!(s, tree.root_node)

    while !isempty(s)
        curr = pop!(s)
        curr_lbs = pop!(lb_s)
        curr_ubs = pop!(ub_s)

        if typeof(curr) == LEAFNODE
            push!(lbs, curr_lbs)
            push!(ubs, curr_ubs)
            push!(probs, curr.prob)
        else
            # Traverse tree and keep track of bounds
            dim = curr.dim
            split = curr.split
            # Go left, upper bounds will change
            left_ubs = copy(curr_ubs)
            left_ubs[dim] = split

            push!(lb_s, curr_lbs)
            push!(ub_s, left_ubs)
            push!(s, curr.left)

            # Go right, lower bounds will change
            right_lbs = copy(curr_lbs)
            right_lbs[dim] = split

            push!(lb_s, right_lbs)
            push!(ub_s, curr_ubs)
            push!(s, curr.right)
        end
    end

    return lbs, ubs, probs
end

function get_overlapping_nodes(root_node, lbs, ubs)
    curr_node_list = Vector{LEAFNODE}()
    overlapping_nodes_helper(root_node, lbs, ubs, curr_node_list)
    return curr_node_list
end

function overlapping_nodes_helper(curr_node, lbs, ubs, curr_node_list)
    # Base case
    if typeof(curr_node) == LEAFNODE
        push!(curr_node_list, curr_node)
    else
        # Check if can prune either side
        dim = curr_node.dim
        split = curr_node.split
        if split > ubs[dim]
            # Can prune the right half
            overlapping_nodes_helper(curr_node.left, lbs, ubs, curr_node_list)
        elseif split < lbs[dim]
            # Can prune the left half
            overlapping_nodes_helper(curr_node.right, lbs, ubs, curr_node_list)
        else
            overlapping_nodes_helper(curr_node.left, lbs, ubs, curr_node_list)
            overlapping_nodes_helper(curr_node.right, lbs, ubs, curr_node_list)
        end
    end
end

function split!(s, lb_s, ub_s, curr_node, curr_lbs, curr_ubs, dim)
    split = (curr_lbs[dim] + curr_ubs[dim]) / 2

    curr_node.split = split
    curr_node.dim = dim

    curr_node.left = kdnode(min_control = curr_node.min_control, max_control = curr_node.max_control)
    curr_node.right = kdnode(min_control = curr_node.min_control, max_control = curr_node.max_control)

    # Get new bounds
    # Go left, upper bounds will change
    left_lbs = copy(curr_lbs)
    left_ubs = copy(curr_ubs)
    left_ubs[dim] = split
    # Go right, lower bounds will change
    right_lbs = copy(curr_lbs)
    right_lbs[dim] = split
    right_ubs = copy(curr_ubs)

    # Add everything to the stacks
    push!(s, curr_node.left)
    push!(lb_s, curr_lbs)
    push!(ub_s, left_ubs)
    push!(s, curr_node.right)
    push!(lb_s, right_lbs)
    push!(ub_s, curr_ubs)
end

function split_specific_dims!(s, lb_s, ub_s, curr_node, curr_lbs, curr_ubs, dims)
    split!(s, lb_s, ub_s, curr_node, curr_lbs, curr_ubs, dims[1])
    for i = 2:length(dims)
        nodes = []
        lbs = []
        ubs = []
        for j = 1:2^(i-1)
            push!(nodes, pop!(s))
            push!(lbs, pop!(lb_s))
            push!(ubs, pop!(ub_s))
        end
        for j = 1:2^(i-1)
            split!(s, lb_s, ub_s, nodes[j], lbs[j], ubs[j], dims[i])
        end
    end
end

function create_tree(min_widths, lbs, ubs)
    root_node = kdnode()

    # Stacks to go through tree
    lb_s = Stack{Vector{Float64}}()
    ub_s = Stack{Vector{Float64}}()
    s = Stack{Union{LEAFNODE, KDNODE}}()

    push!(lb_s, lbs)
    push!(ub_s, ubs)
    push!(s, root_node)

    while !isempty(s)
        curr = pop!(s)
        curr_lbs = pop!(lb_s)
        curr_ubs = pop!(ub_s)

        # println("lbs: ", curr_lbs)
        # println("ubs: ", curr_ubs)

        # Figure out which dimensions need splitting (if any)
        curr_widths = curr_ubs .- curr_lbs
        # println("widths: ", curr_widths)
        # println()
        dims_to_split = findall(curr_widths .> min_widths)

        # Split them
        if length(dims_to_split) > 0
            split_specific_dims!(s, lb_s, ub_s, curr, curr_lbs, curr_ubs, dims_to_split)
        end
    end

    # Traverse and remove extra nodes (this is pretty hacky)
    s = Stack{Union{LEAFNODE, KDNODE}}()
    push!(s, root_node)

    while !isempty(s)
        curr = pop!(s)

        if typeof(curr.left.left) == LEAFNODE
            curr.left = leafnode()
        else
            push!(s, curr.left)
        end

        if typeof(curr.right.right) == LEAFNODE
            curr.right = leafnode()
        else
            push!(s, curr.right)
        end
    end

    return KDTREE(lbs, ubs, root_node)
end

function Base.:(==)(k1::KDNODE, k2::KDNODE)
    return k1.dim == k2.dim && k1.split == k2.split && k1.left == k2.left && k1.right == k2.right &&
                k1.min_control == k2.min_control && k1.max_control == k2.max_control
end

function Base.:(==)(l1::LEAFNODE, l2::LEAFNODE)
    return l1.min_control == l2.min_control && l1.max_control == l2.max_control && l1.prob == l2.prob
end

function Base.:(==)(t1::KDTREE, t2::KDTREE)
    if (t1.lbs != t2.lbs) || (t1.ubs != t2.ubs)
        return false
    else
        equal = true

        s1 = Stack{Union{LEAFNODE, KDNODE}}()
        s2 = Stack{Union{LEAFNODE, KDNODE}}()

        push!(s1, t1.root_node)
        push!(s2, t2.root_node)

        while !isempty(s1)
            curr1 = pop!(s1)
            curr2 = pop!(s2)

            equal = curr1 == curr2

            if !equal
                return false
            else
                if typeof(curr1) != LEAFNODE
                    push!(s1, curr1.left)
                    push!(s1, curr1.right)
                    push!(s2, curr2.left)
                    push!(s2, curr2.right)
                end
            end
        end
    end
    return equal
end

function Base.copy(l::LEAFNODE)
    return LEAFNODE(l.min_control, l.max_control, l.prob)
end

function Base.copy(tree::KDTREE)
    root = tree.root_node
    if typeof(root) == LEAFNODE
        new_tree = copy(root)
    else
        new_tree = kdnode(dim = root.dim, split = root.split,
                                min_control = root.min_control, max_control = root.max_control)

        s_old = Stack{Union{LEAFNODE, KDNODE}}()
        s_new = Stack{Union{LEAFNODE, KDNODE}}()

        push!(s_old, tree.root_node)
        push!(s_new, new_tree)

        while !isempty(s_old)
            curr_old = pop!(s_old)
            curr_new = pop!(s_new)

            old_left = curr_old.left
            old_right = curr_old.right

            left_leaf = typeof(old_left) == LEAFNODE
            right_leaf = typeof(old_right) == LEAFNODE

            if left_leaf && right_leaf
                curr_new.left = copy(old_left)
                curr_new.right = copy(old_right)
            elseif left_leaf
                curr_new.left = copy(old_left)
                curr_new.right = kdnode(dim = old_right.dim, split = old_right.split,
                                min_control = old_right.min_control, max_control = old_right.max_control)
                push!(s_old, old_right)
                push!(s_new, curr_new.right)
            elseif right_leaf
                curr_new.right = copy(old_right)
                curr_new.left = kdnode(dim = old_left.dim, split = old_left.split,
                                min_control = old_left.min_control, max_control = old_left.max_control)
                push!(s_old, old_left)
                push!(s_new, curr_new.left)
            else
                curr_new.left = kdnode(dim = old_left.dim, split = old_left.split,
                                min_control = old_left.min_control, max_control = old_left.max_control)
                push!(s_old, old_left)
                push!(s_new, curr_new.left)
                curr_new.right = kdnode(dim = old_right.dim, split = old_right.split,
                                min_control = old_right.min_control, max_control = old_right.max_control)
                push!(s_old, old_right)
                push!(s_new, curr_new.right)
            end
        end
    end
    return KDTREE(copy(tree.lbs), copy(tree.ubs), new_tree)
end

max_stars_repo_name: sisl/VerifyGAN

max_issues_repo_name: sisl/VerifyGAN

max_forks_repo_name: sisl/VerifyGAN


text:
module TestGeneralised
using Test

using Diversity
using Diversity.ShortNames

metacom = Metacommunity([1 2; 1 2]./6)

# Basic checks for the diversity() function
@testset "diversity()" begin
    @test diversity([metacommunityDiversity], [ρ̄], metacom, 1)[!,:diversity] ≈ [1.0]
    @test diversity([subcommunityDiversity], [ρ̄], metacom, 1)[!,:diversity] ≈ [1.0, 1.0]
end

pops = rand(5, 4)
pops = pops / sum(pops)
meta = Metacommunity(pops)
@testset "all-in-one norm/raw_alpha/beta/rho/gamma_div()" begin
    ra = α(meta)
    na = ᾱ(meta)
    rb = β(meta)
    nb = β̄(meta)
    rr = ρ(meta)
    nr = ρ̄(meta)
    g  = Γ(meta)
    # Alpha diversities
    @test raw_sub_alpha(meta, 0)[!,:diversity] ≈ subdiv(ra, 0)[!,:diversity]
    @test raw_meta_alpha(meta, 1)[!,:diversity] ≈ metadiv(ra, 1)[!,:diversity]
    @test norm_sub_alpha(meta, 2)[!,:diversity] ≈ subdiv(na, 2)[!,:diversity]
    @test norm_meta_alpha(meta, 3)[!,:diversity] ≈ metadiv(na, 3)[!,:diversity]
    # Beta diversities
    @test raw_sub_beta(meta, 0.0)[!,:diversity] ≈ subdiv(rb, 0.0)[!,:diversity]
    @test raw_meta_beta(meta, 1.0)[!,:diversity] ≈ metadiv(rb, 1.0)[!,:diversity]
    @test norm_sub_beta(meta, 2.0)[!,:diversity] ≈ subdiv(nb, 2.0)[!,:diversity]
    @test norm_meta_beta(meta, 3.0)[!,:diversity] ≈ metadiv(nb, 3.0)[!,:diversity]

    @test raw_sub_rho(meta, Inf)[!,:diversity] ≈ subdiv(rr, Inf)[!,:diversity]
    @test raw_meta_rho(meta, Inf)[!,:diversity] ≈ metadiv(rr, Inf)[!,:diversity]
    @test norm_sub_rho(meta, Inf)[!,:diversity] ≈ subdiv(nr, Inf)[!,:diversity]
    @test norm_meta_rho(meta, Inf)[!,:diversity] ≈ metadiv(nr, Inf)[!,:diversity]

    # Gamma diversities
    @test sub_gamma(meta, 0.5)[!,:diversity] ≈ subdiv(g, 0.5)[!,:diversity]
    @test meta_gamma(meta, 1.5)[!,:diversity] ≈ metadiv(g, 1.5)[!,:diversity]

end

end