text:
# Initialising simulation conditions

using Distributions: Normal
using Random: rand!
using LinearAlgebra: normalize, dot, norm

# Constants (SI units)
const kB = 1.38064852e-23 #Boltzmann
const c = 3e8 # Speed of light (m/s)
const ε₀ = 8.854e-12 # Vacuum permittivity
const h̄ = 6.626e-34 / (2 * π) # Reduced Planck
const g = 9.81 # Gravity
const a₀ = 5.28e-11 # Bohr radius

# ATOM INITIALISATION

# Define a neutral atom species
struct AtomSpecies
    m :: Float64    # Atomic mass (kg)
    aₛ :: Float64   # Scattering length (m)
    σ :: Float64    # Scattering cross-section (m^2)
    α :: Float64    # Real part of polarizability
end

const Rb87 = AtomSpecies(1.454660e-25, 98a₀, 8*pi * ( 98a₀ )^2,
                         1.1e-38)

# Dipole trap constant
kappa(species::AtomSpecies) = species.α / (2 * ε₀ * c)

# Initialise N (test) particles with position and velocity components
# distributed uniformly in the range [-size, size] and [-v, v], respectively.
function uniform_cloud(N, size, v)
    positions = 2.0 * size * (rand(Float64, 3, N) .- 0.5)
    velocities = 2.0 * v * (rand(Float64, 3, N) .- 0.5)
    return positions, velocities
end

# Initialise Boltzmann-distributed velocities at temperature
# T for N particles, each of mass m.
boltzmann_velocities(N, m, T) = reshape(rand( Normal(0, sqrt(kB * T / m)), 3*N), 3, N)

# Initialise Boltzmann-distributed positions in a harmonic trap at
# temperature T. Generates N particles, each of mass m.
function harmonic_boltzmann_positions(N, m, T, ωx, ωy, ωz)
    coeff = sqrt(kB * T / m)
    positions = zeros(Float64, 3, N)
    rand!( Normal(0, coeff / ωx), view(positions, 1, :) )
    rand!( Normal(0, coeff / ωy), view(positions, 2, :) )
    rand!( Normal(0, coeff / ωz), view(positions, 3, :) )
    return positions
end

# FIELD INITIALISATION

# Convert values to time-parametrised functions
function time_parametrize(value)
    f(_) = value
    f() = value
    return f
end

# Allow functions to be parametrised (for convenience)
function time_parametrize(func::Function)
    f(t) = func(t)
    function f()
        @warn "Time-dependent function given no argument. Assuming t = 0."
        return f(0)
    end
    return f
end

# Parametrise multiple functions/values at once (for convenience)
time_parametrize(args...) = [time_parametrize(arg) for arg in args]

# Harmonic potential (time-dependent)
struct HarmonicField
    # Time-parametrised trap frequencies
    ωx :: Function
    ωy :: Function
    ωz :: Function
    centre :: Function
    # Initialise
    HarmonicField(ωx, ωy, ωz, centre = [0.0,0,0]) = (
        new( time_parametrize(ωx, ωy, ωz, centre)... )
    )
end

# Gaussian beam (time-dependent)
struct GaussianBeam
    focus :: Function
    direction :: Function
    power :: Function
    waist :: Function
    wavelength :: Function

    function GaussianBeam(args...)
        focus, direction, power, waist, wavelength = time_parametrize(args...)
        unit_dir(t) = normalize(direction(t))
        return new(focus, unit_dir, power, waist, wavelength)
    end
end

# E.g. gravity
struct UniformField
    strength :: Function # Vector acceleration (m/s^2)
    origin :: Function
    UniformField(str, origin = [0,0,0]) = new( time_parametrize(str, origin)... )
end

const gravity = UniformField([0, -g, 0])

# Acceleration (constant field)
function acceleration(field::UniformField, positions, species, t, output)
    # 'species' parameter required to conform to correct call signature.
    Nt = size(positions, 2)
    strx, stry, strz = field.strength(t)
    Threads.@threads for atom in 1:Nt
        output[1, atom] += strx
        output[2, atom] += stry
        output[3, atom] += strz
    end
    return nothing
end

# Potential (constant field)
function potential(field::UniformField, positions, species, t, output)
    # 'species' parameter required to conform to correct call signature.
    Nt = size(positions, 2)
    strx, stry, strz = field.strength(t)
    ox, oy, oz = field.origin(t)
    m = species.m
    Threads.@threads for atom in 1:Nt
        px, py, pz = view(positions, :, atom)
        output[atom] -= m * (strx * (px - ox) + stry * (py - oy) + strz * (pz - oz))
    end
    return nothing
end

# Overload accel/potential by generating new output array when none provided.
acceleration(f, p, s, t) = acceleration(f, p, s, t, zeros(Float64, size(p)))
potential(f, p, s, t) = potential(f, p, s, t, zeros(Float64, size(p, 2)))
# Initialise a potential or acceleration function to
potential(f) = (args...) -> potential(f, args...) #pos, spec, t, out
acceleration(f) = (args...) -> acceleration(f, args...)

# Acceleration (harmonic field)
function acceleration(field::HarmonicField, positions, species, t, output)
    # 'species' parameter required to conform to correct call signature.
    Nt = size(positions, 2)
    cx, cy, cz = field.centre(t)
    ωx2, ωy2, ωz2 = field.ωx(t)^2, field.ωy(t)^2, field.ωz(t)^2
    Threads.@threads for atom in 1:Nt
        output[1, atom] -= ωx2 * (positions[1, atom] - cx)
        output[2, atom] -= ωy2 * (positions[2, atom] - cy)
        output[3, atom] -= ωz2 * (positions[3, atom] - cz)
    end
    return nothing
end

# Potential (harmonic field)
function potential(field::HarmonicField, positions, species, t, output)
    Nt = size(positions, 2)
    m = species.m
    cx, cy, cz = field.centre(t)
    ωx2, ωy2, ωz2 = field.ωx(t)^2, field.ωy(t)^2, field.ωz(t)^2
    Threads.@threads for atom in 1:Nt
        x, y, z = view(positions, :, atom)
        output[atom] += 0.5 * m * (ωx2 * (x-cx)^2 + ωy2 * (y-cy)^2 + ωz2 * (z-cz)^2)
    end
    return nothing
end

# Get parameters of a Gaussian beam
parameters(b::GaussianBeam, t::Float64) = ( b.focus(t), b.direction(t), b.power(t),
                                   b.waist(t), b.wavelength(t) )

# Acceleration (Gaussian beam)
function acceleration(field::GaussianBeam, positions, species, t, output)
    foc, dir, P, w₀, _ = parameters(field, t)
    κ = kappa(species)
    N = size(positions, 2)

    fx, fy, fz = foc
    ux, uy, uz = dir

    r²coeff = - 2 / w₀^2
    acoeff = 8 * P * κ / (π * species.m * w₀^4)
    Threads.@threads for atom in 1:N
        px, py, pz = view(positions, :, atom)
        dx = px - fx # Displacement from beam focus
        dy = py - fy
        dz = pz - fz
        z = dx*ux + dy*uy + dz*uz # Axial coordinate
        rx = dx - (z * ux) # Radial vector
        ry = dy - (z * uy)
        rz = dz - (z * uz)
        r² = rx^2 + ry^2 + rz^2
        a = (-acoeff * exp(r²coeff * r²))
        output[1, atom] += a*rx
        output[2, atom] += a*ry
        output[3, atom] += a*rz
    end
    return nothing
end

# Potential (Gaussian beam)
function potential(field::GaussianBeam, positions, species, t, output)
    foc, dir, P, w₀, _ = parameters(field, t)

    fx, fy, fz = foc
    ux, uy, uz = dir

    Nt = size(positions, 2)
    κ = kappa(species)
    r²coeff = - 2 / w₀^2
    Ucoeff = - κ * 2 * P / (π * w₀^2)
    Threads.@threads for atom in 1:Nt
        px, py, pz = view(positions, :, atom)
        dx = px - fx
        dy = py - fy
        dz = pz - fz
        z = dx*ux + dy*uy + dz*uz
        r² = (dx^2 + dy^2 + dz^2) - z^2
        output[atom] += Ucoeff * exp(r²coeff * r²) - Ucoeff
    end
    return nothing
end

# EVAPORATION PROBABILITIES

# An exponential ramp
exponential_ramp(start, stop, τ) = (t) -> stop + (start - stop) * exp(- t / τ)

# A linear ramp
linear_ramp(start, stop, τ) = (t) -> start + (stop - start) * (t / τ)

# Perfect energy-based removal
function energy_evap(εₜ, positions, velocities, conditions,
        potential, t, out = zeros(Float64, size(positions, 2)))

    potential(positions, conditions.species, t, out)
    #out initially contains pe
    m = conditions.species.m
    N = size(positions, 2)
    for atom in 1:N
        #vx, vy, vz = view(velocities, :, atom)
        #ke = 0.5 * m * (vx^2 + vy^2 + vz^2)
        if out[atom] > εₜ
            out[atom] = 1 # Replace PE with evaporation probability
        else
            out[atom] = 0
        end
    end
    return out
end

# Overload for correct type signature and time-parametrisation
function energy_evap(depth, potential)
    εₜ = time_parametrize(depth)
    return (pos, vel, cond, t, o) -> energy_evap(εₜ(t), pos, vel, cond, potential, t, o)
end

# Removal when beyond ellipsoid
function ellipsoid_evap(ωx, ωy, ωz, T, p, centre,
                        species, positions, out = zeros(Float64, size(positions, 2)))

    N = size(positions, 2)

    m = species.m
    σx² = kB * T / (m * ωx^2)
    σy² = kB * T / (m * ωy^2)
    σz² = kB * T / (m * ωz^2)

    bound = -2 * log(p) # Boundary condition
    cx, cy, cz = centre

    for atom in 1:N
        x = positions[1, atom] - cx
        y = positions[2, atom] - cy
        z = positions[3, atom] - cz
        if (x^2 / σx² + y^2 / σy² + z^2 / σz²) > bound
            out[atom] = 1.0
        else
            out[atom] = 0.0
        end
    end

    return out
end

# Overload for correct type signature and time-parametrisation
function ellipsoid_evap(ωx, ωy, ωz, T, p, centre = [0.0,0,0])
    ω_x, ω_y, ω_z, T = time_parametrize(ωx, ωy, ωz, T)
    centre = time_parametrize(centre)
    function evap(pos, vel, cond, t, o)
        return ellipsoid_evap(ω_x(t), ω_y(t), ω_z(t), T(t), p, centre(t), cond.species, pos, o)
    end
    return evap
end

# No evaporation (evap. probability 0)
no_evap(pos, args...) = zeros(Float64, size(pos, 2))

# SIMULATION CONDITIONS
struct SimulationConditions
    species :: AtomSpecies

    F :: Float64 # Initial F-scale
    positions :: Matrix{Float64} # Initial test-particle positions
    velocities :: Matrix{Float64} # Initial test-particle velocities

    acceleration :: Function # Time-parametrised, vectorised acceleration func.
    potential :: Function # Potential (also parametrised/vectorised)

    threebodyloss :: Float64 # 0 # 3-body loss constant (m^6/s, NOT cm^6/s)
    evaporate :: Function # no_evap # Particle loss-function
    τ_bg :: Float64 # Inf # Background loss time constant

    # Initialisation
    function SimulationConditions(spec, F, pos, vel, accel, potential;
            K = 0, evap = no_evap, τ_bg = Inf) # Optional arguments

        # Check that array sizes are 3 x Nt
        num_components, Nt = size(pos)
        size(vel) != (num_components, Nt) && throw(DimensionMismatch(
            "Velocity and position arrays must have the same size.")
        )
        (num_components != 3) && throw(ArgumentError(
            "Position and velocity must have three components.")
        )

        return new(spec, F, pos, vel, accel, potential, K, evap, τ_bg)
    end
end


text:
# Multi-dimensional FFTs based on the 1d FFTs in ctfft.jl

mutable struct MultiDimPlan{T,forward} <: Plan{T}
    p::Vector{CTPlan{T,forward}} # 1d plans along each transformed dimension,
                                 # with p.n == 0 for untransformed dims.
    w::Vector{T} # workspace (for in-place 1d transforms)
    lastdim::Int # last transformed dimension
    pinv::ScaledPlan{T}
    MultiDimPlan{T,forward}(p,w,lastdim) where {T,forward} = new(p,w,lastdim)
end

summary(p::MultiDimPlan{T,forw}) where {T,forw} =
    string(forw ? "forward " : "backward ", length(p.p),
           "-dimensional MultiDimPlan{$T} of size ",
           join(map(P -> P.n == 0 ? "_" : string(P.n), p.p), "x"))

function show(io::IO, p::MultiDimPlan)
    print(io, summary(p), ", via:\n",
          join(map(string, p.p), "\n"))
end

size(p::MultiDimPlan) = ntuple(n -> p.p[n].n, length(p.p))

function plan_inv(p::MultiDimPlan{T,forward}) where {T,forward}
    sz = filter(n -> n != 0, map(P -> P.n, p.p))
    ScaledPlan(MultiDimPlan{T,!forward}(map(invCT, p.p), p.w, p.lastdim),
               isempty(sz) ? 1 : normalization(real(T), sz, 1:length(sz)))
end

function MultiDimPlan(::Type{T}, forward::Bool, region, sz) where T<:Complex
    sregion = sort(Int[d for d in region])
    N = length(sz)
    p = Array{CTPlan{T,forward}}(undef, N)
    i = 0
    for d in sregion
        (d < 0 || d > N) && throw(ArgumentError("invalid dimension $d"))
        d == i && throw(ArgumentError("dimension $d specified twice"))
        for j = i+1:d-1
            p[j] = CTPlan(T,forward) # non-transformed dimensions
        end
        p[d] = CTPlan(T, forward, sz[d])
        i = d
    end
    for j = i+1:N
        p[j] = CTPlan(T,forward) # non-transformed dimensions
    end
    w = Array{T}(undef, length(sregion) <= 1 ? 0 : maximum(sz[sregion[1:end-1]]))
    MultiDimPlan{T,forward}(p, w, i)
end

plan_fft(x::AbstractArray{Complex{Tr}}, region) where Tr<:AbstractFloat =
    MultiDimPlan(Complex{Tr}, true, region, size(x))::MultiDimPlan{Complex{Tr}, true}
plan_bfft(x::AbstractArray{Complex{Tr}}, region) where Tr<:AbstractFloat =
    MultiDimPlan(Complex{Tr}, false, region, size(x))::MultiDimPlan{Complex{Tr}, false}

# recursive execution of a MultiDim plan, starting at dimension d, for
# strided arrays (so that we can use linear indexing):
function applydims(p::MultiDimPlan{T}, d,
                   x::StridedArray{T}, x0, y::StridedArray{T}, y0) where T
    if d == p.lastdim
        if d == ndims(x)
            applystep(p.p[d], x,x0,stride(x,d), y,y0,stride(y,d), 1)
        else
            applydim(p, d, d+1, x,x0, y,y0)
        end
    else
        sx = stride(x,d)
        sy = stride(y,d)
        y0_ = y0
        for i = 1:size(x,d)
            applydims(p, d+1, x,x0, y,y0)
            x0 += sx
            y0 += sy
        end
        if p.p[d].n != 0
            applydim(p, d, d+1, y,y0_)
        end
    end
end
# apply p to dimension d of y, in-place, looping over dimensions >= k
function applydim(p::MultiDimPlan{T}, d, k, y::StridedArray{T}, y0) where T
    sy_k = stride(y,k)
    if k == ndims(y)
        P = p.p[d]
        w = p.w
        sy = stride(y,d)
        ny = size(y,d)
        for i = 1:size(y,k)
            applystep(P, y,y0,sy, w,1,1, 1)
            jy = y0
            @inbounds for j = 1:ny
                y[jy] = w[j]
                jy += sy
            end
            y0 += sy_k
        end
    else
        for i = 1:size(y,k)
            applydim(p, d, k+1, y, y0)
            y0 += sy_k
        end
    end
end
# as above, but out-of-place
function applydim(p::MultiDimPlan{T}, d, k,
                  x::StridedArray{T}, x0,
                  y::StridedArray{T}, y0) where T
    sx_k = stride(x,k)
    sy_k = stride(y,k)
    if k == ndims(y)
        P = p.p[d]
        sx = stride(x,d)
        sy = stride(y,d)
        for i = 1:size(y,k)
            applystep(P, x,x0,sx, y,y0,sy, 1)
            x0 += sx_k
            y0 += sy_k
        end
    else
        for i = 1:size(y,k)
            applydim(p, d, k+1, x,x0, y,y0)
            x0 += sx_k
            y0 += sy_k
        end
    end
end

function mul!(y::StridedArray{T},
                  p::MultiDimPlan{T}, x::StridedArray{T}) where T
    N = ndims(x)
    ndims(y) != N && throw(BoundsError())
    for i = 1:N
        P = p.p[i]
        P.n == 0 || (P.n == size(x,i) == size(y,i)) || throw(BoundsError())
    end
    if N > 0 && p.lastdim > 0
        applydims(p, 1, x,1, y,1)
    else
        copyto!(y, x)
    end
    return y
end

*(p::MultiDimPlan{T}, x::StridedArray{T}) where T = mul!(similar(x), p, x)

ext: jl


text:
"""
    NoncentralF(ν1, ν2, λ)
"""
struct NoncentralF{T<:Real} <: ContinuousUnivariateDistribution
    ν1::T
    ν2::T
    λ::T

    function NoncentralF{T}(ν1::T, ν2::T, λ::T) where T
        @check_args(NoncentralF, ν1 > zero(T) && ν2 > zero(T))
        @check_args(NoncentralF, λ >= zero(T))
        new{T}(ν1, ν2, λ)
    end
end

NoncentralF(ν1::T, ν2::T, λ::T) where {T<:Real} = NoncentralF{T}(ν1, ν2, λ)
NoncentralF(ν1::Real, ν2::Real, λ::Real) = NoncentralF(promote(ν1, ν2, λ)...)
NoncentralF(ν1::Integer, ν2::Integer, λ::Integer) = NoncentralF(Float64(ν1), Float64(ν2), Float64(λ))

@distr_support NoncentralF 0.0 Inf

#### Conversions

function convert(::Type{NoncentralF{T}}, ν1::S, ν2::S, λ::S) where {T <: Real, S <: Real}
    NoncentralF(T(ν1), T(ν2), T(λ))
end
function convert(::Type{NoncentralF{T}}, d::NoncentralF{S}) where {T <: Real, S <: Real}
    NoncentralF(T(d.ν1), T(d.ν2), T(d.λ))
end

### Parameters

params(d::NoncentralF) = (d.ν1, d.ν2, d.λ)
@inline partype(d::NoncentralF{T}) where {T<:Real} = T

### Statistics

function mean(d::NoncentralF{T}) where T<:Real
    d.ν2 > 2 ? d.ν2 / (d.ν2 - 2) * (d.ν1 + d.λ) / d.ν1 : T(NaN)
end

var(d::NoncentralF{T}) where {T<:Real} = d.ν2 > 4 ? 2d.ν2^2 *
               ((d.ν1 + d.λ)^2 + (d.ν2 - 2)*(d.ν1 + 2d.λ)) /
               (d.ν1 * (d.ν2 - 2)^2 * (d.ν2 - 4)) : T(NaN)

### Evaluation & Sampling

@_delegate_statsfuns NoncentralF nfdist ν1 ν2 λ

function rand(rng::AbstractRNG, d::NoncentralF)
    r1 = rand(rng, NoncentralChisq(d.ν1,d.λ)) / d.ν1
    r2 = rand(rng, Chisq(d.ν2)) / d.ν2
    r1 / r2
end

# TODO: remove RFunctions dependency once NoncentralChisq has its removed
@rand_rdist(NoncentralF)
function rand(d::NoncentralF)
    r1 = rand(NoncentralChisq(d.ν1,d.λ)) / d.ν1
    r2 = rand(Chisq(d.ν2)) / d.ν2
    r1 / r2
end


text:
using ApproxFun, BandedMatrices, BlockBandedMatrices, SparseArrays, LinearAlgebra, LazyArrays, BlockArrays, MultivariateOrthogonalPolynomials, Test
using SparseVolterraExamples

#############################
## Various tests for integration Qy and expansion Ey operator consistency.
#############################
@testset "Testing that E_x from P01 works correctly" begin
    g = Fun(x -> exp(x+3x^2), Jacobi(0,1,0..1));
    N = ncoefficients(g);
    G = Fun(JacobiTriangle(), triExgenP01(N) * g.coefficients)
    @test G(0.1,0.2) ≈ g(0.1);
    #########################################################################
    f = Fun(x -> exp(x+3x^2), Legendre(0..1));
    N = ncoefficients(f);
    C = (I : Legendre(0..1) → Jacobi(0,1,0..1))[1:N,1:N];
    F = Fun(JacobiTriangle(), triExgenP01(N) * (C * f.coefficients));
    @test F(0.1,0.2) ≈ f(0.1);
end

@testset "Testing that Q and E_y from P01 work correctly" begin
    # define a function which we want to integrate from 0 to 0.9
    # the expected result for this function is approx. 5.24914
    f = Fun(y -> exp(y+3*y^2), Jacobi(0,1,0..1));
    Ey=triEygenP01(length(f.coefficients));
    F = Fun(JacobiTriangle(), Ey*f.coefficients);
    @test F(0.1,0.2) ≈ f(0.2);
    @test F(0.1,0.3) ≈ f(0.3);
    #########################################################################
    # define a function which we want to integrate from 0 to 0.9
    # the expected result for this function is approx. 5.24914
    f = Fun(y -> exp(y+3*y^2), Jacobi(0,1,0..1));
    N = ncoefficients(f);
    # option 1: use direct ApproxFun + Q generator. This does not use E.
    @test triNoKernelDirect(f,0.9) ≈ sum(Fun(x -> f(x), 0..0.9));
    # option 2: generate Q and E separately and apply them to f.coefficients
    Q=triQgen(N);
    Ey=triEygenP01(length(f.coefficients));
    @test Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*f.coefficients)(0.1) ≈
                sum(Fun(x -> f(x), 0..0.9));
    # we can check that Ey does what it should like this:
    @test pad(Fun((x,y) -> f(y), JacobiTriangle()).coefficients,length(Ey*f.coefficients)) ≈ Ey*f.coefficients;
    #########################################################################
    # again for a different function from 0 to 0.5, expected result here is approx. 0.0917462
    g = Fun(y -> sin(8*y)+y^3-tan(y), Jacobi(0,1,0..1));
    N = ncoefficients(g);

    # option 1: use direct ApproxFun + Q generator. This does not use E.
    triNoKernelDirect(g,0.5);

    # option 2: generate Q and E separately and apply them to f.coefficients
    Q=triQgen(N);
    Ey=triEygenP01(N);
    @test Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*g.coefficients)(0.5) ≈
                sum(Fun(x -> g(x), 0..0.5));

    # we can check that Ey does what it should like this:
    @test pad(Fun((x,y) -> g(y), JacobiTriangle()).coefficients,length(Ey*g.coefficients)) ≈ Ey*g.coefficients;
end

@testset "Testing that Q and E_y from P00 work correctly" begin
    # define a function which we want to integrate from 0 to 0.9
    # the expected result for this function is approx. 5.24914
    f = Fun(y -> exp(y+3*y^2), Legendre(0..1));
    Ey=triEygenP00(length(f.coefficients));
    F = Fun(JacobiTriangle(), Ey*f.coefficients);
    @test F(0.1,0.2) ≈ f(0.2);
    @test F(0.1,0.3) ≈ f(0.3);
    #########################################################################
    # define a function which we want to integrate from 0 to 0.9
    # the expected result for this function is approx. 5.24914
    f = Fun(y -> exp(y+3*y^2), Legendre(0..1));
    N = ncoefficients(f);
    # option 1: use direct ApproxFun + Q generator. This does not use R.
    @test triNoKernelDirect(f,0.9) ≈ sum(Fun(x -> f(x), 0..0.9));
    # option 2: generate Q and E separately and apply them to f.coefficients
    Q=triQgen(N);
    Ey=triEygenP00(length(f.coefficients));
    @test Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*f.coefficients)(0.1) ≈
                sum(Fun(x -> f(x), 0..0.9));
    # we can check that Ey does what it should like this:
    @test pad(Fun((x,y) -> f(y), JacobiTriangle()).coefficients,length(Ey*f.coefficients)) ≈ Ey*f.coefficients;
    #########################################################################
    # again for a different function from 0 to 0.5, expected result here is approx. 0.0917462
    g = Fun(y -> sin(8*y)+y^3-tan(y), Legendre(0..1));
    N = ncoefficients(g);

    # option 1: use direct ApproxFun + Q generator. This does not use R.
    triNoKernelDirect(g,0.5);

    # option 2: generate Q and R separately and apply them to f.coefficients
    Q=triQgen(N);
    Ey=triEygenP00(N);
    @test Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*g.coefficients)(0.5) ≈
                sum(Fun(x -> g(x), 0..0.5));

    # we can check that Ey does what it should like this:
    @test pad(Fun((x,y) -> g(y), JacobiTriangle()).coefficients,length(Ey*g.coefficients)) ≈ Ey*g.coefficients;
end

#############################
## Various tests for the interactions of Qy,Ey with the Jacobi operators
#############################

@testset "Testing that Q*E*f integrates correctly" begin
    S = Jacobi(0,1,0..1);
    f = Fun(x -> 1, S);
    N = length(f.coefficients);
    Q = triQgen(N);
    Ey = triEygenP01(N);
    QEf = Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*f.coefficients);
    @test triNoKernelDirect(f,1.) ≈ QEf(0);
    #########################################################################
    S = Jacobi(0,1,0..1);
    f = Fun(x -> x, S);
    N = length(f.coefficients);
    Q = triQgen(N);
    Ey = triEygenP01(N);
    QEf = Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*f.coefficients);
    @test triNoKernelDirect(f,1.) ≈ QEf(0);
    #########################################################################
    S = Jacobi(0,1,0..1);
    f = Fun(x -> 1,);
    xf = Fun(x -> x, S);
    N = length(f.coefficients);
    y = Fun(identity, S);
    Jy = Multiplication(y, S);
    Q = triQgen(N);
    Ey = triEygenP01(N);
    QEf = Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*f.coefficients);
    QExf = Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*Jy[Block.(1:N), Block.(1:N)]*f.coefficients);
    @test triNoKernelDirect(f,1.) ≈ QEf(0);
    #########################################################################
    S = Jacobi(0,1,0..1);
    f = Fun(x -> 1+exp(3*x^2), S);
    xf = Fun(x -> x*(1+exp(3*x^2)),S);
    N = length(f.coefficients);
    y = Fun(identity, S);
    Jy = Multiplication(y, S);
    Q = triQgen(N);
    Ey = triEygenP01(N);
    QEf = Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*f.coefficients);
    QExf = Fun(JacobiWeight(0,1, Jacobi(0,1, 0..1)),(1/2)*Q*Ey*Jy[Block.(1:N), Block.(1:N)]*f.coefficients);
    @test triNoKernelDirect(f,1.) ≈ QEf(0);
end

@testset "Testing if expected commutation relations hold" begin
    # we test if the direct Q*EyP01 diagonal matrix is computed correctly
    N=50;
    @test triQEygenP01(N) ≈ triQgen(N)*triEygenP01(N);
    #########################################################################
    # we test expected commutation relations of R and Y on the triangle
    # once with a function and one without function
    N = 50;
    S = Jacobi(0,1,0..1);
    T = JacobiTriangle();
    f = Fun(x->2+x+exp(x^5+1), S);
    E = sparse(triEygenP01(N));
    y = Fun(identity, S);
    J = sparse(Multiplication(y, S)[Block.(1:N), Block.(1:N)]);
    x, y = Fun(identity, T);
    Y = sparse(Multiplication(y, T)[Block.(1:N), Block.(1:N)]);
    @test Y*E ≈ E*J;
    @test Y*E*pad(f.coefficients,N) ≈ E*(J*pad(f.coefficients,N));
    #########################################################################
    # we test expected commutation relations of Q and X
    N = 50;
    S = Jacobi(0,1,0..1);
    T = JacobiTriangle();
    Q = sparse(triQgen(N));
    x = Fun(identity, S);
    J = sparse(Multiplication(x, S)[Block.(1:N), Block.(1:N)]);
    x, y = Fun(identity, T);
    X = sparse(Multiplication(x, T)[Block.(1:N), Block.(1:N)]);
    @test Q*X ≈ J*Q
end

@testset "Selective testing of monomial kernel basis consistency" begin
    k = 81;
    N = 30;
    basisvec = triJQEJbasislexP01(k,N,false);
    basisgrid = triJQEJbasisgridgenP01(k,k,N);
    # first couple
    @test basisvec[1] ≈ basisgrid[1][1];
    @test basisvec[2] ≈ basisgrid[1][2];
    @test basisvec[3] ≈ basisgrid[2][1];
    @test basisvec[4] ≈ basisgrid[2][2];
    # square number sequence elements
    @test basisvec[9] ≈ basisgrid[3][3];
    @test basisvec[16] ≈ basisgrid[4][4];
    @test basisvec[25] ≈ basisgrid[5][5];
    @test basisvec[36] ≈ basisgrid[6][6];
    @test basisvec[81] ≈ basisgrid[9][9];
    # some pseudo random picks
    @test basisvec[12] ≈ basisgrid[3][4];
    @test basisvec[27] ≈ basisgrid[2][6];
    @test basisvec[80] ≈ basisgrid[9][1];
end


text:
using Symbolics
using LinearAlgebra
using OrdinaryDiffEq: Tsit5, solve, ODEProblem, ODEFunction, remake
using ForwardDiff: jacobian!

"""
    function ambient_forcing(f::ODEProblem, z, tau, Frand)

    Solves a given Ambient Forcing ODEProblem with initial condition z,
    end time tau and ambient vector Frand.

"""
function ambient_forcing(afoprob::ODEProblem, z, tau, Frand)
    !(typeof(afoprob.f.f) <: AmbientForcingODE) && error("Expected an AmbientForcingODE.")
    return solve(remake(afoprob, u0=z, tspan=(0.0, tau), p=Frand),
        Tsit5(), save_everystep=false, save_start=false)[end]
end

"""
    function ambient_forcing(f::ODEProblem,  Frand)

    Solves a given Ambient Forcing ODEProblem with ambient vector Frand.

"""
function ambient_forcing(afoprob::ODEProblem, Frand)
    !(typeof(afoprob.f.f) <: AmbientForcingODE) && error("Expected an AmbientForcingODE.")
    return solve(remake(afoprob, p=Frand),
        Tsit5(), save_everystep=false, save_start=false)[end]
end

"""
    constraint_equations(f::ODEFunction)
Returns the constraint equations from an ODEFunction f used in DifferentialEquations.jl.
f must be in Mass Matrix form meaning: M ẋ = f(x), with M diagonal.
f should be inplace.
"""
function constraint_equations(f::ODEFunction, p=nothing)
    M = f.mass_matrix
    M == I && error("There are no constraints in the system!")
    M != Diagonal(M) && error("The constraints are not diagonal.")
    cidx = findall(diag(M) .== 0)
    g(x) = (dx = similar(x);
    f(dx, x, p, 0.0);
    dx[cidx])
    return g
end

"""
    constrained_jac_from_f(f::ODEFunction, dim, const_idx)

    Returns a mutating functions (_out,u) that writes
    the Jacobian of the constraint equations of f at u into
    the Matrix _out. Uses Symbolics.jl
"""
function constrained_jac_from_f(f::ODEFunction, dim, cidx)
    @variables usym[1:dim]
    @variables _dusym[1:dim]
    dusym = collect(_dusym)
    f(dusym, collect(usym), nothing, 0.0) # mutate dusym
    duconst = simplify.(dusym[cidx])
    # This builds a jacobian (could be sparse?)
    # [!] We should offer an option for parameters
    symjac = Symbolics.jacobian(duconst, collect(usym))
    return Symbolics.build_function(symjac, usym, expression=Val{false})[2]
end

"""
    constrained_jac_from_f_fd(f::ODEFunction)

    Returns a mutating functions (_out,u) that writes
    the Jacobian of the constraint equations of f at u into
    the Matrix _out. Uses ForwardDiff.jl
"""
function constrained_jac_from_fd(f::ODEFunction)
    # [!] We should offer an option for parameters
    g = constraint_equations(f)
    return (_out, u) -> jacobian!(_out, g, u)
end

"""
    function ambient_forcing_problem(f::ODEFunction, z, tau, Frand)

    Returns an AmbientForcingProblem for a given constraint ODE function f.
    z is an initial condition that fulfills the constraints, tau is the
    integration period for the ambient forcing algorithm, Frand is a random
    direction from the ambient space.

    The AmbientForcingODE describes the evolution of trajectories in
    the constant vector field Frand projected onto the tangential bundle
    of the constraint manifold.

    Throws an error if f has no constraints in
    mass_matrix form or if the mass_matrix is not diagonal.
"""
function ambient_forcing_problem(f::ODEFunction, z, tau, Frand; method = :Symbolics)
    # Check for consistent constraints
    M = f.mass_matrix
    M == I && error("There are no constraints in the system!")
    M != Diagonal(M) && error("The constraints are not diagonal.")
    dim = size(M, 1)
    cidx = findall(diag(M) .== 0)
    cdim = length(cidx)
    if method == :Symbolics
        fjac = constrained_jac_from_f(f, dim, cidx)
    elseif method == :ForwardDiff
        fjac = constrained_jac_from_fd(f)
    else
        error("Invalid differentiation method specified: $(method).")
    end
    # Initialize the buffers for the Jacobian of the constraints J
    # and the intermediate vectors
    J = similar(Frand, cdim, dim)
    invJJT = similar(Frand, cdim, cdim)
    buff1 = similar(Frand, cdim)
    buff2 = similar(Frand, cdim)
    nafo = AmbientForcingODE(fjac, J, invJJT, buff1, buff2)
    prob = ODEProblem(nafo, z, (0.0, tau), Frand)
    return prob
end

"""
    struct AmbientForcingODE{T,S}

    fjac is a function returning the Jacobian matrix of the constraint equations
    at a point u_0. J, invJJT, buff1, buff2 are buffers needed for the projection.

"""
struct AmbientForcingODE{T,S}
    fjac::T
    J::Matrix{S}
    invJJT::Matrix{S}
    buff1::Vector{S}
    buff2::Vector{S}
end

"""
    function (afo::AmbientForcingODE)(du, u, Frand, t)

    Computes the projection of Frand onto the nullspace of
    the Jacobian afo.J at u by evaluating

    `(I - J^T * inv(J * J^T) * J) * Frand`

    and storing it in du.

    Since manifolds have the same dimension everywhere the inverse
    in this product should exist.
"""
function (afo::AmbientForcingODE)(du, u, Frand, t)
    # Evaluate Jacobian at u and save into J
    afo.fjac(afo.J, u) # add (p,t)?
    afo.invJJT .= inv(afo.J * transpose(afo.J))
    # Allocation free matrix multiply
    mul!(afo.buff1, afo.J, Frand)
    mul!(afo.buff2, afo.invJJT, afo.buff1)
    # Next two lines are du = Frand - J^T * buff2
    du .= Frand
    mul!(du, transpose(afo.J), afo.buff2, -1.0, 1)
    return nothing
end


text:
using PkgBenchmark
using CategoricalArrays

@benchgroup "isequal(A, v::String)" begin
    function sumequals(A::AbstractArray, v::Any)
        n = 0
        @inbounds for x in A
            n += isequal(x, v)
        end
        n
    end

    ca = CategoricalArray(repeat(string.('A':'J'), outer=1000))
    nca = CategoricalArray(repeat([missing; string.('A':'J')], outer=1000))
    @bench "CategoricalArray{String}" sumequals(ca, "D")
    @bench "CategoricalArray{Union{String, Missing}}" sumequals(nca, "D")
end

@benchgroup "isequal(A, v::CategoricalValue)" begin
    function sumequals(A::AbstractArray, v::CategoricalValue)
        n = 0
        @inbounds for x in A
            n += isequal(x, v)
        end
        n
    end

    ca = CategoricalArray(repeat(string.('A':'J'), outer=1000))
    nca = CategoricalArray(repeat([missing; string.('A':'J')], outer=1000))
    @bench "CategoricalArray{String}" sumequals(ca, ca[1])
    @bench "CategoricalArray{Union{String, Missing}}" sumequals(nca, nca[1])
end

# With many levels, checking whether some levels have been dropped can be very slow (#93)
@benchgroup "CategoricalArray{String} with many levels" begin
    a = rand([@sprintf("id%010d", k) for k in 1:1000], 10000)
    @bench "CategoricalArray(::Vector{String})" CategoricalArray(a)

    a = rand([@sprintf("id%010d", k) for k in 1:1000], 10000)
    ca = CategoricalArray(a)
    levs = levels(ca)
    @bench "levels! with original levels" levels!(ca, levs)

    levs = reverse(levels(ca))
    @bench "levels! with resorted levels" levels!(ca, levs)

    levs = [levels(ca); [@sprintf("id2%010d", k) for k in 1:1000]]
    @bench "levels! with many additional levels" levels!(ca, levs)
end


text:
#!/usr/bin/env julia

using Printf

background_color = "black"

function gen_digit(color_name, color, digit)
	hex_digit = @sprintf("%x", digit)
	file_name = "$(color_name)_$hex_digit"
	open("$file_name.svg", "w") do of
		p(args...) = println(of, args...)
		p("<?xml version=\"1.0\" encoding=\"UTF-8\"?>")
		p("<!DOCTYPE svg>")
		p("<svg version=\"1.1\" width=\"32\" height=\"64\">")
		l(d) = p("<path fill=\"$color\" d=\"$d\"/>")
		if digit in [0, 2, 3, 5, 6, 7, 8, 9, 10, 12, 14, 15] # segm a
			l("M5 3  L8 0  L23 0  L26 3  L23 6  L8 6  Z")
		end
		if digit in [0, 1, 2, 3, 4, 7, 8, 9, 10, 13] # segm b
			l("M28 5  L31 8  L31 27  L28 30  L25 27  L25 8  Z")
		end
		if digit in [0, 1, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13] # segm c
			l("M28 57  L31 54  L31 35  L28 32  L25 35  L25 54  Z")
		end
		if digit in [0, 2, 3, 5, 6, 8, 9, 11, 12, 13, 14] # segm d
			l("M5 59  L8 62  L23 62  L26 59  L23 56  L8 56  Z")
		end
		if digit in [0, 2, 6, 8, 10, 11, 12, 13, 14, 15] # segm e
			l("M3 57  L6 54  L6 35  L3 32  L0 35  L0 54  Z")
		end
		if digit in [0, 4, 5, 6, 8, 9, 10, 11, 12, 14, 15] # segm f
			l("M3 5  L6 8  L6 27  L3 30  L0 27  L0 8  Z")
		end
		if digit in [2, 3, 4, 5, 6, 8, 9, 10, 11, 13, 14, 15] # segm g
			l("M5 31  L8 28  L23 28  L26 31  L23 34  L8 34  Z")
		end
		p("</svg>")
	end

	run(Cmd([
		"inkscape",
		"--without-gui",
		"--export-width=32",
		"--export-height=64",
		"--export-background=$background_color",
		"$file_name.svg",
		"--export-png=$file_name.png",
	]))
	run(`rm $file_name.svg`)
	#run(`eog $file_name.png`)

	return file_name
end

function gen_digit_set(color_name, color)
	for d in 0:15
		gen_digit(color_name, color, d)
	end
end

gen_digit_set("green", "#00CC00")
gen_digit_set("red", "#FF0000")


text:
@testset "474.ones-and-zeroes.jl" begin
    @test find_max_form(["10", "0001", "111001", "1", "0"], 5, 3) == 4
    @test find_max_form(["10", "0", "1"], 1, 1) == 2
end


text:
@pyimport networkx as nx
@pyimport networkx.algorithms as nx_algo

"""Number of vertices in the curve"""
nv(c::Curve) = length(c.points)

"""Number of edges in the curve"""
ne(c::Curve) = length(c.segments)

"""Map edge indices to the 2 vertex indices connected to it"""
function edge_vertex_connectivity(c::Curve)
    # Make it a dict for compatibility
    mapping = Dict{Int, Vector{Int}}()
    for (edge, vertices) in enumerate(c.segments)
        mapping[edge] = collect(vertices)
    end
    mapping
end

# FIXME
"""
Compute the inverse mapping of when single K maps to several Vs. So
one V maps to several Ks.
"""
function invert{K, V}(mapping::Dict{K, Vector{V}})
    i_mapping = Dict{V, Set{K}}()
    for pair in mapping
        value = pair.first
        for key in pair.second
            haskey(i_mapping, key) ? push!(i_mapping[key], value) : push!(i_mapping, Pair(key, Set(value)))
        end
    end
    Dict(k => collect(v) for (k, v) in i_mapping)
end

"""Map of vertex indices to at least one edge that is connected to it"""
vertex_edge_connectivity(c::Curve) = invert(edge_vertex_connectivity(c))

# For several properties it will be useful to rely on representation of
# the curve as graph
function Graph(c::Curve)
    g = Graph(nv(c))
    for edge in c.segments
        add_edge!(g, edge)
    end
    # Graph is not MultiGraph!
    @assert ne(c) == ne(g)
    g
end

"""A list of nodes [indices in vertices(graph)] connected to only one edge"""
function boundary_vertices(g::Graph)
    degrees = degree(g)
    find(x -> x == 1, degrees)
end

boundary_vertices(c::Curve) = boundary_vertices(Graph(c))

"""A list of nodes [indices in vertice(graph)] connected to more than 2 edges"""
function branching_vertices(g::Graph)
    degrees = degree(g)
    find(x -> x > 2, degrees)
end

branching_vertices(c::Curve) = branching_vertices(Graph(c))

"""An array holding for each vertex[index] a number of edges connected to it"""
degree(c::Curve) = degree(Graph(c))

"""Represent degree of the curve as VertexFunction"""
degree_f(c::Curve) = VertexFunction(c, degree(c))

"""
We are interested in branching of the curve. A branch is a collection
of edges bounded by 2 vertices which are terminal, meaning that the vertex
is either a branching vertex or a boundary vertex. Here we pair the terminal
vertices with the edges connected to them.
"""
function terminal_map(c::Curve)
    degrees = degree(c)
    # Avoid building the graph twice
    terminals = find(x -> x == 1 || x > 2, degrees)

    v2e = vertex_edge_connectivity(c)
    Dict(t => v2e[t] for t in terminals)
end

"""
A branch is a collection of vertices where the first and last vertices
are `terminal` vertices. Here we make a map of branches
(first, last) => [[edges of the branch]] (where edges as given by vertex
to edge map. NOTE the value datatype is Vector{Vector{Int}} because in
general there can be multiple branches between the vertices. In such case
the reduced curve (each branch is an edge) cannot be represented as a Graph
(but MultiGraph is okay)
"""
function branch_map(c::Curve)
    tmap = terminal_map(c::Curve)

    v2e = vertex_edge_connectivity(c)
    e2v = edge_vertex_connectivity(c)

    terminals = Dict(t => length(tmap[t]) for t in keys(tmap))
    # The idea is to pich a terminal, follow some edge connected to
    # it until we hit a terminal. We keep removing the starting/ending
    # edges so that the at some point nnzkeys is empty
    bmap = Dict{Tuple{Int, Int}, Vector{Vector{Int}}}()
    while !isempty(terminals)
        # Pick som terminal
        start = first(keys(terminals))
        # And we will follow some edge connected to it; leaving no trail
        edge = pop!(tmap[start])
        terminals[start] -= 1
        terminals[start] == 0 && delete!(terminals, start)

        branch = [edge]
        next_v = start
        while true
            v0, v1 = e2v[edge]
            # Pick a different vertex then which lead us here
            next_v = (v0 == next_v) ? v1 : v0
            # Terminal vertex?
            next_v ∈ keys(terminals) && break

            # Otherwise this is a vertex with 2 edges. Next edge is different
            # Then current
            e0, e1 = v2e[next_v]
            edge = (e0 == edge) ? e1 : e0
            push!(branch, edge)
        end
        # If the vertex is terminal remove the edge from the end so that
        # we don't follow the path back
        deleteat!(tmap[next_v], findfirst(x -> x == edge, tmap[next_v]))
        # This reduces connections of the terminal
        terminals[next_v] -= 1
        terminals[next_v] == 0 && delete!(terminals, next_v)

        # We make a rule that key is an ordered pair
        key = (start < next_v) ? (start, next_v) : (next_v, start)
        # And adjust the path if we flipped
        key != (start, next_v) && reverse!(branch)

        (key in keys(bmap))? push!(bmap[key], branch) : push!(bmap, Pair(key, [branch]))
    end
    bmap
end

"""Branch count"""
function num_branches(c::Curve, bmap::Dict{Tuple{Int, Int}, Vector{Vector{Int}}})
    sum((length(v) for v in values(bmap)))
end

num_branches(c::Curve) = num_branches(c, branch_map(c))

"""
Loops (as collection of edges) in the curve (with precomputed branch map)
NOTE: We produce independent basis not the `minimal` basis by which I mean
that

 ___
| /|
|/_| is not decomposed into 2 triangles but maybe one tri + a square.
"""
function loop_basis!(c::Curve, bmap::Dict{Tuple{Int, Int}, Vector{Vector{Int}}})
    loops = Vector{Vector{Int}}()

    # Self loops (start, start) can be added right-away
    paths = collect(keys(bmap))
    closed = find(p -> p[1] == p[2], paths)
    for c in reverse!(closed)
        for loop in bmap[paths[c]]
            push!(loops, loop)
        end
        delete!(bmap, paths[c])
    end

    # To reduce to graph we build loops from multi edges
    paths = collect(keys(bmap))
    while any(length(v) > 1 for v in values(bmap))
        # Get a multiedge
        path = paths[findfirst(p -> length(bmap[p]) > 1, paths)]
        # This could be a pop, but I will remove the one which is longest
        sort!(bmap[path], by=length)
        there = pop!(bmap[path])
        # Each removed edge formed a loop with some of the remaining
        # This is an okay question
        back = bmap[path][end]
        # A loop there - back
        push!(loops, append!(there, back))
    end

    # NOTE: LightGraphs does not have cycle basis so we reach out to python
    py_graph = nx.Graph()
    for (v0, v1) in keys(bmap)
        py"$(py_graph).add_edge($(v0), $(v1))"
    end

    # Here in each row is
    cycle_basis = nx_algo.cycle_basis(py_graph)
    for row in 1:size(cycle_basis, 1)
        vertices = cycle_basis[row, :]

        # [1, 2, 3, 4] = (1, 2), (2, 3), (3, 4), (4, 1)
        loop = Vector{Int}()
        for (start, stop) in zip(vertices, vcat(vertices[2:end], vertices[1]))
            # We have sorted keys
            path = (start < stop) ? (start, stop) : (stop, start)
            @assert length(bmap[path]) == 1
            append!(loop, first(bmap[path]))
        end
        push!(loops, loop)
    end
    loops
end

"""Loops in the curve (without precomputed branch map)"""
loop_basis(c::Curve) = loop_basis!(c, branch_map(c))

"""Number of loops"""
function num_loops(c::Curve, bmap::Dict{Tuple{Int, Int}, Vector{Vector{Int}}})
    # The number of loops in the (vertex)graph is the same as the number
    # of zeros in the spectrum of the laplacian matrix of the dual(edge) graph.
    # This definition is potentialy very expensive to use so the here is
    # to first buld a reduced graph where each branch in the orginal
    # graph is a signel edges => smaller graph. However, in the reduction
    # process the graph might become a MultiGraph so we first reduce
    # to Graph case

    #  __
    # |  |
    # |__|
    nloops = 0
    paths = collect(keys(bmap))
    # The case where the situation is obcious from keys
    closed = find(p -> p[1] == p[2], paths)
    for c in reverse!(closed)
        nloops += length(bmap[paths[c]])
        deleteat!(paths, c)
    end

    #  __ B
    # | /|
    # |/_|
    # A
    multi_edges = filter(p -> length(bmap[p]) > 1, paths)
    # Count multiplicity of multi edges
    nloops += isempty(multi_edges) ? 0 : sum(length(bmap[edge])-1 for edge in multi_edges)

    # NOTE: LightGraphs does not have cycle basis so we reach out to python
    py_graph = nx.Graph()
    for (v0, v1) in paths
        py"$(py_graph).add_edge($(v0), $(v1))"
    end
    # Here in each row is one cycle basis
    nloops += size(nx_algo.cycle_basis(py_graph), 1)
    nloops
end

num_loops(c::Curve) = num_loops(c, branch_map(c))

"""
Graph from the mesh has as edges the line elements. In dual graph the
edges are vertices. Zeros of the Laplace(DualGraph(curve)) count the
loops
"""
function DualGraph(g::Graph)
    vertex_neighbors = LightGraphs.SimpleGraphs.adj(g)

    dual_graph = Graph(ne(g))
    for (v, neighbors) in enumerate(vertex_neighbors)
        for n in neighbors
            add_edge!(dual_graph, (n, v))
        end
    end
    dual_graph
end

DualGraph(c::Curve) = DualGraph(Graph(c))

"""
If the graph is split into branches how many colors are needed to
colorthem so that no 2 neighboring branches have the same color.
"""
function num_colors(c::Curve, bmap::Dict{Tuple{Int, Int}, Vector{Vector{Int}}})
    # Build a reduced representation where each branch is a vertex
    # a graph. We code the branch as bmap[first(code)][last(code)]
    codes = Vector{Tuple{Tuple{Int, Int}, Int}}()
    for key in keys(bmap)
        for vertex in 1:length(bmap[key])
            push!(codes, (key, vertex))
        end
    end

    # Edges
    g = Graph(length(codes))
    for (v0, code0) in enumerate(codes)
        e0, e1 = first(code0)  # Vertices of the branch
        for (v1, code1) in enumerate(codes[(v0+1):end])
            if (e0 in first(code1) || e1 in first(code1))
                add_edge!(g, (v0, v0+v1))
            end
        end
    end

    colors = greedy_color(g)
    colors.num_colors
end

num_colors(c::Curve) = num_colors(c, branch_map(c))

"""
If the graph is split into branches how many colors are needed to
colorthem so that no 2 neighboring branches have the same color.
Produce a color function for the graph (EdgeFunction).
"""
function color_f(curve::Curve, bmap::Dict{Tuple{Int, Int}, Vector{Vector{Int}}})
    # Build a reduced representation where each branch is a vertex
    # a graph. We code the branch as bmap[first(code)][last(code)]
    codes = Vector{Tuple{Tuple{Int, Int}, Int}}()
    for key in keys(bmap)
        for vertex in 1:length(bmap[key])
            push!(codes, (key, vertex))
        end
    end

    # Edges
    g = Graph(length(codes))
    for (v0, code0) in enumerate(codes)
        e0, e1 = first(code0)  # Vertices of the branch
        for (v1, code1) in enumerate(codes[(v0+1):end])
            if (e0 in first(code1) || e1 in first(code1))
                add_edge!(g, (v0, v0+v1))
            end
        end
    end
    color_map = greedy_color(g)

    color_data = Vector{Int}(ne(curve))
    for (branch, color) in enumerate(color_map.colors)
        key, index = codes[branch]
        for edge in bmap[key][index]
            color_data[edge] = color
        end
    end
    EdgeFunction(curve, color_data)
end

color_f(c::Curve) = color_f(c, branch_map(c))

#using StaticArrays
segments = [(SVector(0., 0.), SVector(1., 0)),
            (SVector(1., 0.), SVector(1., -1)),
            (SVector(1., -1.), SVector(0., -1)),
            (SVector(0., -1.), SVector(0., 0)),
            (SVector(0., 0.), SVector(-1., 0)),
            (SVector(-1., 0.), SVector(-1., 1)),
            (SVector(-1., 1.), SVector(0., 1)),
            (SVector(0., 1.), SVector(0., 0))]
c = Curve(segments)

segments = [(SVector(0., 0.), SVector(1., 0)),
            (SVector(1., 0.), SVector(1., -1)),
            (SVector(1., -1.), SVector(0., -1)),
            (SVector(0., -1.), SVector(0., 0)),
            (SVector(0., 0.), SVector(-1., 0)),
            (SVector(-1., 0.), SVector(-1., 1)),
            (SVector(-1., 1.), SVector(0., 1)),
            (SVector(0., 1.), SVector(0., 0)),
            (SVector(-1., 0), SVector(0., 1)),
            (SVector(0., -1.), SVector(1., 0))]
cc = Curve(segments)

segments = [(SVector(0., 0.), SVector(1., 0)),
            (SVector(1., 0.), SVector(1., -1)),
            (SVector(1., -1.), SVector(0., -1)),
            (SVector(0., -1.), SVector(0., 0)),
            (SVector(0., 0.), SVector(-1., 0)),
            (SVector(-1., 0.), SVector(-1., 1)),
            (SVector(-1., 1.), SVector(0., 1)),
            (SVector(0., 1.), SVector(0., 0)),
            (SVector(-1., 0), SVector(0., 1)),
            (SVector(0., -1.), SVector(1., 0)),
            (SVector(1., -1.), SVector(1., -2)),
            (SVector(1., -2.), SVector(2., -2)),
            (SVector(2., -2.), SVector(2., -1)),
            (SVector(2., -1.), SVector(1., -1))]
ccc = Curve(segments)

segments = [(SVector(0., 0.), SVector(1., 0)),
            (SVector(1., 0.), SVector(1., 1)),
            (SVector(1., 1.), SVector(0., 1)),
            (SVector(0., 1.), SVector(0., 0)),
            (SVector(0., 0.), SVector(1., 1.))]
cccc = Curve(segments)


text:
module ImaginePlots

using ImagineInterface
using AxisArrays, Plots, UnitfulRecipes, Unitful
import Plots: plot

include("plot.jl")

end # module

max_stars_repo_name: HolyLab/ImaginePlots

max_issues_repo_name: HolyLab/ImaginePlots

max_forks_repo_name: HolyLab/ImaginePlots


text:
# Set up model for testing
using ModelObjectsLH, ModelParams

import ModelParams.has_pvector

mdl = ModelParams;

mutable struct Obj1 <: ModelObject
    objId :: ObjectId
    x :: Float64
    y :: Vector{Float64}
    z :: Array{Float64,2}
    pvec :: ParamVector
end

# function Obj1(x, y, z)
#     return Obj1(x, y, z, ParamVector(ObjectId(:pv1)))
# end

function init_obj1(objId)
    # objId = ObjectId(:obj1);
    px = Param(:x, "x obj1", "x1", 11.1, 9.9, 1.1, 99.9, true);
    # Important to have vector of length 1 as test case
    valueY = fill(1.1, 1);
    py = Param(:y, "y obj1", "y1", valueY, valueY .+ 1.0,
        valueY .- 5.0, valueY .+ 5.0, true);
    valueZ = [3.3 4.4; 5.5 7.6];
    pz = Param(:z, "z obj1", "z1", valueZ, valueZ .+ 1.0,
        valueZ .- 5.0, valueZ .+ 5.0, false);
    pvector = ParamVector(objId, [px, py, pz])
    o1 = Obj1(objId, px.value, py.value, pz.value, pvector);
    ModelParams.set_own_values_from_pvec!(o1, true);
    return o1
end

mutable struct Obj3Switches <: ModelSwitches
    pvec :: ParamVector
end

function init_obj3_switches(objId :: ObjectId)
    # objId = ObjectId(:obj3);
    px = Param(:x, "x", "x", 0.5, 0.6, 0.0, 1.0, true);
    yV = [1.0, 2.0];
    py = Param(:y, "y", "y", yV, yV .+ 1.0, [0.0, 0.0], [9.0, 9.0], false);
    v = [1.0 2.0 3.0; 1.1 2.1 3.1];
    isCalM = [true false true; false true false];
    pca = CalArray(:ca, "CalArray", "CalArray", v, v .+ 0.1, v .- 2.0, v .+ 2.0,
        isCalM);
    pvec = ParamVector(objId, [px, py, pca]);
    return Obj3Switches(pvec)
end

ModelObjectsLH.get_object_id(switches :: Obj3Switches) = get_object_id(switches.pvec);

mutable struct Obj3 <: ModelObject
    objId :: ObjectId
    switches :: Obj3Switches
    x :: Float64
    y :: Vector{Float64}
    ca :: Matrix{Float64}
end

ModelParams.get_pvector(o :: Obj3) = o.switches.pvec;

function init_obj3(objId :: ObjectId)
    switches = init_obj3_switches(objId);
    ca = retrieve(switches.pvec, :ca);
    return Obj3(get_object_id(switches), switches, 0.5, [1.0, 2.0], default_value(ca));
end

mutable struct Obj2 <: ModelObject
    objId :: ObjectId
    a :: Float64
    y :: Float64
    b :: Array{Float64,2}
    bvec :: Vector{Float64}
    o3 :: Obj3
    pvec :: ParamVector
end

# function Obj2(a, y, b)
#     return Obj2(a, y, b, ParamVector(ObjectId(:pv1)))
# end

function init_obj2(objId)
    # objId = ObjectId(:obj2);
    pa = Param(:a, "a obj2", "a2", 12.1, 7.9, -1.1, 49.9, true);
    valueY = 9.4;
    py = Param(:y, "y obj2", "y2", valueY, valueY .+ 1.0,
        valueY .- 5.0, valueY .+ 5.0, true);
    valueB = 2.0 .+ [3.3 4.4; 5.5 7.6];
    pb = Param(:b, "b obj2", "b2", valueB, valueB .+ 1.0,
        valueB .- 5.0, valueB .+ 5.0, true);
    pBvec = mdl.make_test_bvector(:bvec; isCalibrated = true, increasing = :increasing);
    pvector = ParamVector(objId, [pa, py, pb, pBvec]);
    obj3 = init_obj3(make_child_id(objId, :obj3));
    o2 = Obj2(objId, pa.value, py.value, pb.value, ModelParams.value(pBvec), obj3, pvector);
    ModelParams.set_own_values_from_pvec!(o2, true);
    return o2
end

mutable struct TestModel <: ModelObject
    objId :: ObjectId
    o1 :: Obj1
    o2 :: Obj2
    a :: Float64
    y :: Float64
end

has_pvector(o :: TestModel) = false;

function init_test_model()
    objName = ObjectId(:testModel, "Test model");
    o1 = init_obj1(ModelParams.make_child_id(objName, :o1, "Child object 1"));
    o2 = init_obj2(ModelParams.make_child_id(objName, :o2, "Child object 2"));
    m = TestModel(objName, o1, o2, 9.87, 87.73);
    sync_values!(m);
    return m
end

# ------------------


text:
const LICENSE_DIR = normpath(joinpath(@__DIR__, "..", "licenses"))
const LICENSES = Dict(
    "MIT" => "MIT \"Expat\" License",
    "BSD2" => "Simplified \"2-clause\" BSD License",
    "BSD3" => "Modified \"3-clause\" BSD License",
    "ISC" => "Internet Systems Consortium License",
    "ASL" => "Apache License, Version 2.0",
    "MPL" => "Mozilla Public License, Version 2.0",
    "GPL-2.0+" => "GNU Public License, Version 2.0+",
    "GPL-3.0+" => "GNU Public License, Version 3.0+",
    "LGPL-2.1+" => "Lesser GNU Public License, Version 2.1+",
    "LGPL-3.0+" => "Lesser GNU Public License, Version 3.0+",
    "EUPL-1.2+" => "European Union Public Licence, Version 1.2+",
)

"""
    available_licenses([io::IO]) -> Nothing

Print the names of all available licenses.
"""
available_licenses(io::IO) = print(io, join(("$k: $v" for (k, v) in LICENSES), "\n"))
available_licenses() = available_licenses(stdout)

"""
    show_license([io::IO], license::AbstractString) -> Nothing

Print the text of `license`. Errors if the license is not found.
"""
show_license(io::IO, license::AbstractString) = print(io, read_license(license))
show_license(license::AbstractString) = show_license(stdout, license)

"""
    read_license(license::AbstractString) -> String

Returns the contents of `license`. Errors if the license is not found. Use
[`available_licenses`](@ref) to view available licenses.
"""
function read_license(license::AbstractString)
    path = joinpath(LICENSE_DIR, license)
    if isfile(path)
        return string(readchomp(path))
    else
        throw(ArgumentError("License '$license' is not available"))
    end
end


text:
struct StaticFlagNode{NBranches,T,Value<:Cell{T}}
    nchildren::Int
    children::NTuple{NBranches,Int}
    flag::OneWayObservable{Bool}
    value::Value
end

StaticFlagNode{NBranches,T,Value}(
    nchildren,
    children,
    flag,
) where {NBranches,T,Value<:Cell{T}} =
    StaticFlagNode{NBranches,T,Value}(nchildren, children, flag, Value())

function concrete(::Type{StaticFlagNode{NBranches,T}}) where {NBranches,T}
    if Base.issingletontype(T)
        return StaticFlagNode{NBranches,T,SingletonCell{T}}
    else
        return StaticFlagNode{NBranches,T,MutableCell{T}}
    end
end

struct StaticTreeBarrier{NArrive,NDepart,T,Value,Op} <:
       SyncBarriers.StaticTreeBarrier{NArrive,NDepart,T}
    n::Int
    local_sense::Vector{Bool}  # TODO: pad
    arrives::Vector{StaticFlagNode{NArrive,T,Value}}
    departs::Vector{StaticFlagNode{NDepart,Nothing,SingletonCell{Nothing}}}
    op::Op
end

StaticTreeBarrier{NArrive,NDepart}(n::Integer) where {NArrive,NDepart} =
    StaticTreeBarrier{NArrive,NDepart,Nothing}(right, n)

function StaticTreeBarrier{NArrive,NDepart,T}(op, n::Integer) where {NArrive,NDepart,T}
    @argcheck n > 0
    arrives = Vector{concrete(StaticFlagNode{NArrive,T})}(undef, n)
    departs = Vector{concrete(StaticFlagNode{NDepart,Nothing})}(undef, n)
    static_tree!(arrives)
    static_tree!(departs)
    local_sense = [false for _ in 1:n]
    return StaticTreeBarrier(n, local_sense, arrives, departs, op)
end

function static_tree_depth(B, n)
    len(d) = (B^d - 1) ÷ (B - 1)
    d = ceil(Int, log(B, (B - 1) * n + 1))
    @assert len(d - 1) <= n <= len(d)
    return d
end

function static_tree!(nodes::AbstractVector{<:StaticFlagNode{NBranches}}) where {NBranches}
    i = _static_tree!(nodes, 1, 1, static_tree_depth(NBranches, length(nodes)))
    @assert i == length(nodes) + 1
end

# Filling nodes in depth-first order so that the reduction does not require commutativity.
function _static_tree!(
    nodes::AbstractVector{<:Node},
    i,
    depth,
    maxdepth,
) where {NBranches,Node<:StaticFlagNode{NBranches}}
    children = ntuple(_ -> firstindex(nodes) - 1, Val(NBranches))
    nchildren = 0
    if depth < maxdepth
        for k in 1:NBranches
            i <= length(nodes) - depth || break
            i = _static_tree!(nodes, i, depth + 1, maxdepth)
            nchildren += 1
            children = Base.setindex(children, i - 1, k)
        end
    end
    nodes[i] = Node(nchildren, children, OneWayObservable{Bool}(false))
    return i + 1
end

function foldr_children(op, acc, nodes::AbstractVector{<:StaticFlagNode}, i)
    node = nodes[i]
    for k in node.nchildren:-1:1
        acc = op(nodes[node.children[k]], acc)
    end
    return acc
end

# The iteration order of foldr_children and foreach_child_flag are different so
# that the communication is done in FIFO manner. (TODO: check if it matters)
function foreach_child_flag(f::F, nodes, i) where {F}
    node = nodes[i]
    for k in 1:node.nchildren
        f(nodes[node.children[k]].flag)
    end
    return
end

SyncBarriers.cycle!(
    handle::BarrierHandle{<:StaticTreeBarrier{<:Any,<:Any,Nothing}},
    spin::Union{Nothing,Integer} = nothing,
) = SyncBarriers.reduce!(handle, nothing, spin)

function SyncBarriers.reduce!(
    handle::BarrierHandle{<:StaticTreeBarrier{<:Any,<:Any,T}},
    value,
    spin::Union{Nothing,Integer} = nothing,
) where {T}
    value = convert(T, value)
    i = handle.i
    barrier = handle.barrier

    acc = foldr_children(value, barrier.arrives, i) do node, value
        waitif(!=(true), node.flag, spin)
        node.flag.value[] = false  # for next episode
        barrier.op(node.value[], value)
    end
    barrier.arrives[i].value[] = acc
    barrier.arrives[i].flag[] = true

    sense = barrier.local_sense[i] = !barrier.local_sense[i]
    root_node = length(barrier)
    if i != root_node
        waitif(!=(sense), barrier.departs[i].flag, spin)
    end
    foreach_child_flag(barrier.departs, i) do flag
        flag[] = sense
    end

    return barrier.arrives[root_node].value[]
end


text:
export AnyPlatform

"""
    AnyPlatform()

A special platform to be used to build platform-independent tarballs, like those
containing only header files.  [`FileProduct`](@ref) is the only product type
allowed with this platform.
"""
struct AnyPlatform <: Platform end
abi_agnostic(p::AnyPlatform) = p
Pkg.BinaryPlatforms.platform_name(::AnyPlatform) = "AnyPlatform"
Pkg.BinaryPlatforms.triplet(::AnyPlatform) = "any"
# Fallback on x86_64-linux-musl, but this shouldn't really matter in practice.
Pkg.BinaryPlatforms.arch(::AnyPlatform) = arch(Linux(:x86_64, libc=:musl))
Base.show(io::IO, ::AnyPlatform) = print(io, "AnyPlatform()")


text:
#
# macros.jl -
#
# Implementation of macros to simplify calling functions of the Spinnaker C
# SDK.
#
#------------------------------------------------------------------------------

"""
    @unchecked_call(func, argtypes, args...) -> err

calls function `func` in the Spinnaker C SDK with arguments `args...` of types
`argtypes` and returns the result `err`.  The function to call may be specified
as a string or as a symbol.  Example:

    @unchecked_call(:spinSystemGetLibraryVersion,
                    (SystemHandle, Ptr{LibraryVersion},),
                    handle(system), ref)

is equivalent to (`lib` is the constant string the path to the Spinnaker
dynamic library and `Cint` is the type of the result returned by all functions
of the Spinnaker C SDK):

    err = ccall((:spinSystemGetLibraryVersion, lib), Cint,
                (SystemHandle, Ptr{LibraryVersion},),
                handle(system), ref)

or

   err = @ccall lib.spinSystemGetLibraryVersion(
       handle(system)::SystemHandle,
       ref::Ptr{LibraryVersion}
   )::Cint

"""
macro unchecked_call(func, args...)
    esc(_unchecked_call_expr(func, args...))
end

"""
    @checked_call(func, argtypes, args...)

calls function `func` in the Spinnaker C SDK with arguments `args...` of types
`argtypes` throwing an exception if the value returned by the function
indicates an error.  The function to call may be specified as a string or as a
symbol.  Example:

    @checked_call(:spinSystemGetLibraryVersion,
                  (SystemHandle, Ptr{LibraryVersion},),
                  handle(system), ref)

is equivalent to (`lib` is the constant string the path to the Spinnaker
dynamic library and `Cint` is the type of the result returned by all functions
of the Spinnaker C SDK):

    let err = ccall((:spinSystemGetLibraryVersion, lib), Cint,
                    (SystemHandle, Ptr{LibraryVersion},),
                    handle(system), ref)
        _check(err, :spinSystemGetLibraryVersion)
    end

"""
macro checked_call(func, args...)
    esc(_checked_call_expr(func, args...))
end

# Yield an expression equivalent to a symbol.
_quote_expr(obj::QuoteNode) = _quote_expr(obj.value)
_quote_expr(str::AbstractString) = _quote_expr(Symbol(str))
_quote_expr(sym::Symbol) = Expr(:quote, sym)

# Yield the expression to call a function of the SDK and return its result.
_unchecked_call_expr(func, args...) =
    Expr(:call, :ccall, Expr(:tuple, _quote_expr(func), :lib), :Cint, args...)

# Yield the expression to call a function of the SDK and check its result.
_checked_call_expr(func, args...) =
    Expr(:call, :_check, _unchecked_call_expr(func, args...), _quote_expr(func))


text:
# TODO: the stdlib function findmax(arr, dim) should do this now
function indvalmax(a::Matrix{T}, dim::Integer=2) where T
    out_size = dim == 2 ? size(a, 1) : size(a, 2)
    out_v = Array(T, out_size)
    out_i = Array(Int64, out_size)

    if dim == 2
        for i=1:out_size
            out_v[i], out_i[i] = findmax(a[i, :])
        end
    elseif dim == 1
        for i=1:out_size
            out_v[i], out_i[i] = findmax(a[:, i])
        end
    else
        error("dim must be 1 or 2. Received $dim")
    end
    return out_v, out_i
end

function valmax(v, f, P, delta)
    n, m = size(f)
    v, x = indvalmax(f + reshape(delta .* (P * v), n, m), 2)
    return v, x
end

function valpol(x,f,P)
    n, m = size(f)
    ind = n * x + (1-n:0)
    fstar = f[ind]
    pstar = P[ind, :]
    return pstar, fstar, ind
end

function expandg(g)
    # Only need if I supply "transfunc". Not doing so
    n, m = size(g)
    P = sparse(1:n*m, g(:), 1, n*m, n)
    return P
end

function diagmult(a::Vector{T}, b::Matrix{T}) where T <: Real
    n = length(a)
    return sparse(1:n, 1:n, a, n, n)*b
end

function ddpsolve(model, v_init=zeros(size(model["reward"], 1)))
    maxit = 100
    tol = sqrt(eps())

    delta = model["discount"]
    f = model["reward"]

    # else block L53-62
    P = model["transprob"]
    n, m = size(f)

    if ndims(P) == 3 && size(P) == (m, n, n)
        P = permutedims(P, [2, 1, 3])
        P = reshape(P, m*n, n)
    elseif ndims(P) == 2 && size(P) == (m*n, n)
    else
        error("P has wrong size")  # Todo make this meaningful
    end

    if length(delta) == 1
        delta = fill(delta, m*n)
    end

    algorithm = "newton"
    x = zeros(n)
    v = v_init

    pstar = Array(Float64, n, n)

    for it = 1:maxit
        vold = v
        xold = x
        v, x = valmax(v, f, P, delta)
        pstar, fstar, ind = valpol(x, f, P)
        v = (speye(n)-diagmult(delta[ind], pstar)) \ fstar
        err = norm(v - vold)
        println("it: $it\terr: $err")
        if x == xold
            break
        end
    end
    return v, x, pstar
end


text:
using Documenter, Generators

DocMeta.setdocmeta!(Generators, :DocTestSetup, :(using Generators), recursive = true)
makedocs(sitename="fred.com", modules=[Generators], format = Documenter.HTML())


text:
using ArnoldiMethod
using ArnoldiMethodTransformations
using LinearAlgebra
using SparseArrays
using Statistics
using Test

N = 20
# construct fixed eval matrix in random basis

@testset "generalized eigenvalue problem: " begin
    for i ∈ 1:500

        A = rand(N,N)
        B = N*I + rand(ComplexF64,N,N)
        # find eigenpairs closest to 5.001 (cannot be 5 as algorithm is unstable if σ is exactly an eval)
        decomp, hist = partialschur(A,B,5.001)
        # get evecs
        λ, v = partialeigen(decomp,5.001)
        @test norm(A*v-B*v*diagm(0=>λ))/prod(size(v)) ≤ sqrt(eps(float(N)))N

        As = sparse(A)
        Bs = sparse(B)
        # find eigenpairs closest to 5.001 (cannot be 5 as algorithm is unstable if σ is exactly an eval)
        decomp, hist = partialschur(As,Bs,5.001)
        # get evecs
        λ, v = partialeigen(decomp,5.001)
        @test norm(A*v-Bs*v*diagm(0=>λ))/prod(size(v)) ≤ sqrt(eps(float(N)))N
    end
end


text:
"""
# p44

Analysis of elastic rigid-jointed frames using a 2-node
Frame structural element and Line finite elements
in 2 or 3 dimensions.

### Constructors
```julia
p44(data)
```
### Arguments
```julia
* `data::Dict{Symbol, Any}` : Dictionary containing all input data
```

### Dictionary keys
```julia
* struc_el::StructuralElement                          : Type of  structural fin_el
* support::Array{Tuple{Int,Array{Int,1}},1}        : Fixed-displacements vector
* loaded_nodes::Array{Tuple{Int,Array{Float64,1}},1} : Node load vector
* properties::Vector{Float64}                          : Material properties
* x_coords::FloatRange{Float64}                        : x coordinate vector
* y_coords::FloatRange{Float64}                        : y coordinate vector
* g_num::Array{Int,2}                                : Element node connections
* fixed_freedoms::Array{Tuple{Vector{Int}}           : Fixed freedoms
```

### Optional additional dictionary keys
```julia
* etype::Vector{Int}                                 : Element material vector
* penalty::Float64                                     : Penalty for fixed freedoms
* z_coords::FloatRange{Float64}                        : z coordinate vector
* eq_nodal_forces_and_moments                          : Equivalent nodal loads
```

### Return values
```julia
* (jfem, dis_df, fm_df)        : Tuple of jFem, dis_df and fm_df
                                 where:
                                    jfem::jFem    : Computational result type
                                    dis_df        : Displacement data table
                                    fm_df         : Forces and moments data table
```

### Related help
```julia
?StructuralElement  : Help on structural elements
?Beam               : Help on a Beam structural fin_el
?FiniteElement      : Help on finite element types
```
"""
function p44(data::Dict{Symbol, Any})

  # Parse & check FEdict data

  if :struc_el in keys(data)
    struc_el = data[:struc_el]
  else
    println("No fin_el type specified.")
    return
  end

  ndim = struc_el.ndim
  nst = struc_el.nst

  # Add radial stress
  if ndim == 3 && typeof(struc_el) !== Frame && struc_el.axisymmetric
    nst = 4
  end

  fin_el = struc_el.fin_el
  @assert typeof(fin_el) <: FiniteElement

  if typeof(fin_el) == Line
    if typeof(struc_el) == Frame
      nels = struc_el.nels
      nn = struc_el.nn
    else
      (nels, nn) = mesh_size(fin_el, struc_el.nxe)
    end
  elseif typeof(fin_el) == Triangle || typeof(fin_el) == Quadrilateral
    (nels, nn) = mesh_size(fin_el, struc_el.nxe, struc_el.nye)
  elseif typeof(fin_el) == Hexahedron
    (nels, nn) = mesh_size(fin_el, struc_el.nxe, struc_el.nye, struc_el.nze)
  else
    println("$(typeof(fin_el)) is not a known finite element.")
    return
  end

  nodof = Int(ndim == 2 ? 3 : 6)      # Degrees of freedom per node
  ndof = fin_el.nod * nodof             # Degrees of freedom per fin_el

  # Update penalty if specified in FEdict

  penalty = 1e20
  if :penalty in keys(data)
    penalty = data[:penalty]
  end

  # Allocate all arrays

  # Start with arrays to be initialized from FEdict

  if :properties in keys(data)
    prop = zeros(size(data[:properties], 1), size(data[:properties], 2))
    for i in 1:size(data[:properties], 1)
      prop[i, :] = data[:properties][i, :]
    end
  else
    println("No :properties key found in FEdict")
  end

  nf = ones(Int, nodof, nn)

  if :support in keys(data)
    for i in 1:size(data[:support], 1)
      nf[:, data[:support][i][1]] = data[:support][i][2]
    end
  end

  x_coords = zeros(nn)
  if :x_coords in keys(data)
    x_coords = data[:x_coords]
  end

  y_coords = zeros(nn)
  if :y_coords in keys(data)
    y_coords = data[:y_coords]
  end

  z_coords = zeros(nn)
  if :z_coords in keys(data)
    z_coords = data[:z_coords]
  end

  etype = ones(Int, nels)
  if :etype in keys(data)
    etype = data[:etype]
  end

  g_num = zeros(Int, fin_el.nod, nels)
  if :g_num in keys(data)
    g_num = data[:g_num]
  end

  gamma = zeros(nels)
  if :gamma in keys(data)
    gamma = data[:gamma]
  end

  # All other arrays

  points = zeros(struc_el.nip, ndim)
  g = zeros(Int, ndof)
  g_coord = zeros(ndim,nn)
  fun = zeros(fin_el.nod)
  coord = zeros(fin_el.nod, ndim)
  jac = zeros(ndim, ndim)
  der = zeros(ndim, fin_el.nod)
  deriv = zeros(ndim, fin_el.nod)
  bee = zeros(nst,ndof)
  km = zeros(ndof, ndof)
  mm = zeros(ndof, ndof)
  gm = zeros(ndof, ndof)
  kg = zeros(ndof, ndof)
  eld = zeros(ndof)
  weights = zeros(struc_el.nip)
  g_g = zeros(Int, ndof, nels)
  num = zeros(Int, fin_el.nod)
  actions = zeros(ndof, nels)
  displacements = zeros(size(nf, 1), ndim)
  gc = ones(ndim, ndim)
  dee = zeros(nst,nst)
  sigma = zeros(nst)
  axial = zeros(nels)

  # Set global coordinates

  g_coord[1,:] = data[:x_coords]
  if ndim > 1
    g_coord[2,:] = data[:y_coords]
  end
  if ndim > 2
    g_coord[3,:] = data[:z_coords]
  end

  formnf!(nodof, nn, nf)
  neq = maximum(nf)

  ell = zeros(nels)
  if :x_coords in keys(data)
    for i in 1:length(data[:x_coords])-1
      ell[i] = data[:x_coords][i+1] - data[:x_coords][i]
    end
  end

  for i in 1:nels
    num = g_num[:, i]
    num_to_g!(num, nf, g)
    g_g[:, i] = g
  end

  println("There are $(neq) equations.")

  gsm = spzeros(neq, neq)
  for i in 1:nels
    num = g_num[:, i]
    coord = g_coord[:, num]'
    km = rigid_jointed!(km, prop, gamma, etype, i, coord)
    g = g_g[:, i]
    fsparm!(gsm, g, km)
  end

  loads = OffsetArray(zeros(neq+1), 0:neq)
  if :loaded_nodes in keys(data)
    for i in 1:size(data[:loaded_nodes], 1)
      loads[nf[:, data[:loaded_nodes][i][1]]] = data[:loaded_nodes][i][2]
    end
  end

  fixed_freedoms = 0
  if :fixed_freedoms in keys(data)
    fixed_freedoms = size(data[:fixed_freedoms], 1)
  end
  no = zeros(Int, fixed_freedoms)
  node = zeros(Int, fixed_freedoms)
  sense = zeros(Int, fixed_freedoms)
  value = zeros(Float64, fixed_freedoms)
  if fixed_freedoms > 0
    for i in 1:fixed_freedoms
      node[i] = data[:fixed_freedoms][i][1]
      sense[i] = data[:fixed_freedoms][i][2]
      no[i] = nf[sense[i], node[i]]
      value[i] = data[:fixed_freedoms][i][3]
      gsm[no[i], no[i]] += penalty
      loads[no[i]] = gsm[no[i], no[i]] .* value[i]
    end
  end

  # Compute Cholesky factored global stiffness matrix for
  # future re-use. If re-use is not appropriate the loads
  # can be computed directly using gsm:
  #   loads[1:neq] = gsm \ loads[1:neq]

  cfgsm = cholesky(gsm)
  loads[1:neq] = cfgsm \ loads[1:neq]
  println()

  displacements = zeros(size(nf))
  for i in 1:size(displacements, 1)
    for j in 1:size(displacements, 2)
      if nf[i, j] > 0
        displacements[i,j] = loads[nf[i, j]]
      end
    end
  end

  loads[0] = 0.0
  for i in 1:nels
    num = g_num[:, i]
    coord = g_coord[:, num]'
    g = g_g[:, i]
    eld = loads[g]
    km = rigid_jointed!(km, prop, gamma, etype, i, coord)
    actions[:, i] = km * eld
  end

  if ndim == 2
    dis_df = DataFrame(
      x_translation = displacements[1, :],
      y_translation = displacements[2, :],
      rotation = displacements[3, :]
    )
    fm_df = DataFrame(
      x1_Force = actions[1, :],
      y1_Force = actions[2, :],
      z1_Moment = actions[3, :],
      x2_Force = actions[4, :],
      y2_Force = actions[5, :],
      z2_Moment = actions[6, :]
    )
  elseif ndim == 3
    dis_df = DataFrame(
      x_translation = displacements[1, :],
      y_translation = displacements[2, :],
      z_translation = displacements[3, :],
      x_rotation = displacements[4, :],
      y_rotation = displacements[5, :],
      z_rotation = displacements[6, :]
    )
    fm_df = DataFrame(
      x1_Force = actions[1, :],
      y1_Force = actions[2, :],
      z1_Force = actions[3, :],
      x1_Moment = actions[4, :],
      y1_Moment = actions[5, :],
      z1_Moment = actions[6, :],
      x2_Force = actions[7, :],
      y2_Force = actions[8, :],
      z2_Force = actions[9, :],
      x2_Moment = actions[10, :],
      y2_Moment = actions[11, :],
      z2_Moment = actions[12, :]
    )
  end

  # Correct element forces and moments for equivalent nodal
  # forces and moments introduced for loading between nodes
  if :eq_nodal_forces_and_moments in keys(data)
    k = data[:struc_el].fin_el.nod * data[:struc_el].fin_el.nodof
    eqfm = data[:eq_nodal_forces_and_moments]
    for t in eqfm
      vals = permutedims(Vector(fm_df[t[1], :]))
      for i in 1:k
        fm_df[t[1], i] = round.(vals[i] - t[2][i], digits=2)
      end
    end
  end

  fem = jFEM(struc_el, fin_el, ndim, nels, nst, ndof, nn, nodof, neq, penalty,
    etype, g, g_g, g_num, nf, no, node, num, sense, actions,
    bee, coord, gamma, dee, der, deriv, displacements, eld, fun, gc,
    g_coord, jac, km, mm, gm, cfgsm, loads, points, prop, sigma, value,
    weights, x_coords, y_coords, z_coords, axial)

  (fem, dis_df, fm_df)
end

max_stars_repo_path: src/4 Static Equilibrium/p44.jl

max_issues_repo_path: src/4 Static Equilibrium/p44.jl

max_forks_repo_path: src/4 Static Equilibrium/p44.jl


text:
const callsites = Dict{Int, Set{Symbol}}() # Optimizer id => jitted function names
const freshcallsites = Dict{Int, Set{Symbol}}() # Call sites found during the current round, not yet jitted

function get_callsites!(optimizerid)
    return get!(callsites, optimizerid) do
        Set{Symbol}()
    end
end

function get_freshcallsites!(optimizerid)
    return get!(freshcallsites, optimizerid) do
        Set{Symbol}()
    end
end

function exploreexpr(calledfn, argname)
    fnstr = string(calledfn) # Poor man's quote
    argnamestr = string(argname)
    return quote
        Catwalk.log_callsite(Catwalk.explorer(jitctx), Symbol($fnstr), Symbol($argnamestr))
    end
end

abstract type Explorer end

struct NoExplorer{TOptimizerId} <: Explorer end # Allows only manual registering
function NoExplorer(optimizerid)
    get_callsites!(optimizerid)
    return NoExplorer{optimizerid}()
end

struct BasicExplorer{TOptimizerId} <: Explorer end
function BasicExplorer(optimizerid)
    get_callsites!(optimizerid)
    return BasicExplorer{optimizerid}()
end

optimizerid(::Type{BasicExplorer{TOptimizerId}}) where TOptimizerId = TOptimizerId
optimizerid(::Type{NoExplorer{TOptimizerId}}) where TOptimizerId = TOptimizerId

function register_callsite!(explorer, fnsym)
    push!(callsites[optimizerid(typeof(explorer))], fnsym)
    return nothing
end

log_callsite(::Type{NoExplorer{T}}, calledfn, argname) where T = true

function log_callsite(explorer::Type{BasicExplorer{TOptimizerId}}, calledfn, argname) where TOptimizerId
    calledfns = callsites[TOptimizerId]
    if calledfn in calledfns
        return true
    else
        freshfns = get_freshcallsites!(TOptimizerId)
        if !(calledfn in freshfns)
            @debug "Found call site `$calledfn`, optimizing argument `$argname`."
            push!(freshfns, calledfn)
        end
    end
    return false
end

step!(::NoExplorer) = nothing

function step!(explorer::BasicExplorer{TOptimizerId}) where TOptimizerId
    empty!(freshcallsites[TOptimizerId])
end


text:
"""
Calculate Jacobians for elements

    rs_jacobi(cells, points)

Isosceles right triangle element

X = [x, y] = λ¹V¹ + λ²V² + λ³V³

λs are linear:
λ¹ =  -(r+s)/2
λ² = (r+1)/2
λ³ = (s+1)/2

Jacobian:
[xr xs
 yr ys]

Xᵣ = -V¹/2 + V²/2
Xₛ = -V¹/2 + V³/2

"""
function rs_jacobi(cells, points)
    ncell = size(cells, 1)
    J = [
        begin
            xr, yr = (points[cells[i, 2], 1:2] - points[cells[i, 1], 1:2]) ./ 2
            xs, ys = (points[cells[i, 3], 1:2] - points[cells[i, 1], 1:2]) ./ 2
            [xr xs; yr ys]
        end for i = 1:ncell
    ]

    return J
end

"""
    rs_jacobi(r, s, vertices)

Quadrilateral element

```
4       3
|-------|
|       |
|       |
|-------|
1       2
```

X = λ¹V¹ + λ²V² + λ³V³ + λ⁴V⁴

λs are bilinear rectangle shape functions:
λ¹ = (r-1)(s-1)/4
λ² = (r+1)(1-s)/4
λ³ = (r+1)(s+1)/4
λ⁴ = (1-r)(s+1)/4

Jacobian:
Xᵣ = (s-1)V¹/4 + (1-s)V²/4 + (s+1)V³/4 - (s+1)V⁴/4
Xₛ = (r-1)V¹/4 - (r+1)V²/4 + (r+1)V³/4 + (1-r)V⁴/4

Unlike linear simplex elements,
J varies from point to point within an element for a general linear quadrilateral.
As a special case, the Jacobian matrix is a constant for each element in rectangular mesh.

"""
function rs_jacobi(r::T, s::T, vertices::T1) where {T<:Real,T1<:AbstractMatrix}
    xr, yr = @. (s - 1.0) * vertices[1, :] / 4 +
       (1.0 - s) * vertices[2, :] / 4 +
       (s + 1.0) * vertices[3, :] / 4 - (s + 1.0) * vertices[4, :] / 4
    xs, ys = @. (r - 1.0) * vertices[1, :] / 4 - (r + 1.0) * vertices[2, :] / 4 +
       (r + 1.0) * vertices[3, :] / 4 +
       (1.0 - r) * vertices[4, :] / 4

    J = [xr xs; yr ys]

    return J
end

rs_jacobi(r::T, s::T, vertices::T1) where {T<:AbstractVector,T1<:AbstractMatrix} =
    [rs_jacobi(r[i], s[i], vertices) for i in eachindex(r)]

rs_jacobi(r::T, s::T, vertices::T1) where {T<:AbstractMatrix,T1<:AbstractMatrix} =
    [rs_jacobi(r[i], s[i], vertices) for i in axes(r, 1), j in axes(s, 2)]

rs_jacobi(r, s, vertices::AbstractArray{T,4}) where {T<:AbstractFloat} = [
    rs_jacobi(r, s, @view vertices[i, j, :, :]) for i in axes(vertices, 1),
    j in axes(vertices, 2)
]

# syntax sugar for inner points with same samplings in x and y
rs_jacobi(r::T, vertices::T1) where {T<:AbstractVector,T1<:AbstractMatrix} =
    [rs_jacobi(r[i], r[j], vertices) for i in eachindex(r), j in eachindex(r)]

rs_jacobi(r, vertices::AbstractArray{T,4}) where {T<:AbstractFloat} = [
    rs_jacobi(r, @view vertices[i, j, :, :]) for i in axes(vertices, 1),
    j in axes(vertices, 2)
]


text:
"""
    ag_sum(x)
Aggregate through `+` a vector into a single scalar.
"""
ag_sum(x) = sum(x)

"""
    ag_count_positive(x)
Count the number of strictly positive elements of `x`.
"""
ag_count_positive(x) = count(y -> y > 0.0, x)

"""
    aggregation_layer()
Generate the layer of aggregations of the ICN. The operations are mutually exclusive, that is only one will be selected.
"""
function aggregation_layer()
    aggregations = LittleDict{Symbol, Function}(
        :sum => ag_sum,
        :count_positive => ag_count_positive,
    )

    return Layer(aggregations, true)
end


text:
using Plots
using DelimitedFiles
using LinearAlgebra
using ProgressMeter
using LsqFit
using KernelDensity
using Statistics
using LaTeXStrings
using StatsBase
using Measures
using JLD2
using PyCall
stats = pyimport("scipy.stats")
include("../Analysis/AlphaShapes/AlphaShapes.jl")
include("../Analysis/analysis-utils.jl")
# this is usually an "overnighter"
using ArgParse

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table! s begin
        "-T"
            arg_type = Float64
			default = 10.			# seconds of simulated data
    end
    return parse_args(s)
end

args = parse_commandline()
@info args

model_seconds = args["T"]

@info "Simulating $model_seconds seconds of model data each iteration"

"""
A wrapper to run the model and collect the trajectories
"""
function Sim(N,alpha,tau,v,mu,mur=1.,dt=1.0/900.0)
    cmd = `../CUDAABP --initial-packing-fraction 0.5 -N $N -T $model_seconds -a 1 -b 1 -alpha $alpha -tau $tau -dt $dt -mu $mu -mur $mur -v $v -Dr 2.34 --random-seed $(Int(floor(time()))) -silent 1`
    run(cmd)
    T = readdlm("trajectories.txt",',')
    #h = Int(floor(size(T,1)/2))
    d_data = T[:,1:2]
    return d_data
end

d = Sim(10,1,1,1,30,1.);

len = Int(ceil(size(d,1) / 10))

runs = 3

Ps = [50,100,200]

alpha = -1.0987074891648896
tau = 19.593664471442796
DR = 2.34
v0 = 13.19*2.
mu = 31.624014716770823

bins_density = Float64.(logspace(1e-4,6.0,100))
bins_time = 100

Data = Dict()
@showprogress for p in 1:length(Ps)
    @info Ps[p]
    pdfs_model = zeros(runs,length(bins_density))
    Ts = zeros(runs,len,Ps[p],2)
    for r in 1:runs
        D = Sim(Ps[p],alpha,tau,v0,mu,1.);
        for i in 1:size(Ts,2)
            for j in 1:size(Ts,3)
                Ts[r,i,j,:] = D[(i-1)*Ps[p]+j,:]
            end
        end
        d = readdlm("trajectories.txt",',')[:,4]
        h = Int(floor(size(d,1)/2.))
        kde = stats.gaussian_kde(d[h:end])
        pdfs_model[r,:] = kde(bins_density)
    end
    Data["N=$(Ps[p])"] = Dict("pdfs"=>pdfs_model,"Trajectories"=>Ts)
end
# write to this dir
JLD2.@save "N-pd.jld" Data

max_stars_repo_path: Model Analysis/get-model-data-N.jl

max_stars_repo_name: harveydevereux/CUDA-Whirligigs

max_issues_repo_path: Model Analysis/get-model-data-N.jl

max_issues_repo_name: harveydevereux/CUDA-Whirligigs

max_forks_repo_path: Model Analysis/get-model-data-N.jl

max_forks_repo_name: harveydevereux/CUDA-Whirligigs


text:
using LinearAlgebra;
using Printf;
using Random;
using BenchmarkTools;

("." in LOAD_PATH) ? [] : push!(LOAD_PATH, ".");

using BivMatFun;

Random.seed!(123);

function run_test()
    n = 512;

    A = randn(n, n);
    # A = A / opnorm(A);
    A = A + one(A);

    B = randn(n, n);
    # B = B / opnorm(B);
    B = B + one(B);

    C = randn(n, n);

    A = convert(Matrix{ComplexF64}, A);
    B = convert(Matrix{ComplexF64}, B);
    C = convert(Matrix{ComplexF64}, C);

    f = (x,y,i,j) -> 1 ./ (x-y) ./ abs(x+y);
    # f = (x,y,i,j) -> 1 ./ (x-y) ./ abs(x+y);
    # f = (x,y,i,j) -> x - y;

    method = BivMatFun.DiagNoHp

    # Profile.clear();
    @btime fun2m($f, $A, $B, $C, delta = $0.05, method = $method, parallel = $true);
    # @time F2, info2 = fun2m(f, A, B, C, delta = 0.05, method = method, parallel = true);
    # @btime fun2m($f, $A, $B, $C, 0.05, $method);
    # Juno.profiler();

    @btime BivMatFun.diag_fun($f, $A, $B, $C);

    # @btime BivMatFun.diag_fun($f, $A, $B, $C);
    res = norm(X - F);
    # res = 0.0;

    @printf("Residue: %e, nblocks = %d, %d\n", res, info.nblocksA, info.nblocksB);
end

run_test();

max_stars_repo_name: numpi/BivMatFun

max_issues_repo_name: numpi/BivMatFun

max_forks_repo_name: numpi/BivMatFun


text:
import JSON
function getConfig()
    open("../config.json","r") do io
        dicttext = read(io, String)
        config = JSON.parse(dicttext)
        close(io)
        return config
    end
end

# function save_obj(obj, name)
#     stringdata = JSON.json(obj)
#     open(name + ".json", "w") do io
#         write(io, stringdata)
#     end
# end

# def load_obj(name):
#     with open(name + '.json', 'r') as f:
#         return json.load(f)

max_stars_repo_name: TestSubjector/HybridMeshfreeSolver

max_issues_repo_name: TestSubjector/HybridMeshfreeSolver

max_forks_repo_name: TestSubjector/HybridMeshfreeSolver


text:
export bitarray, basis, packbits, bfloat, bfloat_r, bint, bint_r, flip
export anyone, allone, bmask, baddrs, readbit, setbit, controller
export swapbits, ismatch, neg, breflect, btruncate
export rand_bitstr64

"""
    bitarray(v::Vector, [nbits::Int]) -> BitArray
    bitarray(v::Int, nbits::Int) -> BitArray
    bitarray(nbits::Int) -> Function

Construct BitArray from an integer vector, if nbits not supplied, it is 64.
If an integer is supplied, it returns a function mapping a Vector/Int to bitarray.
"""
function bitarray(v::Vector{T}, nbits::Int)::BitArray{2} where {T<:Number}
    ba = BitArray(undef, 0, 0)
    ba.len = 64 * length(v)
    ba.chunks = UInt64.(v)
    ba.dims = (64, length(v))
    view(ba, 1:nbits, :)
end

function bitarray(v::Vector{T})::BitArray{2} where {T<:Union{UInt64,Int64}}
    ba = BitArray(undef, 0, 0)
    ba.len = 64 * length(v)
    ba.chunks = reinterpret(UInt64, v)
    ba.dims = (64, length(v))
    ba
end

bitarray(v::Number, nbits::Int)::BitArray{1} = vec(bitarray([v], nbits))
bitarray(nbits::Int) = x -> bitarray(x, nbits)

"""
    basis([IntType], nbits::Int) -> UnitRange{IntType}
    basis([IntType], state::AbstractArray) -> UnitRange{IntType}

Returns the UnitRange for basis in Hilbert Space of nbits qubits.
If an array is supplied, it will return a basis having the same size
with the first diemension of array.
"""
basis(arg::Union{Int,AbstractArray}) = basis(Int, arg)
basis(::Type{T}, nbits::Int) where {T<:Integer} = UnitRange{T}(0, 1 << nbits - 1)
basis(::Type{T}, state::AbstractArray) where {T<:Integer} = UnitRange{T}(0, size(state, 1) - 1)

"""
    packbits(arr::AbstractArray) -> AbstractArray

pack bits to integers, usually take a BitArray as input.
"""
packbits(arr::AbstractVector) = _packbits(arr)[]
packbits(arr::AbstractArray) = _packbits(arr)
_packbits(arr) =
    selectdim(sum(mapslices(x -> x .* (1 .<< (0:size(arr, 1)-1)), arr, dims = 1), dims = 1), 1, 1)

"""
    btruncate(b, n)

Truncate bits `b` to given length `n`.
"""
btruncate(b::Integer, n) = b & (1 << n - 1)

"""
    bfloat(b::Integer; nbits::Int=bit_length(b)) -> Float64

float view, with current bit numbering.
See also [`bfloat_r`](@ref).

Ref: [wiki: bit numbering](https://en.wikipedia.org/wiki/Bit_numbering)
"""
bfloat(b::Integer; nbits::Int = bit_length(b)) = breflect(b; nbits = nbits) / (1 << nbits)

"""
    bfloat_r(b::Integer; nbits::Int=bit_length(b)) -> Float64

float view, with reversed bit numbering. See also [`bfloat`](@ref).
"""
bfloat_r(b::Integer; nbits::Int) = b / (1 << nbits)

"""
    bint(b; nbits=nothing) -> Int

integer view, with LSB 0 bit numbering.
See also [wiki: bit numbering](https://en.wikipedia.org/wiki/Bit_numbering)
"""
bint(b::Integer; nbits = nothing) = b
bint(x::Float64; nbits::Int) = breflect(bint_r(x, nbits = nbits); nbits = nbits)

"""
    bint_r(b; nbits::Int) -> Integer

integer read in inverse order.
"""
bint_r(b::Integer; nbits::Int) = breflect(b; nbits = nbits)
bint_r(x::Float64; nbits::Int) = Int(round(x * (1 << nbits)))

"""
    bmask(::Type{T}) where T <: Integer -> zero(T)
    bmask([T::Type], positions::Int...) -> T
    bmask([T::Type], range::UnitRange{Int}) -> T

Return an integer mask of type `T` where `1` is the position masked according to
`positions` or `range`. Directly use `T` will return an empty mask `0`.
"""
function bmask end

bmask(args...) = bmask(Int, args...)
bmask(::Type{T}) where {T<:Integer} = zero(T)
bmask(::Type{T}, positions::Int...) where {T<:Integer} = bmask(T, positions)
bmask(::Type{T}, itr) where {T<:Integer} =
    isempty(itr) ? 0 : reduce(+, one(T) << (b - 1) for b in itr)

@inline function bmask(::Type{T}, range::UnitRange{Int})::T where {T<:Integer}
    ((one(T) << (range.stop - range.start + 1)) - one(T)) << (range.start - 1)
end

# NOTE: we have to return a vector here since this is calculated in runtime

"""
    baddrs(b::Integer) -> Vector

get the locations of nonzeros bits, i.e. the inverse operation of bmask.
"""
function baddrs(b::Integer)
    locs = Vector{Int}(undef, count_ones(b))
    k = 1
    for i in 1:bit_length(b)
        if readbit(b, i) == 1
            locs[k] = i
            k += 1
        end
    end
    return locs
end

"""
    readbit(x, loc...)

Read the bit config at given location.
"""
readbit(x::T, loc::Int) where {T<:Integer} = (x >> (loc - 1)) & one(T)

@inline function readbit(x::T, bits::Int...) where {T<:Integer}
    res = zero(T)
    for (i, loc) in enumerate(bits)
        res += readbit(x, loc) << (i - 1)
    end
    return res
end

"""
    anyone(index::Integer, mask::Integer) -> Bool

Return `true` if any masked position of index is 1.

# Example

`true` if any masked positions is 1.

```jldoctest
julia> anyone(0b1011, 0b1001)
true

julia> anyone(0b1011, 0b1100)
true

julia> anyone(0b1011, 0b0100)
false
```
"""
anyone(index::T, mask::T) where {T<:Integer} = (index & mask) != zero(T)

"""
    allone(index::Integer, mask::Integer) -> Bool

Return `true` if all masked position of index is 1.

# Example

`true` if all masked positions are 1.

```jldoctest
julia> allone(0b1011, 0b1011)
true

julia> allone(0b1011, 0b1001)
true

julia> allone(0b1011, 0b0100)
false
```
"""
allone(index::T, mask::T) where {T<:Integer} = (index & mask) == mask

"""
    ismatch(index::Integer, mask::Integer, target::Integer) -> Bool

Return `true` if bits at positions masked by `mask` equal to `1` are equal to `target`.

## Example

```julia
julia> n = 0b11001; mask = 0b10100; target = 0b10000;

julia> ismatch(n, mask, target)
true
```
"""
ismatch(index::T, mask::T, target::T) where {T<:Integer} = (index & mask) == target

"""
    setbit(index::Integer, mask::Integer) -> Integer

set the bit at masked position to 1.

# Example

```jldoctest
julia> setbit(0b1011, 0b1100) |> bit(len=4)
1111 (15)

julia> setbit(0b1011, 0b0100) |> bit(len=4)
1111 (15)

julia> setbit(0b1011, 0b0000) |> bit(len=4)
1011 (11)
```
"""
setbit(index::T, mask::T) where {T<:Integer} = index | mask

"""
    flip(index::Integer, mask::Integer) -> Integer

Return an Integer with bits at masked position flipped.

# Example

```jldoctest
julia> flip(0b1011, 0b1011) |> bit(len=4)
0000 (0)
```
"""
flip(index::T, mask::T) where {T<:Integer} = index ⊻ mask

"""
    neg(index::Integer, nbits::Int) -> Integer

Return an integer with all bits flipped (with total number of bit `nbits`).

# Example

```jldoctest
julia> neg(0b1111, 4) |> bit(len=4)
0000 (0)

julia> neg(0b0111, 4) |> bit(len=4)
1000 (8)
```
"""
neg(index::T, nbits::Int) where {T<:Integer} = bmask(T, 1:nbits) ⊻ index

"""
    swapbits(n::Integer, mask_ij::Integer) -> Integer
    swapbits(n::Integer, i::Int, j::Int) -> Integer

Return an integer with bits at `i` and `j` flipped.

# Example

```jldoctest
julia> swapbits(0b1011, 0b1100) == 0b0111
true
```

!!! tip

    locations `i` and `j` specified by mask could be faster when [`bmask`](@ref)
    is not straight forward but known by constant.

!!! warning

    `mask_ij` should only contain two `1`, `swapbits` will not check it, use at
    your own risk.
"""
swapbits(b::T, i::Int, j::Int) where {T<:Integer} = swapbits(b, bmask(T, i, j))

@inline function swapbits(b::T, mask::T) where {T<:Integer}
    bm = b & mask
    if bm != 0 && bm != mask
        b ⊻= mask
    end
    return b
end

"""
    breflect(b::Integer[, masks::Vector{Integer}]; nbits) -> Integer

Return left-right reflected integer.

# Example

Reflect the order of bits.

```jldoctest
julia> breflect(0b1011; nbits=4) == 0b1101
true
```
"""
function breflect end

@inline function breflect(b::Integer; nbits::Int)
    @simd for i in 1:nbits÷2
        b = swapbits(b, i, nbits - i + 1)
    end
    return b
end

@inline function breflect(b::T, masks::AbstractVector{T}; nbits::Int)::T where {T<:Integer}
    @simd for m in masks
        b = swapbits(b, m)
    end
    return b
end

# TODO: use a more general trait if possible
const IntIterator{T} = Union{NTuple{<:Any,T},Vector{T},T,UnitRange{T}} where {T<:Integer}

"""
    controller(cbits, cvals) -> Function

Return a function that checks whether a basis at `cbits` takes specific value `cvals`.
"""
function controller(cbits::IntIterator{Int}, cvals::IntIterator{Int})
    do_mask = bmask(cbits)
    target =
        length(cvals) == 0 ? 0 :
        mapreduce(xy -> (xy[2] == 1 ? 1 << (xy[1] - 1) : 0), |, zip(cbits, cvals))
    return b -> ismatch(b, do_mask, target)
end

function rand_bitstr64(nbit::Int)
    T = BitStr64{nbit}
    rand(typemin(T):typemax(T))
end
Base.rand(::Type{T}) where {T<:BitStr} = rand(typemin(T):typemax(T))


text:
using Mamba,Stan,Distributions,Gadfly,DataFrames,Compat

# Note: This example, due to Chris Fisher, only runs on the latest development
#       version of CmdStan & Stan, i.e. the version post (Cmd)Stan 2.14.0.
#       Daniel (Lee, from the Stan team) has confirmed this.

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  #############################
  #    Load Stan Model        #
  #############################

  include("ModelFunctions.jl")

  Nchains = 4
  subjParm = ["theta"]
  hyperParm = ["kappa", "omega"]
  subjVal = Any[[.5,3]]
  hyperVal = [.5,3]

  Nsubj = 10
  Ntrials = 10
  Nobs = Nsubj*Ntrials
  kappa = .7
  omega = 5

  y,SubjIdx = GenerateData(kappa, omega, Nsubj, Ntrials)

  global init, stanmodel, rc, sim
  initpars = InitializeChain(Nsubj, Nchains, subjParm, hyperParm, subjVal, hyperVal)

  observeddata = [
      Dict(
        "Nobs" => Nobs,
        "Nsubj" => Nsubj,
        "SubjIdx"=>SubjIdx,
        "y" => y),
  ]

  ###########################
  #     Run Chains          #
  ###########################

  f = open("model.stan","r")
  TestModel = read(f, String)
  close(f)

  stanmodel = Stanmodel(
    Sample(num_samples = 1000, num_warmup = 200, adapt = Stan.Adapt(delta=0.88)),
    nchains=Nchains, num_samples=1000, thin=1, name="TestModel", model=TestModel
  )

  @time rc, sim = stan(stanmodel, observeddata, ProjDir, init=initpars,
    diagnostics=false)
  rc == 0 && describe(sim)
end


text:
module Mandelbrot
export mandelbrot_image, mandelbrot

using Base.Threads

function mandelbrot(c::Complex, max_iter::Int)
    z = 0.0 + 0.0im
    for i in 1:max_iter
        z = z^2 + c
        if abs(z) > 2
            return max_iter - i
        end
    end
    return 0
end

function mandelbrot_image(width, height, max_iter = 1000; rmin = -2, rmax = 1, imin = -1, imax = 1)
    A = [
        Complex((j/width * (rmax - rmin) + rmin), (i/height * (imax - imin) + imin))
        for i in 1:height, j in 1:width
    ]

    set = Array{Int}(undef, height, width)

    @time @inbounds @threads for i in 1:height
        for j in 1:width
            set[i, j] = mandelbrot(A[i, j], max_iter)
        end
    end

    return set
end

end